[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 ↗

(On Diff #513834)

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 ↗

(On Diff #513834)

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/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