[mlir] Add a `GlobalSymbol` trait.

[fmorac]

Summary

For an explanation of these patches see D154153.

Commit message:
Currently there's no mechanism to indicate that a particular symbol should be
treated as a global symbol, this patch adds this trait and modifies the LLVM
translation mechanism to be aware of the trait and convert the operation
alongside all the other LLVM globals.

[fmorac]

Rebasing.

[mehdi_amini]

What does it mean for a symbol to have "global scope"?

[fmorac]

In D154098#4464065, @mehdi_amini wrote:

What does it mean for a symbol to have "global scope"?

I don't have a good definition, but here's how I see the definition of global scope: if there were a Module trait, then a global scope Op means having a parent with the Module trait.

[mehdi_amini]

I don't understand what this is intended to model at the moment unfortunately, do you have an example of why this is needed?

[fmorac]

In D154098#4464847, @mehdi_amini wrote:

I don't understand what this is intended to model at the moment unfortunately, do you have an example of why this is needed?

Yes, so I use it in the gpu.binary op in D154132. The reason I need it there is because the binary embedding should (ideally) be performed before converting the gpu.launch_func op, but without this trait the gpu.launch_func op would be converted first because it's inside a function body: https://github.com/llvm/llvm-project/blob/main/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp#L1411-L1446 .

For me it also makes sense as it allows translating non MLIR LLVM global ops in a given order, ie translating global ops before function bodies.
As to why I view the gpu.binary op as global op and not a binary string, it's because the embedding can be performed in multiple ways. The default MLIR way is just adding a string, however Clang Offload annotations embeds it as a string coupled with other annotations, a CUDA embedding would be the fat binary + the kernel stubs.

[fmorac]

Rebasing.

[krzysz00]

mlir/include/mlir/IR/SymbolTable.h

453

I don't think "global scope" is the right wording here, given what this trait is doing.

Heck, the name is one I'm iffy on, though I don't have any clearer suggestions.

But, in any case, how's this:

"Trait for operations that define symbols which, within the context of a module, define global values, whether variable or constant. These operations must also implement the Symbol trait and cannot produce results. Examples of such operations include llvm.global, memref.global, and gpu.binary.

This trait is not meant for operations that define functions, modules, or other values that have more complex internal structure."

And then, perhaps, applying the GlobalSymbol

[fmorac]

Changed the trait's explanation.

[mehdi_amini]

In D154098#4464983, @fmorac wrote:

In D154098#4464847, @mehdi_amini wrote:

I don't understand what this is intended to model at the moment unfortunately, do you have an example of why this is needed?

Yes, so I use it in the gpu.binary op in D154132.

So I there that you're adding it to the op, but can you also point me where it is used then?

[fmorac]

In D154098#4526990, @mehdi_amini wrote:

So I there that you're adding it to the op, but can you also point me where it is used then?

The only places that it's being used in C++ code is in this patch, Lines 811 & 1416 in ModuleTranslation.cpp. The sole purpose of the trait right now is allowing certain symbols to be translated alongside LLVM::GlobalOp and before translating all non-LLVM ops. Thus you only have to attach it to the trait list and the modifications in this patch take care of the rest.

[mehdi_amini]

In D154098#4527561, @fmorac wrote:

In D154098#4526990, @mehdi_amini wrote:

So I there that you're adding it to the op, but can you also point me where it is used then?

The only places that it's being used in C++ code is in this patch, Lines 811 & 1416 in ModuleTranslation.cpp. The sole purpose of the trait right now is allowing certain symbols to be translated alongside LLVM::GlobalOp and before translating all non-LLVM ops. Thus you only have to attach it to the trait list and the modifications in this patch take care of the rest.

OK I see. I think @ftynse should chime in here about why this is needed: the translation code does not say why we phase things such that we first process "globals" and then functions. Is this because we need to have definitions before references? I would think that we could have a two-pass model: first translate the declaration for every global (that includes functions in LLVM terminology I believe), and then convert the definitions.

