[mlir] Add internal linkage to `memref.global`

[guraypp]

Summary

The internal is missing, so this work adds that.

[guraypp]

The motivation behind adding internal is to enable demotion of global shared memory variable in the kernel (maybe there are other ways to do it)

Here is the input MLIR code:

memref.global "private" @gv1 : memref<1xi64,3>
func.func @get_gv0_memref() {
  %0 = memref.get_global @gv1 : memref<1xi64,3>
  %c0 = arith.constant 0 : index
  %c4 = arith.constant 4 : i64
  memref.store %c4, %0[%c0] : memref<1xi64,3>
  return
}

When generating PTX code for memref.global "private", we get:

.shared .align 8 .b8 gv1[8];
.visible .func get_gv0_memref()
{
        .reg .b64       %rd<2>;
        mov.u64         %rd1, 4;
        st.shared.u64   [gv1], %rd1;
        ret;
}

In the case where the linkage was memref.global "internal", the PTX code would look like this:

.visible .func get_gv0_memref()
{
        .reg .b64       %rd<2>;
        // demoted variable
        .shared .align 8 .b8 gv1[8];
        mov.u64         %rd1, 4;
        st.shared.u64   [gv1], %rd1;
        ret;
}

These examples demonstrate the internal, specifically the demotion of the variable gv1 to shared memory in the PTX code.

[Dinistro]

LLVM linkage and MLIR symbol visibility are orthogonal concepts and should not be confused. For example, MLIR func.func declarations always have a private visibility, while their corresponding LLVM IR linkage can be external.

mlir/include/mlir/IR/SymbolTable.h

88

I do not understand this comment, could you expand and explicitly state the difference between internal and private?

[guraypp]

Thanks for the feedback. As I mentioned, I am not sure this is the right way. Any ideas are welcome!

My goal is to create
memref.global "internal" @gv1 : memref<1xi64,3>

And lower into:
llvm.mlir.global internal @gv1() {addr_space = 3 : i32} : !llvm.array<1 x i64>

Currently, memref.global only supports : private, public, nested

[Dinistro]

I'm not really familiar with PTX, but do you know why it treats internal and private differently? Their semantic difference is minimal and shouldn't be relevant for PTX, I assume.
Is this just an oversight in the PTX lowering, or are you aware of any semantic difference there?

My goal is to create
memref.global "internal" @gv1 : memref<1xi64,3>

And lower into:
llvm.mlir.global internal @gv1() {addr_space = 3 : i32} : !llvm.array<1 x i64>

One could also add a flag to the MemRefToLLVM conversion that ensures that private visibility is converted into internal linkage. Maybe it even makes sense to always emit internal linkage for these memrefs.

Currently, memref.global only supports : private, public, nested

Again, this is symbol visibility, which is an MLIR internal construct and really only specifies where this symbol is visible inside the IR.

[guraypp]

Thanks for feedback. See the linkage types (internal vs private) in llvm https://llvm.org/docs/LangRef.html#linkage-types

I could get rid of Internal in MLIR symbols if it is not desired. In this case, I will need to have code like below in MemRefToLLVM.cpp. I am not sure which way is better.

LLVM::Linkage linkage;
if (global.isPublic()) {
  linkage = LLVM::Linkage::External;
} else {
  if (type.getMemorySpaceAsInt() == 3) {
    linkage = LLVM::Linkage::Internal;
  } else {
    linkage = LLVM::Linkage::Private;
  }
}

[Dinistro]

In D154074#4462817, @guraypp wrote:

Thanks for feedback. See the linkage types (internal vs private) in llvm https://llvm.org/docs/LangRef.html#linkage-types

The only difference is the visibility of a symbol in an ELF file, how does this impact PTX code generations behaviour? I do not assume that PTX cares about such things.
Especially considering that private doesn't expose a symbol in the symbol table, so why can this not be treated as internal global can.

I could get rid of Internal in MLIR symbols if it is not desired. In this case, I will need to have code like below in MemRefToLLVM.cpp. I am not sure which way is better.

LLVM::Linkage linkage;
if (global.isPublic()) {
  linkage = LLVM::Linkage::External;
} else {
  if (type.getMemorySpaceAsInt() == 3) {
    linkage = LLVM::Linkage::Internal;
  } else {
    linkage = LLVM::Linkage::Private;
  }
}

