[mlir] Add global_load and global_store ops to ml_program.

[stellaraccident]

Summary

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

[stellaraccident]

Note: design discussion continuing on discourse: https://discourse.llvm.org/t/rfc-introduce-ml-program-dialect-and-top-level-ops-proposal-v2/60907/51

[Mogball]

*breathes a sigh of relief that no ODS bugs were stepped on*

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

162

It is possible to specify memory effects on symbol reference resources specifically, but ODS doesn't generate that code for you. Doing so should provide the benefit of being able to reorder/CSE certain reads and writes.

184

I don't have context on this but does disallowing nested symbol references make sense for this dialect?

246

Ditto with the symbol reference resource.

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

178

add braces

229

add braces

233

add braces

[stellaraccident]

In D126230#3535583, @Mogball wrote:

*breathes a sigh of relief that no ODS bugs were stepped on*

Haha - yes, this one was much easier! Thank you for all of your help on the previous patch.

Since the conversation on design is still happening on discourse, I'll leave this until that resolves before applying your suggestions. Thank you for the review!

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

162

Oh, thank you for the tip -- I had wondered about this in the past! I can go reading, but do you have any reference to look at?

[rriddle]

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

162

https://github.com/llvm/llvm-project/blob/2a5d5078d5dadb1aaa4487c89a92c91a98e2ed12/mlir/test/mlir-tblgen/op-side-effects.td#L17

[Mogball]

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

162

Oh, ODS does the work for you. I had never seen examples in TableGen before :P

[stellaraccident]

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

162

This is nice. Thanks!

184

Not a good reason. Changed.

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

178

I am still confused on when to brace and when to not. But taking your word for it.

[stellaraccident]

Comments

[stellaraccident]

Add chains.

[jpienaar]

Nice, mostly looks good, clarification question (and sorry I jumped to the code while you may have been making update in discourse/my mind is clear after long weekend so may be asking things already discussed)

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

414

Is this a pure compilation side type (e.g., same as shape.witness) or does it have a runtime value too?

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

19

You can elide this (replace { ... } with just ; and it behaves the same).

24

You can just elide description if nothing useful :)

If this is meant to be compile time only, this could be useful part of description though.

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

24

ah token and chains ... Mmm, bike shedding: could we have this be something like "after" instead? I'm not even sure if you need to have the mapping there for the tokens produced, that is captured in the type isn't it? Oh but not in the pretty print type. So it would need to be something like

%token4, %token5 = ml_program.global_store @global_mutable_undef = %0 : tensor<?xi32> after (%token1, %token3) : (!ml_program.token, !ml_program.token)

if not captured in mapping. But what does it mean for a store to produce multiple tokens? I'm not sure what it means from an ordering point of view. Why not always make it produce one token?

%after_token = ml_program.global_store @global_mutable_undef = %0 after (%token1, %token3) : tensor<?xi32>

[stellaraccident]

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

414

I don't think we care/specify: I have seen implementations that just use this for compile time support, and I have seen implementations that transform/surface such things across boundaries to "externalize" the ordering constraint as needed.

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

24

Oh, I left it there because I meant to write the description. Sorry!

[stellaraccident]

Comments

[stellaraccident]

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

24

Ok, made the result be one optional token (I meant to do that but got turned around when undoing the earlier tablegen experiment).

I fiddled with it some more. Here are some examples:

%token1 = ml_program.token
%0, %token2 = ml_program.global_load @global_mutable_undef
    after (none -> !ml_program.token) : tensor<?xi32>
%token3 = ml_program.global_store @global_mutable_undef = %0
    after (%token1, %token2 -> !ml_program.token) : tensor<?xi32>
ml_program.global_store @global_mutable_undef = %0
    after (%token3) : tensor<?xi32>

I don't love all of these but I think they convey the meaning in the expected common cases.

Some considerations:

• It would be nice if the after clause was at the end, but this actually creates an ambiguity with a dialect called after I think :/
• I think that having the producing token be optional is right, since it makes it statically clear when this is a "chaining" op vs not. Practically, it also makes the printed form without tokens nice.
• I'm still not sold on after. Maybe chaining, ordered-by or some such thing?

