[mlir] Add global_load and global_store ops to ml_program.

* Adds simple, non-atomic, non-volatile, non-synchronized direct load/store ops.

Differential Revision: https://reviews.llvm.org/D126230

Author: stellaraccident

mlir/include/mlir/Dialect/MLProgram/IR/CMakeLists.txt

@@ -8,3 +8,10 @@ mlir_tablegen(MLProgramAttributes.cpp.inc -gen-attrdef-defs)
 add_public_tablegen_target(MLIRMLProgramAttributesIncGen)
 add_dependencies(mlir-headers MLIRMLProgramAttributesIncGen)
 add_mlir_doc(MLProgramAttributes MLProgramAttributes Dialects/ -gen-attrdef-doc)
+
+set(LLVM_TARGET_DEFINITIONS MLProgramTypes.td)
+mlir_tablegen(MLProgramTypes.h.inc -gen-typedef-decls)
+mlir_tablegen(MLProgramTypes.cpp.inc -gen-typedef-defs)
+add_public_tablegen_target(MLIRMLProgramTypesIncGen)
+add_dependencies(mlir-headers MLIRMLProgramTypesIncGen)
+add_mlir_doc(MLProgramTypes MLProgramTypes Dialects/ -gen-typedef-doc)

mlir/include/mlir/Dialect/MLProgram/IR/MLProgram.h

@@ -9,6 +9,7 @@
 #define MLIR_DIALECT_MLPROGRAM_IR_MLPROGRAM_H_
 
 #include "mlir/Dialect/MLProgram/IR/MLProgramAttributes.h"
+#include "mlir/Dialect/MLProgram/IR/MLProgramTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/OpDefinition.h"

mlir/include/mlir/Dialect/MLProgram/IR/MLProgramBase.td

@@ -28,6 +28,7 @@ def MLProgram_Dialect : Dialect {
   }];
 
   let useDefaultAttributePrinterParser = 1;
+  let useDefaultTypePrinterParser = 1;
   let emitAccessorPrefix = kEmitAccessorPrefix_Prefixed;
 }
 

mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td

@@ -10,6 +10,7 @@
 #define MLPROGRAM_OPS
 
 include "mlir/Dialect/MLProgram/IR/MLProgramBase.td"
