[mlir] Split MLProgram global load and store to Graph variants

[jpienaar]

Summary

• Split ops into X_graph variants as discussed;
• Remove tokens from non-Graph region variants and rely on side-effect modelling there while removing side-effect modelling from Graph variants and relying on explicit ordering there;
• Make tokens required to be produced by Graph variants - but kept explicit token type specification given previous discussion on this potentially being configurable in future;

This results in duplicating some code. I considered adding helper
functions but decided against adding an abstraction there early given
size of duplication and creating accidental coupling.

[stellaraccident]

Thanks for doing this!

[mehdi_amini]

Graph variants should have exactly the same side effects.

[jpienaar]

Updated to same side-effects

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

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.
+    without additional consideration to evaluation order constraints. See
+    `global_load_graph` for op which allows for explicit ordering constraints.
 
     Example:
 
     ```mlir
     %0 = ml_program.global_load @foobar : tensor<?xi32>
     ```
   }];
 
   let arguments = (ins
-    Arg<SymbolRefAttr, "", [MemRead]>:$global,
-    Variadic<MLProgram_TokenType>:$consumeTokens
+    Arg<SymbolRefAttr, "", [MemRead]>:$global
   );
   let results = (outs
-    AnyType:$result,
-    Optional<MLProgram_TokenType>:$produceToken
+    AnyType:$result
   );
 
   let assemblyFormat = [{
-    $global `` custom<TokenOrdering>($consumeTokens, type($produceToken)) `:` type($result) attr-dict
+    $global `:` type($result) attr-dict
   }];
 
   let extraClassDeclaration = [{
     /// Gets the corresponding GlobalOp (or nullptr).
     GlobalOp getGlobalOp(SymbolTableCollection &symbolTable);
   }];
 }
 