[Mogball]

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

24

I'm not crazy about none. If there will be at most one token, what about

%0, %token1 = ml_program.global_load @foo after (%token1) : tensor<...>
%token2 = ml_program.global_store @bar = %0 after(%token1, %token0) : tensor<...>

39

drop braces

mlir/test/Dialect/MLProgram/ops.mlir

45

extra //

[Mogball]

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

24

My edit didn't go through... it's meant to be after (() -> !ml_program.token) or maybe some other way of indicating the optional token since the type itself isn't needed.

[stellaraccident]

Comments.

[stellaraccident]

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

24

Switched to the '()' means nil form -- have a look. I had thought about some other way to avoid spelling out the return token type (since it is buildable and we don't technically need it). However, I expect that in the not too distant future, for interop, we're going to want to widen the allowable TokenTypes (since existing graph based dialects all have one and there is little practical reason to restrict them, if we can come up with some common API).

39

Ok, I promise I'm trying and I've read this part of the style guide like 15 times -- and missing something. This case and the if directly above seem to satisfy this statement by my reading: "Maintainability is harmed if the body of an if ends with a (directly or indirectly) nested if statement with no else. Braces on the outer if would help to avoid running into a “dangling else” situation."

(ftr - I really don't care much but am trying to comply and get the feel for how this codebase wants it so you all don't need to keep leaving these comments for me :) )

[stellaraccident]

Bazel updates.

[jpienaar]

Looks good thanks!

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

24

Yes esp if we do "before -> after" then after as keyword is not accurate. I'm good with chaining, ordered-by or other words convenying it is a scheduling constraint.

Re: ambiguity of after at end, if optional & op in dialect called after without return value is a thing. Also good point on potential different return types ... One could have the return type at end (rather than operand type) and avoid ambiguity

ml_program.global_store @global_mutable_undef(%0 : tensor<?xi32>) after(%token3, %token4) : !ml_program.token

and

ml_program.global_store @global_mutable_undef(%0 : tensor<?xi32>) : None

or

ml_program.global_store @global_mutable_undef after(%token3, %token4 : !ml_program.token) = %0 : tensor<?xi32>

But I'm not going to block about any of these, yours reads well enough and we can iterate.

[stellaraccident]

Tweak based on comments.

[stellaraccident]

In D126230#3550833, @jpienaar wrote:

Looks good thanks!

I ended up trying to use an "ordered-by" keyword but parseOptionalKeyword silently will not parse (presumably because of the dash). I ended up spelling it "ordering" but am open to change the spelling. Also added some comments to the custom directive functions.

[Mogball]

Just one, otherwise braces LGTM

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

39

For me, it's been a journey of trial and getting nitted by River :P

It also occasionally changes vs. older parts of the codebase...

53

drop braces.

The rule here is something like nested indents one line each are okay, i.e.

if (foo)
  for (a in b)
    while (bar)
       if (c)
          do_something()

is okay

[stellaraccident]

Remove braces.

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

 set(LLVM_TARGET_DEFINITIONS MLProgramOps.td)
 add_mlir_dialect(MLProgramOps ml_program)
 add_mlir_doc(MLProgramOps MLProgramOps Dialects/ -gen-dialect-doc)
 
 set(LLVM_TARGET_DEFINITIONS MLProgramAttributes.td)
 mlir_tablegen(MLProgramAttributes.h.inc -gen-attrdef-decls)
 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

 //===- MLProgram.h - MLProgram dialect ----------------------------*- 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_MLPROGRAM_H_
 #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"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/RegionKindInterface.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"

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

let description = [{
     memory model components appropriate for such an abstract level of detail.
 
     This dialect is under active development, and while stability is an
     eventual goal, it is not guaranteed at this juncture. Given the early state,
     it is recommended to inquire further prior to using this dialect.
   }];
 
   let useDefaultAttributePrinterParser = 1;
