GPUToLLVMConversion.cpp

Go to the documentation of this file.

//===- ConvertLaunchFuncToGpuRuntimeCalls.cpp - MLIR GPU lowering passes --===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements a pass to convert gpu.launch_func op into a sequence of
// GPU runtime calls. As most of GPU runtimes does not have a stable published
// ABI, this pass uses a slim runtime layer that builds on top of the public
// API from GPU runtime headers.
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 
#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
#include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
#include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
#include "mlir/Conversion/ConvertToLLVM/ToLLVMInterface.h"
#include "mlir/Conversion/ConvertToLLVM/ToLLVMPass.h"
#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"
#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
#include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
#include "mlir/Dialect/Async/IR/Async.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/GPU/Transforms/Passes.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypes.h"
 
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/FormatVariadic.h"
 
#define DEBUG_TYPE "gpu-to-llvm"
 
namespace mlir {
#define GEN_PASS_DEF_GPUTOLLVMCONVERSIONPASS
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir
 
using namespace mlir;
 
static constexpr const char *kGpuBinaryStorageSuffix = "_gpubin_cst";
 
namespace {
class GpuToLLVMConversionPass
    : public impl::GpuToLLVMConversionPassBase<GpuToLLVMConversionPass> {
public:
  using Base::Base;
  void getDependentDialects(DialectRegistry &registry) const final {
    Base::getDependentDialects(registry);
    registerConvertToLLVMDependentDialectLoading(registry);
  }
  // Run the dialect converter on the module.
  void runOnOperation() override;
};
 
template <typename OpTy>
class ConvertOpToGpuRuntimeCallPattern : public ConvertOpToLLVMPattern<OpTy> {
public:
  explicit ConvertOpToGpuRuntimeCallPattern(
      const LLVMTypeConverter &typeConverter)
      : ConvertOpToLLVMPattern<OpTy>(typeConverter) {}
 
protected:
  Value getNumElements(ConversionPatternRewriter &rewriter, Location loc,
                       MemRefType type, MemRefDescriptor desc) const {
    Type indexType = ConvertToLLVMPattern::getIndexType();
    return type.hasStaticShape()
               ? ConvertToLLVMPattern::createIndexAttrConstant(
                     rewriter, loc, indexType, type.getNumElements())
               // For identity maps (verified by caller), the number of
               // elements is stride[0] * size[0].
               : rewriter.create<LLVM::MulOp>(loc,
                                              desc.stride(rewriter, loc, 0),
                                              desc.size(rewriter, loc, 0));
  }
 
  MLIRContext *context = &this->getTypeConverter()->getContext();
 
  Type llvmVoidType = LLVM::LLVMVoidType::get(context);
  LLVM::LLVMPointerType llvmPointerType = LLVM::LLVMPointerType::get(context);
  Type llvmInt8Type = IntegerType::get(context, 8);
  Type llvmInt16Type = IntegerType::get(context, 16);
  Type llvmInt32Type = IntegerType::get(context, 32);
  Type llvmInt64Type = IntegerType::get(context, 64);
  Type llvmFloat32Type = Float32Type::get(context);
  Type llvmIntPtrType = IntegerType::get(
      context, this->getTypeConverter()->getPointerBitwidth(0));
 
