[mlir] Translate global initializers after creating all LLVM IR globals

[jeanPerier]

Summary

In case an operation in a global initializer region refers to another
global variable defined afterwards in the module of itself, translation
to LLVM IR was currently crashing because it could not find the LLVM IR global
when going through the initializer block.

To solve this problem, split global conversion to LLVM IR into two passes. A
first pass that creates LLVM IR global variables, and a second one that converts
the initializer, if any, and adds it to the llvm global.

[ftynse]

Thanks!

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

369

Nit: I would prefer a name that makes it clearer what this function does, e.g. shouldDropGlobalInitializer. The current want reads like it performs the action of dropping

[jeanPerier]

Rename dropGlobalInitializer to shouldDropGlobalInitializer.

[jeanPerier]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

369

Thanks, that's a better name.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 }
 
 /// A helper method to get the single Block in an operation honoring LLVM's
 /// module requirements.
 static Block &getModuleBody(Operation *module) {
   return module->getRegion(0).front();
 }
 
+/// A helper method to decide if a constant must not be set as a global variable
+/// initializer.
+static bool shouldDropGlobalInitializer(llvm::GlobalValue::LinkageTypes linkage,
+                                        llvm::Constant *cst) {
+  return (linkage == llvm::GlobalVariable::ExternalLinkage &&
+          isa<llvm::UndefValue>(cst)) ||
+         linkage == llvm::GlobalVariable::ExternalWeakLinkage;
+}
+
 /// Create named global variables that correspond to llvm.mlir.global
 /// definitions.
 LogicalResult ModuleTranslation::convertGlobals() {
   for (auto op : getModuleBody(mlirModule).getOps<LLVM::GlobalOp>()) {
     llvm::Type *type = convertType(op.getType());
     llvm::Constant *cst = llvm::UndefValue::get(type);
     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 = op.getValueOrNull().dyn_cast_or_null<StringAttr>()) {
         cst = llvm::ConstantDataArray::getString(
             llvmModule->getContext(), strAttr.getValue(), /*AddNull=*/false);
         type = cst->getType();
       } else if (!(cst = getLLVMConstant(type, op.getValueOrNull(), op.getLoc(),
                                          *this))) {
         return failure();
       }
-    } else if (Block *initializer = op.getInitializerBlock()) {
-      llvm::IRBuilder<> builder(llvmModule->getContext());
-      for (auto &op : initializer->without_terminator()) {
-        if (failed(convertOperation(op, builder)) ||
-            !isa<llvm::Constant>(lookupValue(op.getResult(0))))
-          return emitError(op.getLoc(), "unemittable constant value");
-      }
-      ReturnOp ret = cast<ReturnOp>(initializer->getTerminator());
-      cst = cast<llvm::Constant>(lookupValue(ret.getOperand(0)));
     }
 
     auto linkage = convertLinkageToLLVM(op.linkage());
-    bool anyExternalLinkage =
-        ((linkage == llvm::GlobalVariable::ExternalLinkage &&
-          isa<llvm::UndefValue>(cst)) ||
-         linkage == llvm::GlobalVariable::ExternalWeakLinkage);
     auto addrSpace = op.addr_space();
     auto *var = new llvm::GlobalVariable(
         *llvmModule, type, op.constant(), linkage,
-        anyExternalLinkage ? nullptr : cst, op.sym_name(),
+        shouldDropGlobalInitializer(linkage, cst) ? nullptr : cst,
+        op.sym_name(),
         /*InsertBefore=*/nullptr, llvm::GlobalValue::NotThreadLocal, addrSpace);
 
     globalsMapping.try_emplace(op, var);
   }
 
+  // Convert global variable bodies. This is done after all global variables
+  // have been created in LLVM IR because a global body may refer to another
+  // global or itself. So all global variables need to be mapped first.
+  for (auto op : getModuleBody(mlirModule).getOps<LLVM::GlobalOp>()) {
+    if (Block *initializer = op.getInitializerBlock()) {
+      llvm::IRBuilder<> builder(llvmModule->getContext());
+      for (auto &op : initializer->without_terminator()) {
+        if (failed(convertOperation(op, builder)) ||
+            !isa<llvm::Constant>(lookupValue(op.getResult(0))))
+          return emitError(op.getLoc(), "unemittable constant value");
+      }
+      ReturnOp ret = cast<ReturnOp>(initializer->getTerminator());
+      llvm::Constant *cst =
+          cast<llvm::Constant>(lookupValue(ret.getOperand(0)));
+      auto *global = cast<llvm::GlobalVariable>(lookupGlobal(op));
+      if (!shouldDropGlobalInitializer(global->getLinkage(), cst))
+        global->setInitializer(cst);
+    }
+  }
+
   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

mlir/test/Target/LLVMIR/llvmir.mlir

 
 llvm.mlir.global internal constant @taker_of_address() : !llvm.ptr<func<void ()>> {
   %0 = llvm.mlir.addressof @address_taken : !llvm.ptr<func<void ()>>
   llvm.return %0 : !llvm.ptr<func<void ()>>
 }
 
 // -----
 
+// CHECK: @forward_use_of_address = linkonce global float* @address_declared_after_use
+llvm.mlir.global linkonce @forward_use_of_address() : !llvm.ptr<f32> {
+  %0 = llvm.mlir.addressof @address_declared_after_use : !llvm.ptr<f32>
+  llvm.return %0 : !llvm.ptr<f32>
+}
+
+llvm.mlir.global linkonce @address_declared_after_use() : f32
+
+// -----
+
+// CHECK: @take_self_address = linkonce global { i32, i32* } {{.*}} { i32, i32* }* @take_self_address
+llvm.mlir.global linkonce @take_self_address() : !llvm.struct<(i32, !llvm.ptr<i32>)> {
+  %z32 = llvm.mlir.constant(0 : i32) : i32
+  %0 = llvm.mlir.undef : !llvm.struct<(i32, !llvm.ptr<i32>)>
+  %1 = llvm.mlir.addressof @take_self_address : !llvm.ptr<!llvm.struct<(i32, !llvm.ptr<i32>)>>
+  %2 = llvm.getelementptr %1[%z32, %z32] : (!llvm.ptr<!llvm.struct<(i32, !llvm.ptr<i32>)>>, i32, i32) -> !llvm.ptr<i32>
+  %3 = llvm.insertvalue %z32, %0[0 : i32] : !llvm.struct<(i32, !llvm.ptr<i32>)>
+  %4 = llvm.insertvalue %2, %3[1 : i32] : !llvm.struct<(i32, !llvm.ptr<i32>)>
+  llvm.return %4 : !llvm.struct<(i32, !llvm.ptr<i32>)>
+}
+
+// -----
+
 // Check that branch weight attributes are exported properly as metadata.
 llvm.func @cond_br_weights(%cond : i1, %arg0 : i32,  %arg1 : i32) -> i32 {
   // CHECK: !prof ![[NODE:[0-9]+]]
   llvm.cond_br %cond weights(dense<[5, 10]> : vector<2xi32>), ^bb1, ^bb2
 ^bb1:  // pred: ^bb0
   llvm.return %arg0 : i32
 ^bb2:  // pred: ^bb0
   llvm.return %arg1 : i32