I'm not a fan of encoding target specific behaviour into general purpose passes. Address spaces are fully target specific, so treating 3 somehow differently is not what we want.
For almost all intends and purposes, private and internal behave the same way, so we might really just lower private visibility to internal linkage, if PTX cannot deal with it. Ideally, this is done only when a certain pass option is set, to ensure that we do not mess up other existing flows.

Alternatively, the linkage itself could be added to memref.global, but that would leak LLVM internals up into higher level dialects that shouldn't really care about things like linkage.

[qcolombet]

Talked with @guraypp offline. I think this is a limitation in the PTX backend.

[Dinistro]

In D154074#4471645, @qcolombet wrote:

Talked with @guraypp offline. I think this is a limitation in the PTX backend.

Is it possible to fix this in the PTX backend then? A limitation in this backend should not be the reason to extend the symbol visibility that potentially affects all of MLIR.

[qcolombet]

In D154074#4471755, @Dinistro wrote:

In D154074#4471645, @qcolombet wrote:

Talked with @guraypp offline. I think this is a limitation in the PTX backend.

Is it possible to fix this in the PTX backend then? A limitation in this backend should not be the reason to extend the symbol visibility that potentially affects all of MLIR.

That’s what I asked @guraypp to check.
Sorry for not being clearer earlier.

[qcolombet]

