[mlir][sparse][gpu] sparse GPU code generator pipeline setup

[aartbik]

rebased with main

[aartbik]

rebased with main

[aartbik]

rebased with main

[mehdi_amini]

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

106

It's a bit weird to me that a macro will change the behavior here, can you add comment / documentation around the behavior?
We can have gpuCodegen == true but MLIR_GPU_TO_CUBIN_PASS_ENABLE == 0?

mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h

PassOptions::Option<bool> armSVE{
            "dialect"),
       init(false)};
   PassOptions::Option<bool> x86Vector{
       *this, "enable-x86vector",
       desc("Enables the use of X86Vector dialect while lowering the vector "
            "dialect"),
       init(false)};
 
+  /// These options are used to enable GPU code generation.
+  PassOptions::Option<std::string> gpuTriple{*this, "gpu-triple",
+                                             desc("GPU target triple")};
+  PassOptions::Option<std::string> gpuChip{*this, "gpu-chip",
+                                           desc("GPU target architecture")};
+  PassOptions::Option<std::string> gpuFeatures{*this, "gpu-features",
+                                               desc("GPU target features")};
+
   /// Projects out the options for `createSparsificationPass`.
   SparsificationOptions sparsificationOptions() const {
     return SparsificationOptions(parallelization, enableIndexReduction);
   }
 
   /// Projects out the options for `createSparseTensorConversionPass`.
   SparseTensorConversionOptions sparseTensorConversionOptions() const {
     return SparseTensorConversionOptions(

mlir/lib/Dialect/SparseTensor/Pipelines/CMakeLists.txt

 add_mlir_dialect_library(MLIRSparseTensorPipelines
   SparseTensorPipelines.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
 
   LINK_LIBS PUBLIC
   MLIRArithTransforms
   MLIRAffineToStandard
   MLIRBufferizationTransforms
   MLIRComplexToLLVM
   MLIRComplexToLibm
   MLIRComplexToStandard
   MLIRFuncTransforms
+  MLIRGPUToNVVMTransforms
+  MLIRGPUTransforms
   MLIRLinalgTransforms
   MLIRMathToLibm
   MLIRMathToLLVM
   MLIRMemRefToLLVM
   MLIRPass
   MLIRReconcileUnrealizedCasts
   MLIRSCFToControlFlow
   MLIRSparseTensorDialect
   MLIRSparseTensorTransforms
   MLIRTensorTransforms
   MLIRVectorToLLVM
   MLIRVectorTransforms
 )

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

 //===- SparseTensorPipelines.cpp - Pipelines for sparse tensor code -------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SparseTensor/Pipelines/Passes.h"
 
+#include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/Passes.h"
 #include "mlir/Dialect/Arith/Transforms/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/GPU/Transforms/Passes.h"
+#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
pm.addPass(createSparsificationAndBufferizationPass(
       options.enableBufferInitialization, options.vectorLength,
       /*enableVLAVectorization=*/options.armSVE,
       /*enableSIMDIndex32=*/options.force32BitVectorIndices));
   if (options.testBufferizationAnalysisOnly)
     return;
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   pm.addNestedPass<func::FuncOp>(
       mlir::bufferization::createFinalizingBufferizePass());
+
+  // GPU code generation.
+  const bool gpuCodegen = options.gpuTriple.hasValue();
+  if (gpuCodegen) {
+    pm.addPass(createSparseGPUCodegenPass());
+    pm.addNestedPass<gpu::GPUModuleOp>(createStripDebugInfoPass());
+    pm.addNestedPass<gpu::GPUModuleOp>(createConvertSCFToCFPass());
+    pm.addNestedPass<gpu::GPUModuleOp>(createLowerGpuOpsToNVVMOpsPass());
+  }
+
   // TODO(springerm): Add sparse support to the BufferDeallocation pass and add
   // it to this pipeline.
   pm.addNestedPass<func::FuncOp>(createConvertLinalgToLoopsPass());
   pm.addNestedPass<func::FuncOp>(createConvertVectorToSCFPass());
   pm.addNestedPass<func::FuncOp>(createConvertSCFToCFPass());
   pm.addPass(memref::createExpandStridedMetadataPass());
   pm.addPass(createLowerAffinePass());
   pm.addPass(createConvertVectorToLLVMPass(options.lowerVectorToLLVMOptions()));
   pm.addPass(createFinalizeMemRefToLLVMConversionPass());
   pm.addNestedPass<func::FuncOp>(createConvertComplexToStandardPass());
-  pm.addNestedPass<mlir::func::FuncOp>(mlir::arith::createArithExpandOpsPass());
+  pm.addNestedPass<func::FuncOp>(arith::createArithExpandOpsPass());
   pm.addNestedPass<func::FuncOp>(createConvertMathToLLVMPass());
   pm.addPass(createConvertMathToLibmPass());
   pm.addPass(createConvertComplexToLibmPass());
   // Repeat convert-vector-to-llvm.
   pm.addPass(createConvertVectorToLLVMPass(options.lowerVectorToLLVMOptions()));
   pm.addPass(createConvertComplexToLLVMPass());
   pm.addPass(createConvertVectorToLLVMPass(options.lowerVectorToLLVMOptions()));
   pm.addPass(createConvertFuncToLLVMPass());
+
+  // Finalize GPU code generation.
+  if (gpuCodegen) {
+#if MLIR_GPU_TO_CUBIN_PASS_ENABLE
+    pm.addNestedPass<gpu::GPUModuleOp>(createGpuSerializeToCubinPass(
+        options.gpuTriple, options.gpuChip, options.gpuFeatures));
+#endif
+    pm.addPass(createGpuToLLVMConversionPass());
+  }
+
   pm.addPass(createReconcileUnrealizedCastsPass());
 }
 
 //===----------------------------------------------------------------------===//
 // Pipeline registration.
 //===----------------------------------------------------------------------===//
 
 void mlir::sparse_tensor::registerSparseTensorPipelines() {
   PassPipelineRegistration<SparseCompilerOptions>(
       "sparse-compiler",
       "The standard pipeline for taking sparsity-agnostic IR using the"
       " sparse-tensor type, and lowering it to LLVM IR with concrete"
       " representations and algorithms for sparse tensors.",
       buildSparseCompiler);
 }

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

cc_library(
     ],
 )
 
 cc_library(
     name = "SparseTensorPipelines",
     srcs = glob(["lib/Dialect/SparseTensor/Pipelines/*.cpp"]),
     hdrs = ["include/mlir/Dialect/SparseTensor/Pipelines/Passes.h"],
     includes = ["include"],
+    local_defines = if_cuda_available(["MLIR_GPU_TO_CUBIN_PASS_ENABLE"]),
     deps = [
         ":ArithTransforms",
         ":BufferizationTransforms",
         ":ConversionPasses",
         ":FuncDialect",
         ":FuncTransforms",
+        ":GPUDialect",
+        ":GPUToNVVMTransforms",
+        ":GPUTransforms",
         ":LinalgTransforms",
         ":MemRefTransforms",
+        ":NVVMDialect",
         ":Pass",
+        ":SerializeToCubin",
         ":SparseTensorDialect",
         ":SparseTensorTransforms",
         ":TensorTransforms",
         ":Transforms",
         ":VectorToLLVM",
         ":VectorTransforms",
     ],
 )