[mlir][GPU] make the index bitwidth configurable during the GPU to LLVM conversion

[gysit]

Summary

The patch exposes the LLVM lowering options of the GPU to LLVM conversion pass. Additionally, it introduces the logic needed to cast index type parameters passed to the GPU runtime wrapper if the bitwidth of the index type is set to less than the 64-bits used by the GPU runtime wrapper.

For the review: This is a rather intrusive change and it may interfere with the entire data layout refactoring. If you prefer I can maintain this as a local change.

[csigg]

Hi Tobias!

I'm curious, wouldn't it be easier to compile the runtime wrappers for 32bit instead of casting the arguments?

[mehdi_amini]

mlir/include/mlir/Conversion/Passes.td

127

You're adding many options that I don't find tests for at the moment?

[gysit]

In D95160#2579798, @csigg wrote:

Hi Tobias!

I'm curious, wouldn't it be easier to compile the runtime wrappers for 32bit instead of casting the arguments?

Hi Christian,

My idea was to keep things isolated in the pass and not touch the runtime wrapper. I didn't think about 32-bit compilation. It seems like a clean solution assuming all the rocm / cuda libraries work together with a 32-bit binary? An alternative could be to use preprocessor directives in the wrapper. I am not a big fan of that solution but it may be cleaner / require less code (there are multiple wrappers though).

I suggest we abandon this revision for the time being (we can apply the patch locally in our project). Without the casting it basically boils down to adding the llvm lowering options which should probably be replaced by data layout attributes anyways?

mlir/include/mlir/Conversion/Passes.td

127

That is a good point! I test only the option I actually use and added the other ones for completeness. I should have added and tested only the ones I use instead.