[gysit]

mlir/include/mlir/IR/SymbolTable.h

460

Is the trait marking all non-function global objects (= GlobalVariable in LLVM terminology https://llvm.org/doxygen/classllvm_1_1GlobalObject.html)? If yes, I would consider naming the trait GlobalVariable or GlobalVariableLike?

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

810

Should the existing global variable operations have the trait as well?

[fmorac]

In D154098#4529658, @mehdi_amini wrote:

OK I see. I think @ftynse should chime in here about why this is needed: the translation code does not say why we phase things such that we first process "globals" and then functions. Is this because we need to have definitions before references? I would think that we could have a two-pass model: first translate the declaration for every global (that includes functions in LLVM terminology I believe), and then convert the definitions.

In the future I'd argue we should even consider introducing interfaces for translation, like having LLVMFuncTranslationInterface with methods for translating the signature and the body instead on only looking at LLVMFuncOp, that would allow non-LLVM dialects to have the same treatment as the LLVM dialect, instead of being relegated to be translated at the end, see: https://github.com/llvm/llvm-project/blob/main/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp#L1383-L1395 . But that's out of scope for this patch.

mlir/include/mlir/IR/SymbolTable.h

460

I avoided LLVM names because the trait could be applicable in non-LLVM contexts, however the name is open to change and I like GlobalVariableLike or GlobalSymbolLike.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

810

That's a more difficult conversation , because LLVM globals are already handled a certain way, but I'd argue yes provided there are other refactors -see my comment non-inline comment.

[mehdi_amini]

To be clear: my previous comment is trying to find an alternative solution, as I'm not convinced by this trait as it is at the moment.

[fmorac]

In D154098#4533672, @mehdi_amini wrote:

To be clear: my previous comment is trying to find an alternative solution, as I'm not convinced by this trait as it is at the moment.

Oh ok, I see what you mean, LLVM translation has the following pattern (ModuleTranslation.cpp):

if (failed(translator.convertFunctions()))
  return nullptr;

// Convert other top-level operations if possible.
llvm::IRBuilder<> llvmBuilder(llvmContext);
for (Operation &o : getModuleBody(module).getOperations()) {
  if (!isa<LLVM::LLVMFuncOp, LLVM::GlobalOp, LLVM::GlobalCtorsOp,
           LLVM::GlobalDtorsOp, LLVM::MetadataOp, LLVM::ComdatOp>(&o) &&
      !o.hasTrait<OpTrait::IsTerminator>() &&
      failed(translator.convertOperation(o, llvmBuilder))) {
    return nullptr;
  }
}

If we change it to:

// Convert other top-level operations if possible.
llvm::IRBuilder<> llvmBuilder(llvmContext);
for (Operation &o : getModuleBody(module).getOperations()) {
  if (!isa<LLVM::LLVMFuncOp, LLVM::GlobalOp, LLVM::GlobalCtorsOp,
           LLVM::GlobalDtorsOp, LLVM::MetadataOp, LLVM::ComdatOp>(&o) &&
      !o.hasTrait<OpTrait::IsTerminator>() &&
      failed(translator.convertOperation(o, llvmBuilder))) {
    return nullptr;
  }
}

if (failed(translator.convertFunctions()))
  return nullptr;

The issue for gpu.binary is solved, however I also don't know if that breaks something else (I think it shouldn't, as no op should depend on the body of a function).

[gysit]

Can you explain why the gpu.binary operation needs to be translated before the gpu.launch? Couldn't the gpu.launch conversion introduce an empty global variable that then is finalized once you translate the gpu.binary operation? The latter could then use getOrInsertGlobal to get and modify the global introduced during the conversion of the gpu.launch. This would make the translation robust against changing the order in which things are converted.

The issue for gpu.binary is solved, however I also don't know if that breaks something else (I think it shouldn't, as no op should depend on the body of a function).

I think I prefer moving the function conversion after the conversion of the non-LLVM dialect globals over introducing the trait, but I am unsure if this has undesired side-effects (I would assume no).