  FunctionCallBuilder moduleLoadCallBuilder = {
      "mgpuModuleLoad",
      llvmPointerType /* void *module */,
      {llvmPointerType /* void *cubin */, llvmInt64Type /* size_t size */}};
  FunctionCallBuilder moduleUnloadCallBuilder = {
      "mgpuModuleUnload", llvmVoidType, {llvmPointerType /* void *module */}};
  FunctionCallBuilder moduleGetFunctionCallBuilder = {
      "mgpuModuleGetFunction",
      llvmPointerType /* void *function */,
      {
          llvmPointerType, /* void *module */
          llvmPointerType  /* char *name   */
      }};
  FunctionCallBuilder launchKernelCallBuilder = {
      "mgpuLaunchKernel",
      llvmVoidType,
      {
          llvmPointerType, /* void* f */
          llvmIntPtrType,  /* intptr_t gridXDim */
          llvmIntPtrType,  /* intptr_t gridyDim */
          llvmIntPtrType,  /* intptr_t gridZDim */
          llvmIntPtrType,  /* intptr_t blockXDim */
          llvmIntPtrType,  /* intptr_t blockYDim */
          llvmIntPtrType,  /* intptr_t blockZDim */
          llvmInt32Type,   /* unsigned int sharedMemBytes */
          llvmPointerType, /* void *hstream */
          llvmPointerType, /* void **kernelParams */
          llvmPointerType, /* void **extra */
          llvmInt64Type    /* size_t paramsCount */
      }};
  FunctionCallBuilder streamCreateCallBuilder = {
      "mgpuStreamCreate", llvmPointerType /* void *stream */, {}};
  FunctionCallBuilder streamDestroyCallBuilder = {
      "mgpuStreamDestroy", llvmVoidType, {llvmPointerType /* void *stream */}};
  FunctionCallBuilder streamSynchronizeCallBuilder = {
      "mgpuStreamSynchronize",
      llvmVoidType,
      {llvmPointerType /* void *stream */}};
  FunctionCallBuilder streamWaitEventCallBuilder = {
      "mgpuStreamWaitEvent",
      llvmVoidType,
      {llvmPointerType /* void *stream */, llvmPointerType /* void *event */}};
  FunctionCallBuilder eventCreateCallBuilder = {
      "mgpuEventCreate", llvmPointerType /* void *event */, {}};
  FunctionCallBuilder eventDestroyCallBuilder = {
      "mgpuEventDestroy", llvmVoidType, {llvmPointerType /* void *event */}};
  FunctionCallBuilder eventSynchronizeCallBuilder = {
      "mgpuEventSynchronize",
      llvmVoidType,
      {llvmPointerType /* void *event */}};
  FunctionCallBuilder eventRecordCallBuilder = {
      "mgpuEventRecord",
      llvmVoidType,
      {llvmPointerType /* void *event */, llvmPointerType /* void *stream */}};
  FunctionCallBuilder hostRegisterCallBuilder = {
      "mgpuMemHostRegisterMemRef",
      llvmVoidType,
      {llvmIntPtrType /* intptr_t rank */,
       llvmPointerType /* void *memrefDesc */,
       llvmIntPtrType /* intptr_t elementSizeBytes */}};
  FunctionCallBuilder hostUnregisterCallBuilder = {
      "mgpuMemHostUnregisterMemRef",
      llvmVoidType,
      {llvmIntPtrType /* intptr_t rank */,
       llvmPointerType /* void *memrefDesc */,
       llvmIntPtrType /* intptr_t elementSizeBytes */}};
  FunctionCallBuilder allocCallBuilder = {
      "mgpuMemAlloc",
      llvmPointerType /* void * */,
      {llvmIntPtrType /* intptr_t sizeBytes */,
       llvmPointerType /* void *stream */,
       llvmInt8Type /* bool isHostShared */}};
  FunctionCallBuilder deallocCallBuilder = {
      "mgpuMemFree",
      llvmVoidType,
      {llvmPointerType /* void *ptr */, llvmPointerType /* void *stream */}};
  FunctionCallBuilder memcpyCallBuilder = {
      "mgpuMemcpy",
      llvmVoidType,
      {llvmPointerType /* void *dst */, llvmPointerType /* void *src */,
       llvmIntPtrType /* intptr_t sizeBytes */,
       llvmPointerType /* void *stream */}};
  FunctionCallBuilder memset16CallBuilder = {
      "mgpuMemset16",
      llvmVoidType,
      {llvmPointerType /* void *dst */,
       llvmInt16Type /* unsigned short value */,
       llvmIntPtrType /* intptr_t sizeBytes */,
       llvmPointerType /* void *stream */}};
  FunctionCallBuilder memset32CallBuilder = {
      "mgpuMemset32",
      llvmVoidType,
      {llvmPointerType /* void *dst */, llvmInt32Type /* unsigned int value */,
       llvmIntPtrType /* intptr_t sizeBytes */,
       llvmPointerType /* void *stream */}};
  FunctionCallBuilder setDefaultDeviceCallBuilder = {
      "mgpuSetDefaultDevice",
      llvmVoidType,
      {llvmInt32Type /* uint32_t devIndex */}};
  FunctionCallBuilder createDnVecCallBuilder = {
      "mgpuCreateDnVec",
      llvmPointerType,
      {llvmIntPtrType, llvmPointerType, llvmInt32Type,
       llvmPointerType /* void *stream */}};
  FunctionCallBuilder destroyDnVecCallBuilder = {
      "mgpuDestroyDnVec",
      llvmVoidType,
      {llvmPointerType, llvmPointerType /* void *stream */}};
  FunctionCallBuilder createDnMatCallBuilder = {
      "mgpuCreateDnMat",
      llvmPointerType,
      {llvmIntPtrType, llvmIntPtrType, llvmPointerType, llvmInt32Type,
       llvmPointerType /* void *stream */}};
  FunctionCallBuilder destroyDnMatCallBuilder = {
      "mgpuDestroyDnMat",
      llvmVoidType,
      {llvmPointerType, llvmPointerType /* void *stream */}};
  FunctionCallBuilder createCooCallBuilder = {
      "mgpuCreateCoo",
      llvmPointerType,
      {llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, llvmPointerType,
       llvmPointerType, llvmPointerType, llvmInt32Type, llvmInt32Type,
       llvmPointerType /* void *stream */}};
  FunctionCallBuilder createCooAoSCallBuilder = {
      "mgpuCreateCooAoS", // deprecated in cuSPARSE 11.2
      llvmPointerType,
      {llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, llvmPointerType,
       llvmPointerType, llvmInt32Type, llvmInt32Type,
       llvmPointerType /* void *stream */}};
  FunctionCallBuilder createCsrCallBuilder = {
      "mgpuCreateCsr",
      llvmPointerType,
      {llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, llvmPointerType,
       llvmPointerType, llvmPointerType, llvmInt32Type, llvmInt32Type,
       llvmInt32Type, llvmPointerType /* void *stream */}};
  FunctionCallBuilder createCscCallBuilder = {
      "mgpuCreateCsc",
      llvmPointerType,
      {llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, llvmPointerType,
       llvmPointerType, llvmPointerType, llvmInt32Type, llvmInt32Type,
       llvmInt32Type, llvmPointerType /* void *stream */}};
  FunctionCallBuilder createBsrCallBuilder = {
      "mgpuCreateBsr",
      llvmPointerType,
      {llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, llvmIntPtrType,
       llvmIntPtrType, llvmPointerType, llvmPointerType, llvmPointerType,
       llvmInt32Type, llvmInt32Type, llvmInt32Type,
       llvmPointerType /* void *stream */}};
  FunctionCallBuilder destroySpMatCallBuilder = {
      "mgpuDestroySpMat",
      llvmVoidType,
      {llvmPointerType, llvmPointerType /* void *stream */}};
  FunctionCallBuilder spMVBufferSizeCallBuilder = {
      "mgpuSpMVBufferSize",
      llvmIntPtrType,
      {llvmInt32Type, llvmPointerType, llvmPointerType, llvmPointerType,
       llvmInt32Type, llvmPointerType /* void *stream */}};
  FunctionCallBuilder spMVCallBuilder = {
      "mgpuSpMV",
      llvmVoidType,
      {llvmInt32Type, llvmPointerType, llvmPointerType, llvmPointerType,
       llvmInt32Type, llvmPointerType, llvmPointerType /* void *stream */}};
  FunctionCallBuilder createSpMMBufferSizeCallBuilder = {
      "mgpuSpMMBufferSize",
      llvmIntPtrType,
      {llvmInt32Type, llvmInt32Type, llvmPointerType, llvmPointerType,
       llvmPointerType, llvmInt32Type, llvmPointerType /* void *stream */}};
  FunctionCallBuilder createSpMMCallBuilder = {
      "mgpuSpMM",
      llvmVoidType,
      {llvmInt32Type, llvmInt32Type, llvmPointerType, llvmPointerType,
       llvmPointerType, llvmInt32Type, llvmPointerType,
       llvmPointerType /* void *stream */}};
  FunctionCallBuilder createSDDMMBufferSizeCallBuilder = {
      "mgpuSDDMMBufferSize",
      llvmIntPtrType,
      {llvmInt32Type, llvmInt32Type, llvmPointerType, llvmPointerType,
       llvmPointerType, llvmInt32Type, llvmPointerType /* void *stream */}};
  FunctionCallBuilder createSDDMMCallBuilder = {
      "mgpuSDDMM",
      llvmVoidType,
      {llvmInt32Type, llvmInt32Type, llvmPointerType, llvmPointerType,
       llvmPointerType, llvmInt32Type, llvmPointerType,
       llvmPointerType /* void *stream */}};
  FunctionCallBuilder createLtDnMatCallBuilder = {
      "mgpuCreateCuSparseLtDnMat",
      llvmVoidType,
      {llvmPointerType, llvmIntPtrType, llvmIntPtrType, llvmPointerType,
       llvmInt32Type, llvmPointerType /* void *stream */}};
  FunctionCallBuilder destroyCuSparseLtSpMatBuilder = {
      "mgpuDestroyCuSparseLtSpMat",
      llvmVoidType,
      {llvmPointerType, llvmPointerType /* void *stream */}};
  FunctionCallBuilder destroyCuSparseLtDnMatBuilder = {
      "mgpuDestroyCuSparseLtDnMat",
      llvmVoidType,
      {llvmPointerType, llvmPointerType /* void *stream */}};
  FunctionCallBuilder create2To4SpMatCallBuilder = {
      "mgpuCusparseLtCreate2To4SpMat",
      llvmVoidType,
      {llvmPointerType, llvmIntPtrType, llvmIntPtrType, llvmPointerType,
       llvmInt32Type, llvmPointerType /* void *stream */}};
  FunctionCallBuilder createCuSparseLtSpMMBufferSizeBuilder = {
      "mgpuCuSparseLtSpMMBufferSize",
      llvmVoidType,
      {llvmPointerType, llvmInt32Type, llvmInt32Type, llvmPointerType,
       llvmPointerType, llvmPointerType, llvmInt32Type, llvmInt32Type,
       llvmPointerType /*void *stream*/}};
  FunctionCallBuilder createCuSparseLtSpMMBuilder = {
      "mgpuCuSparseLtSpMM",
      llvmVoidType,
      {llvmPointerType, llvmPointerType, llvmPointerType, llvmPointerType,
       llvmPointerType, llvmPointerType, llvmPointerType /*void *stream*/}};
  FunctionCallBuilder createSpGEMMCreateDescrBuilder = {
      "mgpuSpGEMMCreateDescr",
      llvmPointerType,
      {llvmPointerType /*void *stream*/}};
  FunctionCallBuilder createSpGEMMDestroyDescrBuilder = {
      "mgpuSpGEMMDestroyDescr",
      llvmVoidType,
      {llvmPointerType /*s*/, llvmPointerType /*void *stream*/}};
  FunctionCallBuilder createSpGEMMWorkEstimationBuilder = {
      "mgpuSpGEMMWorkEstimation",
      llvmIntPtrType,
      {llvmPointerType /*s*/, llvmInt32Type /*ma*/, llvmInt32Type /*mb*/,
       llvmPointerType /*a*/, llvmPointerType /*b*/, llvmPointerType /*c*/,
       llvmInt32Type /*ctp*/, llvmIntPtrType /*bs*/, llvmPointerType /*buf*/,
       llvmPointerType /*void *stream*/}};
  FunctionCallBuilder createSpGEMMComputeBuilder = {
      "mgpuSpGEMMCompute",
      llvmIntPtrType,
      {llvmPointerType /*s*/, llvmInt32Type /*ma*/, llvmInt32Type /*mb*/,
       llvmPointerType /*a*/, llvmPointerType /*b*/, llvmPointerType /*c*/,
       llvmInt32Type /*ctp*/, llvmIntPtrType /*bs*/, llvmPointerType /*buf*/,
       llvmPointerType /*void *stream*/}};
  FunctionCallBuilder createSpGEMMCopyBuilder = {
      "mgpuSpGEMMCopy",
      llvmVoidType,
      {llvmPointerType /*s*/, llvmInt32Type /*ma*/, llvmInt32Type /*mb*/,
       llvmPointerType /*a*/, llvmPointerType /*b*/, llvmPointerType /*c*/,
       llvmInt32Type /*ctp*/, llvmPointerType /*void *stream*/}};
  FunctionCallBuilder createSpMatGetSizeBuilder = {
      "mgpuSpMatGetSize",
      llvmVoidType,
      {llvmPointerType /*mc*/, llvmPointerType /*rc*/, llvmPointerType /*cc*/,
       llvmPointerType /*nc*/, llvmPointerType /*void *stream*/}};
  FunctionCallBuilder createSetCsrPointersBuilder = {
      "mgpuSetCsrPointers",
      llvmVoidType,
      {llvmPointerType /*spmat*/, llvmPointerType /*pos*/,
       llvmPointerType /*crd*/, llvmPointerType /*val*/,
       llvmPointerType /*void *stream*/}};
};
 
/// A rewrite pattern to convert gpu.host_register operations into a GPU runtime
/// call. Currently it supports CUDA and ROCm (HIP).
class ConvertHostRegisterOpToGpuRuntimeCallPattern
    : public ConvertOpToGpuRuntimeCallPattern<gpu::HostRegisterOp> {
public:
  ConvertHostRegisterOpToGpuRuntimeCallPattern(
      const LLVMTypeConverter &typeConverter)
      : ConvertOpToGpuRuntimeCallPattern<gpu::HostRegisterOp>(typeConverter) {}
 
private:
  LogicalResult
  matchAndRewrite(gpu::HostRegisterOp hostRegisterOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
class ConvertHostUnregisterOpToGpuRuntimeCallPattern
    : public ConvertOpToGpuRuntimeCallPattern<gpu::HostUnregisterOp> {
public:
  ConvertHostUnregisterOpToGpuRuntimeCallPattern(
      const LLVMTypeConverter &typeConverter)
      : ConvertOpToGpuRuntimeCallPattern<gpu::HostUnregisterOp>(typeConverter) {
  }
 
private:
  LogicalResult
  matchAndRewrite(gpu::HostUnregisterOp hostUnregisterOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
/// A rewrite pattern to convert gpu.alloc operations into a GPU runtime
/// call. Currently it supports CUDA and ROCm (HIP).
class ConvertAllocOpToGpuRuntimeCallPattern
    : public ConvertOpToGpuRuntimeCallPattern<gpu::AllocOp> {
public:
  ConvertAllocOpToGpuRuntimeCallPattern(const LLVMTypeConverter &typeConverter)
      : ConvertOpToGpuRuntimeCallPattern<gpu::AllocOp>(typeConverter) {}
 
private:
  LogicalResult
  matchAndRewrite(gpu::AllocOp allocOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
/// A rewrite pattern to convert gpu.dealloc operations into a GPU runtime
/// call. Currently it supports CUDA and ROCm (HIP).
class ConvertDeallocOpToGpuRuntimeCallPattern
    : public ConvertOpToGpuRuntimeCallPattern<gpu::DeallocOp> {
public:
  ConvertDeallocOpToGpuRuntimeCallPattern(
      const LLVMTypeConverter &typeConverter)
      : ConvertOpToGpuRuntimeCallPattern<gpu::DeallocOp>(typeConverter) {}
 
private:
  LogicalResult
  matchAndRewrite(gpu::DeallocOp deallocOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
class ConvertAsyncYieldToGpuRuntimeCallPattern
    : public ConvertOpToGpuRuntimeCallPattern<async::YieldOp> {
public:
  ConvertAsyncYieldToGpuRuntimeCallPattern(
      const LLVMTypeConverter &typeConverter)
      : ConvertOpToGpuRuntimeCallPattern<async::YieldOp>(typeConverter) {}
 
private:
  LogicalResult
  matchAndRewrite(async::YieldOp yieldOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
/// A rewrite pattern to convert gpu.wait operations into a GPU runtime
/// call. Currently it supports CUDA and ROCm (HIP).
class ConvertWaitOpToGpuRuntimeCallPattern
    : public ConvertOpToGpuRuntimeCallPattern<gpu::WaitOp> {
public:
  ConvertWaitOpToGpuRuntimeCallPattern(const LLVMTypeConverter &typeConverter)
      : ConvertOpToGpuRuntimeCallPattern<gpu::WaitOp>(typeConverter) {}
 
private:
  LogicalResult
  matchAndRewrite(gpu::WaitOp waitOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
/// A rewrite pattern to convert gpu.wait async operations into a GPU runtime
/// call. Currently it supports CUDA and ROCm (HIP).
class ConvertWaitAsyncOpToGpuRuntimeCallPattern
    : public ConvertOpToGpuRuntimeCallPattern<gpu::WaitOp> {
public:
  ConvertWaitAsyncOpToGpuRuntimeCallPattern(
      const LLVMTypeConverter &typeConverter)
      : ConvertOpToGpuRuntimeCallPattern<gpu::WaitOp>(typeConverter) {}
 
private:
  LogicalResult
  matchAndRewrite(gpu::WaitOp waitOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
/// A rewrite patter to convert gpu.launch_func operations into a sequence of
/// GPU runtime calls. Currently it supports CUDA and ROCm (HIP).
///
/// In essence, a gpu.launch_func operations gets compiled into the following
/// sequence of runtime calls:
///
/// * moduleLoad        -- loads the module given the cubin / hsaco data
/// * moduleGetFunction -- gets a handle to the actual kernel function
/// * getStreamHelper   -- initializes a new compute stream on GPU
/// * launchKernel      -- launches the kernel on a stream
/// * streamSynchronize -- waits for operations on the stream to finish
///
/// Intermediate data structures are allocated on the stack.
class ConvertLaunchFuncOpToGpuRuntimeCallPattern
    : public ConvertOpToGpuRuntimeCallPattern<gpu::LaunchFuncOp> {
public:
  ConvertLaunchFuncOpToGpuRuntimeCallPattern(
      const LLVMTypeConverter &typeConverter, StringRef gpuBinaryAnnotation,
      bool kernelBarePtrCallConv, SymbolTable *cachedModuleTable)
      : ConvertOpToGpuRuntimeCallPattern<gpu::LaunchFuncOp>(typeConverter),
        gpuBinaryAnnotation(gpuBinaryAnnotation),
        kernelBarePtrCallConv(kernelBarePtrCallConv),
        cachedModuleTable(cachedModuleTable) {}
 
private:
  Value generateParamsArray(gpu::LaunchFuncOp launchOp, OpAdaptor adaptor,
                            OpBuilder &builder) const;
  Value generateKernelNameConstant(StringRef moduleName, StringRef name,
                                   Location loc, OpBuilder &builder) const;
 
  LogicalResult
  matchAndRewrite(gpu::LaunchFuncOp launchOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
 
  llvm::SmallString<32> gpuBinaryAnnotation;
  bool kernelBarePtrCallConv;
  SymbolTable *cachedModuleTable;
};
 
class EraseGpuModuleOpPattern : public OpRewritePattern<gpu::GPUModuleOp> {
  using OpRewritePattern<gpu::GPUModuleOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(gpu::GPUModuleOp op,
                                PatternRewriter &rewriter) const override {
    // GPU kernel modules are no longer necessary since we have a global
    // constant with the CUBIN, or HSACO data.
    rewriter.eraseOp(op);
    return success();
  }
};
 
/// A rewrite pattern to convert gpu.memcpy operations into a GPU runtime
/// call. Currently it supports CUDA and ROCm (HIP).
class ConvertMemcpyOpToGpuRuntimeCallPattern
    : public ConvertOpToGpuRuntimeCallPattern<gpu::MemcpyOp> {
public:
  ConvertMemcpyOpToGpuRuntimeCallPattern(const LLVMTypeConverter &typeConverter)
      : ConvertOpToGpuRuntimeCallPattern<gpu::MemcpyOp>(typeConverter) {}
 
private:
  LogicalResult
  matchAndRewrite(gpu::MemcpyOp memcpyOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
/// A rewrite pattern to convert gpu.memset operations into a GPU runtime
/// call. Currently it supports CUDA and ROCm (HIP).
class ConvertMemsetOpToGpuRuntimeCallPattern
    : public ConvertOpToGpuRuntimeCallPattern<gpu::MemsetOp> {
public:
  ConvertMemsetOpToGpuRuntimeCallPattern(const LLVMTypeConverter &typeConverter)
      : ConvertOpToGpuRuntimeCallPattern<gpu::MemsetOp>(typeConverter) {}
 
private:
  LogicalResult
  matchAndRewrite(gpu::MemsetOp memsetOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
/// A rewrite pattern to convert gpu.set_default_device to a GPU runtime call.
/// Currently supports CUDA and ROCm (HIP)
class ConvertSetDefaultDeviceOpToGpuRuntimeCallPattern
    : public ConvertOpToGpuRuntimeCallPattern<gpu::SetDefaultDeviceOp> {
public:
  ConvertSetDefaultDeviceOpToGpuRuntimeCallPattern(
      const LLVMTypeConverter &typeConverter)
      : ConvertOpToGpuRuntimeCallPattern<gpu::SetDefaultDeviceOp>(
            typeConverter) {}
 
  LogicalResult
  matchAndRewrite(gpu::SetDefaultDeviceOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
 
/// Generic rewriting rule for operation on sparse matrices.
/// Currently supports CUDA (by means of cuSparse and cuSparseLt).
#define DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(op_name)                \
  class Convert##op_name##ToGpuRuntimeCallPattern                              \
      : public ConvertOpToGpuRuntimeCallPattern<gpu::op_name> {                \
  public:                                                                      \
    Convert##op_name##ToGpuRuntimeCallPattern(                                 \
        const LLVMTypeConverter &typeConverter)                                \
        : ConvertOpToGpuRuntimeCallPattern<gpu::op_name>(typeConverter) {}     \
                                                                               \
  private:                                                                     \
    LogicalResult                                                              \
    matchAndRewrite(gpu::op_name op, OpAdaptor adaptor,                        \
                    ConversionPatternRewriter &rewriter) const override;       \
  };
 
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(CreateDnTensorOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(DestroyDnTensorOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(CreateCooOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(CreateCooAoSOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(CreateCsrOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(CreateCscOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(CreateBsrOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(Create2To4SpMatOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(DestroySpMatOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(SpMVBufferSizeOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(SpMVOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(SpMMBufferSizeOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(SDDMMBufferSizeOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(SpMMOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(SDDMMOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(SpGEMMCreateDescrOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(SpGEMMDestroyDescrOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(SpGEMMWorkEstimationOrComputeOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(SpGEMMCopyOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(SpMatGetSizeOp)
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(SetCsrPointersOp)
 
} // namespace
 
void GpuToLLVMConversionPass::runOnOperation() {
  MLIRContext *context = &getContext();
  SymbolTable symbolTable = SymbolTable(getOperation());
  LowerToLLVMOptions options(context);
  options.useBarePtrCallConv = hostBarePtrCallConv;
  RewritePatternSet patterns(context);
  ConversionTarget target(*context);
  target.addLegalDialect<LLVM::LLVMDialect>();
  LLVMTypeConverter converter(context, options);
 
  // Populate all patterns from all dialects that implement the
  // `ConvertToLLVMPatternInterface` interface.
  for (Dialect *dialect : context->getLoadedDialects()) {
    auto iface = dyn_cast<ConvertToLLVMPatternInterface>(dialect);
    if (!iface)
      continue;
    iface->populateConvertToLLVMConversionPatterns(target, converter, patterns);
  }
 
  // Preserve GPU modules if they have target attributes.
  target.addDynamicallyLegalOp<gpu::GPUModuleOp>(
      [](gpu::GPUModuleOp module) -> bool {
        return module.getTargetsAttr() != nullptr;
      });
  // Accept as legal LaunchFuncOps if they refer to GPU Modules with targets and
  // the operands have been lowered.
  target.addDynamicallyLegalOp<gpu::LaunchFuncOp>(
      [&](gpu::LaunchFuncOp op) -> bool {
        auto module =
            symbolTable.lookup<gpu::GPUModuleOp>(op.getKernelModuleName());
        return converter.isLegal(op->getOperandTypes()) &&
               converter.isLegal(op->getResultTypes()) &&
               (module && module.getTargetsAttr() &&
                !module.getTargetsAttr().empty());
      });
 
  // These aren't covered by the ConvertToLLVMPatternInterface right now.
  populateVectorToLLVMConversionPatterns(converter, patterns);
  populateFinalizeMemRefToLLVMConversionPatterns(converter, patterns);
  populateAsyncStructuralTypeConversionsAndLegality(converter, patterns,
                                                    target);
  populateGpuToLLVMConversionPatterns(converter, patterns, gpuBinaryAnnotation,
                                      kernelBarePtrCallConv, &symbolTable);
 
  if (failed(
          applyPartialConversion(getOperation(), target, std::move(patterns))))
    signalPassFailure();
}
 
LLVM::CallOp FunctionCallBuilder::create(Location loc, OpBuilder &builder,
                                         ArrayRef<Value> arguments) const {
  auto module = builder.getBlock()->getParent()->getParentOfType<ModuleOp>();
  auto function = [&] {
    if (auto function = module.lookupSymbol<LLVM::LLVMFuncOp>(functionName))
      return function;
    return OpBuilder::atBlockEnd(module.getBody())
        .create<LLVM::LLVMFuncOp>(loc, functionName, functionType);
  }();
  return builder.create<LLVM::CallOp>(loc, function, arguments);
}
 
// Corresponding to cusparseIndexType_t defined in cusparse.h.
static int32_t getCuSparseIndexTypeFrom(Type type) {
  if (type.isInteger(16))
    return 1; // CUSPARSE_INDEX_16U
  if (type.isInteger(32))
    return 2; // CUSPARSE_INDEX_32I
  return 3;   // CUSPARSE_INDEX_64I
}
 
static int32_t getCuSparseLtDataTypeFrom(Type type) {
  if (type.isF16())
    return 0; // CUSPARSE_COMPUTE_16F,
  if (type.isInteger(32))
    return 1; // CUSPARSE_COMPUTE_32I
  llvm_unreachable("unsupported type");
  // TODO: add support to TF32
}
 
// Corresponding to cudaDataType_t defined in CUDA library_types.h.
static int32_t getCuSparseDataTypeFrom(Type type) {
  if (llvm::isa<ComplexType>(type)) {
    // get the element type
    auto elementType = type.cast<ComplexType>().getElementType();
    if (elementType.isBF16())
      return 15; // CUDA_C_16BF
    if (elementType.isF16())
      return 6; // CUDA_C_16F
    if (elementType.isF32())
      return 4; // CUDA_C_32F
    if (elementType.isF64())
      return 5; // CUDA_C_64F
    if (elementType.isInteger(8))
      return 7; // CUDA_C_8I
    if (elementType.isInteger(16))
      return 21; // CUDA_C_16I
    if (elementType.isInteger(32))
      return 11; // CUDA_C_32I
  }
  if (type.isBF16())
    return 14; // CUDA_R_16BF
  if (type.isF16())
    return 2; // CUDA_R_16F
  if (type.isF32())
    return 0; // CUDA_R_32F
  if (type.isF64())
    return 1; // CUDA_R_64F
  if (type.isInteger(8))
    return 3; // CUDA_R_8I
  if (type.isInteger(16))
    return 20; // CUDA_R_16I
  if (type.isInteger(32))
    return 10; // CUDA_R_32I
 
  llvm_unreachable("unsupported element type");
}
 
static gpu::Prune2To4SpMatFlag get2To4PruneFlag(Value spMat) {
  return spMat.getDefiningOp<gpu::Create2To4SpMatOp>().getPruneFlag();
}
 
// TODO:  We may want a run-time (of the mlir compiler) disablement/warning:
// cusparseLt currently won't work for cuda architecture <8.0 and will trigger a
// runtime (of the CUDA program) error , but it might be great if we could at
// least output a warning when we found the target architecture is <8.0 and the
// user still wants to use cusparseLt. to make sure when lowering gpu sparse
// dialect to llvm calls, the cusparselt calls are disabled for cuda
// architecture <8.0
static bool is2To4Sparsity(Value spMat) {
  if (auto op = spMat.getDefiningOp<gpu::Create2To4SpMatOp>())
    return true;
  if (auto op = spMat.getDefiningOp<gpu::CreateCooOp>())
    return false;
  if (auto op = spMat.getDefiningOp<gpu::CreateCooAoSOp>())
    return false;
  if (auto op = spMat.getDefiningOp<gpu::CreateCsrOp>())
    return false;
  if (auto op = spMat.getDefiningOp<gpu::CreateCscOp>())
    return false;
  if (auto op = spMat.getDefiningOp<gpu::CreateBsrOp>())
    return false;
  // Print the spMat defining op
  spMat.getDefiningOp()->print(llvm::errs());
  llvm_unreachable("cannot find spmat def");
}
 
static bool isSpMMCusparseLtOp(Value op) {
  for (Operation *user : op.getUsers()) {
    auto spmmOp = dyn_cast<gpu::SpMMOp>(user);
    // If the other operator is 50% sparsity then we should use cusparseLt
    if (!spmmOp)
      continue;
    if (is2To4Sparsity(spmmOp.getSpmatA()))
      return true;
  }
  return false;
}
 
// 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());
      }))
    return rewriter.notifyMatchFailure(
        op, "Cannot convert if operands aren't of LLVM type.");
  return success();
}
 
static LogicalResult
isAsyncWithOneDependency(ConversionPatternRewriter &rewriter,
                         gpu::AsyncOpInterface op) {
  if (op.getAsyncDependencies().size() != 1)
    return rewriter.notifyMatchFailure(
        op, "Can only convert with exactly one async dependency.");
 
  if (!op.getAsyncToken())
    return rewriter.notifyMatchFailure(op, "Can convert only async version.");
 
  return success();
}
 
LogicalResult ConvertHostRegisterOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::HostRegisterOp hostRegisterOp, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  auto *op = hostRegisterOp.getOperation();
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)))
    return failure();
 
  Location loc = op->getLoc();
 
  auto memRefType = hostRegisterOp.getValue().getType();
  auto elementType = cast<UnrankedMemRefType>(memRefType).getElementType();
  auto elementSize = getSizeInBytes(loc, elementType, rewriter);
 
  auto arguments = getTypeConverter()->promoteOperands(
      loc, op->getOperands(), adaptor.getOperands(), rewriter);
  arguments.push_back(elementSize);
  hostRegisterCallBuilder.create(loc, rewriter, arguments);
 
  rewriter.eraseOp(op);
  return success();
}
 
LogicalResult ConvertHostUnregisterOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::HostUnregisterOp hostUnregisterOp, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Operation *op = hostUnregisterOp.getOperation();
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)))
    return failure();
 
  Location loc = op->getLoc();
 
  auto memRefType = hostUnregisterOp.getValue().getType();
  auto elementType = cast<UnrankedMemRefType>(memRefType).getElementType();
  auto elementSize = getSizeInBytes(loc, elementType, rewriter);
 
  auto arguments = getTypeConverter()->promoteOperands(
      loc, op->getOperands(), adaptor.getOperands(), rewriter);
  arguments.push_back(elementSize);
  hostUnregisterCallBuilder.create(loc, rewriter, arguments);
 
  rewriter.eraseOp(op);
  return success();
}
 
LogicalResult ConvertAllocOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::AllocOp allocOp, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
 
  MemRefType memRefType = allocOp.getType();
 
  if (failed(areAllLLVMTypes(allocOp, adaptor.getOperands(), rewriter)) ||
      !isConvertibleAndHasIdentityMaps(memRefType))
    return failure();
 
  auto loc = allocOp.getLoc();
 
  bool isShared = allocOp.getHostShared();
 
  if (isShared && allocOp.getAsyncToken())
    return rewriter.notifyMatchFailure(
        allocOp, "Host Shared allocation cannot be done async");
  if (!isShared && failed(isAsyncWithOneDependency(rewriter, allocOp)))
    return failure();
 
  // Get shape of the memref as values: static sizes are constant
  // values and dynamic sizes are passed to 'alloc' as operands.
  SmallVector<Value, 4> shape;
  SmallVector<Value, 4> strides;
  Value sizeBytes;
  getMemRefDescriptorSizes(loc, memRefType, adaptor.getDynamicSizes(), rewriter,
                           shape, strides, sizeBytes);
 
  // Allocate the underlying buffer and store a pointer to it in the MemRef
  // descriptor.
  auto nullPtr = rewriter.create<mlir::LLVM::ZeroOp>(loc, llvmPointerType);
  Value stream = adaptor.getAsyncDependencies().empty()
                     ? nullPtr
                     : adaptor.getAsyncDependencies().front();
 
  auto isHostShared = rewriter.create<mlir::LLVM::ConstantOp>(
      loc, llvmInt8Type, rewriter.getI8IntegerAttr(isShared));
 
  Value allocatedPtr =
      allocCallBuilder.create(loc, rewriter, {sizeBytes, stream, isHostShared})
          .getResult();
 
  // No alignment.
  Value alignedPtr = allocatedPtr;
 
  // Create the MemRef descriptor.
  auto memRefDescriptor = this->createMemRefDescriptor(
      loc, memRefType, allocatedPtr, alignedPtr, shape, strides, rewriter);
 
  if (allocOp.getAsyncToken()) {
    // Async alloc: make dependent ops use the same stream.
    rewriter.replaceOp(allocOp, {memRefDescriptor, stream});
  } else {
    rewriter.replaceOp(allocOp, {memRefDescriptor});
  }
 
  return success();
}
 
LogicalResult ConvertDeallocOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::DeallocOp deallocOp, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(deallocOp, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, deallocOp)))
    return failure();
 
  Location loc = deallocOp.getLoc();
 
  Value pointer =
      MemRefDescriptor(adaptor.getMemref()).allocatedPtr(rewriter, loc);
  Value stream = adaptor.getAsyncDependencies().front();
  deallocCallBuilder.create(loc, rewriter, {pointer, stream});
 
  rewriter.replaceOp(deallocOp, {stream});
  return success();
}
 
static bool isGpuAsyncTokenType(Value value) {
  return isa<gpu::AsyncTokenType>(value.getType());
}
 
// Converts !gpu.async.token operands of `async.yield` to runtime calls. The
// !gpu.async.token are lowered to stream within the async.execute region, but
// are passed as events between them. For each !gpu.async.token operand, we
// create an event and record it on the stream.
LogicalResult ConvertAsyncYieldToGpuRuntimeCallPattern::matchAndRewrite(
    async::YieldOp yieldOp, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (llvm::none_of(yieldOp.getOperands(), isGpuAsyncTokenType))
    return rewriter.notifyMatchFailure(yieldOp, "no gpu async token operand");
 
  Location loc = yieldOp.getLoc();
  SmallVector<Value, 4> newOperands(adaptor.getOperands());
  llvm::SmallDenseSet<Value> streams;
  for (auto &operand : yieldOp->getOpOperands()) {
    if (!isGpuAsyncTokenType(operand.get()))
      continue;
    auto idx = operand.getOperandNumber();
    auto stream = adaptor.getOperands()[idx];
    auto event = eventCreateCallBuilder.create(loc, rewriter, {}).getResult();
    eventRecordCallBuilder.create(loc, rewriter, {event, stream});
    newOperands[idx] = event;
    streams.insert(stream);
  }
  for (auto stream : streams)
    streamDestroyCallBuilder.create(loc, rewriter, {stream});
 
  rewriter.modifyOpInPlace(yieldOp, [&] { yieldOp->setOperands(newOperands); });
  return success();
}
 
// Returns whether `value` is the result of an LLVM::CallOp to `functionName`.
static bool isDefinedByCallTo(Value value, StringRef functionName) {
  assert(isa<LLVM::LLVMPointerType>(value.getType()));
  if (auto defOp = value.getDefiningOp<LLVM::CallOp>())
    return defOp.getCallee()->equals(functionName);
  return false;
}
 
// Converts `gpu.wait` to runtime calls. The converted op synchronizes the host
// with the stream/event operands. The operands are destroyed. That is, it
// assumes that it is not used afterwards or elsewhere. Otherwise we will get a
// runtime error. Eventually, we should guarantee this property.
LogicalResult ConvertWaitOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::WaitOp waitOp, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (waitOp.getAsyncToken())
    return rewriter.notifyMatchFailure(waitOp, "Cannot convert async op.");
 
  Location loc = waitOp.getLoc();
 
  for (auto operand : adaptor.getOperands()) {
    if (isDefinedByCallTo(operand, streamCreateCallBuilder.functionName)) {
      // The converted operand's definition created a stream.
      streamSynchronizeCallBuilder.create(loc, rewriter, {operand});
      streamDestroyCallBuilder.create(loc, rewriter, {operand});
    } else {
      // Otherwise the converted operand is an event. This assumes that we use
      // events in control flow code as well.
      eventSynchronizeCallBuilder.create(loc, rewriter, {operand});
      eventDestroyCallBuilder.create(loc, rewriter, {operand});
    }
  }
 
  rewriter.eraseOp(waitOp);
  return success();
}
 
// Converts `gpu.wait async` to runtime calls. The converted op creates a new
// stream that is synchronized with stream/event operands. The operands are
// destroyed. That is, it assumes that it is not used afterwards or elsewhere.
// Otherwise we will get a runtime error. Eventually, we should guarantee this
// property.
LogicalResult ConvertWaitAsyncOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::WaitOp waitOp, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (!waitOp.getAsyncToken())
    return rewriter.notifyMatchFailure(waitOp, "Can only convert async op.");
 
  Location loc = waitOp.getLoc();
 
  auto insertionPoint = rewriter.saveInsertionPoint();
  SmallVector<Value, 1> events;
  for (auto pair :
       llvm::zip(waitOp.getAsyncDependencies(), adaptor.getOperands())) {
    auto operand = std::get<1>(pair);
    if (isDefinedByCallTo(operand, streamCreateCallBuilder.functionName)) {
      // The converted operand's definition created a stream. Insert an event
      // into the stream just after the last use of the original token operand.
      auto *defOp = std::get<0>(pair).getDefiningOp();
      rewriter.setInsertionPointAfter(defOp);
      auto event = eventCreateCallBuilder.create(loc, rewriter, {}).getResult();
      eventRecordCallBuilder.create(loc, rewriter, {event, operand});
      events.push_back(event);
    } else {
      // Otherwise the converted operand is an event. This assumes that we use
      // events in control flow code as well.
      events.push_back(operand);
    }
  }
  rewriter.restoreInsertionPoint(insertionPoint);
  auto stream = streamCreateCallBuilder.create(loc, rewriter, {}).getResult();
  for (auto event : events)
    streamWaitEventCallBuilder.create(loc, rewriter, {stream, event});
  for (auto event : events)
    eventDestroyCallBuilder.create(loc, rewriter, {event});
  rewriter.replaceOp(waitOp, {stream});
 
  return success();
}
 
// Creates a struct containing all kernel parameters on the stack and returns
// an array of type-erased pointers to the fields of the struct. The array can
// then be passed to the CUDA / ROCm (HIP) kernel launch calls.
// The generated code is essentially as follows:
//
// %struct = alloca(sizeof(struct { Parameters... }))
// %array = alloca(NumParameters * sizeof(void *))
// for (i : [0, NumParameters))
//   %fieldPtr = llvm.getelementptr %struct[0, i]
//   llvm.store parameters[i], %fieldPtr
//   %elementPtr = llvm.getelementptr %array[i]
//   llvm.store %fieldPtr, %elementPtr
// return %array
Value ConvertLaunchFuncOpToGpuRuntimeCallPattern::generateParamsArray(
    gpu::LaunchFuncOp launchOp, OpAdaptor adaptor, OpBuilder &builder) const {
  auto loc = launchOp.getLoc();
  auto numKernelOperands = launchOp.getNumKernelOperands();
  // Note: If `useBarePtrCallConv` is set in the type converter's options,
  // the value of `kernelBarePtrCallConv` will be ignored.
  SmallVector<Value, 4> arguments = getTypeConverter()->promoteOperands(
      loc, launchOp.getOperands().take_back(numKernelOperands),
      adaptor.getOperands().take_back(numKernelOperands), builder,
      /*useBarePtrCallConv=*/kernelBarePtrCallConv);
  auto numArguments = arguments.size();
  SmallVector<Type, 4> argumentTypes;
  argumentTypes.reserve(numArguments);
  for (auto argument : arguments)
    argumentTypes.push_back(argument.getType());
  auto structType = LLVM::LLVMStructType::getNewIdentified(context, StringRef(),
                                                           argumentTypes);
  auto one = builder.create<LLVM::ConstantOp>(loc, llvmInt32Type, 1);
  auto structPtr =
      builder.create<LLVM::AllocaOp>(loc, llvmPointerType, structType, one,
                                     /*alignment=*/0);
  auto arraySize =
      builder.create<LLVM::ConstantOp>(loc, llvmInt32Type, numArguments);
  auto arrayPtr = builder.create<LLVM::AllocaOp>(
      loc, llvmPointerType, llvmPointerType, arraySize, /*alignment=*/0);
  for (const auto &en : llvm::enumerate(arguments)) {
    const auto index = static_cast<int32_t>(en.index());
    Value fieldPtr =
        builder.create<LLVM::GEPOp>(loc, llvmPointerType, structType, structPtr,
                                    ArrayRef<LLVM::GEPArg>{0, index});
    builder.create<LLVM::StoreOp>(loc, en.value(), fieldPtr);
    auto elementPtr =
        builder.create<LLVM::GEPOp>(loc, llvmPointerType, llvmPointerType,
                                    arrayPtr, ArrayRef<LLVM::GEPArg>{index});
    builder.create<LLVM::StoreOp>(loc, fieldPtr, 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.
// The code is essentially:
//
// llvm.global constant @kernel_name("function_name\00")
// func(...) {
//   %0 = llvm.addressof @kernel_name
//   %1 = llvm.constant (0 : index)
//   %2 = llvm.getelementptr %0[%1, %1] : !llvm<"i8*">
// }
Value ConvertLaunchFuncOpToGpuRuntimeCallPattern::generateKernelNameConstant(
    StringRef moduleName, StringRef name, Location loc,
    OpBuilder &builder) const {
  // Make sure the trailing zero is included in the constant.
  std::vector<char> kernelName(name.begin(), name.end());
  kernelName.push_back('\0');
 
  std::string globalName =
      std::string(llvm::formatv("{0}_{1}_kernel_name", moduleName, name));
  return LLVM::createGlobalString(
      loc, builder, globalName, StringRef(kernelName.data(), kernelName.size()),
      LLVM::Linkage::Internal);
}
 
// Emits LLVM IR to launch a kernel function. Expects the module that contains
// the compiled kernel function as a cubin in the 'nvvm.cubin' attribute, or a
// hsaco in the 'rocdl.hsaco' attribute of the kernel function in the IR.
//
// %0 = call %binarygetter
// %1 = call %moduleLoad(%0)
// %2 = <see generateKernelNameConstant>
// %3 = call %moduleGetFunction(%1, %2)
// %4 = call %streamCreate()
// %5 = <see generateParamsArray>
// call %launchKernel(%3, <launchOp operands 0..5>, 0, %4, %5, nullptr)
// call %streamSynchronize(%4)
// call %streamDestroy(%4)
// call %moduleUnload(%1)
//
// If the op is async, the stream corresponds to the (single) async dependency
// as well as the async token the op produces.
LogicalResult ConvertLaunchFuncOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::LaunchFuncOp launchOp, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(launchOp, adaptor.getOperands(), rewriter)))
    return failure();
 
  if (launchOp.getAsyncDependencies().size() > 1)
    return rewriter.notifyMatchFailure(
        launchOp, "Cannot convert with more than one async dependency.");
 
  // Fail when the synchronous version of the op has async dependencies. The
  // lowering destroys the stream, and we do not want to check that there is no
  // use of the stream after this op.
  if (!launchOp.getAsyncToken() && !launchOp.getAsyncDependencies().empty())
    return rewriter.notifyMatchFailure(
        launchOp, "Cannot convert non-async op with async dependencies.");
 
  Location loc = launchOp.getLoc();
 
  // Create an LLVM global with CUBIN extracted from the kernel annotation and
  // obtain a pointer to the first byte in it.
  gpu::GPUModuleOp kernelModule;
  if (cachedModuleTable)
    kernelModule = cachedModuleTable->lookup<gpu::GPUModuleOp>(
        launchOp.getKernelModuleName());
  else
    kernelModule = SymbolTable::lookupNearestSymbolFrom<gpu::GPUModuleOp>(
        launchOp, launchOp.getKernelModuleName());
  assert(kernelModule && "expected a kernel module");
 
  // If the module has Targets then just update the op operands.
  if (ArrayAttr targets = kernelModule.getTargetsAttr()) {
    Value stream = Value();
    if (!adaptor.getAsyncDependencies().empty())
      stream = adaptor.getAsyncDependencies().front();
    // If the async keyword is present and there are no dependencies, then a
    // stream must be created to pass to subsequent operations.
    else if (launchOp.getAsyncToken())
      stream = streamCreateCallBuilder.create(loc, rewriter, {}).getResult();
 
    // Lower the kernel operands to match kernel parameters.
    // Note: If `useBarePtrCallConv` is set in the type converter's options,
    // the value of `kernelBarePtrCallConv` will be ignored.
    SmallVector<Value, 4> arguments = getTypeConverter()->promoteOperands(
        loc, launchOp.getKernelOperands(), adaptor.getKernelOperands(),
        rewriter, /*useBarePtrCallConv=*/kernelBarePtrCallConv);
 
    std::optional<gpu::KernelDim3> clusterSize = std::nullopt;
    if (launchOp.hasClusterSize()) {
      clusterSize =
          gpu::KernelDim3{adaptor.getClusterSizeX(), adaptor.getClusterSizeY(),
                          adaptor.getClusterSizeZ()};
    }
    rewriter.create<gpu::LaunchFuncOp>(
        launchOp.getLoc(), launchOp.getKernelAttr(),
        gpu::KernelDim3{adaptor.getGridSizeX(), adaptor.getGridSizeY(),
                        adaptor.getGridSizeZ()},
        gpu::KernelDim3{adaptor.getBlockSizeX(), adaptor.getBlockSizeY(),
                        adaptor.getBlockSizeZ()},
        adaptor.getDynamicSharedMemorySize(), arguments, stream, clusterSize);
    if (launchOp.getAsyncToken())
      rewriter.replaceOp(launchOp, {stream});
    else
      rewriter.eraseOp(launchOp);
    return success();
  }
 
  auto binaryAttr =
      kernelModule->getAttrOfType<StringAttr>(gpuBinaryAnnotation);
  if (!binaryAttr) {
    kernelModule.emitOpError()
        << "missing " << gpuBinaryAnnotation << " attribute";
    return failure();
  }
 
  SmallString<128> nameBuffer(kernelModule.getName());
  nameBuffer.append(kGpuBinaryStorageSuffix);
  Value data =
      LLVM::createGlobalString(loc, rewriter, nameBuffer.str(),
                               binaryAttr.getValue(), LLVM::Linkage::Internal);
 
  // Pass the binary size. SPIRV requires binary size.
  auto gpuBlob = binaryAttr.getValue();
  auto gpuBlobSize = rewriter.create<mlir::LLVM::ConstantOp>(
      loc, llvmInt64Type,
      mlir::IntegerAttr::get(llvmInt64Type,
                             static_cast<int64_t>(gpuBlob.size())));
 
  auto module =
      moduleLoadCallBuilder.create(loc, rewriter, {data, gpuBlobSize});
 
  // Pass the count of the parameters to runtime wrappers
  auto paramsCount = rewriter.create<mlir::LLVM::ConstantOp>(
      loc, llvmInt64Type,
      mlir::IntegerAttr::get(
          llvmInt64Type,
          static_cast<int64_t>(launchOp.getNumKernelOperands())));
 
  // Get the function from the module. The name corresponds to the name of
  // the kernel function.
  auto kernelName = generateKernelNameConstant(
      launchOp.getKernelModuleName().getValue(),
      launchOp.getKernelName().getValue(), loc, rewriter);
  auto function = moduleGetFunctionCallBuilder.create(
      loc, rewriter, {module.getResult(), kernelName});
  Value zero = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type, 0);
  Value stream =
      adaptor.getAsyncDependencies().empty()
          ? streamCreateCallBuilder.create(loc, rewriter, {}).getResult()
          : adaptor.getAsyncDependencies().front();
  // Create array of pointers to kernel arguments.
  auto kernelParams = generateParamsArray(launchOp, adaptor, rewriter);
  auto nullpointer = rewriter.create<LLVM::ZeroOp>(loc, llvmPointerType);
  Value dynamicSharedMemorySize = launchOp.getDynamicSharedMemorySize()
                                      ? launchOp.getDynamicSharedMemorySize()
                                      : zero;
  launchKernelCallBuilder.create(
      loc, rewriter,
      {function.getResult(), adaptor.getGridSizeX(), adaptor.getGridSizeY(),
       adaptor.getGridSizeZ(), adaptor.getBlockSizeX(), adaptor.getBlockSizeY(),
       adaptor.getBlockSizeZ(), dynamicSharedMemorySize, stream, kernelParams,
       /*extra=*/nullpointer, paramsCount});
 
  if (launchOp.getAsyncToken()) {
    // 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.
    streamSynchronizeCallBuilder.create(loc, rewriter, stream);
    streamDestroyCallBuilder.create(loc, rewriter, stream);
    rewriter.eraseOp(launchOp);
  }
  moduleUnloadCallBuilder.create(loc, rewriter, module.getResult());
 
  return success();
}
 
static Value bitAndAddrspaceCast(Location loc,
                                 ConversionPatternRewriter &rewriter,
                                 LLVM::LLVMPointerType destinationType,
                                 Value sourcePtr,
                                 const LLVMTypeConverter &typeConverter) {
  auto sourceTy = cast<LLVM::LLVMPointerType>(sourcePtr.getType());
  if (destinationType.getAddressSpace() != sourceTy.getAddressSpace())
    sourcePtr = rewriter.create<LLVM::AddrSpaceCastOp>(
        loc,
        LLVM::LLVMPointerType::get(rewriter.getContext(),
                                   destinationType.getAddressSpace()),
        sourcePtr);
  return sourcePtr;
}
 
LogicalResult ConvertMemcpyOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::MemcpyOp memcpyOp, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  auto memRefType = cast<MemRefType>(memcpyOp.getSrc().getType());
 
  if (failed(areAllLLVMTypes(memcpyOp, adaptor.getOperands(), rewriter)) ||
      !isConvertibleAndHasIdentityMaps(memRefType) ||
      failed(isAsyncWithOneDependency(rewriter, memcpyOp)))
    return failure();
 
  auto loc = memcpyOp.getLoc();
 
  MemRefDescriptor srcDesc(adaptor.getSrc());
  Value numElements = getNumElements(rewriter, loc, memRefType, srcDesc);
 
  Type elementPtrType = getElementPtrType(memRefType);
  Value nullPtr = rewriter.create<LLVM::ZeroOp>(loc, elementPtrType);
  Value gepPtr = rewriter.create<LLVM::GEPOp>(
      loc, elementPtrType,
      typeConverter->convertType(memRefType.getElementType()), nullPtr,
      numElements);
  auto sizeBytes =
      rewriter.create<LLVM::PtrToIntOp>(loc, getIndexType(), gepPtr);
 
  auto src = bitAndAddrspaceCast(loc, rewriter, llvmPointerType,
                                 srcDesc.alignedPtr(rewriter, loc),
                                 *getTypeConverter());
  auto dst = bitAndAddrspaceCast(
      loc, rewriter, llvmPointerType,
      MemRefDescriptor(adaptor.getDst()).alignedPtr(rewriter, loc),
      *getTypeConverter());
 
  auto stream = adaptor.getAsyncDependencies().front();
  memcpyCallBuilder.create(loc, rewriter, {dst, src, sizeBytes, stream});
 
  rewriter.replaceOp(memcpyOp, {stream});
 
  return success();
}
 
LogicalResult ConvertMemsetOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::MemsetOp memsetOp, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  auto memRefType = cast<MemRefType>(memsetOp.getDst().getType());
 
  if (failed(areAllLLVMTypes(memsetOp, adaptor.getOperands(), rewriter)) ||
      !isConvertibleAndHasIdentityMaps(memRefType) ||
      failed(isAsyncWithOneDependency(rewriter, memsetOp)))
    return failure();
 
  auto loc = memsetOp.getLoc();
 
  Type valueType = adaptor.getValue().getType();
  unsigned bitWidth = valueType.getIntOrFloatBitWidth();
  // Ints and floats of 16 or 32 bit width are allowed.
  if (!valueType.isIntOrFloat() || (bitWidth != 16 && bitWidth != 32)) {
    return rewriter.notifyMatchFailure(
        memsetOp, "value must be a 16 or 32 bit int or float");
  }
 
  unsigned valueTypeWidth = valueType.getIntOrFloatBitWidth();
  Type bitCastType = valueTypeWidth == 32 ? llvmInt32Type : llvmInt16Type;
 
  MemRefDescriptor dstDesc(adaptor.getDst());
  Value numElements = getNumElements(rewriter, loc, memRefType, dstDesc);
 
  auto value =
      rewriter.create<LLVM::BitcastOp>(loc, bitCastType, adaptor.getValue());
  auto dst = bitAndAddrspaceCast(loc, rewriter, llvmPointerType,
                                 dstDesc.alignedPtr(rewriter, loc),
                                 *getTypeConverter());
 
  auto stream = adaptor.getAsyncDependencies().front();
  FunctionCallBuilder builder =
      valueTypeWidth == 32 ? memset32CallBuilder : memset16CallBuilder;
  builder.create(loc, rewriter, {dst, value, numElements, stream});
 
  rewriter.replaceOp(memsetOp, {stream});
  return success();
}
 
LogicalResult ConvertSetDefaultDeviceOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::SetDefaultDeviceOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Location loc = op.getLoc();
  auto call = setDefaultDeviceCallBuilder.create(loc, rewriter,
                                                 {adaptor.getDevIndex()});
  rewriter.replaceOp(op, call);
  return success();
}
 
template <typename T>
static Value genConstInt32From(OpBuilder &builder, Location loc, T tValue) {
  Type llvmInt32Type = builder.getIntegerType(32);
  return builder.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                          static_cast<int32_t>(tValue));
}
 
template <typename T>
static Value genConstFloat32From(OpBuilder &builder, Location loc, T tValue) {
  Type llvmFloat32Type = builder.getF32Type();
  return builder.create<LLVM::ConstantOp>(
      loc, llvmFloat32Type,
      builder.getF32FloatAttr(static_cast<float>(tValue)));
}
 
LogicalResult ConvertCreateDnTensorOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::CreateDnTensorOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto stream = adaptor.getAsyncDependencies().front();
  Value pTensor =
      MemRefDescriptor(adaptor.getMemref()).allocatedPtr(rewriter, loc);
  Type dType = op.getMemref().getType().getElementType();
  auto dtp = genConstInt32From(rewriter, loc, getCuSparseDataTypeFrom(dType));
 
  SmallVector<Value, 4> dims;
  for (Value dim : adaptor.getDims()) {
    dims.push_back(dim);
  }
 
  Value handle;
  // TODO: For now, we track the use of the handle and lower it to cusparse /
  // cusparseLt accordingly. If in a block, both cusparse and cusparseLt are
  // used, we require two separate Creation ops to be the correct logic. In
  // future, we may add support to using one handle in sparse tensor / GPU
  // dialect in both cusparse and cusparseLt. use the cusparseLt create call if
  // the dnmat is used with spmat with 2:4 sparsity
  if (dims.size() == 2) {
    if (isSpMMCusparseLtOp(op.getDnTensor())) {
      auto handleSz = rewriter.create<LLVM::ConstantOp>(
          loc, getIndexType(), rewriter.getIndexAttr(11032));
      handle = rewriter.create<LLVM::AllocaOp>(
          loc, llvmPointerType, llvmInt8Type, handleSz, /*alignment=*/16);
      handle = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, handle);
 
      createLtDnMatCallBuilder
          .create(loc, rewriter,
                  {handle, dims[0], dims[1], pTensor, dtp, stream})
          .getResult();
    } else {
      handle =
          createDnMatCallBuilder
              .create(loc, rewriter, {dims[0], dims[1], pTensor, dtp, stream})
              .getResult();
    }
  } else {
    assert(dims.size() == 1 && "Only 1D and 2D tensors are supported");
    handle = createDnVecCallBuilder
                 .create(loc, rewriter, {dims[0], pTensor, dtp, stream})
                 .getResult();
  }
  rewriter.replaceOp(op, {handle, stream});
  return success();
}
 
LogicalResult ConvertDestroyDnTensorOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::DestroyDnTensorOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto stream = adaptor.getAsyncDependencies().front();
  auto definingOp = op.getDnTensor().getDefiningOp<gpu::CreateDnTensorOp>();
  SmallVector<Value, 4> dims;
  for (Value dim : definingOp.getDims()) {
    dims.push_back(dim);
  }
  if (dims.size() == 2) {
    // Use the cusparseLt destroy call if the dnmat is used with spmat with
    // 2:4 sparsity
    if (isSpMMCusparseLtOp(op.getDnTensor())) {
      destroyCuSparseLtDnMatBuilder.create(loc, rewriter,
                                           {adaptor.getDnTensor(), stream});
    } else {
      destroyDnMatCallBuilder.create(loc, rewriter,
                                     {adaptor.getDnTensor(), stream});
    }
  } else {
    assert(dims.size() == 1 && "Only 1D and 2D tensors are supported");
    destroyDnVecCallBuilder.create(loc, rewriter,
                                   {adaptor.getDnTensor(), stream});
  }
  rewriter.replaceOp(op, {stream});
  return success();
}
 
LogicalResult ConvertCreateCooOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::CreateCooOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto stream = adaptor.getAsyncDependencies().front();
  Value pRowIdxs =
      MemRefDescriptor(adaptor.getRowIdxs()).allocatedPtr(rewriter, loc);
  Value pColIdxs =
      MemRefDescriptor(adaptor.getColIdxs()).allocatedPtr(rewriter, loc);
  Value pValues =
      MemRefDescriptor(adaptor.getValues()).allocatedPtr(rewriter, loc);
  Type iType =
      llvm::cast<MemRefType>(op.getColIdxs().getType()).getElementType();
  Type dType =
      llvm::cast<MemRefType>(op.getValues().getType()).getElementType();
  auto itp = genConstInt32From(rewriter, loc, getCuSparseIndexTypeFrom(iType));
  auto dtp = genConstInt32From(rewriter, loc, getCuSparseDataTypeFrom(dType));
  auto handle =
      createCooCallBuilder
          .create(loc, rewriter,
                  {adaptor.getRows(), adaptor.getCols(), adaptor.getNnz(),
                   pRowIdxs, pColIdxs, pValues, itp, dtp, stream})
          .getResult();
  rewriter.replaceOp(op, {handle, stream});
  return success();
}
 
LogicalResult ConvertCreateCooAoSOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::CreateCooAoSOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto stream = adaptor.getAsyncDependencies().front();
  Value pIdxs = MemRefDescriptor(adaptor.getIdxs()).allocatedPtr(rewriter, loc);
  Value pValues =
      MemRefDescriptor(adaptor.getValues()).allocatedPtr(rewriter, loc);
  Type iType = llvm::cast<MemRefType>(op.getIdxs().getType()).getElementType();
  Type dType =
      llvm::cast<MemRefType>(op.getValues().getType()).getElementType();
  auto itp = genConstInt32From(rewriter, loc, getCuSparseIndexTypeFrom(iType));
  auto dtp = genConstInt32From(rewriter, loc, getCuSparseDataTypeFrom(dType));
  auto handle =
      createCooAoSCallBuilder
          .create(loc, rewriter,
                  {adaptor.getRows(), adaptor.getCols(), adaptor.getNnz(),
                   pIdxs, pValues, itp, dtp, stream})
          .getResult();
  rewriter.replaceOp(op, {handle, stream});
  return success();
}
 
LogicalResult ConvertCreateCsrOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::CreateCsrOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto stream = adaptor.getAsyncDependencies().front();
  Value pRowPos =
      MemRefDescriptor(adaptor.getRowPos()).allocatedPtr(rewriter, loc);
  Value pColIdxs =
      MemRefDescriptor(adaptor.getColIdxs()).allocatedPtr(rewriter, loc);
  Value pValues =
      MemRefDescriptor(adaptor.getValues()).allocatedPtr(rewriter, loc);
  Type pType =
      llvm::cast<MemRefType>(op.getRowPos().getType()).getElementType();
  Type iType =
      llvm::cast<MemRefType>(op.getColIdxs().getType()).getElementType();
  Type dType =
      llvm::cast<MemRefType>(op.getValues().getType()).getElementType();
  auto ptp = genConstInt32From(rewriter, loc, getCuSparseIndexTypeFrom(pType));
  auto itp = genConstInt32From(rewriter, loc, getCuSparseIndexTypeFrom(iType));
  auto dtp = genConstInt32From(rewriter, loc, getCuSparseDataTypeFrom(dType));
  auto handle =
      createCsrCallBuilder
          .create(loc, rewriter,
                  {adaptor.getRows(), adaptor.getCols(), adaptor.getNnz(),
                   pRowPos, pColIdxs, pValues, ptp, itp, dtp, stream})
          .getResult();
  rewriter.replaceOp(op, {handle, stream});
  return success();
}
 
LogicalResult ConvertCreate2To4SpMatOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::Create2To4SpMatOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto stream = adaptor.getAsyncDependencies().front();
  Value pMat =
      MemRefDescriptor(adaptor.getMemref()).allocatedPtr(rewriter, loc);
  Type dType =
      llvm::cast<MemRefType>(op.getMemref().getType()).getElementType();
  auto dtp = genConstInt32From(rewriter, loc, getCuSparseDataTypeFrom(dType));
 
  // CUDA runner asserts the size is 44104 bytes.
  auto handleSz = rewriter.create<LLVM::ConstantOp>(
      loc, getIndexType(), rewriter.getIndexAttr(44104));
  Value handle = rewriter.create<LLVM::AllocaOp>(
      loc, llvmPointerType, llvmInt8Type, handleSz, /*alignment=*/16);
  handle = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, handle);
 
  create2To4SpMatCallBuilder
      .create(loc, rewriter,
              {handle, adaptor.getRows(), adaptor.getCols(), pMat, dtp, stream})
      .getResult();
  rewriter.replaceOp(op, {handle, stream});
  return success();
}
 
LogicalResult ConvertDestroySpMatOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::DestroySpMatOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto stream = adaptor.getAsyncDependencies().front();
  // Use the cusparseLt destroy call if the spmat is 2:4 sparsity
  if (is2To4Sparsity(op.getSpmat())) {
    destroyCuSparseLtSpMatBuilder.create(loc, rewriter,
                                         {adaptor.getSpmat(), stream});
 
  } else {
    destroySpMatCallBuilder.create(loc, rewriter, {adaptor.getSpmat(), stream});
  }
  rewriter.replaceOp(op, {stream});
  return success();
}
 
LogicalResult ConvertSpMVBufferSizeOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::SpMVBufferSizeOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto modeA = genConstInt32From(rewriter, loc, op.getModeA());
  auto computeType = genConstInt32From(
      rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
  auto stream = adaptor.getAsyncDependencies().front();
  auto bufferSize = spMVBufferSizeCallBuilder
                        .create(loc, rewriter,
                                {modeA, adaptor.getSpmatA(), adaptor.getDnX(),
                                 adaptor.getDnY(), computeType, stream})
                        .getResult();
  rewriter.replaceOp(op, {bufferSize, stream});
  return success();
}
 
LogicalResult ConvertSpMVOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::SpMVOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
  auto computeType = genConstInt32From(
      rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
  auto stream = adaptor.getAsyncDependencies().front();
  Value pBuf =
      MemRefDescriptor(adaptor.getBuffer()).allocatedPtr(rewriter, loc);
  spMVCallBuilder.create(loc, rewriter,
                         {modeA, adaptor.getSpmatA(), adaptor.getDnX(),
                          adaptor.getDnY(), computeType, pBuf, stream});
  rewriter.replaceOp(op, {stream});
  return success();
}
 
LogicalResult ConvertSpMMBufferSizeOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::SpMMBufferSizeOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
  auto modeB = genConstInt32From(rewriter, loc, adaptor.getModeB());
  auto stream = adaptor.getAsyncDependencies().front();
  Value bufferSize;
  if (is2To4Sparsity(op.getSpmatA())) {
    auto pruneFlag =
        genConstInt32From(rewriter, loc, get2To4PruneFlag(op.getSpmatA()));
    auto computeType = genConstInt32From(
        rewriter, loc, getCuSparseLtDataTypeFrom(adaptor.getComputeType()));
    auto three = rewriter.create<LLVM::ConstantOp>(loc, getIndexType(),
                                                   rewriter.getIndexAttr(3));
    auto bufferSize = rewriter.create<LLVM::AllocaOp>(
        loc, llvmPointerType, llvmPointerType, three, /*alignment=*/16);
    createCuSparseLtSpMMBufferSizeBuilder
        .create(loc, rewriter,
                {bufferSize, modeA, modeB, adaptor.getSpmatA(),
                 adaptor.getDnmatB(), adaptor.getDnmatC(), computeType,
                 pruneFlag, stream})
        .getResult();
 
    auto bufferSizePtr1 = rewriter.create<LLVM::GEPOp>(
        loc, llvmPointerType, llvmPointerType, bufferSize,
        ValueRange{rewriter.create<LLVM::ConstantOp>(
            loc, getIndexType(), rewriter.getIndexAttr(1))});
    auto bufferSizePtr2 = rewriter.create<LLVM::GEPOp>(
        loc, llvmPointerType, llvmPointerType, bufferSize,
        ValueRange{rewriter.create<LLVM::ConstantOp>(
            loc, getIndexType(), rewriter.getIndexAttr(2))});
    auto bufferSize0 =
        rewriter.create<LLVM::LoadOp>(loc, llvmInt64Type, bufferSize);
    auto bufferSize1 =
        rewriter.create<LLVM::LoadOp>(loc, llvmInt64Type, bufferSizePtr1);
    auto bufferSize2 =
        rewriter.create<LLVM::LoadOp>(loc, llvmInt64Type, bufferSizePtr2);
 
    rewriter.replaceOp(op, {bufferSize0, bufferSize1, bufferSize2, stream});
  } else {
    auto computeType = genConstInt32From(
        rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
    bufferSize =
        createSpMMBufferSizeCallBuilder
            .create(loc, rewriter,
                    {modeA, modeB, adaptor.getSpmatA(), adaptor.getDnmatB(),
                     adaptor.getDnmatC(), computeType, stream})
            .getResult();
    rewriter.replaceOp(op, {bufferSize, stream});
  }
  return success();
}
 
LogicalResult ConvertSDDMMBufferSizeOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::SDDMMBufferSizeOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
  auto modeB = genConstInt32From(rewriter, loc, adaptor.getModeB());
  auto computeType = genConstInt32From(
      rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
  auto stream = adaptor.getAsyncDependencies().front();
  auto bufferSize =
      createSDDMMBufferSizeCallBuilder
          .create(loc, rewriter,
                  {modeA, modeB, adaptor.getDnmatA(), adaptor.getDnmatB(),
                   adaptor.getSpmatC(), computeType, stream})
          .getResult();
  rewriter.replaceOp(op, {bufferSize, stream});
  return success();
}
 
LogicalResult ConvertSpMMOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::SpMMOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
  auto modeB = genConstInt32From(rewriter, loc, adaptor.getModeB());
  auto computeType = genConstInt32From(
      rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
 
  auto stream = adaptor.getAsyncDependencies().front();
 
  // Lower to cusparseLt if applicable
  if (is2To4Sparsity(op.getSpmatA())) {
    SmallVector<Value> pBufs;
    for (Value buffer : adaptor.getBuffers()) {
      Value pBuf = MemRefDescriptor(buffer).allocatedPtr(rewriter, loc);
      pBufs.push_back(pBuf);
    }
    createCuSparseLtSpMMBuilder.create(
        loc, rewriter,
        {adaptor.getSpmatA(), adaptor.getDnmatB(), adaptor.getDnmatC(),
         pBufs[0], pBufs[1], pBufs[2], stream});
  } else {
    Value pBuf = MemRefDescriptor(adaptor.getBuffers().front())
                     .allocatedPtr(rewriter, loc);
    createSpMMCallBuilder.create(loc, rewriter,
                                 {modeA, modeB, adaptor.getSpmatA(),
                                  adaptor.getDnmatB(), adaptor.getDnmatC(),
                                  computeType, pBuf, stream});
  }
  rewriter.replaceOp(op, {stream});
  return success();
}
 
template <typename T>
static void addOpaquePointerConversion(LLVMTypeConverter &converter) {
  converter.addConversion([&converter](T) -> Type {
    return LLVM::LLVMPointerType::get(&converter.getContext());
  });
}
 
LogicalResult ConvertSDDMMOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::SDDMMOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto computeType = genConstInt32From(
      rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
  auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
  auto modeB = genConstInt32From(rewriter, loc, adaptor.getModeB());
  auto stream = adaptor.getAsyncDependencies().front();
  Value pBuf =
      MemRefDescriptor(adaptor.getBuffer()).allocatedPtr(rewriter, loc);
  createSDDMMCallBuilder.create(loc, rewriter,
                                {modeA, modeB, adaptor.getDnmatA(),
                                 adaptor.getDnmatB(), adaptor.getSpmatC(),
                                 computeType, pBuf, stream});
  rewriter.replaceOp(op, {stream});
  return success();
}
 
LogicalResult
ConvertSpGEMMCreateDescrOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::SpGEMMCreateDescrOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto stream = adaptor.getAsyncDependencies().front();
  Value descr = createSpGEMMCreateDescrBuilder.create(loc, rewriter, {stream})
                    .getResult();
  rewriter.replaceOp(op, {descr, stream});
  return success();
}
 
LogicalResult
ConvertSpGEMMDestroyDescrOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::SpGEMMDestroyDescrOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto stream = adaptor.getAsyncDependencies().front();
  createSpGEMMDestroyDescrBuilder.create(loc, rewriter,
                                         {adaptor.getDesc(), stream});
  rewriter.replaceOp(op, {stream});
  return success();
}
 
LogicalResult
ConvertSpGEMMWorkEstimationOrComputeOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::SpGEMMWorkEstimationOrComputeOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto computeType = genConstInt32From(
      rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
  auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
  auto modeB = genConstInt32From(rewriter, loc, adaptor.getModeB());
  auto stream = adaptor.getAsyncDependencies().front();
 
  Value pBuf =
      MemRefDescriptor(adaptor.getBuffer()).allocatedPtr(rewriter, loc);
  Value bufferSizeNew;
 
  if (adaptor.getKind() ==
      gpu::SpGEMMWorkEstimationOrComputeKind::WORK_ESTIMATION) {
    bufferSizeNew =
        createSpGEMMWorkEstimationBuilder
            .create(loc, rewriter,
                    {adaptor.getDesc(), modeA, modeB, adaptor.getSpmatA(),
                     adaptor.getSpmatB(), adaptor.getSpmatC(), computeType,
                     adaptor.getBufferSz(), pBuf, stream})
            .getResult();
  } else {
    bufferSizeNew =
        createSpGEMMComputeBuilder
            .create(loc, rewriter,
                    {adaptor.getDesc(), modeA, modeB, adaptor.getSpmatA(),
                     adaptor.getSpmatB(), adaptor.getSpmatC(), computeType,
                     adaptor.getBufferSz(), pBuf, stream})
            .getResult();
  }
  rewriter.replaceOp(op, {bufferSizeNew, stream});
  return success();
}
 
LogicalResult ConvertSpGEMMCopyOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::SpGEMMCopyOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto computeType = genConstInt32From(
      rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
  auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
  auto modeB = genConstInt32From(rewriter, loc, adaptor.getModeB());
  auto stream = adaptor.getAsyncDependencies().front();
  createSpGEMMCopyBuilder.create(loc, rewriter,
                                 {adaptor.getDesc(), modeA, modeB,
                                  adaptor.getSpmatA(), adaptor.getSpmatB(),
                                  adaptor.getSpmatC(), computeType, stream});
  rewriter.replaceOp(op, {stream});
  return success();
}
 
LogicalResult ConvertSpMatGetSizeOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::SpMatGetSizeOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto stream = adaptor.getAsyncDependencies().front();
 
  auto three = rewriter.create<LLVM::ConstantOp>(loc, getIndexType(),
                                                 rewriter.getIndexAttr(3));
  auto buffer = rewriter.create<LLVM::AllocaOp>(
      loc, llvmPointerType, llvmInt64Type, three, /*alignment=*/16);
 
  auto rowsPtr = rewriter.create<LLVM::GEPOp>(
      loc, llvmPointerType, llvmPointerType, buffer,
      ValueRange{rewriter.create<LLVM::ConstantOp>(loc, getIndexType(),
                                                   rewriter.getIndexAttr(0))});
  auto colsPtr = rewriter.create<LLVM::GEPOp>(
      loc, llvmPointerType, llvmPointerType, buffer,
      ValueRange{rewriter.create<LLVM::ConstantOp>(loc, getIndexType(),
                                                   rewriter.getIndexAttr(1))});
  auto nnzsPtr = rewriter.create<LLVM::GEPOp>(
      loc, llvmPointerType, llvmPointerType, buffer,
      ValueRange{rewriter.create<LLVM::ConstantOp>(loc, getIndexType(),
                                                   rewriter.getIndexAttr(2))});
  createSpMatGetSizeBuilder.create(
      loc, rewriter, {adaptor.getSpmat(), rowsPtr, colsPtr, nnzsPtr, stream});
  auto rows = rewriter.create<LLVM::LoadOp>(loc, llvmInt64Type, rowsPtr);
  auto cols = rewriter.create<LLVM::LoadOp>(loc, llvmInt64Type, colsPtr);
  auto nnzs = rewriter.create<LLVM::LoadOp>(loc, llvmInt64Type, nnzsPtr);
 
  rewriter.replaceOp(op, {rows, cols, nnzs, stream});
  return success();
}
 
LogicalResult ConvertSetCsrPointersOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::SetCsrPointersOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto stream = adaptor.getAsyncDependencies().front();
  Value pPos =
      MemRefDescriptor(adaptor.getPositions()).allocatedPtr(rewriter, loc);
  Value pCrd =
      MemRefDescriptor(adaptor.getCoordinates()).allocatedPtr(rewriter, loc);
  Value pVal =
      MemRefDescriptor(adaptor.getValues()).allocatedPtr(rewriter, loc);
  createSetCsrPointersBuilder.create(
      loc, rewriter, {adaptor.getSpmat(), pPos, pCrd, pVal, stream});
  rewriter.replaceOp(op, {stream});
  return success();
}
 
LogicalResult ConvertCreateCscOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::CreateCscOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto stream = adaptor.getAsyncDependencies().front();
  Value pColPos =
      MemRefDescriptor(adaptor.getColPos()).allocatedPtr(rewriter, loc);
  Value pRowIdxs =
      MemRefDescriptor(adaptor.getRowIdxs()).allocatedPtr(rewriter, loc);
  Value pValues =
      MemRefDescriptor(adaptor.getValues()).allocatedPtr(rewriter, loc);
  Type pType =
      llvm::cast<MemRefType>(op.getColPos().getType()).getElementType();
  Type iType =
      llvm::cast<MemRefType>(op.getRowIdxs().getType()).getElementType();
  Type dType =
      llvm::cast<MemRefType>(op.getValues().getType()).getElementType();
  auto ptp = genConstInt32From(rewriter, loc, getCuSparseIndexTypeFrom(pType));
  auto itp = genConstInt32From(rewriter, loc, getCuSparseIndexTypeFrom(iType));
  auto dtp = genConstInt32From(rewriter, loc, getCuSparseDataTypeFrom(dType));
  auto handle =
      createCscCallBuilder
          .create(loc, rewriter,
                  {adaptor.getRows(), adaptor.getCols(), adaptor.getNnz(),
                   pColPos, pRowIdxs, pValues, ptp, itp, dtp, stream})
          .getResult();
  rewriter.replaceOp(op, {handle, stream});
  return success();
}
 
LogicalResult ConvertCreateBsrOpToGpuRuntimeCallPattern::matchAndRewrite(
    gpu::CreateBsrOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
      failed(isAsyncWithOneDependency(rewriter, op)))
    return failure();
  Location loc = op.getLoc();
  auto stream = adaptor.getAsyncDependencies().front();
  Value pRowPos =
      MemRefDescriptor(adaptor.getBRowPos()).allocatedPtr(rewriter, loc);
  Value pColIdxs =
      MemRefDescriptor(adaptor.getBColIdxs()).allocatedPtr(rewriter, loc);
  Value pValues =
      MemRefDescriptor(adaptor.getValues()).allocatedPtr(rewriter, loc);
  Type pType =
      llvm::cast<MemRefType>(op.getBRowPos().getType()).getElementType();
  Type iType =
      llvm::cast<MemRefType>(op.getBColIdxs().getType()).getElementType();
  Type dType =
      llvm::cast<MemRefType>(op.getValues().getType()).getElementType();
  auto ptp = genConstInt32From(rewriter, loc, getCuSparseIndexTypeFrom(pType));
  auto itp = genConstInt32From(rewriter, loc, getCuSparseIndexTypeFrom(iType));
  auto dtp = genConstInt32From(rewriter, loc, getCuSparseDataTypeFrom(dType));
  auto handle =
      createBsrCallBuilder
          .create(loc, rewriter,
                  {adaptor.getBrows(), adaptor.getBcols(), adaptor.getBnnz(),
                   adaptor.getRBlockSize(), adaptor.getCBlockSize(), pRowPos,
                   pColIdxs, pValues, ptp, itp, dtp, stream})
          .getResult();
  rewriter.replaceOp(op, {handle, stream});
  return success();
}
 
void mlir::populateGpuToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                               RewritePatternSet &patterns,
                                               StringRef gpuBinaryAnnotation,
                                               bool kernelBarePtrCallConv,
                                               SymbolTable *cachedModuleTable) {
  addOpaquePointerConversion<gpu::AsyncTokenType>(converter);
  addOpaquePointerConversion<gpu::SparseDnTensorHandleType>(converter);
  addOpaquePointerConversion<gpu::SparseSpMatHandleType>(converter);
  addOpaquePointerConversion<gpu::SparseSpGEMMOpHandleType>(converter);
 
  patterns.add<ConvertAllocOpToGpuRuntimeCallPattern,
               ConvertDeallocOpToGpuRuntimeCallPattern,
               ConvertHostRegisterOpToGpuRuntimeCallPattern,
               ConvertHostUnregisterOpToGpuRuntimeCallPattern,
               ConvertMemcpyOpToGpuRuntimeCallPattern,
               ConvertMemsetOpToGpuRuntimeCallPattern,
               ConvertSetDefaultDeviceOpToGpuRuntimeCallPattern,
               ConvertWaitAsyncOpToGpuRuntimeCallPattern,
               ConvertWaitOpToGpuRuntimeCallPattern,
               ConvertAsyncYieldToGpuRuntimeCallPattern,
               ConvertCreateDnTensorOpToGpuRuntimeCallPattern,
               ConvertDestroyDnTensorOpToGpuRuntimeCallPattern,
               ConvertCreateCooOpToGpuRuntimeCallPattern,
               ConvertCreateCooAoSOpToGpuRuntimeCallPattern,
               ConvertCreateCsrOpToGpuRuntimeCallPattern,
               ConvertCreateCscOpToGpuRuntimeCallPattern,
               ConvertCreateBsrOpToGpuRuntimeCallPattern,
               ConvertCreate2To4SpMatOpToGpuRuntimeCallPattern,
               ConvertDestroySpMatOpToGpuRuntimeCallPattern,
               ConvertSpMVBufferSizeOpToGpuRuntimeCallPattern,
               ConvertSpMVOpToGpuRuntimeCallPattern,
               ConvertSpMMBufferSizeOpToGpuRuntimeCallPattern,
               ConvertSDDMMBufferSizeOpToGpuRuntimeCallPattern,
               ConvertSpMMOpToGpuRuntimeCallPattern,
               ConvertSDDMMOpToGpuRuntimeCallPattern,
               ConvertSpGEMMCreateDescrOpToGpuRuntimeCallPattern,
               ConvertSpGEMMDestroyDescrOpToGpuRuntimeCallPattern,
               ConvertSpGEMMWorkEstimationOrComputeOpToGpuRuntimeCallPattern,
               ConvertSpGEMMCopyOpToGpuRuntimeCallPattern,
               ConvertSpMatGetSizeOpToGpuRuntimeCallPattern,
               ConvertSetCsrPointersOpToGpuRuntimeCallPattern>(converter);
  patterns.add<ConvertLaunchFuncOpToGpuRuntimeCallPattern>(
      converter, gpuBinaryAnnotation, kernelBarePtrCallConv, cachedModuleTable);
  patterns.add<EraseGpuModuleOpPattern>(&converter.getContext());
}