[MLIR] Introduce std.global_memref and std.get_global_memref operations.

[jurahul]

Summary

• Add standard dialect operations to define global variables with memref types and to retrieve the memref for to a named global variable
• Extend unit tests to test verification for these operations.

[silvas]

Thanks so much for pushing this forward!

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2013

summary seems out of date.

2032

nit: no space after :.

[silvas]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2015

Can you provide some description here about the intent that the memref descriptor itself is not considered mutable, to simplify global alias analysis.

2026

can provide examples with some commentary for:

1. external global memref.
2. uninitialized global memref.

2038

also add isUnintialized()?

2051

for clarity: "writing to the result memref (such as through a std.store op)"

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2151

nit: newline after this banner comment

2165

is ::mlir needed?

2186

nit: Capitalize first letter

Also, I think a description equivalent to this comment belongs in the let description = on the ODS side. (and then you can remove this comment if desired)

2211

nit: newline after this banner comment

2215

These lookups can be expensive since there is no symbol table build here.

I think the preferred pattern these days is to implement SymbolUserOpInterface which allows for efficient verification https://mlir.llvm.org/docs/Interfaces/#symbolinterfaces

mlir/test/Dialect/Standard/invalid.mlir

268

should private visibility + missing initial_value be a verifier error?

[rriddle]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2152

Can you just use the declarative syntax for this? You just need to use SymbolNameAttr for the sym_name attribute and it will format it the way you want.

[jurahul]

Address review comments.

[jurahul]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2152

I did look into that. For Symbols the documentation says that sym_name should be a string attribute. And with that the declarative syntax would then format it as just a string. I also tried adding @ to the syntax but that's not one of the allowed literal, so I fell back to manual print/parse.

Let me try the SymbolNameAttr.

[jurahul]

Use declarative syntax for global_memref as well.

[jurahul]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2152

SymbolNameAttr works. Thanks!

[silvas]

Thanks Rahul, this looks excellent.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2026

I think the terminology "memref pointing to the global variable" is a bit confusing. Maybe "retrieve the memref descriptor held in this std.global_memref op"?

2084

Can we tighten this to AnyStaticShapeMemRef? I know that the SymbolUserOpInterface implementation does the all the necessary checking, but this field also serves a documentation purpose (e.g. people reading the .td file or in generated documentation).

I also notice that you have some code in verify() that could be removed by this.

[mehdi_amini]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2022–2023

Can you specify the meaning of this? (for example mention that a store to a constant global is UB)

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2156

In general, I like to be more explicit in error messages to avoid having to spawn a debugger (or add more print on demand):

op.emitOpError() << "type should be static shaped memref, but got " << op.type();

(same below)

2183

Actually I'm not entirely convinced that the *declaration* should be public, for example it wouldn't be DCE if it is the case.

mlir/test/Dialect/Standard/ops.mlir

97

I wonder if we could do print more "pretty":

global_memref @foo : memref<2x2xf32>

global_memref @foo : memref<2x2xf32> = [0.0, 1.0]

global_memref constant @foo : memref<2x2xf32> = [0.0, 1.0]

global_memref private constant @foo : memref<2x2xf32> = uninitialized

[jurahul]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2026

The std.global_memref op itself does not hold the memref in my mental model, i.e., its contents are not the memref itself but the buffer. I changed it to "retrieve the memref for the global variable". In my mental model, std.get_global_memref will actually construct the memref by packing say the backing buffer pointer and whatever other things are needed to construct the descriptor.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2183

Looking into other such cases, I did not see any verification like this for either FuncOp or LLVM::GlobalOp.

I think its pretty clear that a definition can have all 3 allowed visibility attributes (public, private, and nested). When it comes to declarations, the intent of a declaration to "import" a symbol that is defined external to the IR. Such an import is legal only if the definition of the symbol in the defining IR is public. That's a linker level check which we cannot do here. The declaration itself can then define the visibility of this "import" in the current IR. Private visibility means that only get_global_memrefs local to the current operation can reference the declaration. Nested means that any operation in the current IR can reference the declared GV. It seems public "import" is problematic since if external IR wants to reference the GV, that reference will be to the public definition of the symbol in the defining IR, and not a declaration in this IR (which does not define the symbol). So declarations themselves cannot have public visibility. So it seems what we want is:

if (op.isExternal() && op.isPublic()) {
    "global variable declarations cannot have public visibility"
}

WDTY? That way, unreferenced GV's will get eliminated correctly by SymbolDCE.

mlir/test/Dialect/Standard/ops.mlir