AFAIK the llvm lowering options are replaced by data layout attributes anyways? I think that is a much nicer solution not only from a testing perspective. I thus suggest to abandon this revision (see my answer to @csigg 's comment).

mlir/include/mlir/Conversion/GPUCommon/GPUCommonPass.h

 //===- GPUCommonPass.h - MLIR GPU runtime support -------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #ifndef MLIR_CONVERSION_GPUCOMMON_GPUCOMMONPASS_H_
 #define MLIR_CONVERSION_GPUCOMMON_GPUCOMMONPASS_H_
 
+#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/IR/Module.h"
 #include <vector>
 
 namespace mlir {
 
 class LLVMTypeConverter;
 class Location;

 
 /// Creates a pass to convert a gpu.launch_func operation into a sequence of
 /// GPU runtime calls.
 ///
 /// This pass does not generate code to call GPU runtime APIs directly but
 /// instead uses a small wrapper library that exports a stable and conveniently
 /// typed ABI on top of GPU runtimes such as CUDA or ROCm (HIP).
 std::unique_ptr<OperationPass<ModuleOp>>
-createGpuToLLVMConversionPass(StringRef gpuBinaryAnnotation = "");
+createGpuToLLVMConversionPass(StringRef gpuBinaryAnnotation = "",
+                              const LowerToLLVMOptions &options =
+                                  LowerToLLVMOptions::getDefaultOptions());
 
 /// Collect a set of patterns to convert from the GPU dialect to LLVM.
 void populateGpuToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                          OwningRewritePatternList &patterns,
                                          StringRef gpuBinaryAnnotation);
 
 /// Creates a pass to convert kernel functions into GPU target object blobs.
 ///

mlir/include/mlir/Conversion/Passes.td

 
 def GpuToLLVMConversionPass : Pass<"gpu-to-llvm", "ModuleOp"> {
   let summary = "Convert GPU dialect to LLVM dialect with GPU runtime calls";
   let constructor = "mlir::createGpuToLLVMConversionPass()";
   let dependentDialects = ["LLVM::LLVMDialect"];
   let options = [
     Option<"gpuBinaryAnnotation", "gpu-binary-annotation", "std::string",
            "", "Annotation attribute string for GPU binary">,
+    Option<"useAlignedAlloc", "use-aligned-alloc", "bool", /*default=*/"false",
+           "Use aligned_alloc in place of malloc for heap allocations">,
+    Option<"useBarePtrCallConv", "use-bare-ptr-memref-call-conv", "bool",
+           /*default=*/"false",
+           "Replace FuncOp's MemRef arguments with bare pointers to the MemRef "
+           "element types">,
+    Option<"emitCWrappers", "emit-c-wrappers", "bool", /*default=*/"false",
+           "Emit wrappers for C-compatible pointer-to-struct memref "
+           "descriptors">,
+    Option<"indexBitwidth", "index-bitwidth", "unsigned",
+           /*default=kDeriveIndexBitwidthFromDataLayout*/"0",
+           "Bitwidth of the index type, 0 to use size of machine word">,
+    Option<"dataLayout", "data-layout", "std::string",
+           /*default=*/"\"\"",
+           "String description (LLVM format) of the data layout that is "
+           "expected on the produced module">
   ];
 }
 
 def LowerHostCodeToLLVM : Pass<"lower-host-to-llvm", "ModuleOp"> {
   let summary = "Lowers the host module code and `gpu.launch_func` to LLVM";
   let constructor = "mlir::createLowerHostCodeToLLVMPass()";
   let dependentDialects = ["LLVM::LLVMDialect"];
 }

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

 
 static constexpr const char *kGpuBinaryStorageSuffix = "_gpubin_cst";
 
 namespace {
 
 class GpuToLLVMConversionPass
     : public GpuToLLVMConversionPassBase<GpuToLLVMConversionPass> {
 public:
-  GpuToLLVMConversionPass(StringRef gpuBinaryAnnotation) {
+  GpuToLLVMConversionPass(StringRef gpuBinaryAnnotation,
+                          bool useBarePtrCallConv, bool emitCWrappers,
+                          unsigned indexBitwidth, bool useAlignedAlloc,
+                          const llvm::DataLayout &dataLayout) {
     if (!gpuBinaryAnnotation.empty())
       this->gpuBinaryAnnotation = gpuBinaryAnnotation.str();
+    this->useBarePtrCallConv = useBarePtrCallConv;
+    this->emitCWrappers = emitCWrappers;
+    this->indexBitwidth = indexBitwidth;
+    this->useAlignedAlloc = useAlignedAlloc;
+    this->dataLayout = dataLayout.getStringRepresentation();
   }
 
   // Run the dialect converter on the module.
   void runOnOperation() override;
 };
 
+/// Helper class to build cast operations that adapt the bitwidth of index and
+/// size arguments to match the target function parameters.
+class IndexCastBuilder {
+public:
+  IndexCastBuilder(unsigned indexBitwidth) : indexBitwidth(indexBitwidth) {}
+  Value create(Location loc, OpBuilder &builder, Value value,
+               Type paramType) const;
+
+private:
+  unsigned indexBitwidth;
+};
+
 class FunctionCallBuilder {
 public:
   FunctionCallBuilder(StringRef functionName, Type returnType,
                       ArrayRef<Type> argumentTypes)
       : functionName(functionName),
         functionType(LLVM::LLVMFunctionType::get(returnType, argumentTypes)) {}
   LLVM::CallOp create(Location loc, OpBuilder &builder,
-                      ArrayRef<Value> arguments) const;
+                      ArrayRef<Value> arguments,
+                      const IndexCastBuilder *indexCastBuilder = nullptr) const;
 
 private:
   StringRef functionName;
   LLVM::LLVMFunctionType functionType;
 };
 
 template <typename OpTy>
 class ConvertOpToGpuRuntimeCallPattern : public ConvertOpToLLVMPattern<OpTy> {
Type llvmPointerType =
       LLVM::LLVMPointerType::get(IntegerType::get(context, 8));
   Type llvmPointerPointerType = LLVM::LLVMPointerType::get(llvmPointerType);
   Type llvmInt8Type = IntegerType::get(context, 8);
   Type llvmInt32Type = IntegerType::get(context, 32);
   Type llvmInt64Type = IntegerType::get(context, 64);
   Type llvmIntPtrType = IntegerType::get(
       context, this->getTypeConverter()->getPointerBitwidth(0));
 
+  IndexCastBuilder indexCastBuilder = {
+      this->getTypeConverter()->getIndexTypeBitwidth()};
+
   FunctionCallBuilder moduleLoadCallBuilder = {
       "mgpuModuleLoad",
       llvmPointerType /* void *module */,
       {llvmPointerType /* void *cubin */}};
   FunctionCallBuilder moduleUnloadCallBuilder = {
       "mgpuModuleUnload", llvmVoidType, {llvmPointerType /* void *module */}};
   FunctionCallBuilder moduleGetFunctionCallBuilder = {
       "mgpuModuleGetFunction",

 private:
   LogicalResult
   matchAndRewrite(gpu::MemcpyOp memcpyOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 } // namespace
 
 void GpuToLLVMConversionPass::runOnOperation() {
-  LLVMTypeConverter converter(&getContext());
+  LowerToLLVMOptions options = {useBarePtrCallConv, emitCWrappers,
+                                indexBitwidth, useAlignedAlloc,
+                                llvm::DataLayout(this->dataLayout)};
+  LLVMTypeConverter converter(&getContext(), options);
   OwningRewritePatternList patterns;
   populateStdToLLVMConversionPatterns(converter, patterns);
   populateGpuToLLVMConversionPatterns(converter, patterns, gpuBinaryAnnotation);
 
   LLVMConversionTarget target(getContext());
   if (failed(
           applyPartialConversion(getOperation(), target, std::move(patterns))))
     signalPassFailure();
 }
 
-LLVM::CallOp FunctionCallBuilder::create(Location loc, OpBuilder &builder,
-                                         ArrayRef<Value> arguments) const {
+Value IndexCastBuilder::create(Location loc, OpBuilder &builder, Value value,
+                               Type paramType) const {
+  // Only cast arguments of index or integer arguments if their bitwidth is
+  // lower than the bitwidth of the target function parameter.
+  if ((value.getType().isIndex() ||
+       value.getType().isSignlessInteger(indexBitwidth)) &&
+      paramType.getIntOrFloatBitWidth() > indexBitwidth) {
+    return builder.create<LLVM::ZExtOp>(loc, paramType, value);
+  }
+  return value;
+}
+
+LLVM::CallOp
+FunctionCallBuilder::create(Location loc, OpBuilder &builder,
+                            ArrayRef<Value> arguments,
+                            const IndexCastBuilder *indexCastBuilder) const {
   auto module = builder.getBlock()->getParent()->getParentOfType<ModuleOp>();
   auto function = [&] {
     if (auto function = module.lookupSymbol<LLVM::LLVMFuncOp>(functionName))
       return function;
     return OpBuilder(module.getBody()->getTerminator())
         .create<LLVM::LLVMFuncOp>(loc, functionName, functionType);
   }();
+  // Optionally cast the index arguments to extend the bitwidth of index or
+  // integer arguments to match the bitwidth of the function parameters.
+  SmallVector<Value, 4> castedArguments(arguments.begin(), arguments.end());
+  if (indexCastBuilder) {
+    castedArguments.reserve(arguments.size());
+    for (auto en : llvm::enumerate(arguments)) {
+      // Get the function parameter type.
+      auto paramType = const_cast<LLVM::LLVMFunctionType &>(functionType)
+                           .getParamType(en.index());
+      castedArguments[en.index()] =
+          indexCastBuilder->create(loc, builder, en.value(), paramType);
+    }
+  }
   return builder.create<LLVM::CallOp>(
       loc, const_cast<LLVM::LLVMFunctionType &>(functionType).getReturnType(),
-      builder.getSymbolRefAttr(function), arguments);
+      builder.getSymbolRefAttr(function), castedArguments);
 }
 
 // Returns whether all operands are of LLVM type.
 static LogicalResult areAllLLVMTypes(Operation *op, ValueRange operands,
                                      ConversionPatternRewriter &rewriter) {
   if (!llvm::all_of(operands, [](Value value) {
         return LLVM::isCompatibleType(value.getType());
       }))
LogicalResult ConvertHostRegisterOpToGpuRuntimeCallPattern::matchAndRewrite(
 
   auto memRefType = hostRegisterOp.value().getType();
   auto elementType = memRefType.cast<UnrankedMemRefType>().getElementType();
   auto elementSize = getSizeInBytes(loc, elementType, rewriter);
 
   auto arguments = getTypeConverter()->promoteOperands(loc, op->getOperands(),
                                                        operands, rewriter);
   arguments.push_back(elementSize);
-  hostRegisterCallBuilder.create(loc, rewriter, arguments);
+  hostRegisterCallBuilder.create(loc, rewriter, arguments, &indexCastBuilder);
 
   rewriter.eraseOp(op);
   return success();
 }
 
 LogicalResult ConvertAllocOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::AllocOp allocOp, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
LogicalResult ConvertAllocOpToGpuRuntimeCallPattern::matchAndRewrite(
   getMemRefDescriptorSizes(loc, memRefType, adaptor.dynamicSizes(), rewriter,
                            shape, strides, sizeBytes);
 
   // Allocate the underlying buffer and store a pointer to it in the MemRef
   // descriptor.
   Type elementPtrType = this->getElementPtrType(memRefType);
   auto stream = adaptor.asyncDependencies().front();
   Value allocatedPtr =
-      allocCallBuilder.create(loc, rewriter, {sizeBytes, stream}).getResult(0);
+      allocCallBuilder
+          .create(loc, rewriter, {sizeBytes, stream}, &indexCastBuilder)
+          .getResult(0);
   allocatedPtr =
       rewriter.create<LLVM::BitcastOp>(loc, elementPtrType, allocatedPtr);
 
   // No alignment.
   Value alignedPtr = allocatedPtr;
 
   // Create the MemRef descriptor.
   auto memRefDescriptor = this->createMemRefDescriptor(
LogicalResult ConvertLaunchFuncOpToGpuRuntimeCallPattern::matchAndRewrite(
   auto kernelParams = generateParamsArray(launchOp, operands, rewriter);
   auto nullpointer = rewriter.create<LLVM::NullOp>(loc, llvmPointerPointerType);
   launchKernelCallBuilder.create(loc, rewriter,
                                  {function.getResult(0), launchOp.gridSizeX(),
                                   launchOp.gridSizeY(), launchOp.gridSizeZ(),
                                   launchOp.blockSizeX(), launchOp.blockSizeY(),
                                   launchOp.blockSizeZ(),
                                   /*sharedMemBytes=*/zero, stream, kernelParams,
-                                  /*extra=*/nullpointer});
+                                  /*extra=*/nullpointer},
+                                 &indexCastBuilder);
 
   if (launchOp.asyncToken()) {
     // Async launch: make dependent ops use the same stream.
     rewriter.replaceOp(launchOp, {stream});
   } else {
     // Synchronize with host and destroy stream. This must be the stream created
     // above (with no other uses) because we check that the synchronous version
     // does not have any async dependencies.
LogicalResult ConvertMemcpyOpToGpuRuntimeCallPattern::matchAndRewrite(
 
   auto src = rewriter.create<LLVM::BitcastOp>(
       loc, llvmPointerType, srcDesc.alignedPtr(rewriter, loc));
   auto dst = rewriter.create<LLVM::BitcastOp>(
       loc, llvmPointerType,
       MemRefDescriptor(adaptor.dst()).alignedPtr(rewriter, loc));
 
   auto stream = adaptor.asyncDependencies().front();
-  memcpyCallBuilder.create(loc, rewriter, {dst, src, sizeBytes, stream});
+  memcpyCallBuilder.create(loc, rewriter, {dst, src, sizeBytes, stream},
+                           &indexCastBuilder);
 
   rewriter.replaceOp(memcpyOp, {stream});
 
   return success();
 }
 
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
-mlir::createGpuToLLVMConversionPass(StringRef gpuBinaryAnnotation) {
-  return std::make_unique<GpuToLLVMConversionPass>(gpuBinaryAnnotation);
+mlir::createGpuToLLVMConversionPass(StringRef gpuBinaryAnnotation,
+                                    const LowerToLLVMOptions &options) {
+  return std::make_unique<GpuToLLVMConversionPass>(
+      gpuBinaryAnnotation, options.useBarePtrCallConv, options.emitCWrappers,
+      options.indexBitwidth, options.useAlignedAlloc, options.dataLayout);
 }
 
 void mlir::populateGpuToLLVMConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
     StringRef gpuBinaryAnnotation) {
   converter.addConversion(
       [context = &converter.getContext()](gpu::AsyncTokenType type) -> Type {
         return LLVM::LLVMPointerType::get(IntegerType::get(context, 8));

mlir/test/Conversion/GPUCommon/lower-alloc-to-gpu-runtime-calls.mlir

 // RUN: mlir-opt %s --gpu-to-llvm | FileCheck %s
+// RUN: mlir-opt %s --gpu-to-llvm="index-bitwidth=32" | FileCheck %s --check-prefix=CHECK32
 
 module attributes {gpu.container_module} {
   // CHECK-LABEL: llvm.func @main
   // CHECK-SAME: %[[size:.*]]: i64
   func @main(%size : index) {
     // CHECK: %[[stream:.*]] = llvm.call @mgpuStreamCreate()
     %0 = gpu.wait async
     // CHECK: %[[gep:.*]] = llvm.getelementptr {{.*}}[%[[size]]]
     // CHECK: %[[size_bytes:.*]] = llvm.ptrtoint %[[gep]]
+    // CHECK32: %[[size_bytes:.*]] = llvm.ptrtoint
+    // CHECK32: {{%.*}} = llvm.zext %[[size_bytes:.*]] : i32 to i64
     // CHECK: llvm.call @mgpuMemAlloc(%[[size_bytes]], %[[stream]])
     %1, %2 = gpu.alloc async [%0] (%size) : memref<?xf32>
     // CHECK: %[[float_ptr:.*]] = llvm.extractvalue {{.*}}[0]
     // CHECK: %[[void_ptr:.*]] = llvm.bitcast %[[float_ptr]]
     // CHECK: llvm.call @mgpuMemFree(%[[void_ptr]], %[[stream]])
     %3 = gpu.dealloc async [%2] %1 : memref<?xf32>
     // CHECK: llvm.call @mgpuStreamSynchronize(%[[stream]])
     // CHECK: llvm.call @mgpuStreamDestroy(%[[stream]])
     gpu.wait [%3]
     return
   }
 }

mlir/test/Conversion/GPUCommon/lower-launch-func-to-gpu-runtime-calls.mlir

 // RUN: mlir-opt %s --gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" | FileCheck %s
+// RUN: mlir-opt %s --gpu-to-llvm="gpu-binary-annotation=nvvm.cubin index-bitwidth=32" | FileCheck %s --check-prefix=CHECK32
 // RUN: mlir-opt %s --gpu-to-llvm="gpu-binary-annotation=rocdl.hsaco" | FileCheck %s --check-prefix=ROCDL
 
 module attributes {gpu.container_module} {
 
   // CHECK: llvm.mlir.global internal constant @[[KERNEL_NAME:.*]]("kernel\00")
   // CHECK: llvm.mlir.global internal constant @[[GLOBAL:.*]]("CUBIN")
   // ROCDL: llvm.mlir.global internal constant @[[GLOBAL:.*]]("HSACO")
 
func @foo(%buffer: memref<?xf32>) {
     gpu.launch_func @kernel_module::@kernel
         blocks in (%c8, %c8, %c8)
         threads in (%c8, %c8, %c8)
         args(%c32 : i32, %buffer : memref<?xf32>)
     return
   }
 
   // CHECK: [[C8:%.*]] = llvm.mlir.constant(8 : index) : i64
+  // CHECK32: [[C8:%.*]] = llvm.mlir.constant(8 : index) : i32
+  // CHECK32: {{%.*}} = llvm.zext [[C8]] : i32 to i64
   // CHECK: [[ADDRESSOF:%.*]] = llvm.mlir.addressof @[[GLOBAL]]
   // CHECK: [[C0:%.*]] = llvm.mlir.constant(0 : index)
   // CHECK: [[BINARY:%.*]] = llvm.getelementptr [[ADDRESSOF]]{{\[}}[[C0]], [[C0]]]
   // CHECK-SAME: -> !llvm.ptr<i8>
 
   // CHECK: [[MODULE:%.*]] = llvm.call @mgpuModuleLoad([[BINARY]])
   // CHECK: [[FUNC:%.*]] = llvm.call @mgpuModuleGetFunction([[MODULE]], {{.*}})
 

mlir/test/Conversion/GPUCommon/lower-memcpy-to-gpu-runtime-calls.mlir

 // RUN: mlir-opt %s --gpu-to-llvm | FileCheck %s
+// RUN: mlir-opt %s --gpu-to-llvm="index-bitwidth=32" | FileCheck %s --check-prefix=CHECK32
 
 module attributes {gpu.container_module} {
 
   // CHECK: func @foo
   func @foo(%dst : memref<7xf32, 1>, %src : memref<7xf32>) {
     // CHECK: %[[t0:.*]] = llvm.call @mgpuStreamCreate
     %t0 = gpu.wait async
     // CHECK: %[[size_bytes:.*]] = llvm.ptrtoint
+    // CHECK32: %[[size_bytes:.*]] = llvm.ptrtoint
+    // CHECK32: {{%.*}} = llvm.zext %[[size_bytes:.*]] : i32 to i64
     // CHECK: %[[src:.*]] = llvm.bitcast
     // CHECK: %[[dst:.*]] = llvm.bitcast
     // CHECK: llvm.call @mgpuMemcpy(%[[dst]], %[[src]], %[[size_bytes]], %[[t0]])
     %t1 = gpu.memcpy async [%t0] %dst, %src : memref<7xf32, 1>, memref<7xf32>
     // CHECK: llvm.call @mgpuStreamSynchronize(%[[t0]])
     // CHECK: llvm.call @mgpuStreamDestroy(%[[t0]])
     gpu.wait [%t1]
     return
   }
 }