I've sent a patch for review to fix the NVPTX backend limitation https://reviews.llvm.org/D154507

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let description = [{
     memref.global "private" @x : memref<2xf32> = dense<0.0,2.0>
 
     // Private variable with an initial value and an alignment (power of 2).
     memref.global "private" @x : memref<2xf32> = dense<0.0,2.0> {alignment = 64}
 
     // Declaration of an external variable.
     memref.global "private" @y : memref<4xi32>
 
+    // Declaration of an internal variable.
+    memref.global "internal" @y : memref<4xi32>
+
     // Uninitialized externally visible variable.
     memref.global @z : memref<3xf16> = uninitialized
 
     // Externally visible constant variable.
     memref.global constant @c : memref<2xi32> = dense<1, 4>
     ```
   }];
 

mlir/include/mlir/IR/SymbolInterfaces.td

InterfaceMethod<"Gets the visibility of this symbol.",
       }]
     >,
     InterfaceMethod<"Returns true if this symbol has nested visibility.",
       "bool", "isNested", (ins),  [{}],
       /*defaultImplementation=*/[{
         return getVisibility() == mlir::SymbolTable::Visibility::Nested;
       }]
     >,
+    InterfaceMethod<"Returns true if this symbol has internal visibility.",
+      "bool", "isInternal", (ins),  [{}],
+      /*defaultImplementation=*/[{
+        return getVisibility() == mlir::SymbolTable::Visibility::Internal;
+      }]
+    >,
     InterfaceMethod<"Returns true if this symbol has private visibility.",
       "bool", "isPrivate", (ins),  [{}],
       /*defaultImplementation=*/[{
         return getVisibility() == mlir::SymbolTable::Visibility::Private;
       }]
     >,
     InterfaceMethod<"Returns true if this symbol has public visibility.",
       "bool", "isPublic", (ins),  [{}],

mlir/include/mlir/IR/SymbolTable.h

enum class Visibility {
     /// to the operations within the current symbol table.
     Private,
 
     /// The symbol is visible to the current IR, which may include operations in
     /// symbol tables above the one that owns the current symbol. `Nested`
     /// visibility allows for referencing a symbol outside of its current symbol
     /// table, while retaining the ability to observe all uses.
     Nested,
+
+    /// The symbol is public and may be referenced internally
+    /// to the visible references in the IR.
+    Internal,
   };
 
   /// Returns the name of the given symbol operation, aborting if no symbol is
   /// present.
   static StringAttr getSymbolName(Operation *symbol);
 
   /// Sets the name of the given symbol operation.
   static void setSymbolName(Operation *symbol, StringAttr name);

mlir/include/mlir/TableGen/Class.h

protected:
   MethodSignature methodSignature;
   /// The body of the method, if it has one.
   MethodBody methodBody;
   /// Deprecation message if the method is deprecated.
   std::optional<std::string> deprecationMessage;
 };
 
 /// This enum describes C++ inheritance visibility.
-enum class Visibility { Public, Protected, Private };
+enum class Visibility { Public, Protected, Private, Internal };
 
 /// Write "public", "protected", or "private".
 llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
                               mlir::tblgen::Visibility visibility);
 
 // Class for holding an op's constructor method for C++ code emission.
 class Constructor : public Method {
 public:

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

struct GlobalMemrefOpLowering
   matchAndRewrite(memref::GlobalOp global, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     MemRefType type = global.getType();
     if (!isConvertibleAndHasIdentityMaps(type))
       return failure();
 
     Type arrayTy = convertGlobalMemrefTypeToLLVM(type, *getTypeConverter());
 
-    LLVM::Linkage linkage =
-        global.isPublic() ? LLVM::Linkage::External : LLVM::Linkage::Private;
+    LLVM::Linkage linkage = global.isPublic()     ? LLVM::Linkage::External
+                            : global.isInternal() ? LLVM::Linkage::Internal
+                                                  : LLVM::Linkage::Private;
 
     Attribute initialValue = nullptr;
     if (!global.isExternal() && !global.isUninitialized()) {
       auto elementsAttr = llvm::cast<ElementsAttr>(*global.getInitialValue());
       initialValue = elementsAttr;
 
       // For scalar memrefs, the global variable created is of the element type,
       // so unpack the elements attribute to extract the value.

mlir/lib/IR/SymbolTable.cpp

SymbolTable::Visibility SymbolTable::getSymbolVisibility(Operation *symbol) {
   StringAttr vis = symbol->getAttrOfType<StringAttr>(getVisibilityAttrName());
   if (!vis)
     return Visibility::Public;
 
   // Otherwise, switch on the string value.
   return StringSwitch<Visibility>(vis.getValue())
       .Case("private", Visibility::Private)
       .Case("nested", Visibility::Nested)
-      .Case("public", Visibility::Public);
+      .Case("public", Visibility::Public)
+      .Case("internal", Visibility::Internal);
 }
 /// Sets the visibility of the given symbol operation.
 void SymbolTable::setSymbolVisibility(Operation *symbol, Visibility vis) {
   MLIRContext *ctx = symbol->getContext();
 
   // If the visibility is public, just drop the attribute as this is the
   // default.
   if (vis == Visibility::Public) {
     symbol->removeAttr(StringAttr::get(ctx, getVisibilityAttrName()));
     return;
   }
 
   // Otherwise, update the attribute.
-  assert((vis == Visibility::Private || vis == Visibility::Nested) &&
+  assert((vis == Visibility::Private || vis == Visibility::Nested ||
+          vis == Visibility::Internal) &&
          "unknown symbol visibility kind");
 
-  StringRef visName = vis == Visibility::Private ? "private" : "nested";
+  StringRef visName = vis == Visibility::Private  ? "private"
+                      : vis == Visibility::Nested ? "nested"
+                                                  : "internal";
   symbol->setAttr(getVisibilityAttrName(), StringAttr::get(ctx, visName));
 }
 
 /// Returns the nearest symbol table from a given operation `from`. Returns
 /// nullptr if no valid parent symbol table could be found.
 Operation *SymbolTable::getNearestSymbolTable(Operation *from) {
   assert(from && "expected valid operation");
   if (isPotentiallyUnknownSymbolTable(from))
raw_ostream &mlir::operator<<(raw_ostream &os,
                               SymbolTable::Visibility visibility) {
   switch (visibility) {
   case SymbolTable::Visibility::Public:
     return os << "public";
   case SymbolTable::Visibility::Private:
     return os << "private";
   case SymbolTable::Visibility::Nested:
     return os << "nested";
+  case SymbolTable::Visibility::Internal:
+    return os << "internal";
   }
   llvm_unreachable("Unexpected visibility");
 }
 
 //===----------------------------------------------------------------------===//
 // SymbolTable Trait Types
 //===----------------------------------------------------------------------===//
 
LogicalResult detail::verifySymbol(Operation *op) {
   // Verify the visibility attribute.
   if (Attribute vis = op->getAttr(mlir::SymbolTable::getVisibilityAttrName())) {
     StringAttr visStrAttr = llvm::dyn_cast<StringAttr>(vis);
     if (!visStrAttr)
       return op->emitOpError() << "requires visibility attribute '"
                                << mlir::SymbolTable::getVisibilityAttrName()
                                << "' to be a string attribute, but got " << vis;
 
-    if (!llvm::is_contained(ArrayRef<StringRef>{"public", "private", "nested"},
-                            visStrAttr.getValue()))
+    if (!llvm::is_contained(
+            ArrayRef<StringRef>{"public", "private", "nested", "internal"},
       return op->emitOpError()
              << "visibility expected to be one of [\"public\", \"private\", "
-                "\"nested\"], but got "
+                "\"nested\", \"internal\"], but got "
              << visStrAttr;
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Symbol Use Lists
 //===----------------------------------------------------------------------===//

mlir/lib/TableGen/Class.cpp

 raw_ostream &operator<<(raw_ostream &os, Visibility visibility) {
   switch (visibility) {
   case Visibility::Public:
     return os << "public";
   case Visibility::Protected:
     return os << "protected";
   case Visibility::Private:
     return os << "private";
+  case Visibility::Internal:
+    return os << "internal";
   }
   return os;
 }
 } // namespace tblgen
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//
 // ParentClass definitions

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

 // CHECK-LABEL: func @store_non_temporal(
 func.func @store_non_temporal(%input : memref<32xf32, affine_map<(d0) -> (d0)>>, %output : memref<32xf32, affine_map<(d0) -> (d0)>>) {
   %1 = arith.constant 7 : index
   %2 = memref.load %input[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>>
   // CHECK: llvm.store %{{.*}}, %{{.*}}  {nontemporal} : f32, !llvm.ptr
   memref.store %2, %output[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>>
   func.return
 }
+
+// -----
+
+// CHECK:    llvm.mlir.global internal @gv1() {addr_space = 3 : i32} : !llvm.array<1 x i64>
+memref.global "internal" @gv1 : memref<1xi64,3>
+
+func.func @get_gv0_memref() {
+  %0 = memref.get_global @gv1 : memref<1xi64,3>
+  %c0 = arith.constant 0 : index
+  %c4 = arith.constant 4 : i64
+  memref.store %c4, %0[%c0] : memref<1xi64,3>
+  return
+}

mlir/test/Dialect/MemRef/invalid.mlir

 
 // constant qualifier must be after visibility.
 // expected-error @+1 {{expected valid '@'-identifier for symbol name}}
 memref.global constant "private" @foo : memref<2x2xf32>  = "foo"
 
 
 // -----
 
-// expected-error @+1 {{op visibility expected to be one of ["public", "private", "nested"], but got "priate"}}
+// expected-error @+1 {{op visibility expected to be one of ["public", "private", "nested", "internal"], but got "priate"}}
 memref.global "priate" constant @memref5 : memref<2xf32>  = uninitialized
 
 // -----
 
 func.func @nonexistent_global_memref() {
   // expected-error @+1 {{'gv' does not reference a valid global memref}}
   %0 = memref.get_global @gv : memref<3xf32>
   return

mlir/test/IR/traits.mlir

 
 // -----
 
 // expected-error@+1 {{invalid properties {sym_name = "foo_2", sym_visibility} for op test.symbol: Invalid attribute `sym_visibility` in property conversion: unit}}
 "test.symbol"() <{sym_name = "foo_2", sym_visibility}> : () -> ()
 
 // -----
 
-// expected-error@+1 {{visibility expected to be one of ["public", "private", "nested"]}}
+// expected-error@+1 {{visibility expected to be one of ["public", "private", "nested", "internal"]}}
 "test.symbol"() {sym_name = "foo_2", sym_visibility = "foo"} : () -> ()
 
 // -----
 
 "test.symbol"() {sym_name = "foo_3", sym_visibility = "nested"} : () -> ()
 "test.symbol"() {sym_name = "foo_4", sym_visibility = "private"} : () -> ()
 "test.symbol"() {sym_name = "foo_5", sym_visibility = "public"} : () -> ()
 "test.symbol"() {sym_name = "foo_6"} : () -> ()