97

I don't know if this is possible with the declarative syntax (specifically sym_visibility is optional attribute). I'll try custom printer/parser here.

[jurahul]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2183

Just to be super clear: That way, unreferenced GV declarations will get eliminated correctly by SymbolDCE.

[mehdi_amini]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2183

WDTY? That way, unreferenced GV's will get eliminated correctly by SymbolDCE.

Yes this is exactly what I had in mind in term of reasoning.

[jurahul]

Add pretty printing/parsing for global_memrefs. Also removed the external vs visibility checks.

[jurahul]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2183

Based on offline discussion, we decided to leave this as a TODO and start a discussion about this on the MLIR discourse since Functions and other symbol operations also likely need such a check.

mlir/test/Dialect/Standard/ops.mlir

97

I have implemented the pretty format using a custom parser and printer. It's not exactly like above, it still prints dense<[...]> but elides the tensor type as that can be derived from the memref type. To be able to do this (hide the tensor type), I had to repeat some of the checks in the Op verification (like verifying that the type is a static shaped memref) so that I can call the getTensorTypeFromMemRefType on a valid memref type during parsing.

[mehdi_amini]

LGTM overall, nice syntax! Thanks

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2252

Can you get coverage for the error cases?

2306

Nit: remove trivial braces (and below)

mlir/test/Dialect/Standard/ops.mlir

97

Nice result!

[rriddle]

mlir/test/Dialect/Standard/ops.mlir

97

Can you please move back to the declarative form and use a custom directive for the initializer. A majority of the syntax can be declarative.

https://mlir.llvm.org/docs/OpDefinitions/#custom-directives

[rriddle]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2055

The hasValue here is unnecessary.

mlir/include/mlir/IR/SymbolTable.h

76 ↗

(On Diff #301807)

Why are these necessary? The visibility attribute is already a StringAttr. If you use the declarative assembly form these shouldn't be necessary.

[jurahul]

mlir/test/Dialect/Standard/ops.mlir

97

Thanks. If this can still be supported with a declarative form, I am all for it. Let me try that.

[jurahul]

Use declarative syntax with custom directives.

[jurahul]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2252

When we switch to declarative assembly syntax, it seems there is no equivalent check. Seems something which needs to be added to the code generated by ODS framework? I'll file a bug for this.

[jurahul]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2252

Filed: https://bugs.llvm.org/show_bug.cgi?id=48025

[jurahul]

Windows build failure is an unrelated issue in RISC-V backend. Linux failure is clang-tidy warnings for the parse helper functions to use camelCase variable names for arguments. I think it makes sense to keep these names matching the names in the operations.

[rriddle]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2159

Why is this necessary? I would expect you can just print the string attribute normally.

2176

You can format this in the declarative syntax like:

(`constant` $constant)?

There is builtin support for formatting UnitAttrs this way, i.e. the UnitAttr won't be in the textual format.

[silvas]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2176

I think it is with a ^

(`constant` $constant^)?

https://mlir.llvm.org/docs/OpDefinitions/#unit-attributes

[jurahul]

Review feedback

[jurahul]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2159

StringAttr would print normally as a string with quotes I think, whereas we want it to be printed without quotes. But I think that might be ok if it lets us the declarative syntax more. I tried that and it needed some fixes to the parser, but seems to be working.

2176

Thanks. I figured that one out. Also, since the sym_visibility is optional, that also needed the anchor in the declarative syntax.

[jurahul]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2058

I tried $sym_name : $type here so that I could use some more generated code, but looks like the custom syntax does not allow binding a value twice, and I do need the parsed type in custom<GlobalMemrefOpTypeAndInitialValue>. So I think this is the best we can squeeze out of the custom syntax.

[rriddle]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2058

Can you use type_ref for that?

[jurahul]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2058

The custom syntax does not allow it on attributes, I get: "invalid argument to 'type' directive". ODS doc says:

type_ref ( input )
input must be either an operand or result variable , the operands directive, or the results directive.

[rriddle]

Thanks Rahul, looks great.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2058

Ah yeah, forgot this was a TypeAttr.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2156

This variable should be camelCase.

2169

Same here.

2182

nit: Can you switch to early exit here?

[jurahul]

Address final comments

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let extraClassDeclaration = [{
     /// the operation.
     static bool areCastCompatible(Type a, Type b);
   }];
 
   let hasFolder = 0;
 }
 
 //===----------------------------------------------------------------------===//