+  let useDefaultTypePrinterParser = 1;
   let emitAccessorPrefix = kEmitAccessorPrefix_Prefixed;
 }
 
 #endif // MLPROGRAM_BASE

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

 //===- MLProgramOps.td - Structural ML Program Ops ---------*- 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_OPS
 #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"
 include "mlir/IR/FunctionInterfaces.td"
 include "mlir/IR/RegionKindInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 
 class MLProgram_Op<string mnemonic, list<Trait> traits = []> :
let assemblyFormat = [{
     custom<TypedInitialValue>($type, $value)
     attr-dict
   }];
 
   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 attr-dict `` custom<TokenOrdering>($consumeTokens, type($produceToken)) `:` type($result)
+  }];
+
+  let extraClassDeclaration = [{
+    /// Gets the corresponding GlobalOp (or nullptr).
+    GlobalOp getGlobalOp(SymbolTableCollection &symbolTable);
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // GlobalLoadConstOp
 //===----------------------------------------------------------------------===//
 
 def MLProgram_GlobalLoadConstOp : MLProgram_Op<"global_load_const", [
     NoSideEffect,
     DeclareOpInterfaceMethods<SymbolUserOpInterface>
   ]> {
   let summary = "Direct load a constant value from a global";
   let description = [{
     Loads a constant (immutable) value from a global directly by symbol.
 
     This op is only legal for globals that are not mutable and exists because
     such a load can be considered to have no side effects.
 
     Example:
 
     ```mlir
     %0 = ml_program.global_load_const @foobar : tensor<?xi32>
     ```
   }];
 
   let arguments = (ins
-    FlatSymbolRefAttr:$global
+    SymbolRefAttr:$global
   );
   let results = (outs
     AnyType:$result
   );
 
   let assemblyFormat = [{
     $global attr-dict `:` type($result)
   }];
 
   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 = [{
+    /// Gets the corresponding GlobalOp (or nullptr).
+    GlobalOp getGlobalOp(SymbolTableCollection &symbolTable);
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // SubgraphOp
 //===----------------------------------------------------------------------===//
 
 def MLProgram_SubgraphOp : MLProgram_Op<"subgraph", [
     CallableOpInterface, FunctionOpInterface, HasOnlyGraphRegion,
     IsolatedFromAbove, RegionKindInterface, SingleBlock, Symbol
   ]> {
   let summary = "An function containing a single `Graph` region";
def MLProgram_ReturnOp : MLProgram_Op<"return", [
   let builders = [OpBuilder<(ins), [{
     build($_builder, $_state, llvm::None);
   }]>];
 
   let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
   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

+//===- 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

+//===- 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

 add_mlir_dialect_library(MLIRMLProgram
   MLProgramOps.cpp
   MLProgramDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MLProgram
 
   DEPENDS
   MLIRMLProgramOpsIncGen
   MLIRMLProgramAttributesIncGen
+  MLIRMLProgramTypesIncGen
 
   LINK_LIBS PUBLIC
   MLIRDialect
   MLIRInferTypeOpInterface
   MLIRIR
   )

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

 
 //===----------------------------------------------------------------------===//
 /// Tablegen Definitions
 //===----------------------------------------------------------------------===//
 
 #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 {
   using OpAsmDialectInterface::OpAsmDialectInterface;
 
   AliasResult getAlias(Attribute attr, raw_ostream &os) const override {
     if (attr.isa<ExternAttr>()) {
       os << "extern";
       return AliasResult::OverridableAlias;
     }
     return AliasResult::NoAlias;
   }
 };
 } // namespace
 
 void ml_program::MLProgramDialect::initialize() {
 #define GET_ATTRDEF_LIST
   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

 
 using namespace mlir;
 using namespace mlir::ml_program;
 
 //===----------------------------------------------------------------------===//
 // Custom asm helpers
 //===----------------------------------------------------------------------===//
 
+static ParseResult parseTokenOrdering(
+    OpAsmParser &parser,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &consumeTokens,
+    Type &produceTokenType) {
+  if (failed(parser.parseOptionalKeyword("after")) ||
+      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 << " after (";
+  if (consumeTokens.empty())
+    p << "()";
+  else
+    p.printOperands(consumeTokens);
+  if (produceTokenType) {
+    p << " -> ";
+    p.printType(produceTokenType);
+  }
+  p << ")";
+}
+
 /// some.op custom<TypeOrAttr>($type, $attr)
 ///
 /// Uninitialized:
 ///   some.op : tensor<3xi32>
 /// Initialized to narrower type than op:
 ///   some.op (dense<0> : tensor<3xi32>) : tensor<?xi32>
 static ParseResult parseTypedInitialValue(OpAsmParser &parser,
                                           TypeAttr &typeAttr, Attribute &attr) {

 
 LogicalResult GlobalOp::verify() {
   if (!getIsMutable() && !getValue())
     return emitOpError() << "immutable global must have an initial value";
   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
 //===----------------------------------------------------------------------===//
 
 GlobalOp GlobalLoadConstOp::getGlobalOp(SymbolTableCollection &symbolTable) {
   return symbolTable.lookupNearestSymbolFrom<GlobalOp>(
       getOperation()->getParentOp(), getGlobalAttr());
 }
 
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());
+}
+
+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
 //===----------------------------------------------------------------------===//
 
 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

 }
 
 // -----
 // expected-error @+1 {{immutable global must have an initial value}}
 ml_program.global private @const : i32
 
 // -----
 ml_program.func @undef_global() -> i32 {
-  // expected-error @+1 {{undefined global: nothere}}
+  // expected-error @+1 {{undefined global: @nothere}}
   %0 = ml_program.global_load_const @nothere : i32
   ml_program.return %0 : i32
 }
 
 // -----
 ml_program.global private mutable @var : i32
 ml_program.func @mutable_const_load() -> i32 {
-  // expected-error @+1 {{op cannot load as const from mutable global var}}
+  // expected-error @+1 {{op cannot load as const from mutable global @var}}
   %0 = ml_program.global_load_const @var : i32
   ml_program.return %0 : i32
 }
 
 // -----
 ml_program.global private @var(42 : i64) : i64
 ml_program.func @const_load_type_mismatch() -> i32 {
   // expected-error @+1 {{cannot load from global typed 'i64' as 'i32'}}
   %0 = ml_program.global_load_const @var : i32
   ml_program.return %0 : i32
 }
+
+// -----
+ml_program.func @load_undef() -> i32 {
+  // expected-error @+1 {{undefined global: @nothere}}
+  %0 = ml_program.global_load @nothere : i32
+  ml_program.return %0 : i32
+}
+
+// -----
+ml_program.global private mutable @var(42 : i64) : i64
+ml_program.func @load_type_mismatch() -> i32 {
+  // expected-error @+1 {{cannot load from global typed 'i64' as 'i32'}}
+  %0 = ml_program.global_load @var : i32
+  ml_program.return %0 : i32
+}
+
+// -----
+ml_program.func @store_undef(%arg0: i32) {
+  // expected-error @+1 {{undefined global: @nothere}}
+  ml_program.global_store @nothere = %arg0 : i32
+  ml_program.return
+}
+
+// -----
+ml_program.global private mutable @var(42 : i64) : i64
+ml_program.func @store_type_mismatch(%arg0: i32) {
+  // expected-error @+1 {{cannot store to a global typed 'i64' from 'i32'}}
+  ml_program.global_store @var = %arg0 : i32
+  ml_program.return
+}
+
+// -----
+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
+}

mlir/test/Dialect/MLProgram/ops.mlir

 // RUN: mlir-opt %s --allow-unregistered-dialect | mlir-opt --allow-unregistered-dialect | FileCheck %s
 // RUN: mlir-opt %s --allow-unregistered-dialect --mlir-print-op-generic | mlir-opt --allow-unregistered-dialect | FileCheck %s
 
 // CHECK-LABEL: ml_program.func private @extern_func
 ml_program.func private @extern_func(i32) -> i32
 
 // CHECK-LABEL: ml_program.func @defined_func
 ml_program.func @defined_func(%arg0 : i32) -> i32 {
   ml_program.return %arg0 : i32
 }
 
 // CHECK-LABEL: ml_program.subgraph private @extern_subgraph
 ml_program.subgraph private @extern_subgraph(i32) -> i32
 
 // CHECK-LABEL: ml_program.subgraph @compute_subgraph
 ml_program.subgraph @compute_subgraph(%arg0 : i32) -> i32 {
-  %1 = "unregistered.dummy"(%0) : (i32) -> i32
+  %token = ml_program.token
+  %1 = "unregistered.dummy"(%0, %token) : (i32, !ml_program.token) -> i32
   %0 = "unregistered.dummy"(%arg0) : (i32) -> i32
   ml_program.output %0 : i32
 }
 
 // CHECK: ml_program.global private @global_same_type(dense<4> : tensor<4xi32>) : tensor<4xi32>
 ml_program.global private @global_same_type(dense<4> : tensor<4xi32>) : tensor<4xi32>
 
 // CHECK: ml_program.global private mutable @global_mutable_undef : tensor<?xi32>
 ml_program.global private mutable @global_mutable_undef : tensor<?xi32>
 
 // CHECK: ml_program.global private mutable @global_extern(#extern) : tensor<?xi32>
 ml_program.global private mutable @global_extern(#ml_program.extern : tensor<4xi32>) : tensor<?xi32>
+
+// CHECK-LABEL: @global_load_const
+ml_program.func @global_load_const() -> tensor<4xi32> {
+  %0 = ml_program.global_load_const @global_same_type : tensor<4xi32>
+  ml_program.return %0 : tensor<4xi32>
+}
+
+// CHECK-LABEL: @global_load_store
+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
+      after (() -> !ml_program.token) : tensor<?xi32>
+  %token3 = ml_program.global_store @global_mutable_undef = %0
+      after (%token1, %token2 -> !ml_program.token) : tensor<?xi32>
+  ml_program.global_store @global_mutable_undef = %0
+      after (%token3) : tensor<?xi32>
+
+  ml_program.output %0, %token3 : tensor<?xi32>, !ml_program.token
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

 ##---------------------------------------------------------------------------##
 
 td_library(
     name = "MLProgramOpsTdFiles",
     srcs = [
         "include/mlir/Dialect/MLProgram/IR/MLProgramAttributes.td",
         "include/mlir/Dialect/MLProgram/IR/MLProgramBase.td",
         "include/mlir/Dialect/MLProgram/IR/MLProgramOps.td",
+        "include/mlir/Dialect/MLProgram/IR/MLProgramTypes.td",
     ],
     includes = ["include"],
     deps = [
         ":CallInterfacesTdFiles",
         ":ControlFlowInterfacesTdFiles",
         ":FunctionInterfacesTdFiles",
         ":OpBaseTdFiles",
         ":RegionKindInterfaceIncGen",
tbl_outs = [
             "include/mlir/Dialect/MLProgram/IR/MLProgramAttributes.cpp.inc",
         ),
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/MLProgram/IR/MLProgramAttributes.td",
     deps = [":MLProgramOpsTdFiles"],
 )
 
+gentbl_cc_library(
+    name = "MLProgramTypesIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            ["-gen-typedef-decls"],
+            "include/mlir/Dialect/MLProgram/IR/MLProgramTypes.h.inc",
+        ),
+        (
+            ["-gen-typedef-defs"],
+            "include/mlir/Dialect/MLProgram/IR/MLProgramTypes.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/MLProgram/IR/MLProgramTypes.td",
+    deps = [":MLProgramOpsTdFiles"],
+)
+
 cc_library(
     name = "MLProgramDialect",
     srcs = glob([
         "lib/Dialect/MLProgram/IR/*.cpp",
         "lib/Dialect/MLProgram/IR/*.h",
     ]),
     hdrs = glob([
         "include/mlir/Dialect/MLProgram/IR/*.h",
     ]),
     includes = ["include"],
     deps = [
         ":ControlFlowInterfaces",
         ":IR",
         ":MLProgramAttributesIncGen",
         ":MLProgramOpsIncGen",
+        ":MLProgramTypesIncGen",
         ":Pass",
         ":Support",
         "//llvm:Support",
     ],
 )
 
 ##---------------------------------------------------------------------------##
 # Allocation interfaces