+include "mlir/Dialect/MLProgram/IR/MLProgramTypes.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
@@ -152,6 +153,51 @@ def MLProgram_GlobalOp : MLProgram_Op<"global", [
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// GlobalLoadOp
+//===----------------------------------------------------------------------===//
+
+def MLProgram_GlobalLoadOp : MLProgram_Op<"global_load", [
+    DeclareOpInterfaceMethods<SymbolUserOpInterface>
+  ]> {
+  let summary = "Direct load of a mutable value from a global";
+  let description = [{
+    Performs a non-atomic, non-volatile, non-synchronized load from a global
+    that may be mutable.
+
+    It is fully expected that these constraints are not suitable for
+    all situations, and alternative ops should be defined and used for more
+    advanced cases.
+
+    This op is side effecting and may not be valid to use in graph regions
+    without additional consideration to evaluation order constraints.
+
+    Example:
+
+    ```mlir
+    %0 = ml_program.global_load @foobar : tensor<?xi32>
+    ```
+  }];
+
+  let arguments = (ins
+    Arg<SymbolRefAttr, "", [MemRead]>:$global,
+    Variadic<MLProgram_TokenType>:$consumeTokens
+  );
+  let results = (outs
+    AnyType:$result,
+    Optional<MLProgram_TokenType>:$produceToken
+  );
+
+  let assemblyFormat = [{
+    $global `` custom<TokenOrdering>($consumeTokens, type($produceToken)) `:` type($result) attr-dict
+  }];
+
+  let extraClassDeclaration = [{
+    /// Gets the corresponding GlobalOp (or nullptr).
+    GlobalOp getGlobalOp(SymbolTableCollection &symbolTable);
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // GlobalLoadConstOp
 //===----------------------------------------------------------------------===//
@@ -175,14 +221,59 @@ def MLProgram_GlobalLoadConstOp : MLProgram_Op<"global_load_const", [
   }];
 
   let arguments = (ins
-    FlatSymbolRefAttr:$global
+    SymbolRefAttr:$global
   );
   let results = (outs
     AnyType:$result
   );
 
   let assemblyFormat = [{
-    $global attr-dict `:` type($result)
+    $global `:` type($result) attr-dict
+  }];
+
+  let extraClassDeclaration = [{
+    /// Gets the corresponding GlobalOp (or nullptr).
+    GlobalOp getGlobalOp(SymbolTableCollection &symbolTable);
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// GlobalStoreOp
+//===----------------------------------------------------------------------===//
+
+def MLProgram_GlobalStoreOp : MLProgram_Op<"global_store", [
+    DeclareOpInterfaceMethods<SymbolUserOpInterface>
+  ]> {
+  let summary = "Direct store of a value into a mutable global";
+  let description = [{
+    Performs a non-atomic, non-volatile, non-synchronized store to a mutable
+    global.
+
+    It is fully expected that these constraints are not suitable for
+    all situations, and alternative ops should be defined and used for more
+    advanced cases.
+
+    This op is side effecting and may not be valid to use in graph regions
+    without additional consideration to evaluation order constraints.
+
+    Example:
+
+    ```mlir
+    ml_program.global_store @foobar = %0 : tensor<?xi32>
+    ```
+  }];
+
+  let arguments = (ins
+    Arg<SymbolRefAttr, "", [MemWrite]>:$global,
+    AnyType:$value,
+    Variadic<MLProgram_TokenType>:$consumeTokens
+  );
+  let results = (outs
+    Optional<MLProgram_TokenType>:$produceToken
+  );
+
+  let assemblyFormat = [{
+    $global `=` $value `` custom<TokenOrdering>($consumeTokens, type($produceToken)) `:` type($value) attr-dict
   }];
 
   let extraClassDeclaration = [{
@@ -310,4 +401,24 @@ def MLProgram_ReturnOp : MLProgram_Op<"return", [
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// TokenOp
+//===----------------------------------------------------------------------===//
+
+def MLProgram_TokenOp : MLProgram_Op<"token", [
+    NoSideEffect
+  ]> {
+  let summary = "Produces a new token value";
+  let description = [{
+    Token values are used to chain side effecting ops in a graph so as to
+    establish an execution order. This op produces a token.
+  }];
+
+  let results = (outs
+    MLProgram_TokenType:$token
+  );
+
+  let assemblyFormat = "attr-dict";
+}
+
 #endif // MLPROGRAM_OPS

mlir/include/mlir/Dialect/MLProgram/IR/MLProgramTypes.h

@@ -0,0 +1,21 @@
+//===- MLProgramTypes.h - Type Classes --------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_MLPROGRAM_IR_MLPROGRAMTYPES_H_
+#define MLIR_DIALECT_MLPROGRAM_IR_MLPROGRAMTYPES_H_
+
+#include "mlir/IR/Types.h"
+
+//===----------------------------------------------------------------------===//
+// Tablegen Type Declarations
+//===----------------------------------------------------------------------===//
+
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/MLProgram/IR/MLProgramTypes.h.inc"
+
+#endif // MLIR_DIALECT_MLPROGRAM_IR_MLPROGRAMTYPES_H_

mlir/include/mlir/Dialect/MLProgram/IR/MLProgramTypes.td

@@ -0,0 +1,24 @@
+//===- MLProgramTypes.td - Type definitions ----------------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLPROGRAM_TYPES
+#define MLPROGRAM_TYPES
+
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/Dialect/MLProgram/IR/MLProgramBase.td"
+
+class MLProgram_Type<string name, list<Trait> traits = [],
+                     string baseCppClass = "::mlir::Type">
+    : TypeDef<MLProgram_Dialect, name, traits, baseCppClass> {}
+
+def MLProgram_TokenType : MLProgram_Type<"Token"> {
+  let summary = "Token for establishing execution ordering in a graph";
+  let mnemonic = "token";
+}
+
+#endif // MLPROGRAM_TYPES

mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt

@@ -8,6 +8,7 @@ add_mlir_dialect_library(MLIRMLProgram
   DEPENDS
   MLIRMLProgramOpsIncGen
   MLIRMLProgramAttributesIncGen
+  MLIRMLProgramTypesIncGen
 
   LINK_LIBS PUBLIC
   MLIRDialect

mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp

@@ -20,6 +20,8 @@ using namespace mlir::ml_program;
 #include "mlir/Dialect/MLProgram/IR/MLProgramOpsDialect.cpp.inc"
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/MLProgram/IR/MLProgramAttributes.cpp.inc"
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/MLProgram/IR/MLProgramTypes.cpp.inc"
 
 namespace {
 struct MLProgramOpAsmDialectInterface : public OpAsmDialectInterface {
@@ -40,9 +42,16 @@ void ml_program::MLProgramDialect::initialize() {
   addAttributes<
 #include "mlir/Dialect/MLProgram/IR/MLProgramAttributes.cpp.inc"
       >();
+
+#define GET_TYPEDEF_LIST
+  addTypes<
+#include "mlir/Dialect/MLProgram/IR/MLProgramTypes.cpp.inc"
+      >();
+
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/MLProgram/IR/MLProgramOps.cpp.inc"
       >();
+
   addInterfaces<MLProgramOpAsmDialectInterface>();
 }

mlir/lib/Dialect/MLProgram/IR/MLProgramOps.cpp

@@ -17,6 +17,64 @@ using namespace mlir::ml_program;
 // Custom asm helpers
 //===----------------------------------------------------------------------===//
 
+/// Parse and print an ordering clause for a variadic of consuming tokens
+/// and an optional producing token.
+///
+/// Syntax:
+///   ordering(%0, %1 -> !ml_program.token)
+///   ordering(() -> !ml_program.token)
+///   ordering(%0, %1)
+///
+/// If both the consuming and producing token are not present on the op, then
+/// the clause prints nothing.
+static ParseResult parseTokenOrdering(
+    OpAsmParser &parser,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &consumeTokens,
+    Type &produceTokenType) {
+  if (failed(parser.parseOptionalKeyword("ordering")) ||
+      failed(parser.parseLParen()))
+    return success();
+
+  // Parse consuming token list. If there are no consuming tokens, the
+  // '()' null list represents this.
+  if (succeeded(parser.parseOptionalLParen())) {
+    if (failed(parser.parseRParen()))
+      return failure();
+  } else {
+    if (failed(parser.parseOperandList(consumeTokens,
+                                       /*requiredOperandCount=*/-1)))
+      return failure();
+  }
+
+  // Parse optional producer token.
+  if (succeeded(parser.parseOptionalArrow()))
+    if (failed(parser.parseType(produceTokenType)))
+      return failure();
+
+  if (failed(parser.parseRParen()))
+    return failure();
+
+  return success();
+}
+
+static void printTokenOrdering(OpAsmPrinter &p, Operation *op,
+                               OperandRange consumeTokens,
+                               Type produceTokenType) {
+  if (consumeTokens.empty() && !produceTokenType)
+    return;
+
+  p << " ordering(";
+  if (consumeTokens.empty())
+    p << "()";
+  else
+    p.printOperands(consumeTokens);
+  if (produceTokenType) {
+    p << " -> ";
+    p.printType(produceTokenType);
+  }
+  p << ")";
+}
+
 /// some.op custom<TypeOrAttr>($type, $attr)
 ///
 /// Uninitialized:
@@ -111,6 +169,30 @@ LogicalResult GlobalOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// GlobalLoadOp
+//===----------------------------------------------------------------------===//
+
+GlobalOp GlobalLoadOp::getGlobalOp(SymbolTableCollection &symbolTable) {
+  return symbolTable.lookupNearestSymbolFrom<GlobalOp>(
+      getOperation()->getParentOp(), getGlobalAttr());
+}
+
+LogicalResult
+GlobalLoadOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  GlobalOp referrent = getGlobalOp(symbolTable);
+  if (!referrent)
+    return emitOpError() << "undefined global: " << getGlobal();
+
+  if (referrent.getType() != getResult().getType()) {
+    return emitOpError() << "cannot load from global typed "
+                         << referrent.getType() << " as "
+                         << getResult().getType();
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // GlobalLoadConstOp
 //===----------------------------------------------------------------------===//
@@ -138,6 +220,35 @@ GlobalLoadConstOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// GlobalStoreOp
+//===----------------------------------------------------------------------===//
+
+GlobalOp GlobalStoreOp::getGlobalOp(SymbolTableCollection &symbolTable) {
+  return symbolTable.lookupNearestSymbolFrom<GlobalOp>(
+      getOperation()->getParentOp(), getGlobalAttr());
+}
+
+LogicalResult
+GlobalStoreOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  GlobalOp referrent = getGlobalOp(symbolTable);
+  if (!referrent)
+    return emitOpError() << "undefined global: " << getGlobal();
+
+  if (!referrent.getIsMutable()) {
+    return emitOpError() << "cannot store to an immutable global "
+                         << getGlobal();
+  }
+
+  if (referrent.getType() != getValue().getType()) {
+    return emitOpError() << "cannot store to a global typed "
+                         << referrent.getType() << " from "
+                         << getValue().getType();
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // SubgraphOp
 //===----------------------------------------------------------------------===//