def MLProgram_GlobalLoadConstOp : MLProgram_Op<"global_load_const", [
 
   let extraClassDeclaration = [{
     /// Gets the corresponding GlobalOp (or nullptr).
     GlobalOp getGlobalOp(SymbolTableCollection &symbolTable);
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// GlobalLoadGraphOp
+//===----------------------------------------------------------------------===//
+
+def MLProgram_GlobalLoadGraphOp : MLProgram_Op<"global_load_graph", [
+    DeclareOpInterfaceMethods<SymbolUserOpInterface>
+  ]> {
+  let summary = "Direct load of a mutable value from a global in Graph region";
+  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, %cstr = ml_program.global_load_graph @foobar
+      ordering (%token -> !ml_program.token) : tensor<?xi32>
+    ```
+  }];
+
+  let arguments = (ins
+    Arg<SymbolRefAttr, "", [MemRead]>:$global,
+    Variadic<MLProgram_TokenType>:$consumeTokens
+  );
+  let results = (outs
+    AnyType:$result,
+    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);
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // 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.
+    without additional consideration to evaluation order constraints. See
+    `global_store_graph` for op which allows for explicit ordering constraints.
 
     Example:
 
     ```mlir
     ml_program.global_store @foobar = %0 : tensor<?xi32>
     ```
   }];
 
   let arguments = (ins
     Arg<SymbolRefAttr, "", [MemWrite]>:$global,
+    AnyType:$value
+  );
+
+  let assemblyFormat = [{
+    $global `=` $value `:` type($value) attr-dict
+  }];
+
+  let extraClassDeclaration = [{
+    /// Gets the corresponding GlobalOp (or nullptr).
+    GlobalOp getGlobalOp(SymbolTableCollection &symbolTable);
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// GlobalStoreGraphOp
+//===----------------------------------------------------------------------===//
+
+def MLProgram_GlobalStoreGraphOp : MLProgram_Op<"global_store_graph", [
+    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
+    %token = ml_program.global_store @foobar = %0 : tensor<?xi32>
+      ordering (%in_token -> !ml_program.token) : tensor<?xi32>
+    ```
+  }];
+
+  let arguments = (ins
+    Arg<SymbolRefAttr, "", [MemRead]>:$global,
     AnyType:$value,
     Variadic<MLProgram_TokenType>:$consumeTokens
   );
   let results = (outs
-    Optional<MLProgram_TokenType>:$produceToken
+    MLProgram_TokenType:$produceToken
   );
 
   let assemblyFormat = [{
     $global `=` $value `` custom<TokenOrdering>($consumeTokens, type($produceToken)) `:` type($value) attr-dict
   }];
 
   let extraClassDeclaration = [{
     /// Gets the corresponding GlobalOp (or nullptr).

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

 using namespace mlir;
 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.
+/// and an 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();
+  // Parse producer token.
+  if (failed(parser.parseArrow()))
+    return failure();
 
   if (failed(parser.parseRParen()))
     return failure();
 
   return success();
 }
 
 static void printTokenOrdering(OpAsmPrinter &p, Operation *op,
if (referrent.getType() != getResult().getType())
     return emitOpError() << "cannot load from global typed "
                          << referrent.getType() << " as "
                          << getResult().getType();
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// GlobalLoadGraphOp
+//===----------------------------------------------------------------------===//
+
+GlobalOp GlobalLoadGraphOp::getGlobalOp(SymbolTableCollection &symbolTable) {
+  return symbolTable.lookupNearestSymbolFrom<GlobalOp>(
+      getOperation()->getParentOp(), getGlobalAttr());
+}
+
+LogicalResult
+GlobalLoadGraphOp::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();
+}
+
+//===----------------------------------------------------------------------===//
 // GlobalStoreOp
 //===----------------------------------------------------------------------===//
 
 GlobalOp GlobalStoreOp::getGlobalOp(SymbolTableCollection &symbolTable) {
   return symbolTable.lookupNearestSymbolFrom<GlobalOp>(
       getOperation()->getParentOp(), getGlobalAttr());
 }
 
return emitOpError() << "cannot store to a global typed "
                          << referrent.getType() << " from "
                          << getValue().getType();
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// GlobalStoreGraphOp
+//===----------------------------------------------------------------------===//
+
+GlobalOp GlobalStoreGraphOp::getGlobalOp(SymbolTableCollection &symbolTable) {
+  return symbolTable.lookupNearestSymbolFrom<GlobalOp>(
+      getOperation()->getParentOp(), getGlobalAttr());
+}
+
+LogicalResult
+GlobalStoreGraphOp::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
 //===----------------------------------------------------------------------===//
 
 ParseResult SubgraphOp::parse(OpAsmParser &parser, OperationState &result) {
   auto buildFuncType =
       [](Builder &builder, ArrayRef<Type> argTypes, ArrayRef<Type> results,
          function_interface_impl::VariadicFlag,
          std::string &) { return builder.getFunctionType(argTypes, results); };

mlir/test/Dialect/MLProgram/invalid.mlir

 
 // -----
 ml_program.global private @var(42 : i64) : i64
 ml_program.func @store_immutable(%arg0: i64) {
   // expected-error @+1 {{cannot store to an immutable global @var}}
   ml_program.global_store @var = %arg0 : i64
   ml_program.return
 }
+
+// -----
+
+ml_program.global private mutable @global_mutable_undef : tensor<?xi32>
+ml_program.subgraph @global_load_store_tokens() -> (tensor<?xi32>, !ml_program.token) {
+  %token1 = ml_program.token
+  %0, %token2 = ml_program.global_load_graph @global_mutable_undef
+      ordering(() -> !ml_program.token) : tensor<?xi32>
+  %token3 = ml_program.global_store_graph @global_mutable_undef = %0
+  // expected-error @+1 {{expected '->'}}
+      ordering(%token1, %token2) : tensor<?xi32>
+
+  ml_program.output %0, %token3 : tensor<?xi32>, !ml_program.token
+}

mlir/test/Dialect/MLProgram/ops.mlir

ml_program.func @global_load_store() {
   %0 = ml_program.global_load @global_mutable_undef : tensor<?xi32>
   ml_program.global_store @global_mutable_undef = %0 : tensor<?xi32>
   ml_program.return
 }
 
 // CHECK-LABEL: @global_load_store_tokens
 ml_program.subgraph @global_load_store_tokens() -> (tensor<?xi32>, !ml_program.token) {
   %token1 = ml_program.token
-  %0, %token2 = ml_program.global_load @global_mutable_undef
+  %0, %token2 = ml_program.global_load_graph @global_mutable_undef
       ordering(() -> !ml_program.token) : tensor<?xi32>
-  %token3 = ml_program.global_store @global_mutable_undef = %0
+  %token3 = ml_program.global_store_graph @global_mutable_undef = %0
       ordering(%token1, %token2 -> !ml_program.token) : tensor<?xi32>
-  ml_program.global_store @global_mutable_undef = %0
-      ordering(%token3) : tensor<?xi32>
+  %token4 = ml_program.global_store_graph @global_mutable_undef = %0
+      ordering(%token3 -> !ml_program.token) : tensor<?xi32>
 
   ml_program.output %0, %token3 : tensor<?xi32>, !ml_program.token
 }