[mlir][llvm] allow mlir-translate carry custom triple and data layout.

[whchung]

Summary

• Teach mlir-translate to use custom triple and data layout.
• Change convert-to-rocdlir pass to pass AMDGPU-specific triple and target layout string.
• Amend test case to check alloca on non-zero addrspace.

[ftynse]

Thanks, I have been discussing a similar change with @aartbik just today. Please add the test on translation and feel free to land.

As an improvement (fine for a follow-up), could you please expose the triple and data layout as mlir-translate command line options?

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

81

Could you add a test that the LLVM module indeed has the specified triple and target and put it under test/Target ?

[whchung]

Check triple and data layout string.

[whchung]

In D79019#2008214, @ftynse wrote:

Thanks, I have been discussing a similar change with @aartbik just today. Please add the test on translation and feel free to land.

As an improvement (fine for a follow-up), could you please expose the triple and data layout as mlir-translate command line options?

@ftynse I just revised the test. Will submit a follow-up patch for mlir-translate command line options. I don't have authority to land the patch though.

[whchung]

@ftynse I used to submit PRs to MLIR on GitHub. It's actually my first patch on Phabricator so I don't have access push to master. Would you mind help me make the commit? Thanks a lot.

I've got several other patches under review right now. Once I'm more acquainted with the process I'll request write access.

[mehdi_amini]

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

54

I am not convinced by this API: the datalayout is not something we should inject blindly after the fact. This is something that should be defined in MLIR in the first place, otherwise anything done before will make assumption about it.

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

83

Please don't use auto here, use StringRef instead.

Also having these hard-coded here is a bit strange: can we get this from LLVM instead?

[whchung]

Address some review comments.

[whchung]

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

54

@mehdi_amini I'm not quite sure if introducing the concept of triple and data layout in MLIR is the right approach.

Up until now, all dialects (including llvm dialect) do not require any special treatment for target triple or data layout. It's perfect fine to use std.alloca, or llvm.alloca to model allocating on-stack variables in std and llvm dialect. The only thing which is broken at this moment, is when we translate llvm dialect to llvm IR we couldn't get the proper instruction.

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

83

@mehdi_amini I just changed the patch to use StringRef.

On most of the targets on LLVM there is an computeDataLayout function which get you the data layout string given a triple. But none of those functions are exposed as public functions. I want to restrict the scope of the patch to be within MLIR so I hard-coded the string.

[mehdi_amini]

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

54

Up until now, all dialects (including llvm dialect) do not require any special treatment for target triple or data layout.

Well I'm not sure the assumption made in here is always correct, for example on the alignment of everything.
Since LLVM transformations actually need the datalayout to be correct, why can MLIR transformations be correct without a similar concept?

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

83

I expect that a TargetMachine can be created from a Triple and in turn it exposes access to a DataLayout?

There has to be a public way to get this out of LLVM otherwise every frontend would duplicate this information.

[whchung]

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

54

@mehdi_amini alignment is an optional attribute for std.alloca and llvm.alloca. So it's up to the IR builder to figure it out, when an alignment is not specified. And that would correspond to the topic how to properly specify Target Triple and Data Layout. And I think we can move this discussion to the next thread wrt should a TargetMachine be instantiated.

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

83

@mehdi_amini It's definitely possible to create a TargetMachine so we can get Triple and DataLayout. However mlir-translate is designed to produce LLVM IR, not machine ISA, so currently a TargetMachine is not constructed.

Would it be feasible if target triple and data layout be set as command-line options for mlir-translate? This approach was also suggested by @ftynse .

[mehdi_amini]

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

83

Have you looked into how do any frontend that wants to emit LLVM IR retrieve a data layout? Because they have to provide one right?

83

Unless: if we don't provide one and provide only a triple we get the default one for the triple?

[whchung]

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

83

@mehdi_amini

Most of the LLVM applications I know follow the idiom of:

1. create a TargetMachine from a triple
2. use TargetMachine::createDataLayout to fetch the data layout string associated with the triple
3. use Module::setDataLayout

• For clang, llvm::sys::getDefaultTargetTriple() is used if nothing is specified, and a custom -triple command line option could be set for cross compilation. A TargetMachine would be created in the code generation process.

• For XLA, a TargetMachine is also created according to the target.

• Within MLIR, there are several utilities which would also create a TargetMachine : JitRunner, ExecutionEngine, ConvertKernelFuncToCubin. Notice all of these utilities require a TargetMachine because target machine instructions are expected to be produced.

So indeed they create a TargetMachine, specify target triple and data layout, because machine instructions are expected.

On the other hand, the implementation of mlir-translate doesn't employ a TargetMachine. All of the IR coming out from mlir-translate right now, doesn't carry triple or data layout at all. This patch is the first attempt to allow an llvm Module coming out from mlir-translate to carry non-empty triple & data layout.