+// GlobalMemrefOp
+//===----------------------------------------------------------------------===//
+
+def GlobalMemrefOp : Std_Op<"global_memref", [NoSideEffect, Symbol]> {
+  let summary = "declare or define a global memref variable";
+  let description = [{
+    The `global_memref` operation declares or defines a named global variable.
+    The backing memory for the variable is allocated statically and is described
+    by the type of the variable (which should be a statically shaped memref
+    type). The operation is a declaration if no `inital_value` is specified,
+    else it is a definition. The `initial_value` can either be a unit attribute
+    to represent a definition of an uninitialized global variable, or an
+    elements attribute to represent the definition of a global variable with an
+    initial value. The global variable can also be marked constant using the
+    `constant` unit attribute.
+
+    The global variable can be accessed by using the `get_global_memref` to
+    retrieve the memref pointing to the global variable. Note that the memref
+    for such global variable itself is immutable (i.e., get_global_memref for a
+    given global variable will always return the same memref descriptor).
+
+    Example:
+
+    ```mlir
+    // Private variable with an initial value.
+    global_memref @x : memref<2xf32> { sym_visibility = "private",
+                         initial_value = dense<0.0,2.0> : tensor<2xf32> }
+
+    // External variable.
+    global_memref @y : memref<4xi32> { sym_visibility = "public" }
+
+    // Uninitialized externally visible variable.
+    global_memref @z : memref<3xf16> { sym_visibility = "public",
+                         initial_value }
+    ```
+  }];
+
+  let arguments = (ins
+      SymbolNameAttr:$sym_name,
+      TypeAttr:$type,
+      OptionalAttr<AnyAttr>:$initial_value,
+      UnitAttr:$constant
+  );
+  let assemblyFormat = "$sym_name `:` $type attr-dict";
+
+  let extraClassDeclaration = [{
+     bool isExternal() { return !initial_value().hasValue(); }
+     bool isUnitialized() {
+       return initial_value().hasValue() &&
+              initial_value().getValue().isa<UnitAttr>();
+     }
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// GetGlobalMemrefOp
+//===----------------------------------------------------------------------===//
+
+def GetGlobalMemrefOp : Std_Op<"get_global_memref",
+    [NoSideEffect, DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
+  let summary = "get the memref pointing to a global variable";
+  let description = [{
+     The `get_global_memref` operation retrieves the memref pointing to a
+     named global variable. If the global variable is marked constant, writing
+     to the result memref (such as through a `std.store` operation) is
+     undefined.
+
+     Example:
+
+     ```mlir
+     %x = get_global_memref @foo : memref<2xf32>
+     ```
+  }];
+
+  let arguments = (ins FlatSymbolRefAttr:$name);
+  let results = (outs AnyType:$result);
+  let assemblyFormat = "$name `:` type($result) attr-dict";
+
+  // `GetGlobalMemrefOp` is fully verified by its traits.
+  let verifier = ?;
+}
+
+//===----------------------------------------------------------------------===//
 // ImOp
 //===----------------------------------------------------------------------===//
 
 def ImOp : Std_Op<"im",
     [NoSideEffect,
      TypesMatchWith<"complex element type matches result type",
                     "complex", "imaginary",
                     "$_self.cast<ComplexType>().getElementType()">]> {

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

if (auto va = a.dyn_cast<VectorType>())
     if (auto vb = b.dyn_cast<VectorType>())
       return va.getShape().equals(vb.getShape()) &&
              areElementsCastCompatible(va.getElementType(),
                                        vb.getElementType());
   return false;
 }
 
 //===----------------------------------------------------------------------===//
+// Helpers for Tensor[Load|Store]Op, TensorToMemrefOp, and GlobalMemrefOp
+//===----------------------------------------------------------------------===//
+
+static Type getTensorTypeFromMemRefType(Type type) {
+  if (auto memref = type.dyn_cast<MemRefType>())
+    return RankedTensorType::get(memref.getShape(), memref.getElementType());
+  if (auto memref = type.dyn_cast<UnrankedMemRefType>())
+    return UnrankedTensorType::get(memref.getElementType());
+  return NoneType::get(type.getContext());
+}
+
+//===----------------------------------------------------------------------===//
 // AddFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult AddFOp::fold(ArrayRef<Attribute> operands) {
   return constFoldBinaryOp<FloatAttr>(
       operands, [](APFloat a, APFloat b) { return a + b; });
 }
 

 bool FPTruncOp::areCastCompatible(Type a, Type b) {
   if (auto fa = a.dyn_cast<FloatType>())
     if (auto fb = b.dyn_cast<FloatType>())
       return fa.getWidth() > fb.getWidth();
   return areVectorCastSimpleCompatible(a, b, areCastCompatible);
 }
 
 //===----------------------------------------------------------------------===//
+// GlobalMemrefOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(GlobalMemrefOp op) {
+  auto memrefType = op.type().dyn_cast<MemRefType>();
+  if (!memrefType || !memrefType.hasStaticShape())
+    return op.emitOpError("type should be static shaped memref");
+
+  // Verify that the initial value, if present, is either a unit attribute or
+  // an elements attribute.
+  if (op.initial_value().hasValue()) {
+    Attribute initValue = op.initial_value().getValue();
+    if (!initValue.isa<UnitAttr>() && !initValue.isa<ElementsAttr>()) {
+      return op.emitOpError(
+          "initial value should be a unit or elements attribute");
+    }
+
+    // Check that the type of the initial value is compatible with the type of
+    // the global variable.
+    if (initValue.isa<ElementsAttr>()) {
+      Type initType = initValue.getType();
+      Type tensorType = getTensorTypeFromMemRefType(memrefType);
+      if (initType != tensorType) {
+        return op.emitOpError("initial value expected to be of type ")
+               << tensorType << ", but was of type " << initType;
+      }
+    }
+  }
+
+  // External variables should have public visibility.
+  if (op.isExternal() && !op.isPublic()) {
+    return op.emitOpError(
+        "external global variables should have public visibility");
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// GetGlobalMemrefOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+GetGlobalMemrefOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  // Verify that the result type is same as the type of the referenced
+  // global_memref op.
+  auto global =
+      symbolTable.lookupNearestSymbolFrom<GlobalMemrefOp>(*this, nameAttr());
+  if (!global) {
+    return emitOpError("'")
+           << name() << "' does not reference a valid global memref";
+  }
+
+  Type resultType = result().getType();
+  if (global.type() != resultType) {
+    return emitOpError("result type ")
+           << resultType << " does not match type " << global.type()
+           << " of the global memref @" << name();
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // IndexCastOp
 //===----------------------------------------------------------------------===//
 
 // Index cast is applicable from index to integer and backwards.
 bool IndexCastOp::areCastCompatible(Type a, Type b) {
   if (a.isa<ShapedType>() && b.isa<ShapedType>()) {
     auto aShaped = a.cast<ShapedType>();
     auto bShaped = b.cast<ShapedType>();
 
     return (aShaped.getShape() == bShaped.getShape()) &&
            areCastCompatible(aShaped.getElementType(),
                              bShaped.getElementType());

 
 //===----------------------------------------------------------------------===//
 // LoadOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(LoadOp op) {
   if (op.getNumOperands() != 1 + op.getMemRefType().getRank())
     return op.emitOpError("incorrect number of indices for load");
   return success();
 }
 
 OpFoldResult LoadOp::fold(ArrayRef<Attribute> cstOperands) {
   /// load(memrefcast) -> load
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
   return OpFoldResult();
 }
if (aT.getAffineMaps() != bT.getAffineMaps()) {
           return false;
     }
     if (aT.getMemorySpace() != bT.getMemorySpace())
       return false;
 
     // They must have the same rank, and any specified dimensions must match.
     if (aT.getRank() != bT.getRank())
       return false;
 
     for (unsigned i = 0, e = aT.getRank(); i != e; ++i) {
       int64_t aDim = aT.getDimSize(i), bDim = bT.getDimSize(i);
       if (aDim != -1 && bDim != -1 && aDim != bDim)
         return false;
     }
     return true;
   } else {
     if (!aT && !uaT)

 } // namespace
 
 void TensorCastOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<ChainedTensorCast>(context);
 }
 
 //===----------------------------------------------------------------------===//
-// Helpers for Tensor[Load|Store]Op and TensorToMemrefOp
-//===----------------------------------------------------------------------===//
-
-static Type getTensorTypeFromMemRefType(Type type) {
-  if (auto memref = type.dyn_cast<MemRefType>())
-    return RankedTensorType::get(memref.getShape(), memref.getElementType());
-  if (auto memref = type.dyn_cast<UnrankedMemRefType>())
-    return UnrankedTensorType::get(memref.getElementType());
-  return NoneType::get(type.getContext());
-}
-
-//===----------------------------------------------------------------------===//
 // TensorLoadOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult TensorLoadOp::fold(ArrayRef<Attribute>) {
   if (auto tensorToMemref = memref().getDefiningOp<TensorToMemrefOp>())
     return tensorToMemref.tensor();
   return {};
 }

mlir/test/Dialect/Standard/invalid.mlir

 
 // CHECK-LABEL: memref_reshape_result_affine_map_is_not_identity
 func @memref_reshape_result_affine_map_is_not_identity(
         %buf: memref<4x4xf32>, %shape: memref<1xi32>) {
   // expected-error @+1 {{result memref type should have identity affine map}}
   memref_reshape %buf(%shape)
     : (memref<4x4xf32>, memref<1xi32>) -> memref<8xf32, offset: 0, strides: [2]>
 }
+
+// -----
+
+// expected-error @+1 {{type should be static shaped memref}}
+global_memref @foo : i32
+
+// -----
+
+// expected-error @+1 {{type should be static shaped memref}}
+global_memref @foo : memref<*xf32>
+
+// -----
+
+// expected-error @+1 {{type should be static shaped memref}}
+global_memref @foo : memref<?x?xf32>
+
+// -----
+
+// expected-error @+1 {{initial value should be a unit or elements attribute}}
+global_memref @foo : memref<2x2xf32> { initial_value = "foo" }
+
+// -----
+
+// expected-error @+1 {{initial value expected to be of type 'tensor<2x2xf32>', but was of type 'tensor<2xf32>'}}
+global_memref @foo : memref<2x2xf32> { initial_value = dense<[0.0, 1.0]> : tensor<2xf32> }
+
+// -----
+
+// expected-error @+1 {{external global variables should have public visibility}}
+global_memref @foo : memref<2x2xf32> { sym_visibility = "private" }
+
+// -----
+
+func @nonexistent_global_memref() {
+  // expected-error @+1 {{'gv' does not reference a valid global memref}}
+  %0 = get_global_memref @gv : memref<3xf32>
+  return
+}
+
+// -----
+
+func @foo()
+
+func @nonexistent_global_memref() {
+  // expected-error @+1 {{'foo' does not reference a valid global memref}}
+  %0 = get_global_memref @foo : memref<3xf32>
+  return
+}
+
+// -----
+
+global_memref @gv : memref<3xi32>
+
+func @mismatched_types() {
+  // expected-error @+1 {{result type 'memref<3xf32>' does not match type 'memref<3xi32>' of the global memref @gv}}
+  %0 = get_global_memref @gv : memref<3xf32>
+  return
+}
+

mlir/test/Dialect/Standard/ops.mlir

func @memref_reshape(%unranked: memref<*xf32>, %shape1: memref<1xi32>,
   %dyn_vec = memref_reshape %unranked(%shape1)
                : (memref<*xf32>, memref<1xi32>) -> memref<?xf32>
   %dyn_mat = memref_reshape %dyn_vec(%shape2)
                : (memref<?xf32>, memref<2xi32>) -> memref<?x?xf32>
   %new_unranked = memref_reshape %dyn_mat(%shape3)
                : (memref<?x?xf32>, memref<?xi32>) -> memref<*xf32>
   return %new_unranked : memref<*xf32>
 }
+
+// CHECK-LABEL: global_memref @memref0 : memref<2xf32>
+global_memref @memref0 : memref<2xf32>
+
+// CHECK-LABEL: global_memref @memref1 : memref<2xf32>
+// CHECK-SAME: constant
+// CHECK-SAME: initial_value = dense<[0.000000e+00, 1.000000e+00]> : tensor<2xf32>
+global_memref @memref1 : memref<2xf32>
+       { initial_value = dense<[0.0, 1.0]> : tensor<2xf32>, constant }
+
+// CHECK-LABEL: global_memref @memref2 : memref<2xf32>
+// CHECK-SAME: initial_value
+global_memref @memref2 : memref<2xf32>  { initial_value }
+
+// CHECK-LABEL: global_memref @memref3 : memref<2xf32>
+// CHECK-SAME: initial_value
+// CHECK-SAME: sym_visibility = "private"
+global_memref @memref3 : memref<2xf32>  { initial_value, sym_visibility = "private" }
+
+// CHECK-LABEL: func @write_global_memref
+func @write_global_memref() {
+  %0 = get_global_memref @memref0 : memref<2xf32>
+  %1 = constant dense<[1.0, 2.0]> : tensor<2xf32>
+  tensor_store %1, %0 : memref<2xf32>
+  return
+}
+
+// CHECK-LABEL: func @read_global_memref
+func @read_global_memref() {
+  %0 = get_global_memref @memref0 : memref<2xf32>
+  %1 = tensor_load %0 : memref<2xf32>
+  return
+}
+