[mlir] Add GlobalOp, GlobalLoadConstOp to ml_program.

[stellaraccident]

Summary

The approach I took was to define a dialect 'extern' attribute that a GlobalOp can take as a value to signify external linkage. I think this approach should compose well and should also work with wherever the OpaqueElements work goes in the future (since that is just another kind of attribute). I special cased the GlobalOp parser/printer for this case because it is significantly easier on the eyes.

In the discussion, Jeff Niu had proposed an alternative syntax for GlobalOp that I ended up not taking. I did try to implement it but a) I don't think it made anything easier to read in the common case, and b) it made the parsing/printing logic a lot more complicated (I think I would need a completely custom parser/printer to do it well). Please have a look at the common cases where the global type and initial value type match: I don't think how I have it is too bad. The less common cases seem ok to me.

I chose to only implement the direct, constant load op since that is non side effecting and there was still discussion pending on that.

[rriddle]

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

37

42

The mnemonic is generally passed to MLProgram_Attr, can we follow that style here as well?

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

34

nit: Newline before this. Alternatively, just set useDefaultAttributePrinterParser to 1.

https://mlir.llvm.org/docs/DefiningDialects/#default-attributetype-parsers-and-printers

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

121–132

Please use a proper mlir code block.

143

declararative

Also, I don't think we will add anything builtin for this. The assembly format is not intended to have non-generic things inside of it.

171

Same here.

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

39

Please include the generated tablegen file like with operations, we shouldn't name these directly.

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

54

Drop trivial braces here.

94

Drop trivial braces here and below.

135

Drop trivial braces here and throughout the commit.

150–166

You should use the SymbolUserOpInterface trait for verifying symbol uses.

[rriddle]

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

143

I added some thoughts to the issue.

[jpienaar]

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

18