Would it be feasible if target triple and data layout be set as command-line options for mlir-translate, just like what's there for opt? This approach was also suggested by @ftynse .

[ftynse]

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

83

Correctness of transformations in MLIR wrt data layout is a separate topic, that needs to be handled within MLIR and likely requires non-trivial design, so I feel quite strongly about not blocking this diff on it. (1) We currently don't have any guarantee at all, so potentially _all_ transformations for _any_ layout are wrong. This patch does not make it any worse; if anything, it makes debugging easier because you see the layout in some cases. (2) MLIR translation is a front-end from LLVM's point of view. The front-end should specify a layout, how -- it's up to the front-end. This is no different from typing data layout in the textual module, or from doing casts with vector-type builtins in C.

My thinking is that we ultimately need "triple" and "layout" associated with MLIR's incarnation of the LLVM module. E.g., as attributes on the ModuleOp that comes out of *->LLVM dialect conversion. Then mlir-translate can pick it up transparently, and work in both directions. This requires some untangling in the llvm conversion passes, so in the meantime, we can have mlir-translate set the triple and layout.

Certainly, we can think about having a target description with a layout in MLIR itself, but while we are in the process of thinking, progress is necessary, lest we get stuck in the analysis paralysis.

[mehdi_amini]

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

83

On the other hand, the implementation of mlir-translate doesn't employ a TargetMachine. All of the IR coming out from mlir-translate right now, doesn't carry triple or data layout at all. This patch is the first attempt to allow an llvm Module coming out from mlir-translate to carry non-empty triple & data layout.

Yes, but it does so by patching it after the fact, which as far as I know is not correct: the module has been generated with assumptions from a Datalayout (implicitly) and you're actually force-changing it.

The topic of the TargetMachine is independent, and I don't really get your point: if any LLVM frontend requires a TM to be able to set the DataLayout (I don't know if there is an alternative...) then I don't see why mlir-translate shouldn't do the same.

Certainly, we can think about having a target description with a layout in MLIR itself, but while we are in the process of thinking, progress is necessary,

Right now I see this as implementing something deliberately incorrect.
I am not opposed to taking shortcuts, but I'd like these to be clearly identified, for example there should be a tracking bug and big TODO statement:

// TODO(pr1234): adding a datalayout after the fact is incorrect and we need to have it exposed
// from the beginning and propagated during the lowering to LLVM dialect instead.

lest we get stuck in the analysis paralysis.

I don't see any paralysis here, unless you include "no one is working on it"?

[ftynse]

Yes, but it does so by patching it after the fact, which as far as I know is not correct: the module has been generated with assumptions from a Datalayout (implicitly) and you're actually force-changing it.

The topic of the TargetMachine is independent, and I don't really get your point: if any LLVM frontend requires a TM to be able to set the DataLayout (I don't know if there is an alternative...) then I don't see why mlir-translate shouldn't do the same.

That is why I recommended that mlir-translate take the triple and the data layout if necessary, so it can derive the necessary parts of the LLVM module. This is still dangerous, but will start the bottom-up progression of the data layout until we hit the right position for it in the MLIR.

Right now I see this as implementing something deliberately incorrect.

The flow we already have is also incorrect. I see this patch as less problematic as long as it does not claim it corrects the current flow.

I am not opposed to taking shortcuts, but I'd like these to be clearly identified, for example there should be a tracking bug and big TODO statement:

Sure, we agree here, let's do this.

I don't see any paralysis here, unless you include "no one is working on it"?

I abuse the term a bit, but you know how it works in MLIR: any substantial discussion will be repeatedly side-tracked into considering an ever-increasing number of cases (or deciding they are not worth considering) and paralyze the progress here for a significant amount of time because analysis is happening somewhere.

[mehdi_amini]

In D79019#2012531, @ftynse wrote:

Yes, but it does so by patching it after the fact, which as far as I know is not correct: the module has been generated with assumptions from a Datalayout (implicitly) and you're actually force-changing it.

The topic of the TargetMachine is independent, and I don't really get your point: if any LLVM frontend requires a TM to be able to set the DataLayout (I don't know if there is an alternative...) then I don't see why mlir-translate shouldn't do the same.

That is why I recommended that mlir-translate take the triple and the data layout if necessary, so it can derive the necessary parts of the LLVM module.

What do you see being "derived" on the LLVM module by mlir-translate with this information that wouldn't already be materialized in the LLVM dialect?

Right now I see this as implementing something deliberately incorrect.

The flow we already have is also incorrect. I see this patch as less problematic as long as it does not claim it corrects the current flow.

I am not opposed to taking shortcuts, but I'd like these to be clearly identified, for example there should be a tracking bug and big TODO statement:

Sure, we agree here, let's do this.

@whchung how does that sounds?

[ftynse]

What do you see being "derived" on the LLVM module by mlir-translate with this information that wouldn't already be materialized in the LLVM dialect?

Data layout from the triple.

We should then move this earlier in the pipeline, data layout is actually necessary in std->llvm conversion to properly convert index, among others. So you are right that it should be materialized in the LLVM dialect.

[whchung]

Address review comments from @mehdi_amini and @ftynse. Add TODO in the revised interface.

[whchung]

Right now I see this as implementing something deliberately incorrect.

The flow we already have is also incorrect. I see this patch as less problematic as long as it does not claim it corrects the current flow.

I am not opposed to taking shortcuts, but I'd like these to be clearly identified, for example there should be a tracking bug and big TODO statement:

Sure, we agree here, let's do this.

@whchung how does that sounds?

@mehdi_amini Agree. I've revised the patch so it has a TODO statement in the new interface. I'm applying for a new account on Bugzilla so if you could create one that would be very helpful.

@ftynse I agree with your assessment that data layout is best conveyed in std->llvm conversion process. I myself am seeing sub-optimal codes generated due to how index types get produced.

[whchung]

Rebase.

[aartbik]

Is this revision still going anywhere? Last activity was May?

[whchung]

