[mlir][sparse][gpu] generate single module, unique kernel names

[aartbik]

Summary

This fixes a TODO in the first version.

[aartbik]

removed unrelated bazel

[aartbik]

and add related file

[aartbik]

typo

[Peiming]

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

59

Maybe you can make it a static local variable if this is not against LLVM coding convention?

[aartbik]

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

59

Alas, that would make it thread unsafe (since several threads may run this code in parallel on different modules in theory).

[aartbik]

rebased with main

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

 //===----------------------------------------------------------------------===//
 
 /// Marks the given top module as a GPU container module.
 static void markAsGPUContainer(ModuleOp topModule) {
   topModule->setAttr(gpu::GPUDialect::getContainerModuleAttrName(),
                      UnitAttr::get(topModule->getContext()));
 }
 
-/// Constructs a new GPU module (for GPU kernels) inside the given top module.
-static gpu::GPUModuleOp genGPUModule(OpBuilder &builder, ModuleOp topModule,
-                                     StringRef name) {
+/// Constructs a new GPU module (for GPU kernels) inside the given top module,
+/// or returns an existing GPU module if one was built previously.
+static gpu::GPUModuleOp genGPUModule(OpBuilder &builder, ModuleOp topModule) {
+  for (auto op : topModule.getBodyRegion().getOps<gpu::GPUModuleOp>())
+    return op; // existing
   markAsGPUContainer(topModule);
   builder.setInsertionPointToStart(&topModule.getBodyRegion().front());
-  return builder.create<gpu::GPUModuleOp>(topModule->getLoc(), name);
+  return builder.create<gpu::GPUModuleOp>(topModule->getLoc(),
+                                          "sparse_kernels");
 }
 
 /// Constructs a new GPU kernel in the given GPU module.
 static gpu::GPUFuncOp genGPUFunc(OpBuilder &builder, gpu::GPUModuleOp gpuModule,
-                                 StringRef name, SmallVectorImpl<Value> &args) {
+                                 SmallVectorImpl<Value> &args) {
+  // Get a unique kernel name. Not very creative,
+  // but we simply try kernel0, kernel1, etc.
+  unsigned kernelNumber = 0;
+  SmallString<16> kernelName;
+  do {
+    kernelName.clear();
+    ("kernel" + Twine(kernelNumber++)).toStringRef(kernelName);
+  } while (gpuModule.lookupSymbol(kernelName));
+  // Then we insert a new kernel with given arguments into the module.
   builder.setInsertionPointToStart(&gpuModule.getBodyRegion().front());
   SmallVector<Type> argsTp;
   for (unsigned i = 0, e = args.size(); i < e; i++)
     argsTp.push_back(args[i].getType());
   FunctionType type = FunctionType::get(gpuModule->getContext(), argsTp, {});
   auto gpuFunc =
-      builder.create<gpu::GPUFuncOp>(gpuModule->getLoc(), name, type);
+      builder.create<gpu::GPUFuncOp>(gpuModule->getLoc(), kernelName, type);
   gpuFunc->setAttr(gpu::GPUDialect::getKernelFuncAttrName(),
                    builder.getUnitAttr());
   return gpuFunc;
 }
 
 /// Constructs code to launch GPU kernel.
 static void genLaunchGPUFunc(OpBuilder &builder, gpu::GPUFuncOp gpuFunc,
                              SmallVectorImpl<Value> &args,
LogicalResult matchAndRewrite(scf::ParallelOp forallOp,
     Location loc = forallOp->getLoc();
     SmallVector<Value> args;
     for (Value s : scalars)
       args.push_back(s);
     for (Value b : buffers)
       args.push_back(genHostRegisterMemref(rewriter, loc, b));
     auto saveIp = rewriter.saveInsertionPoint();
     // Set up GPU module and construct GPU function.
-    //
-    // TODO: only generate once, avoid name conflict
-    //
     ModuleOp topModule = forallOp->getParentOfType<ModuleOp>();
-    auto gpuModule = genGPUModule(rewriter, topModule, "sparsekernels");
-    auto gpuFunc = genGPUFunc(rewriter, gpuModule, "kernel", args);
+    auto gpuModule = genGPUModule(rewriter, topModule);
+    auto gpuFunc = genGPUFunc(rewriter, gpuModule, args);
     genGPUCode(rewriter, gpuFunc, forallOp, constants, scalars, buffers);
     // Generate code that launches the kernel.
     rewriter.restoreInsertionPoint(saveIp);
     genLaunchGPUFunc(rewriter, gpuFunc, args, numThreads);
     rewriter.eraseOp(forallOp);
     return success();
   }
 

mlir/test/Dialect/SparseTensor/GPU/gpu_combi.mlir

+// RUN: mlir-opt %s --linalg-generalize-named-ops \
+// RUN:             --pre-sparsification-rewrite \
+// RUN:             --sparsification="parallelization-strategy=dense-outer-loop" \
+// RUN:             --sparse-gpu-codegen | FileCheck %s
+
+#CSR = #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>
+
+//
+// CHECK-LABEL: gpu.module @sparse_kernels
+// CHECK-DAG:   gpu.func @kernel0
+// CHECK-DAG:   gpu.func @kernel1
+//
+// CHECK-LABEL: func.func @matmuls
+// CHECK-DAG:   gpu.launch_func @sparse_kernels::@kernel0 blocks
+// CHECK-DAG:   gpu.launch_func @sparse_kernels::@kernel1 blocks
+//
+func.func @matmuls(%A: tensor<1024x8xf64>,
+                   %B: tensor<8x1024xf64, #CSR>,
+       %C: tensor<1024x1024xf64, #CSR>) -> tensor<1024x1024xf64> {
+  %Z = arith.constant dense<0.0> : tensor<1024x1024xf64>
+  %T = linalg.matmul
+      ins(%A, %B: tensor<1024x8xf64>, tensor<8x1024xf64, #CSR>)
+      outs(%Z: tensor<1024x1024xf64>) -> tensor<1024x1024xf64>
+  %D = linalg.matmul
+      ins(%T, %C: tensor<1024x1024xf64>, tensor<1024x1024xf64, #CSR>)
+      outs(%Z: tensor<1024x1024xf64>) -> tensor<1024x1024xf64>
+  return %D : tensor<1024x1024xf64>
+}
+

mlir/test/Dialect/SparseTensor/GPU/gpu_matmul.mlir

 // RUN: mlir-opt %s --linalg-generalize-named-ops \
 // RUN:             --pre-sparsification-rewrite \
 // RUN:             --sparsification="parallelization-strategy=dense-outer-loop" \
 // RUN:             --sparse-gpu-codegen | FileCheck %s
 
 #CSR = #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>
 
 //
 // Compute matrix matrix C = AB
 //
-// CHECK-LABEL: gpu.func @kernel(
+// CHECK-LABEL: gpu.module @sparse_kernels
+// CHECK-LABEL: gpu.func @kernel0(
 // CHECK-SAME:        %[[VAL_0:.*0]]: index,
 // CHECK-SAME:        %[[VAL_1:.*1]]: index,
 // CHECK-SAME:        %[[VAL_2:.*2]]: memref<?xindex>,
 // CHECK-SAME:        %[[VAL_3:.*3]]: memref<?xindex>,
 // CHECK-SAME:        %[[VAL_4:.*4]]: memref<?xf64>,
 // CHECK-SAME:        %[[VAL_5:.*5]]: memref<?x?xf64>,
 // CHECK-SAME:        %[[VAL_6:.*6]]: memref<?x?xf64>) kernel {
 // CHECK:         %[[VAL_7:.*]] = arith.constant 1 : index

 //
 //
 // CHECK-LABEL: func.func @matmul
 // CHECK:       gpu.host_register
 // CHECK:       gpu.host_register
 // CHECK:       gpu.host_register
 // CHECK:       gpu.host_register
 // CHECK:       gpu.host_register
-// CHECK:       gpu.launch_func  @sparsekernels::@kernel blocks
+// CHECK:       gpu.launch_func @sparse_kernels::@kernel0 blocks
 //
 func.func @matmul(%A: tensor<?x?xf64, #CSR>, %B: tensor<?x?xf64>, %C_in: tensor<?x?xf64>) -> tensor<?x?xf64> {
   %C_out = linalg.matmul
       ins(%A, %B: tensor<?x?xf64, #CSR>, tensor<?x?xf64>)
       outs(%C_in: tensor<?x?xf64>) -> tensor<?x?xf64>
   return %C_out : tensor<?x?xf64>
 }

mlir/test/Dialect/SparseTensor/GPU/gpu_matvec.mlir

 // RUN: mlir-opt %s --linalg-generalize-named-ops \
 // RUN:             --pre-sparsification-rewrite \
 // RUN:             --sparsification="parallelization-strategy=dense-outer-loop" \
 // RUN:             --sparse-gpu-codegen | FileCheck %s
 
 #CSR = #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>
 
 //
 // Compute matrix vector y = Ax
 //
-//
-// CHECK:       gpu.func @kernel(
+// CHECK-LABEL: gpu.module @sparse_kernels
+// CHECK:       gpu.func @kernel0(
 // CHECK-SAME:          %[[VAL_0:.*0]]: index,
 // CHECK-SAME:          %[[VAL_1:.*1]]: memref<?xf64>,
 // CHECK-SAME:          %[[VAL_2:.*2]]: memref<?xindex>,
 // CHECK-SAME:          %[[VAL_3:.*3]]: memref<?xindex>,
 // CHECK-SAME:          %[[VAL_4:.*4]]: memref<?xf64>,
 // CHECK-SAME:          %[[VAL_5:.*5]]: memref<?xf64>) kernel {
 // CHECK:         %[[VAL_6:.*]] = arith.constant 1 : index
 // CHECK:         %[[VAL_7:.*]] = gpu.block_id  x

 // CHECK:       }
 //
 // CHECK-LABEL: func.func @matvec
 // CHECK:       gpu.host_register
 // CHECK:       gpu.host_register
 // CHECK:       gpu.host_register
 // CHECK:       gpu.host_register
 // CHECK:       gpu.host_register
-// CHECK:       gpu.launch_func  @sparsekernels::@kernel blocks
+// CHECK:       gpu.launch_func @sparse_kernels::@kernel0 blocks
 //
 func.func @matvec(%A: tensor<?x?xf64, #CSR>, %x: tensor<?xf64>, %y_in: tensor<?xf64>) -> tensor<?xf64> {
   %y_out = linalg.matvec
       ins(%A, %x: tensor<?x?xf64, #CSR>, tensor<?xf64>)
       outs(%y_in: tensor<?xf64>) -> tensor<?xf64>
   return %y_out : tensor<?xf64>
 }

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

tags = [
         "nobuildkite",  # TODO(gcmn): Add support for this target
     ],
 )
 
 cc_binary(
     name = "libmlir_cuda_runtime.so",
     linkshared = True,
     linkstatic = False,
+    deps = [":mlir_cuda_runtime"],
     tags = [
         "manual",  # External dependency
         "nobuildkite",  # TODO(gcmn): Add support for this target
     ],
-    deps = [":mlir_cuda_runtime"],
 )
 
 cc_library(
     name = "VulkanRuntime",
     srcs = ["tools/mlir-vulkan-runner/VulkanRuntime.cpp"],
     hdrs = ["tools/mlir-vulkan-runner/VulkanRuntime.h"],
     tags = [
         "manual",  # External dependency