[mlir][gpu] Fix `gpu.launch_func` arg promotion for LLVM Dialect args.

[pifon2a]

Summary

LaunchFuncOp supports arguments of STD or LLVM types. If the args are already
LLVM types, then memref descriptor would not be flattened correctly.

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

 //   %elementPtr = llvm.getelementptr %array[i]
 //   llvm.store %fieldPtr, %elementPtr
 // return %array
 Value ConvertLaunchFuncOpToGpuRuntimeCallPattern::generateParamsArray(
     gpu::LaunchFuncOp launchOp, ArrayRef<Value> operands,
     OpBuilder &builder) const {
   auto loc = launchOp.getLoc();
   auto numKernelOperands = launchOp.getNumKernelOperands();
-  auto arguments = typeConverter.promoteOperands(
-      loc, launchOp.getOperands().take_back(numKernelOperands),
-      operands.take_back(numKernelOperands), builder);
-  auto numArguments = arguments.size();
-  SmallVector<LLVM::LLVMType, 4> argumentTypes;
-  argumentTypes.reserve(numArguments);
-  for (auto argument : arguments)
-    argumentTypes.push_back(argument.getType().cast<LLVM::LLVMType>());
-  auto structType = LLVM::LLVMType::createStructTy(argumentTypes, StringRef());
-  auto one = builder.create<LLVM::ConstantOp>(loc, llvmInt32Type,
+  auto arguments = operands.take_back(numKernelOperands);
+
+  // Flatten the arguments.
+  SmallVector<Value, 4> flatArgValues;
+  flatArgValues.reserve(arguments.size());
+  for (auto argument : arguments) {
+    auto llvmArgType = argument.getType().cast<LLVM::LLVMType>();
+    if (!llvmArgType.isStructTy()) {
+      flatArgValues.push_back(argument);
+      continue;
+    }
+    for (int32_t j = 0, ej = llvmArgType.getStructNumElements(); j < ej; ++j) {
+      auto elemType = llvmArgType.getStructElementType(j);
+      if (elemType.isArrayTy()) {
+        auto arrayElemType = elemType.getArrayElementType();
+        for (int32_t k = 0, ek = elemType.getArrayNumElements(); k < ek; ++k) {
+          Value elem = builder.create<LLVM::ExtractValueOp>(
+              loc, arrayElemType, argument, builder.getI32ArrayAttr({j, k}));
+          flatArgValues.push_back(elem);
+        }
+        continue;
+      }
+      assert((elemType.isIntegerTy() || elemType.isFloatTy() ||
+              elemType.isDoubleTy() || elemType.isPointerTy()) &&
+             "expected scalar type");
+      Value elem = builder.create<LLVM::ExtractValueOp>(
+          loc, elemType, argument, builder.getI32ArrayAttr(j));
+      flatArgValues.push_back(elem);
+    }
+  }
+
+  // Get types of flattened args.
+  SmallVector<LLVM::LLVMType, 4> flatArgTypes;
+  flatArgTypes.reserve(flatArgValues.size());
+  for (auto argument : flatArgValues) {
+    auto llvmArgType = argument.getType().cast<LLVM::LLVMType>();
+    flatArgTypes.push_back(llvmArgType);
+  }
+
+  // Promote the flattened args to stack.
+  auto structType = LLVM::LLVMType::createStructTy(flatArgTypes, StringRef());
+  Value one = builder.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                                builder.getI32IntegerAttr(1));
-  auto structPtr = builder.create<LLVM::AllocaOp>(
+  Value structPtr = builder.create<LLVM::AllocaOp>(
       loc, structType.getPointerTo(), one, /*alignment=*/0);
-  auto arraySize = builder.create<LLVM::ConstantOp>(
-      loc, llvmInt32Type, builder.getI32IntegerAttr(numArguments));
-  auto arrayPtr = builder.create<LLVM::AllocaOp>(loc, llvmPointerPointerType,
+  Value arraySize = builder.create<LLVM::ConstantOp>(
+      loc, llvmInt32Type, builder.getI32IntegerAttr(flatArgTypes.size()));
+  Value arrayPtr = builder.create<LLVM::AllocaOp>(loc, llvmPointerPointerType,
                                                   arraySize, /*alignment=*/0);
-  auto zero = builder.create<LLVM::ConstantOp>(loc, llvmInt32Type,
+  Value zero = builder.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                                 builder.getI32IntegerAttr(0));
-  for (auto en : llvm::enumerate(arguments)) {
-    auto index = builder.create<LLVM::ConstantOp>(
+  for (auto &en : llvm::enumerate(flatArgValues)) {
+    Value index = builder.create<LLVM::ConstantOp>(
         loc, llvmInt32Type, builder.getI32IntegerAttr(en.index()));
-    auto fieldPtr =
+    Value fieldPtr =
         builder.create<LLVM::GEPOp>(loc, structType.getPointerTo(), structPtr,
-                                    ArrayRef<Value>{zero, index.getResult()});
+                                    ArrayRef<Value>{zero, index});
     builder.create<LLVM::StoreOp>(loc, en.value(), fieldPtr);
-    auto elementPtr = builder.create<LLVM::GEPOp>(loc, llvmPointerPointerType,
-                                                  arrayPtr, index.getResult());
-    auto casted =
+    Value elementPtr = builder.create<LLVM::GEPOp>(loc, llvmPointerPointerType,
+                                                   arrayPtr, index);
+    Value casted =
         builder.create<LLVM::BitcastOp>(loc, llvmPointerType, fieldPtr);
     builder.create<LLVM::StoreOp>(loc, casted, elementPtr);
   }
   return arrayPtr;
 }
 
 // Generates an LLVM IR dialect global that contains the name of the given
 // kernel function as a C string, and returns a pointer to its beginning.

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

 // RUN: mlir-opt -allow-unregistered-dialect %s --gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" | FileCheck %s
 // RUN: mlir-opt -allow-unregistered-dialect %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")
 
   gpu.module @kernel_module attributes {nvvm.cubin = "CUBIN", rocdl.hsaco = "HSACO"} {
     llvm.func @kernel(%arg0: !llvm.float, %arg1: !llvm.ptr<float>) attributes {gpu.kernel} {
       llvm.return
     }
   }
 
-  llvm.func @foo() {
-    %0 = "op"() : () -> (!llvm.float)
-    %1 = "op"() : () -> (!llvm.ptr<float>)
-    %cst = llvm.mlir.constant(8 : index) : !llvm.i64
-
+  llvm.func @foo(%f: !llvm.float, %f_ptr: !llvm.ptr<float>,
+                 %memref : !llvm.struct<
+                   (ptr<i64>, ptr<i64>, i64, array<1 x i64>, array<1 x 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]]) : (!llvm.ptr<i8>) -> !llvm.ptr<i8>
     // CHECK: %[[func:.*]] = llvm.call @mgpuModuleGetFunction(%[[module]], {{.*}}) : (!llvm.ptr<i8>, !llvm.ptr<i8>) -> !llvm.ptr<i8>
     // CHECK: llvm.call @mgpuStreamCreate
-    // CHECK: llvm.call @mgpuLaunchKernel
-    // CHECK: llvm.call @mgpuStreamSynchronize
-    "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1) { kernel = @kernel_module::@kernel }
-        : (!llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.float, !llvm.ptr<float>) -> ()
 
+    // CHECK: llvm.alloca {{.*}} -> !llvm.ptr<struct<"",
+    // CHECK-SAME: (float, ptr<float>, ptr<i64>, ptr<i64>, i64, i64, i64)>>
+
+    // CHECK: llvm.call @mgpuLaunchKernel({{.*}}) : (!llvm.ptr<i8>, !llvm.i64,
+    // CHECK-SAME: !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64,
+    // CHECK-SAME: !llvm.i32, !llvm.ptr<i8>, !llvm.ptr<ptr<i8>>,
+    // CHECK-SAME: !llvm.ptr<ptr<i8>>) -> !llvm.void   
+
+    // CHECK: llvm.call @mgpuStreamSynchronize
+    %cst = llvm.mlir.constant(8 : index) : !llvm.i64
+    "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %f, %f_ptr, %memref) {
+      kernel = @kernel_module::@kernel
+    } : (!llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64,
+         !llvm.float, !llvm.ptr<float>,
+         !llvm.struct<
+           (ptr<i64>, ptr<i64>, i64, array<1 x i64>, array<1 x i64>)
+         >) -> ()
     llvm.return
   }
 
 }