@aartbik The patch is no longer necessary.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

 def LLVM_FRemOp : LLVM_ArithmeticOp<"frem", "CreateFRem">;
 def LLVM_FNegOp : LLVM_UnaryArithmeticOp<"fneg", "CreateFNeg">;
 
 // Memory-related operations.
 def LLVM_AllocaOp :
     LLVM_OneResultOp<"alloca">,
     Arguments<(ins LLVM_Type:$arraySize, OptionalAttr<I64Attr>:$alignment)> {
   string llvmBuilder = [{
+    llvm::Module *module = builder.GetInsertBlock()->getModule();
+    auto allocaAddrSpace = module->getDataLayout().getAllocaAddrSpace();
     auto *alloca = builder.CreateAlloca(
-      $_resultType->getPointerElementType(), $arraySize);
+      $_resultType->getPointerElementType(), allocaAddrSpace, $arraySize);
     if ($alignment.hasValue()) {
       auto align = $alignment.getValue().getZExtValue();
       if (align != 0)
         alloca->setAlignment(llvm::MaybeAlign(align));
     }
     $res = alloca;
   }];
   let builders = [OpBuilder<

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

 /// Implementation class for module translation. Holds a reference to the module
 /// being translated, and the mappings between the original and the translated
 /// functions, basic blocks and values. It is practically easier to hold these
 /// mappings in one class since the conversion of control flow operations
 /// needs to look up block and function mappings.
 class ModuleTranslation {
 public:
   template <typename T = ModuleTranslation>
-  static std::unique_ptr<llvm::Module> translateModule(Operation *m) {
+  static std::unique_ptr<llvm::Module>
+  translateModule(Operation *m, StringRef triple = "",
+                  StringRef dataLayout = "") {
     if (!satisfiesLLVMModule(m))
       return nullptr;
     if (failed(checkSupportedModuleOps(m)))
       return nullptr;
-    auto llvmModule = prepareLLVMModule(m);
+    auto llvmModule = prepareLLVMModule(m, triple, dataLayout);
     if (!llvmModule)
       return nullptr;
 
     LLVM::ensureDistinctSuccessors(m);
 
     T translator(m, std::move(llvmModule));
     if (failed(translator.convertGlobals()))
       return nullptr;
protected:
   ModuleTranslation(Operation *module,
                     std::unique_ptr<llvm::Module> llvmModule);
   virtual ~ModuleTranslation();
 
   virtual LogicalResult convertOperation(Operation &op,
                                          llvm::IRBuilder<> &builder);
   virtual LogicalResult convertOmpOperation(Operation &op,
                                             llvm::IRBuilder<> &builder);
-  static std::unique_ptr<llvm::Module> prepareLLVMModule(Operation *m);
+  static std::unique_ptr<llvm::Module>
+  prepareLLVMModule(Operation *m, StringRef triple = "",
+                    StringRef dayaLayout = "");
 
   /// A helper to look up remapped operands in the value remapping table.
   SmallVector<llvm::Value *, 8> lookupValues(ValueRange values);
 
 private:
   /// Check whether the module contains only supported ops directly in its body.
   static LogicalResult checkSupportedModuleOps(Operation *m);
 

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

#include "mlir/Dialect/LLVMIR/ROCDLConversions.inc"
 
   /// Allow access to the constructor.
   friend LLVM::ModuleTranslation;
 };
 } // namespace
 
 std::unique_ptr<llvm::Module> mlir::translateModuleToROCDLIR(Operation *m) {
   // lower MLIR (with RODL Dialect) to LLVM IR (with ROCDL intrinsics)
-  auto llvmModule =
-      LLVM::ModuleTranslation::translateModule<ModuleTranslation>(m);
+  auto amdgcnTriple = "amdgcn-amd-amdhsa";
+  auto amdgcnDataLayout =
+      "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-i64:64-"
+      "v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:"
+      "1024-v2048:2048-n32:64-S32-A5-ni:7";
+  auto llvmModule = LLVM::ModuleTranslation::translateModule<ModuleTranslation>(
+      m, amdgcnTriple, amdgcnDataLayout);
 
   // foreach GPU kernel
   // 1. Insert AMDGPU_KERNEL calling convention.
   // 2. Insert amdgpu-flat-workgroup-size(1, 1024) attribute.
   for (auto func :
        ModuleTranslation::getModuleBody(m).getOps<LLVM::LLVMFuncOp>()) {
     if (!func.getAttrOfType<UnitAttr>(gpu::GPUDialect::getKernelFuncAttrName()))
       continue;

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

ModuleTranslation::lookupValues(ValueRange values) {
   for (Value v : values) {
     assert(valueMapping.count(v) && "referencing undefined value");
     remapped.push_back(valueMapping.lookup(v));
   }
   return remapped;
 }
 
 std::unique_ptr<llvm::Module>
-ModuleTranslation::prepareLLVMModule(Operation *m) {
+ModuleTranslation::prepareLLVMModule(Operation *m, StringRef triple,
+                                     StringRef dataLayout) {
   auto *dialect = m->getContext()->getRegisteredDialect<LLVM::LLVMDialect>();
   assert(dialect && "LLVM dialect must be registered");
   // Lock the LLVM context as we might create new types here.
   llvm::sys::SmartScopedLock<true> scopedLock(dialect->getLLVMContextMutex());
 
   auto llvmModule = llvm::CloneModule(dialect->getLLVMModule());
   if (!llvmModule)
     return nullptr;
 
   llvm::LLVMContext &llvmContext = llvmModule->getContext();
   llvm::IRBuilder<> builder(llvmContext);
 
+  // Set target triple string.
+  llvmModule->setTargetTriple(triple);
+
+  // Set data layout string.
+  llvmModule->setDataLayout(dataLayout);
+
   // Inject declarations for `malloc` and `free` functions that can be used in
   // memref allocation/deallocation coming from standard ops lowering.
   llvmModule->getOrInsertFunction("malloc", builder.getInt8PtrTy(),
                                   builder.getInt64Ty());
   llvmModule->getOrInsertFunction("free", builder.getVoidTy(),
                                   builder.getInt8PtrTy());
 
   return llvmModule;
 }

mlir/test/Target/rocdl.mlir

 // RUN: mlir-translate -mlir-to-rocdlir %s | FileCheck %s
 
+// CHECK: target datalayout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-ni:7"
+// CHECK-NEXT: target triple = "amdgcn-amd-amdhsa"
+
 llvm.func @rocdl_special_regs() -> !llvm.i32 {
   // CHECK-LABEL: rocdl_special_regs
   // CHECK: call i32 @llvm.amdgcn.workitem.id.x()
   %1 = rocdl.workitem.id.x : !llvm.i32
   // CHECK: call i32 @llvm.amdgcn.workitem.id.y()
   %2 = rocdl.workitem.id.y : !llvm.i32
   // CHECK: call i32 @llvm.amdgcn.workitem.id.z()
   %3 = rocdl.workitem.id.z : !llvm.i32
llvm.func @rocdl_special_regs() -> !llvm.i32 {
   %12 = rocdl.grid.dim.z : !llvm.i64
   llvm.return %1 : !llvm.i32
 }
 
 llvm.func @kernel_func() attributes {gpu.kernel} {
   // CHECK-LABEL: amdgpu_kernel void @kernel_func
   llvm.return
 }
+
+// CHECK-LABEL: @alloca_non_zero_addrspace
+llvm.func @alloca_non_zero_addrspace(%size : !llvm.i64) {
+  // Alignment automatically set by the LLVM IR builder when alignment attribute
+  // is 0.
+  //  CHECK: alloca {{.*}} align 4, addrspace(5)
+  llvm.alloca %size x !llvm.i32 {alignment = 0} : (!llvm.i64) -> (!llvm<"i32 addrspace(5)*">)
+  // CHECK-NEXT: alloca {{.*}} align 8, addrspace(5)
+  llvm.alloca %size x !llvm.i32 {alignment = 8} : (!llvm.i64) -> (!llvm<"i32 addrspace(5)*">)
+  llvm.return
+}