[fmorac]

In D154098#4534270, @gysit wrote:

Can you explain why the gpu.binary operation needs to be translated before the gpu.launch? Couldn't the gpu.launch conversion introduce an empty global variable that then is finalized once you translate the gpu.binary operation? The latter could then use getOrInsertGlobal to get and modify the global introduced during the conversion of the gpu.launch. This would make the translation robust against changing the order in which things are converted.

The issue is that not all gpu.binary ops have to decay into simple strings, they could decay into a string and function stubs in a case of a CUDA like translation. Thus to avoid any potential issues with other kinds of translations it's better to always translate first gpu.binary.

In D154098#4534270, @gysit wrote:

I think I prefer moving the function conversion after the conversion of the non-LLVM dialect globals over introducing the trait, but I am unsure if this has undesired side-effects (I would assume no).

Abandoning in favor of D156284 (I'll wait until all the other patches get through before abandoning).

[gysit]

The issue is that not all gpu.binary ops have to decay into simple strings, they could decay into a string and function stubs in a case of a CUDA like translation. Thus to avoid any potential issues with other kinds of translations it's better to always translate first gpu.binary.

Just out of curiosity, do you have an example in mind where that would be a problem? It sounds at first glance that it doesn't matter too much if gpu.binary lowers to one or more ops?

[fmorac]

In D154098#4535341, @gysit wrote:

The issue is that not all gpu.binary ops have to decay into simple strings, they could decay into a string and function stubs in a case of a CUDA like translation. Thus to avoid any potential issues with other kinds of translations it's better to always translate first gpu.binary.

Just out of curiosity, do you have an example in mind where that would be a problem? It sounds at first glance that it doesn't matter too much if gpu.binary lowers to one or more ops?

Ok, before the example some context, GPU binaries don't have to come from gpu.modules they could come from an external source, eg. a CUDA binary compiled with nvcc.

If the current translation order remains unaltered then gpu.binary will always be translated after gpu.launch_func. Here's an scenario were that's a potential issue:

• Lets say you have an external CUDA binary and you create a gpu.binary out of that external object.
• In this scheme the ObjectManager takes a binary, embeds it into the IR and creates a function stub out of every kernel in the binary.
• When translating gpu.launch_func the ObjectManager searches for the stub and creates the kernel launch call.

If you translate gpu.binary before gpu.launch_func then the above scheme works perfectly and you can verify that each gpu.launch_func is calling a real function because if the stub is not there then there's an error.

You could read the binary almost every time you encounter a launch_func create the stub and verify the kernel launch, however that's a waste of resources.

Also there could be more complicated schemes where is not practical translating the binary after the kernel launch.

[gysit]

Thanks for the explanation. Yeah then reordering the translation steps may make sense in this case!

mlir/include/mlir/IR/SymbolInterfaces.td

let methods = [
     >,
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // Symbol Traits
 //===----------------------------------------------------------------------===//
 
+// Op defines a global symbol.
+def GlobalSymbol : NativeOpTrait<"GlobalSymbol", [Symbol]>;
+
 // Op defines a symbol table.
 def SymbolTable : NativeOpTrait<"SymbolTable">;
 
 #endif // MLIR_IR_SYMBOLINTERFACES

mlir/include/mlir/IR/SymbolTable.h

ParseResult parseOptionalVisibilityKeyword(OpAsmParser &parser,
                                            NamedAttrList &attrs);
 } // namespace impl
 
 } // namespace mlir
 
 /// Include the generated symbol interfaces.
 #include "mlir/IR/SymbolInterfaces.h.inc"
 