This is already default of AttrDef, lets remove until its actually needed to specialize on (which isn't too common, same is true for traits but that is more likely to change quickly so up to in terms of removing that here too)

27

OpaqueExtern - this is not just extern, this is opaque too. One can have an extern that behaves in all ways like the value it represents but is just external to the context. This explicitly doesn't. It doesn't allow for regular interaction in folding nor does it support reproducers (so error case reproducers are no longer self-contained if one uses this attribute). This is actually beyond opaque even (one can have one that doesn't interact in folding but is still reproducer friendly), NrOpaqueExtern. Although for that one could introduce optional param to differentiate between OpaqueExtern kinds.

(error reproducers are in particular very important to me and has been a pain point for many years in systems where they don't exist, opaque part is more about being clear in objective & usage of attribute; having an extern that just allows for cheaper interactions within a larger system, rather than using this in a offline tool, will also be interesting)

27

s/[]// as default

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

129

This syntax reads opposite to what I'd expect (and even feels contradictory with one in #123). Could we iterate on this a bit? Something like

ml_program.global mutable @foobar {#extern : tensor<4xi32>} : tensor<?xi32>

or

ml_program.global mutable @foobar init = {#extern : tensor<4xi32>} : tensor<?xi32>

or

ml_program.global mutable @foobar(#extern : tensor<4xi32>) : tensor<?xi32>

all feel a bit more natural to me

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

151

This is not safe in general. Use verifySymbolUses instead.

[Mogball]

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

129

The last version is the syntax I proposed (in line with llvm.global)

[stellaraccident]

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

129

I'll rework it. I did try it the way requested but was having trouble with the declarative assembly. Now that it is using a custom block, should be able to do it that way.

I don't see the advantage of the alternative suggested but also don't have a strong opinion.

[stellaraccident]

Address NFC comments.

[stellaraccident]

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

18

Thanks. I had a significant challenge finding good examples in tree to copy.

27

I think we may be talking about two different things: you are talking about internal/external to a context, whereas this is indicating external linkage in the traditional sense. I get where you are coming from and can adapt the spelling but just want to be sure we are talking about the same thing.

42

I'm trying to grok what you are asking. I only found a few examples in tree (Test Dialect, Builtin Attributes and Sparse) and I think they are all done as I have done here? Let me know and happy to fix.

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

34

Added newline. Setting useDefaultAttributePrinterParser to 1 does not seem to work on its own. I get an error:

In file included from /home/stella/src/llvm-project/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp:22:
tools/mlir/include/mlir/Dialect/MLProgram/IR/MLProgramAttributes.cpp.inc:99:37: error: out-of-line definition of 'parseAttribute' does not match any declaration in 'mlir::ml_program::MLProgramDialect'
::mlir::Attribute MLProgramDialect::parseAttribute(::mlir::DialectAsmParser &parser,
                                    ^~~~~~~~~~~~~~
tools/mlir/include/mlir/Dialect/MLProgram/IR/MLProgramAttributes.cpp.inc:117:24: error: out-of-line definition of 'printAttribute' does not match any declaration in 'mlir::ml_program::MLProgramDialect'
void MLProgramDialect::printAttribute(::mlir::Attribute attr,

Maybe there is some missing emission of declarations for that switch? It seems like the generated cpp file should mate with the generated header.

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

143

Removed TODO.

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

39

Ah, thanks - I recall going and looking for that thinking it must be there, but I think I must have been looking in the generated header file, not the cpp file (since GET_ATTRDEF_LIST is clearly at the top of the cpp file).

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

135

Sorry about that: My eyes have become uncalibrated to this quirk of LLVM style. Re-read the guidelines and I think I've fixed them all.

150–166

Thank you! I had long wondered the right way to do that and avoid multiple linear scans. Filing an issue in the downstream from which I picked up this bad habbit.

[stellaraccident]

Change global asm format.

[stellaraccident]

Update tests.

[stellaraccident]

Asm format updated - PTAL

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

129

Ok, so trying the obvious:

let assemblyFormat = [{
  custom<SymbolVisibility>($sym_visibility)
  (`mutable` $is_mutable^)?
  $sym_name
  (`(` $value^ `)`)?
  `:`
  $type
  attr-dict
}];

Yields an error:

error: format ambiguity caused by `:` literal found after attribute `value` which does not have a buildable type
    custom<SymbolVisibility>($sym_visibility)

I don't quite understand why this is ambiguous.

129

Ok, I simplified and implemented the last option.

I'm not thrilled with how the #extern attribute prints (it currently prints as an alias) as this is short and load bearing to how to interpret the global (I think it should be inline). If inline, I'd also like it to print without the dialect prefix. I recall there being various discussions about this, but I'm having trouble figuring out the right way to hold it to get what I want. Also, this may be good enough to get started.

[Mogball]

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

42

The self type should not be in the assembly format.

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

129

That sounds like a bug in the assembly format. It should not be ambiguous because $value is followed by a literal.

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

49

If you want to drop the prefix, you can specially check for ExternAttr in both the printer and parser.

mlir/test/Dialect/MLProgram/ops.mlir

29

The type on extern goes away?

[rriddle]

Sorry for the delay, was semi-OOO.

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

42

I meant similar to:

https://github.com/llvm/llvm-project/blob/8bb5b657fec6dc604eeea49e1f26eff4c14084ba/mlir/include/mlir/Dialect/PDL/IR/PDLTypes.td#L23

The docs also have some discussion on this:

https://mlir.llvm.org/docs/AttributesAndTypes/#adding-a-new-attribute-or-type-definition

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

34

Weird, can you file a bug for this?

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

152

Can this be dropped? Or is this comment meant for something?

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

30

Missing comment here?

57

I'll look into moving this somewhere shared in a followup.

135

(No worries at all)

[stellaraccident]

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

42

ODS bug? Changing to let assemblyFormat = ""; produces this error:

error: format is missing reference to parameter: type

^
/home/stella/src/llvm-project/mlir/include/mlir/Dialect/MLProgram/IR/MLProgramAttributes.td:25:1: note: in custom assembly format for this operation
def MLProgram_ExternAttr : MLProgram_Attr<"Extern"> {
^
/home/stella/src/llvm-project/mlir/include/mlir/Dialect/MLProgram/IR/MLProgramAttributes.td:25:1: error: failed to parse assembly format
def MLProgram_ExternAttr : MLProgram_Attr<"Extern"> {

Side note: the lack of actual use of this stuff upstream has made all of the ODS work in this patch incredibly painful and time consuming. Can you point me to a test case that shows proper usage? Otherwise, I can file an issue.

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

34

https://github.com/llvm/llvm-project/issues/55322

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

152

Sorry, that was the result of me trying to figure out what was wrong with the declarative assembly generator. Let me capture in a bug/test op and will drop here.

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

49

I don't care deeply enough about it in this patch to change it. I may followup with something else, but I've spent quite a bit of time tracking through getting this to print/parse and interlocking issues. I would like to checkpoint as I have it here.

mlir/test/Dialect/MLProgram/ops.mlir

29

It is carried in the assembly format. I could not figure out how to get it to print as I expected: #ml_program.extern : tensor<4xi32>. Please point me at a test case that shows proper usage (all I could find were builtins that were all weird in one way or another). Or I can create a test case/issue and iterate on it in a followup.

[stellaraccident]

Comments.

[stellaraccident]

PTAL.

As a side note, this is possibly the single most painful patch I've made. I am not new to defining ops or using ODS, but trying to do it "right" (i.e. in downstreams, we have sometimes just open coded parsers/printers and such) cost hours, hit multiple issues and it was really hard to find non-builtin examples of things done right (the builtins examples help, but they are often "weird" in various ways that it was hard to see through). I did a lot of spelunking through downstreams, each of which represents a point in time and may not have done it right...

I'm not so much complaining as saying that I think it is important that we get more corner case usage of ODS upstream. This dialect should help but we should have more production uses. I've filed this issue: https://github.com/llvm/llvm-project/issues/55322 and will file another (couple) for declarative assembly snags.

[stellaraccident]

filed: https://github.com/llvm/llvm-project/issues/55323

[stellaraccident]

Filed: https://github.com/llvm/llvm-project/issues/55324

[Mogball]

Yeah that's... unfortunate to hear. ODS can be a bag of knives sometimes. I'll take a look at the bugs.

[jpienaar]

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

27

You are correct (i was thinking of extern linkage too, but I wrote here exclusively wrt context indeed). A difference with extern linkage as in C (I think) is that there isn't a generic resolution mechanism or LTO that can be generally done. And that is, as far as I understand, a desirable and required property here: nothing but a specific implementation of a target that can do something with this attribute. It is meant to represent a placeholder/be opaque for all others. No freezing'esqu behavior should happen "accidentally". Is that correct? (And then of course question about leaving space for something with external linkage but generic way of loading that allows for something like generic LTO)

mlir/test/Dialect/MLProgram/ops.mlir

29

I'm almost back, and Jeff is probably back from EuroLLVM then too :-) Then could try and find example here, this form is readable though.

(This part is part fun, part bike shedding to make sure it is easy for folks, especially as this crosses framework boundaries erring on the side of more verbose but trivial to interpret has value. Can help as it sounds like this hasn't been as much fun)

[stellaraccident]

In D124318#3499605, @Mogball wrote:

Yeah that's... unfortunate to hear. ODS can be a bag of knives sometimes. I'll take a look at the bugs.

Ok, you self-type fix worked. Thanks for that!

[stellaraccident]

Rebase and adapt to self-type attribute fix.

[stellaraccident]

Jacques/Jeff, please take a look if you have a chance. I see there is a patch pending to fix the attribute parser declarations, but that will be a minor adaptation. I think the IR format is good enough as I have it here.

[stellaraccident]

Update bazel files.

[jpienaar]

Looks good/we can refine in follow ups, thanks!

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

27

May have forgotten to upload (removing trait default)

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

24

I dont think this function does this one (e.g., it always does colon type at moment). Also s/global/op/

[stellaraccident]

Comments

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

27

I think it is removed? Let me know if I'm seeing something different than you.

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

24

Thanks. Stale comment. Fixed.

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)

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/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/MLProgramAttributes.h

+//===- MLProgramAttributes.h - Attribute 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_MLPROGRAMATTRIBUTES_H_
+#define MLIR_DIALECT_MLPROGRAM_IR_MLPROGRAMATTRIBUTES_H_
+
+#include "mlir/IR/Attributes.h"
+
+//===----------------------------------------------------------------------===//
+// Tablegen Attribute Declarations
+//===----------------------------------------------------------------------===//
+
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/MLProgram/IR/MLProgramAttributes.h.inc"
+
+#endif // MLIR_DIALECT_MLPROGRAM_IR_MLPROGRAMATTRIBUTES_H_

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

+//===- MLProgramAttributed.td - Attr 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_ATTRIBUTES
+#define MLPROGRAM_ATTRIBUTES
+
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/Dialect/MLProgram/IR/MLProgramBase.td"
+
+// Base class for MLProgram dialect attributes.
+class MLProgram_Attr<string name, list<Trait> traits = []>
+    : AttrDef<MLProgram_Dialect, name, traits> {
+  let mnemonic = ?;
+}
+
+//===----------------------------------------------------------------------===//
+// ExternAttr
+//===----------------------------------------------------------------------===//
+
+def MLProgram_ExternAttr : MLProgram_Attr<"Extern"> {
+  let summary = "Value used for a global signalling external resolution";
+  let description = [{
+  When used as the value for a GlobalOp, this indicates that the actual
+  value should be resolved externally in an implementation defined manner.
+  The `sym_name` of the global is the key for locating the value.
+
+  Examples:
+
+  ```mlir
+  extern : tensor<4xi32>
+  ```
+  }];
+
+  let parameters = (ins AttributeSelfTypeParameter<"">:$type);
+  let mnemonic = "extern";
+  let assemblyFormat = "`<` $type `>`";
+}
+
+#endif // MLPROGRAM_ATTRIBUTES

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

let description = [{
     programming model for establishing modules, functions, globals and
     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 extraClassDeclaration = [{
+  private:
+    Attribute parseAttribute(DialectAsmParser& parser, Type type) const override;
+    void printAttribute(Attribute attr, DialectAsmPrinter& p) const override;
+
+  public:
+  }];
+
+  let useDefaultAttributePrinterParser = 1;
   let emitAccessorPrefix = kEmitAccessorPrefix_Prefixed;
 }
 
 #endif // MLPROGRAM_BASE

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

let extraClassDeclaration = [{
 
     bool isDeclaration() { return isExternal(); }
   }];
 
   let hasCustomAssemblyFormat = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// GlobalOp
+//===----------------------------------------------------------------------===//
+
+def MLProgram_GlobalOp : MLProgram_Op<"global", [
+    Symbol
+  ]> {
+  let summary = "Module level declaration of a global variable";
+  let description = [{
+    Declares a named global variable (or constant).
+
+    A global contains a value of a specified type which can be accessed at
+    runtime via appropriate load/store operations. It can be mutable or
+    constant, optionally taking an initial value or declared as
+    extern (in which case, the initial value is found in external storage
+    by symbol name).
+
+    Generally, the type of the global and the type of the initial value
+    will be the same. However, for type hierarchies which can have a more
+    generalized bounding type that can be assigned from a narrow type, this
+    is allowed (but not verified).
+
+    Examples:
+
+    ```mlir
+    // Constant global that stores the same type as its initial value.
+    ml_program.global @foobar = dense<4> : tensor<4xi32>
+
+    // Constant global with an inline initial value of a more specific type.
+    ml_program.global @foobar : tensor<?xi32> = dense<4> : tensor<4xi32>
+
+    // Mutable global with an external initial value of a more specific type.
+    ml_program.global mutable @foobar : tensor<?xi32> = #extern : tensor<4xi32>
+
+    // Mutable global with an undefined initial value.
+    ml_program.global mutable @foobar : tensor<?xi32>
+    ```
+  }];
+
+  let arguments = (ins
+    SymbolNameAttr:$sym_name,
+    TypeAttr:$type,
+    UnitAttr:$is_mutable,
+    OptionalAttr<AnyAttr>:$value,
+    OptionalAttr<StrAttr>:$sym_visibility
+  );
+
+  let assemblyFormat = [{
+    custom<SymbolVisibility>($sym_visibility)
+    (`mutable` $is_mutable^)?
+    $sym_name
+    custom<TypedInitialValue>($type, $value)
+    attr-dict
+  }];
+
+  let hasVerifier = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// 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
+  );
+  let results = (outs
+    AnyType:$result
+  );
+
+  let assemblyFormat = [{
+    $global attr-dict `:` type($result)
+  }];
+
+  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";

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
 
   LINK_LIBS PUBLIC
   MLIRDialect
   MLIRInferTypeOpInterface
   MLIRIR
   )

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

 //===- MLProgramDialect.cpp - MLProgram dialect implementation ------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::ml_program;
 
+//===----------------------------------------------------------------------===//
+/// Tablegen Definitions
+//===----------------------------------------------------------------------===//
+
 #include "mlir/Dialect/MLProgram/IR/MLProgramOpsDialect.cpp.inc"
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/MLProgram/IR/MLProgramAttributes.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"
+      >();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/MLProgram/IR/MLProgramOps.cpp.inc"
       >();
+  addInterfaces<MLProgramOpAsmDialectInterface>();
 }

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

 //===- MLProgramOps.cpp - MLProgram dialect ops implementation ------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/FunctionImplementation.h"
 
 using namespace mlir;
 using namespace mlir::ml_program;
 
 //===----------------------------------------------------------------------===//
+// Custom asm helpers
+//===----------------------------------------------------------------------===//
+
+/// some.op custom<TypeOrAttr>($type, $attr)
+/// ->
+/// some.op : i32
+/// some.op = 42 : i32
+/// some.op : i32 = 42 : index
+/// some.op = extern : i32
+/// some.op : i32 = extern : index
+static ParseResult parseTypedInitialValue(OpAsmParser &parser,
+                                          TypeAttr &typeAttr, Attribute &attr) {
+  auto parseAttr = [&]() -> LogicalResult {
+    // Parse `extern` bareword as an ExternAttr.
+    if (succeeded(parser.parseOptionalKeyword("extern"))) {
+      Type externType;
+      if (failed(parser.parseColonType(externType))) {
+        return parser.emitError(parser.getCurrentLocation()) << "expected type";
+      }
+      attr = ExternAttr::get(parser.getContext(), externType);
+    } else if (failed(parser.parseAttribute(attr))) {
+      return parser.emitError(parser.getCurrentLocation())
+             << "expected attribute";
+    }
+    return success();
+  };
+
+  // Parse `=` value
+  if (succeeded(parser.parseOptionalEqual())) {
+    if (failed(parseAttr()))
+      return failure();
+    typeAttr = TypeAttr::get(attr.getType());
+    return success();
+  }
+
+  // Parse `:` type `=` value
+  Type type;
+  if (failed(parser.parseColonType(type)))
+    return parser.emitError(parser.getCurrentLocation()) << "expected type";
+  typeAttr = TypeAttr::get(type);
+  if (succeeded(parser.parseOptionalEqual())) {
+    if (failed(parseAttr()))
+      return failure();
+  }
+
+  return success();
+}
+
+static void printTypedInitialValue(OpAsmPrinter &p, Operation *op,
+                                   TypeAttr type, Attribute attr) {
+  if (!attr || attr.getType() != type.getValue()) {
+    p << ": ";
+    p.printAttribute(type);
+    p << " ";
+  }
+  if (attr) {
+    p << "= ";
+    if (attr.isa_and_nonnull<ExternAttr>())
+      p << "extern : " << attr.getType();
+    else
+      p.printAttribute(attr);
+  }
+}
+
+/// some.op custom<SymbolVisibility>($sym_visibility) $sym_name
+/// ->
+/// some.op public @foo
+/// some.op private @foo
+static ParseResult parseSymbolVisibility(OpAsmParser &parser,
+                                         StringAttr &symVisibilityAttr) {
+  StringRef symVisibility;
+  parser.parseOptionalKeyword(&symVisibility, {"public", "private", "nested"});
+  if (symVisibility.empty())
+    return parser.emitError(parser.getCurrentLocation())
+           << "expected 'public', 'private', or 'nested'";
+  if (!symVisibility.empty())
+    symVisibilityAttr = parser.getBuilder().getStringAttr(symVisibility);
+  return success();
+}
+
+static void printSymbolVisibility(OpAsmPrinter &p, Operation *op,
+                                  StringAttr symVisibilityAttr) {
+  if (!symVisibilityAttr)
+    p << "public";
+  else
+    p << symVisibilityAttr.getValue();
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/MLProgram/IR/MLProgramOps.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // FuncOp
return function_interface_impl::parseFunctionOp(
       parser, result, /*allowVariadic=*/false, buildFuncType);
 }
 
 void FuncOp::print(OpAsmPrinter &p) {
   function_interface_impl::printFunctionOp(p, *this, /*isVariadic=*/false);
 }
 
 //===----------------------------------------------------------------------===//
+// GlobalOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult GlobalOp::verify() {
+  if (!getIsMutable() && !getValue())
+    return emitOpError() << "immutable global must have an initial value";
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// GlobalLoadConstOp
+//===----------------------------------------------------------------------===//
+
+GlobalOp GlobalLoadConstOp::getGlobalOp(SymbolTableCollection &symbolTable) {
+  return symbolTable.lookupNearestSymbolFrom<GlobalOp>(
+      getOperation()->getParentOp(), getGlobalAttr());
+}
+
+LogicalResult
+GlobalLoadConstOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  GlobalOp referrent = getGlobalOp(symbolTable);
+  if (!referrent)
+    return emitOpError() << "undefined global: " << getGlobal();
+
+  if (referrent.getIsMutable())
+    return emitOpError() << "cannot load as const from mutable global "
+                         << getGlobal();
+
+  if (referrent.getType() != getResult().getType())
+    return emitOpError() << "cannot load from global typed "
+                         << referrent.getType() << " as "
+                         << getResult().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/attrs.mlir

+// RUN: mlir-opt %s --allow-unregistered-dialect | mlir-opt --allow-unregistered-dialect | FileCheck %s
+
+// CHECK: #ml_program.extern<i32>
+"unregistered.attributes"() {
+  value = #ml_program.extern<i32>
+} : () -> ()
+

mlir/test/Dialect/MLProgram/invalid.mlir

ml_program.subgraph @output_arity_match(%arg0 : i32) -> i32 {
   ml_program.output %arg0, %arg0 : i32, i32
 }
 
 // -----
 ml_program.subgraph @output_type_match(%arg0 : i64) -> i32 {
   // expected-error @+1 {{doesn't match function result}}
   ml_program.output %arg0 : i64
 }
+
+// -----
+// 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}}
+  %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}}
+  %0 = ml_program.global_load_const @var : i32
+  ml_program.return %0 : i32
+}
+
+// -----
+ml_program.global private @var = 42 : 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
+}

mlir/test/Dialect/MLProgram/ops.mlir

 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
   %0 = "unregistered.dummy"(%arg0) : (i32) -> i32
   ml_program.output %0 : i32
 }
+
+// CHECK: ml_program.global private @global_same_type = dense<4> : tensor<4xi32>
+ml_program.global private @global_same_type = dense<4> : tensor<4xi32>
+// CHECK: ml_program.global public @global_bounding_type : tensor<?xi32> = dense<4> : tensor<4xi32>
+ml_program.global public @global_bounding_type : tensor<?xi32> = dense<4> : tensor<4xi32>
+// CHECK: ml_program.global private mutable @foobar_mutable_undef : tensor<?xi32>
+ml_program.global private mutable @foobar_mutable_undef : tensor<?xi32>
+// CHECK: ml_program.global private mutable @global_extern = extern : tensor<4xi32>
+ml_program.global private mutable @global_extern = extern : tensor<4xi32>
+// CHECK: ml_program.global private mutable @global_extern_bounding_type : tensor<?xi32> = extern : tensor<4xi32>
+ml_program.global private mutable
+   @global_extern_bounding_type : tensor<?xi32> = extern : tensor<4xi32>