+namespace mlir {
+namespace OpTrait {
+/// This trait is for operations defining a global value within the context of a
+/// module, whether variable or constant. These operations must also implement
+/// the Symbol trait and cannot produce results. Examples of such operations
+/// include `llvm.global`, `memref.global`, and `gpu.binary`. Operations with a
+/// more complex internal structure, such as functions, modules, or other
+/// values, should not use this trait.
+template <typename ConcreteType>
+class GlobalSymbol : public TraitBase<ConcreteType, GlobalSymbol> {
+  static LogicalResult verifyTrait(Operation *op) {
+    static_assert(ConcreteType::template hasTrait<ZeroResults>(),
+                  "expected operation must have zero results");
+    static_assert(ConcreteType::template hasTrait<SymbolOpInterface::Trait>(),
+                  "expected operation must inherit the `Symbol` trait");
+    return success();
+  }
+};
+} // namespace OpTrait
+} // namespace mlir
+
 #endif // MLIR_IR_SYMBOLTABLE_H

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 /// `dsoLocalRequested` is true, otherwise it is left unchanged.
 static void addRuntimePreemptionSpecifier(bool dsoLocalRequested,
                                           llvm::GlobalValue *gv) {
   if (dsoLocalRequested)
     gv->setDSOLocal(true);
 }
 
 /// Create named global variables that correspond to llvm.mlir.global
-/// definitions. Convert llvm.global_ctors and global_dtors ops.
+/// definitions. Convert llvm.global_ctors and global_dtors ops. Finally convert
+/// operations with the `GlobalSymbol` trait.
 LogicalResult ModuleTranslation::convertGlobals() {
   for (auto op : getModuleBody(mlirModule).getOps<LLVM::GlobalOp>()) {
     llvm::Type *type = convertType(op.getType());
     llvm::Constant *cst = nullptr;
     if (op.getValueOrNull()) {
       // String attributes are treated separately because they cannot appear as
       // in-function constants and are thus not supported by getLLVMConstant.
       if (auto strAttr = dyn_cast_or_null<StringAttr>(op.getValueOrNull())) {
for (auto symbolAndPriority : range) {
                      /*Data=*/nullptr);
     }
   }
 
   for (auto op : getModuleBody(mlirModule).getOps<LLVM::GlobalOp>())
     if (failed(convertDialectAttributes(op)))
       return failure();
 
+  // Convert operations having the `GlobalSymbol` trait.
+  {
+    llvm::IRBuilder<> llvmBuilder(llvmModule->getContext());
+    for (Operation &op : getModuleBody(mlirModule).getOperations()) {
+      if (!isa<LLVM::GlobalOp, LLVM::GlobalCtorsOp, LLVM::GlobalDtorsOp>(&op) &&
+          op.hasTrait<OpTrait::GlobalSymbol>() &&
+          failed(convertOperation(op, llvmBuilder))) {
+        return failure();
+      }
+    }
+  }
+
   return success();
 }
 
 /// Attempts to add an attribute identified by `key`, optionally with the given
 /// `value` to LLVM function `llvmFunc`. Reports errors at `loc` if any. If the
 /// attribute has a kind known to LLVM IR, create the attribute of this kind,
 /// otherwise keep it as a string attribute. Performs additional checks for
 /// attributes known to have or not have a value in order to avoid assertions
if (failed(translator.convertFunctions()))
     return nullptr;
 
   // Convert other top-level operations if possible.
   llvm::IRBuilder<> llvmBuilder(llvmContext);
   for (Operation &o : getModuleBody(module).getOperations()) {
     if (!isa<LLVM::LLVMFuncOp, LLVM::GlobalOp, LLVM::GlobalCtorsOp,
              LLVM::GlobalDtorsOp, LLVM::MetadataOp, LLVM::ComdatOp>(&o) &&
         !o.hasTrait<OpTrait::IsTerminator>() &&
+        !o.hasTrait<OpTrait::GlobalSymbol>() &&
         failed(translator.convertOperation(o, llvmBuilder))) {
       return nullptr;
     }
   }
 
   // Convert module itself.
   if (failed(translator.convertOperation(*module, llvmBuilder)))
     return nullptr;
 
   if (llvm::verifyModule(*translator.llvmModule, &llvm::errs()))
     return nullptr;
 
   return std::move(translator.llvmModule);
 }