2020-05-18 13:01:54 -05:00
|
|
|
//===- ConvertLaunchFuncToGpuRuntimeCalls.cpp - MLIR GPU lowering passes --===//
|
2019-06-19 06:22:36 -07:00
|
|
|
//
|
2020-01-26 03:58:30 +00:00
|
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
2019-12-23 09:35:36 -08:00
|
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
|
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
2019-06-19 06:22:36 -07:00
|
|
|
//
|
2019-12-23 09:35:36 -08:00
|
|
|
//===----------------------------------------------------------------------===//
|
2019-06-19 06:22:36 -07:00
|
|
|
//
|
|
|
|
|
// This file implements a pass to convert gpu.launch_func op into a sequence of
|
2020-05-18 13:01:54 -05:00
|
|
|
// 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.
|
2019-06-19 06:22:36 -07:00
|
|
|
//
|
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
|
|
2020-05-18 13:01:54 -05:00
|
|
|
#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
|
2019-06-19 06:22:36 -07:00
|
|
|
|
2022-09-29 11:14:47 -04:00
|
|
|
#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
|
2021-02-03 20:10:51 +01:00
|
|
|
#include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
|
2022-02-06 15:10:03 -08:00
|
|
|
#include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
|
2022-03-01 14:53:41 -08:00
|
|
|
#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"
|
|
|
|
|
#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
|
2021-07-08 18:35:18 +02:00
|
|
|
#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
|
2021-07-07 11:45:27 +02:00
|
|
|
#include "mlir/Conversion/LLVMCommon/Pattern.h"
|
2021-07-08 18:35:18 +02:00
|
|
|
#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
|
2021-03-18 12:59:49 -07:00
|
|
|
#include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
|
2021-02-24 15:02:30 +01:00
|
|
|
#include "mlir/Dialect/Async/IR/Async.h"
|
2022-06-09 21:33:41 +00:00
|
|
|
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
|
|
|
|
|
#include "mlir/Dialect/GPU/Transforms/Passes.h"
|
2019-08-19 11:00:47 -07:00
|
|
|
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
|
2023-01-17 23:20:11 -05:00
|
|
|
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
2019-06-19 06:22:36 -07:00
|
|
|
#include "mlir/IR/Attributes.h"
|
|
|
|
|
#include "mlir/IR/Builders.h"
|
2020-11-19 10:43:12 -08:00
|
|
|
#include "mlir/IR/BuiltinOps.h"
|
2020-12-03 17:22:29 -08:00
|
|
|
#include "mlir/IR/BuiltinTypes.h"
|
2019-06-19 06:22:36 -07:00
|
|
|
|
|
|
|
|
#include "llvm/ADT/STLExtras.h"
|
|
|
|
|
#include "llvm/Support/Error.h"
|
2019-08-20 07:51:32 -07:00
|
|
|
#include "llvm/Support/FormatVariadic.h"
|
2019-06-19 06:22:36 -07:00
|
|
|
|
2022-08-31 10:16:29 +02:00
|
|
|
namespace mlir {
|
|
|
|
|
#define GEN_PASS_DEF_GPUTOLLVMCONVERSIONPASS
|
|
|
|
|
#include "mlir/Conversion/Passes.h.inc"
|
|
|
|
|
} // namespace mlir
|
|
|
|
|
|
2019-06-19 06:22:36 -07:00
|
|
|
using namespace mlir;
|
|
|
|
|
|
2020-05-18 13:01:54 -05:00
|
|
|
static constexpr const char *kGpuBinaryStorageSuffix = "_gpubin_cst";
|
2019-06-19 06:22:36 -07:00
|
|
|
|
|
|
|
|
namespace {
|
|
|
|
|
|
2022-08-30 22:20:36 +02:00
|
|
|
class GpuToLLVMConversionPass
|
2022-08-31 10:16:29 +02:00
|
|
|
: public impl::GpuToLLVMConversionPassBase<GpuToLLVMConversionPass> {
|
2020-08-01 15:06:25 +02:00
|
|
|
public:
|
2023-02-10 20:40:38 +01:00
|
|
|
using Base::Base;
|
2021-03-09 20:19:32 +01:00
|
|
|
|
2020-08-01 15:06:25 +02:00
|
|
|
// Run the dialect converter on the module.
|
|
|
|
|
void runOnOperation() override;
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
template <typename OpTy>
|
|
|
|
|
class ConvertOpToGpuRuntimeCallPattern : public ConvertOpToLLVMPattern<OpTy> {
|
|
|
|
|
public:
|
2023-08-14 08:40:13 +02:00
|
|
|
explicit ConvertOpToGpuRuntimeCallPattern(
|
|
|
|
|
const LLVMTypeConverter &typeConverter)
|
2020-08-01 15:06:25 +02:00
|
|
|
: ConvertOpToLLVMPattern<OpTy>(typeConverter) {}
|
|
|
|
|
|
|
|
|
|
protected:
|
2021-09-04 08:03:33 +02:00
|
|
|
Value getNumElements(ConversionPatternRewriter &rewriter, Location loc,
|
|
|
|
|
MemRefType type, MemRefDescriptor desc) const {
|
2023-08-04 13:27:31 +00:00
|
|
|
Type indexType = ConvertToLLVMPattern::getIndexType();
|
2021-09-04 08:03:33 +02:00
|
|
|
return type.hasStaticShape()
|
2023-08-02 06:23:34 +00:00
|
|
|
? ConvertToLLVMPattern::createIndexAttrConstant(
|
|
|
|
|
rewriter, loc, indexType, type.getNumElements())
|
2021-09-04 08:03:33 +02:00
|
|
|
// 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));
|
|
|
|
|
}
|
|
|
|
|
|
2020-12-03 12:34:26 +01:00
|
|
|
MLIRContext *context = &this->getTypeConverter()->getContext();
|
2020-08-01 15:06:25 +02:00
|
|
|
|
2021-01-05 16:22:53 +01:00
|
|
|
Type llvmVoidType = LLVM::LLVMVoidType::get(context);
|
2023-10-31 09:22:44 +01:00
|
|
|
LLVM::LLVMPointerType llvmPointerType = LLVM::LLVMPointerType::get(context);
|
[mlir] replace LLVMIntegerType with built-in integer type
The LLVM dialect type system has been closed until now, i.e. did not support
types from other dialects inside containers. While this has had obvious
benefits of deriving from a common base class, it has led to some simple types
being almost identical with the built-in types, namely integer and floating
point types. This in turn has led to a lot of larger-scale complexity: simple
types must still be converted, numerous operations that correspond to LLVM IR
intrinsics are replicated to produce versions operating on either LLVM dialect
or built-in types leading to quasi-duplicate dialects, lowering to the LLVM
dialect is essentially required to be one-shot because of type conversion, etc.
In this light, it is reasonable to trade off some local complexity in the
internal implementation of LLVM dialect types for removing larger-scale system
complexity. Previous commits to the LLVM dialect type system have adapted the
API to support types from other dialects.
Replace LLVMIntegerType with the built-in IntegerType plus additional checks
that such types are signless (these are isolated in a utility function that
replaced `isa<LLVMType>` and in the parser). Temporarily keep the possibility
to parse `!llvm.i32` as a synonym for `i32`, but add a deprecation notice.
Reviewed By: mehdi_amini, silvas, antiagainst
Differential Revision: https://reviews.llvm.org/D94178
2021-01-06 16:19:04 +01:00
|
|
|
Type llvmInt8Type = IntegerType::get(context, 8);
|
2023-06-05 12:09:35 +05:30
|
|
|
Type llvmInt16Type = IntegerType::get(context, 16);
|
[mlir] replace LLVMIntegerType with built-in integer type
The LLVM dialect type system has been closed until now, i.e. did not support
types from other dialects inside containers. While this has had obvious
benefits of deriving from a common base class, it has led to some simple types
being almost identical with the built-in types, namely integer and floating
point types. This in turn has led to a lot of larger-scale complexity: simple
types must still be converted, numerous operations that correspond to LLVM IR
intrinsics are replicated to produce versions operating on either LLVM dialect
or built-in types leading to quasi-duplicate dialects, lowering to the LLVM
dialect is essentially required to be one-shot because of type conversion, etc.
In this light, it is reasonable to trade off some local complexity in the
internal implementation of LLVM dialect types for removing larger-scale system
complexity. Previous commits to the LLVM dialect type system have adapted the
API to support types from other dialects.
Replace LLVMIntegerType with the built-in IntegerType plus additional checks
that such types are signless (these are isolated in a utility function that
replaced `isa<LLVMType>` and in the parser). Temporarily keep the possibility
to parse `!llvm.i32` as a synonym for `i32`, but add a deprecation notice.
Reviewed By: mehdi_amini, silvas, antiagainst
Differential Revision: https://reviews.llvm.org/D94178
2021-01-06 16:19:04 +01:00
|
|
|
Type llvmInt32Type = IntegerType::get(context, 32);
|
|
|
|
|
Type llvmInt64Type = IntegerType::get(context, 64);
|
2023-08-01 00:57:47 +00:00
|
|
|
Type llvmFloat32Type = Float32Type::get(context);
|
[mlir] replace LLVMIntegerType with built-in integer type
The LLVM dialect type system has been closed until now, i.e. did not support
types from other dialects inside containers. While this has had obvious
benefits of deriving from a common base class, it has led to some simple types
being almost identical with the built-in types, namely integer and floating
point types. This in turn has led to a lot of larger-scale complexity: simple
types must still be converted, numerous operations that correspond to LLVM IR
intrinsics are replicated to produce versions operating on either LLVM dialect
or built-in types leading to quasi-duplicate dialects, lowering to the LLVM
dialect is essentially required to be one-shot because of type conversion, etc.
In this light, it is reasonable to trade off some local complexity in the
internal implementation of LLVM dialect types for removing larger-scale system
complexity. Previous commits to the LLVM dialect type system have adapted the
API to support types from other dialects.
Replace LLVMIntegerType with the built-in IntegerType plus additional checks
that such types are signless (these are isolated in a utility function that
replaced `isa<LLVMType>` and in the parser). Temporarily keep the possibility
to parse `!llvm.i32` as a synonym for `i32`, but add a deprecation notice.
Reviewed By: mehdi_amini, silvas, antiagainst
Differential Revision: https://reviews.llvm.org/D94178
2021-01-06 16:19:04 +01:00
|
|
|
Type llvmIntPtrType = IntegerType::get(
|
2020-12-03 12:34:26 +01:00
|
|
|
context, this->getTypeConverter()->getPointerBitwidth(0));
|
2020-08-01 15:06:25 +02:00
|
|
|
|
|
|
|
|
FunctionCallBuilder moduleLoadCallBuilder = {
|
|
|
|
|
"mgpuModuleLoad",
|
|
|
|
|
llvmPointerType /* void *module */,
|
2023-09-26 11:27:07 -07:00
|
|
|
{llvmPointerType /* void *cubin */, llvmInt64Type /* size_t size */}};
|
2020-10-29 08:17:27 +01:00
|
|
|
FunctionCallBuilder moduleUnloadCallBuilder = {
|
|
|
|
|
"mgpuModuleUnload", llvmVoidType, {llvmPointerType /* void *module */}};
|
2020-08-01 15:06:25 +02:00
|
|
|
FunctionCallBuilder moduleGetFunctionCallBuilder = {
|
|
|
|
|
"mgpuModuleGetFunction",
|
|
|
|
|
llvmPointerType /* void *function */,
|
|
|
|
|
{
|
|
|
|
|
llvmPointerType, /* void *module */
|
|
|
|
|
llvmPointerType /* char *name */
|
|
|
|
|
}};
|
|
|
|
|
FunctionCallBuilder launchKernelCallBuilder = {
|
|
|
|
|
"mgpuLaunchKernel",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{
|
2023-10-31 09:22:44 +01:00
|
|
|
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 */
|
2020-08-01 15:06:25 +02:00
|
|
|
}};
|
|
|
|
|
FunctionCallBuilder streamCreateCallBuilder = {
|
|
|
|
|
"mgpuStreamCreate", llvmPointerType /* void *stream */, {}};
|
2020-10-08 16:37:44 +02:00
|
|
|
FunctionCallBuilder streamDestroyCallBuilder = {
|
|
|
|
|
"mgpuStreamDestroy", llvmVoidType, {llvmPointerType /* void *stream */}};
|
2020-08-01 15:06:25 +02:00
|
|
|
FunctionCallBuilder streamSynchronizeCallBuilder = {
|
|
|
|
|
"mgpuStreamSynchronize",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmPointerType /* void *stream */}};
|
2020-10-08 16:37:44 +02:00
|
|
|
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 */}};
|
2020-08-10 10:13:57 +02:00
|
|
|
FunctionCallBuilder hostRegisterCallBuilder = {
|
|
|
|
|
"mgpuMemHostRegisterMemRef",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmIntPtrType /* intptr_t rank */,
|
|
|
|
|
llvmPointerType /* void *memrefDesc */,
|
|
|
|
|
llvmIntPtrType /* intptr_t elementSizeBytes */}};
|
2023-04-06 15:07:12 -05:00
|
|
|
FunctionCallBuilder hostUnregisterCallBuilder = {
|
|
|
|
|
"mgpuMemHostUnregisterMemRef",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmIntPtrType /* intptr_t rank */,
|
|
|
|
|
llvmPointerType /* void *memrefDesc */,
|
|
|
|
|
llvmIntPtrType /* intptr_t elementSizeBytes */}};
|
2020-11-24 22:07:34 +01:00
|
|
|
FunctionCallBuilder allocCallBuilder = {
|
|
|
|
|
"mgpuMemAlloc",
|
|
|
|
|
llvmPointerType /* void * */,
|
|
|
|
|
{llvmIntPtrType /* intptr_t sizeBytes */,
|
2023-09-26 15:32:11 -07:00
|
|
|
llvmPointerType /* void *stream */,
|
|
|
|
|
llvmInt8Type /* bool isHostShared */}};
|
2020-11-24 22:07:34 +01:00
|
|
|
FunctionCallBuilder deallocCallBuilder = {
|
|
|
|
|
"mgpuMemFree",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmPointerType /* void *ptr */, llvmPointerType /* void *stream */}};
|
2020-12-22 17:42:59 +01:00
|
|
|
FunctionCallBuilder memcpyCallBuilder = {
|
|
|
|
|
"mgpuMemcpy",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmPointerType /* void *dst */, llvmPointerType /* void *src */,
|
|
|
|
|
llvmIntPtrType /* intptr_t sizeBytes */,
|
|
|
|
|
llvmPointerType /* void *stream */}};
|
2023-06-05 12:09:35 +05:30
|
|
|
FunctionCallBuilder memset16CallBuilder = {
|
|
|
|
|
"mgpuMemset16",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmPointerType /* void *dst */,
|
|
|
|
|
llvmInt16Type /* unsigned short value */,
|
|
|
|
|
llvmIntPtrType /* intptr_t sizeBytes */,
|
|
|
|
|
llvmPointerType /* void *stream */}};
|
|
|
|
|
FunctionCallBuilder memset32CallBuilder = {
|
2021-09-04 08:03:33 +02:00
|
|
|
"mgpuMemset32",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmPointerType /* void *dst */, llvmInt32Type /* unsigned int value */,
|
|
|
|
|
llvmIntPtrType /* intptr_t sizeBytes */,
|
|
|
|
|
llvmPointerType /* void *stream */}};
|
2022-02-15 20:23:44 +00:00
|
|
|
FunctionCallBuilder setDefaultDeviceCallBuilder = {
|
|
|
|
|
"mgpuSetDefaultDevice",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmInt32Type /* uint32_t devIndex */}};
|
2023-05-11 13:56:34 -07:00
|
|
|
FunctionCallBuilder createDnVecCallBuilder = {
|
|
|
|
|
"mgpuCreateDnVec",
|
|
|
|
|
llvmPointerType,
|
|
|
|
|
{llvmIntPtrType, llvmPointerType, llvmInt32Type,
|
|
|
|
|
llvmPointerType /* void *stream */}};
|
|
|
|
|
FunctionCallBuilder destroyDnVecCallBuilder = {
|
|
|
|
|
"mgpuDestroyDnVec",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmPointerType, llvmPointerType /* void *stream */}};
|
2023-05-15 14:55:45 -07:00
|
|
|
FunctionCallBuilder createDnMatCallBuilder = {
|
|
|
|
|
"mgpuCreateDnMat",
|
|
|
|
|
llvmPointerType,
|
|
|
|
|
{llvmIntPtrType, llvmIntPtrType, llvmPointerType, llvmInt32Type,
|
|
|
|
|
llvmPointerType /* void *stream */}};
|
|
|
|
|
FunctionCallBuilder destroyDnMatCallBuilder = {
|
|
|
|
|
"mgpuDestroyDnMat",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmPointerType, llvmPointerType /* void *stream */}};
|
2023-05-11 13:56:34 -07:00
|
|
|
FunctionCallBuilder createCooCallBuilder = {
|
|
|
|
|
"mgpuCreateCoo",
|
|
|
|
|
llvmPointerType,
|
|
|
|
|
{llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, llvmPointerType,
|
|
|
|
|
llvmPointerType, llvmPointerType, llvmInt32Type, llvmInt32Type,
|
|
|
|
|
llvmPointerType /* void *stream */}};
|
2023-06-06 10:56:58 -07:00
|
|
|
FunctionCallBuilder createCooAoSCallBuilder = {
|
|
|
|
|
"mgpuCreateCooAoS", // deprecated in cuSPARSE 11.2
|
|
|
|
|
llvmPointerType,
|
|
|
|
|
{llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, llvmPointerType,
|
|
|
|
|
llvmPointerType, llvmInt32Type, llvmInt32Type,
|
|
|
|
|
llvmPointerType /* void *stream */}};
|
2023-05-11 13:56:34 -07:00
|
|
|
FunctionCallBuilder createCsrCallBuilder = {
|
|
|
|
|
"mgpuCreateCsr",
|
|
|
|
|
llvmPointerType,
|
|
|
|
|
{llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, llvmPointerType,
|
|
|
|
|
llvmPointerType, llvmPointerType, llvmInt32Type, llvmInt32Type,
|
|
|
|
|
llvmInt32Type, llvmPointerType /* void *stream */}};
|
2023-09-27 09:32:25 -07:00
|
|
|
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 */}};
|
2023-05-11 13:56:34 -07:00
|
|
|
FunctionCallBuilder destroySpMatCallBuilder = {
|
|
|
|
|
"mgpuDestroySpMat",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmPointerType, llvmPointerType /* void *stream */}};
|
|
|
|
|
FunctionCallBuilder spMVBufferSizeCallBuilder = {
|
|
|
|
|
"mgpuSpMVBufferSize",
|
|
|
|
|
llvmIntPtrType,
|
2023-06-16 21:38:40 +00:00
|
|
|
{llvmInt32Type, llvmPointerType, llvmPointerType, llvmPointerType,
|
|
|
|
|
llvmInt32Type, llvmPointerType /* void *stream */}};
|
2023-05-11 13:56:34 -07:00
|
|
|
FunctionCallBuilder spMVCallBuilder = {
|
|
|
|
|
"mgpuSpMV",
|
|
|
|
|
llvmVoidType,
|
2023-06-16 21:38:40 +00:00
|
|
|
{llvmInt32Type, llvmPointerType, llvmPointerType, llvmPointerType,
|
|
|
|
|
llvmInt32Type, llvmPointerType, llvmPointerType /* void *stream */}};
|
2023-06-12 23:24:08 +00:00
|
|
|
FunctionCallBuilder createSpMMBufferSizeCallBuilder = {
|
2023-05-15 14:55:45 -07:00
|
|
|
"mgpuSpMMBufferSize",
|
|
|
|
|
llvmIntPtrType,
|
2023-06-16 21:38:40 +00:00
|
|
|
{llvmInt32Type, llvmInt32Type, llvmPointerType, llvmPointerType,
|
|
|
|
|
llvmPointerType, llvmInt32Type, llvmPointerType /* void *stream */}};
|
2023-06-12 23:24:08 +00:00
|
|
|
FunctionCallBuilder createSpMMCallBuilder = {
|
2023-05-15 14:55:45 -07:00
|
|
|
"mgpuSpMM",
|
|
|
|
|
llvmVoidType,
|
2023-06-16 21:38:40 +00:00
|
|
|
{llvmInt32Type, llvmInt32Type, llvmPointerType, llvmPointerType,
|
|
|
|
|
llvmPointerType, llvmInt32Type, llvmPointerType,
|
2023-05-26 02:29:49 +00:00
|
|
|
llvmPointerType /* void *stream */}};
|
2023-06-12 23:24:08 +00:00
|
|
|
FunctionCallBuilder createSDDMMBufferSizeCallBuilder = {
|
2023-05-24 02:21:55 +00:00
|
|
|
"mgpuSDDMMBufferSize",
|
|
|
|
|
llvmIntPtrType,
|
2023-06-16 21:38:40 +00:00
|
|
|
{llvmInt32Type, llvmInt32Type, llvmPointerType, llvmPointerType,
|
|
|
|
|
llvmPointerType, llvmInt32Type, llvmPointerType /* void *stream */}};
|
2023-06-12 23:24:08 +00:00
|
|
|
FunctionCallBuilder createSDDMMCallBuilder = {
|
2023-05-24 02:21:55 +00:00
|
|
|
"mgpuSDDMM",
|
|
|
|
|
llvmVoidType,
|
2023-06-16 21:38:40 +00:00
|
|
|
{llvmInt32Type, llvmInt32Type, llvmPointerType, llvmPointerType,
|
|
|
|
|
llvmPointerType, llvmInt32Type, llvmPointerType,
|
2023-05-24 02:21:55 +00:00
|
|
|
llvmPointerType /* void *stream */}};
|
2023-06-06 23:12:45 +00:00
|
|
|
FunctionCallBuilder createLtDnMatCallBuilder = {
|
|
|
|
|
"mgpuCreateCuSparseLtDnMat",
|
|
|
|
|
llvmVoidType,
|
2023-06-16 21:38:40 +00:00
|
|
|
{llvmPointerType, llvmIntPtrType, llvmIntPtrType, llvmPointerType,
|
|
|
|
|
llvmInt32Type, llvmPointerType /* void *stream */}};
|
2023-06-06 23:12:45 +00:00
|
|
|
FunctionCallBuilder destroyCuSparseLtSpMatBuilder = {
|
|
|
|
|
"mgpuDestroyCuSparseLtSpMat",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmPointerType, llvmPointerType /* void *stream */}};
|
|
|
|
|
FunctionCallBuilder destroyCuSparseLtDnMatBuilder = {
|
|
|
|
|
"mgpuDestroyCuSparseLtDnMat",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmPointerType, llvmPointerType /* void *stream */}};
|
|
|
|
|
FunctionCallBuilder create2To4SpMatCallBuilder = {
|
|
|
|
|
"mgpuCusparseLtCreate2To4SpMat",
|
|
|
|
|
llvmVoidType,
|
2023-06-16 21:38:40 +00:00
|
|
|
{llvmPointerType, llvmIntPtrType, llvmIntPtrType, llvmPointerType,
|
|
|
|
|
llvmInt32Type, llvmPointerType /* void *stream */}};
|
2023-06-12 23:24:08 +00:00
|
|
|
FunctionCallBuilder createCuSparseLtSpMMBufferSizeBuilder = {
|
2023-06-06 23:12:45 +00:00
|
|
|
"mgpuCuSparseLtSpMMBufferSize",
|
|
|
|
|
llvmVoidType,
|
2023-06-16 21:38:40 +00:00
|
|
|
{llvmPointerType, llvmInt32Type, llvmInt32Type, llvmPointerType,
|
2023-07-21 02:03:51 +00:00
|
|
|
llvmPointerType, llvmPointerType, llvmInt32Type, llvmInt32Type,
|
2023-06-06 23:12:45 +00:00
|
|
|
llvmPointerType /*void *stream*/}};
|
2023-06-12 23:24:08 +00:00
|
|
|
FunctionCallBuilder createCuSparseLtSpMMBuilder = {
|
2023-06-06 23:12:45 +00:00
|
|
|
"mgpuCuSparseLtSpMM",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmPointerType, llvmPointerType, llvmPointerType, llvmPointerType,
|
2023-06-16 21:38:40 +00:00
|
|
|
llvmPointerType, llvmPointerType, llvmPointerType /*void *stream*/}};
|
2023-08-14 13:43:29 -07:00
|
|
|
FunctionCallBuilder createSpGEMMCreateDescrBuilder = {
|
|
|
|
|
"mgpuSpGEMMCreateDescr",
|
|
|
|
|
llvmPointerType,
|
|
|
|
|
{llvmPointerType /*void *stream*/}};
|
|
|
|
|
FunctionCallBuilder createSpGEMMDestroyDescrBuilder = {
|
|
|
|
|
"mgpuSpGEMMDestroyDescr",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmPointerType /*s*/, llvmPointerType /*void *stream*/}};
|
2023-08-01 00:57:47 +00:00
|
|
|
FunctionCallBuilder createSpGEMMWorkEstimationBuilder = {
|
|
|
|
|
"mgpuSpGEMMWorkEstimation",
|
|
|
|
|
llvmIntPtrType,
|
|
|
|
|
{llvmPointerType /*s*/, llvmInt32Type /*ma*/, llvmInt32Type /*mb*/,
|
|
|
|
|
llvmPointerType /*a*/, llvmPointerType /*b*/, llvmPointerType /*c*/,
|
2023-08-09 12:19:49 -07:00
|
|
|
llvmInt32Type /*ctp*/, llvmIntPtrType /*bs*/, llvmPointerType /*buf*/,
|
|
|
|
|
llvmPointerType /*void *stream*/}};
|
2023-08-01 00:57:47 +00:00
|
|
|
FunctionCallBuilder createSpGEMMComputeBuilder = {
|
|
|
|
|
"mgpuSpGEMMCompute",
|
|
|
|
|
llvmIntPtrType,
|
|
|
|
|
{llvmPointerType /*s*/, llvmInt32Type /*ma*/, llvmInt32Type /*mb*/,
|
|
|
|
|
llvmPointerType /*a*/, llvmPointerType /*b*/, llvmPointerType /*c*/,
|
2023-08-09 12:19:49 -07:00
|
|
|
llvmInt32Type /*ctp*/, llvmIntPtrType /*bs*/, llvmPointerType /*buf*/,
|
|
|
|
|
llvmPointerType /*void *stream*/}};
|
2023-08-01 00:57:47 +00:00
|
|
|
FunctionCallBuilder createSpGEMMCopyBuilder = {
|
|
|
|
|
"mgpuSpGEMMCopy",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmPointerType /*s*/, llvmInt32Type /*ma*/, llvmInt32Type /*mb*/,
|
|
|
|
|
llvmPointerType /*a*/, llvmPointerType /*b*/, llvmPointerType /*c*/,
|
2023-08-09 12:19:49 -07:00
|
|
|
llvmInt32Type /*ctp*/, llvmPointerType /*void *stream*/}};
|
2023-08-14 13:43:29 -07:00
|
|
|
FunctionCallBuilder createSpMatGetSizeBuilder = {
|
|
|
|
|
"mgpuSpMatGetSize",
|
2023-08-01 00:57:47 +00:00
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmPointerType /*mc*/, llvmPointerType /*rc*/, llvmPointerType /*cc*/,
|
|
|
|
|
llvmPointerType /*nc*/, llvmPointerType /*void *stream*/}};
|
2023-08-10 12:36:15 -07:00
|
|
|
FunctionCallBuilder createSetCsrPointersBuilder = {
|
|
|
|
|
"mgpuSetCsrPointers",
|
|
|
|
|
llvmVoidType,
|
|
|
|
|
{llvmPointerType /*spmat*/, llvmPointerType /*pos*/,
|
|
|
|
|
llvmPointerType /*crd*/, llvmPointerType /*val*/,
|
|
|
|
|
llvmPointerType /*void *stream*/}};
|
2020-08-10 10:13:57 +02:00
|
|
|
};
|
|
|
|
|
|
2020-08-19 15:50:44 +02:00
|
|
|
/// A rewrite pattern to convert gpu.host_register operations into a GPU runtime
|
2020-08-10 10:13:57 +02:00
|
|
|
/// call. Currently it supports CUDA and ROCm (HIP).
|
|
|
|
|
class ConvertHostRegisterOpToGpuRuntimeCallPattern
|
|
|
|
|
: public ConvertOpToGpuRuntimeCallPattern<gpu::HostRegisterOp> {
|
|
|
|
|
public:
|
2023-08-14 08:40:13 +02:00
|
|
|
ConvertHostRegisterOpToGpuRuntimeCallPattern(
|
|
|
|
|
const LLVMTypeConverter &typeConverter)
|
2020-08-10 10:13:57 +02:00
|
|
|
: ConvertOpToGpuRuntimeCallPattern<gpu::HostRegisterOp>(typeConverter) {}
|
|
|
|
|
|
|
|
|
|
private:
|
|
|
|
|
LogicalResult
|
2021-09-24 17:51:20 +00:00
|
|
|
matchAndRewrite(gpu::HostRegisterOp hostRegisterOp, OpAdaptor adaptor,
|
2020-08-10 10:13:57 +02:00
|
|
|
ConversionPatternRewriter &rewriter) const override;
|
2020-08-01 15:06:25 +02:00
|
|
|
};
|
|
|
|
|
|
2023-04-06 15:07:12 -05:00
|
|
|
class ConvertHostUnregisterOpToGpuRuntimeCallPattern
|
|
|
|
|
: public ConvertOpToGpuRuntimeCallPattern<gpu::HostUnregisterOp> {
|
|
|
|
|
public:
|
|
|
|
|
ConvertHostUnregisterOpToGpuRuntimeCallPattern(
|
2023-08-14 08:40:13 +02:00
|
|
|
const LLVMTypeConverter &typeConverter)
|
2023-04-06 15:07:12 -05:00
|
|
|
: ConvertOpToGpuRuntimeCallPattern<gpu::HostUnregisterOp>(typeConverter) {
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
private:
|
|
|
|
|
LogicalResult
|
|
|
|
|
matchAndRewrite(gpu::HostUnregisterOp hostUnregisterOp, OpAdaptor adaptor,
|
|
|
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
|
|
|
};
|
|
|
|
|
|
2020-11-24 22:07:34 +01:00
|
|
|
/// 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:
|
2023-08-14 08:40:13 +02:00
|
|
|
ConvertAllocOpToGpuRuntimeCallPattern(const LLVMTypeConverter &typeConverter)
|
2020-11-24 22:07:34 +01:00
|
|
|
: ConvertOpToGpuRuntimeCallPattern<gpu::AllocOp>(typeConverter) {}
|
|
|
|
|
|
|
|
|
|
private:
|
|
|
|
|
LogicalResult
|
2021-09-24 17:51:20 +00:00
|
|
|
matchAndRewrite(gpu::AllocOp allocOp, OpAdaptor adaptor,
|
2020-11-24 22:07:34 +01:00
|
|
|
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:
|
2023-08-14 08:40:13 +02:00
|
|
|
ConvertDeallocOpToGpuRuntimeCallPattern(
|
|
|
|
|
const LLVMTypeConverter &typeConverter)
|
2020-11-24 22:07:34 +01:00
|
|
|
: ConvertOpToGpuRuntimeCallPattern<gpu::DeallocOp>(typeConverter) {}
|
|
|
|
|
|
|
|
|
|
private:
|
|
|
|
|
LogicalResult
|
2021-09-24 17:51:20 +00:00
|
|
|
matchAndRewrite(gpu::DeallocOp deallocOp, OpAdaptor adaptor,
|
2020-11-24 22:07:34 +01:00
|
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
|
|
|
};
|
|
|
|
|
|
2021-02-24 15:02:30 +01:00
|
|
|
class ConvertAsyncYieldToGpuRuntimeCallPattern
|
|
|
|
|
: public ConvertOpToGpuRuntimeCallPattern<async::YieldOp> {
|
|
|
|
|
public:
|
2023-08-14 08:40:13 +02:00
|
|
|
ConvertAsyncYieldToGpuRuntimeCallPattern(
|
|
|
|
|
const LLVMTypeConverter &typeConverter)
|
2021-02-24 15:02:30 +01:00
|
|
|
: ConvertOpToGpuRuntimeCallPattern<async::YieldOp>(typeConverter) {}
|
|
|
|
|
|
|
|
|
|
private:
|
|
|
|
|
LogicalResult
|
2021-09-24 17:51:20 +00:00
|
|
|
matchAndRewrite(async::YieldOp yieldOp, OpAdaptor adaptor,
|
2021-02-24 15:02:30 +01:00
|
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
|
|
|
};
|
|
|
|
|
|
2020-10-21 08:24:53 +02:00
|
|
|
/// 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:
|
2023-08-14 08:40:13 +02:00
|
|
|
ConvertWaitOpToGpuRuntimeCallPattern(const LLVMTypeConverter &typeConverter)
|
2020-10-21 08:24:53 +02:00
|
|
|
: ConvertOpToGpuRuntimeCallPattern<gpu::WaitOp>(typeConverter) {}
|
|
|
|
|
|
|
|
|
|
private:
|
|
|
|
|
LogicalResult
|
2021-09-24 17:51:20 +00:00
|
|
|
matchAndRewrite(gpu::WaitOp waitOp, OpAdaptor adaptor,
|
2020-10-21 08:24:53 +02:00
|
|
|
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:
|
2023-08-14 08:40:13 +02:00
|
|
|
ConvertWaitAsyncOpToGpuRuntimeCallPattern(
|
|
|
|
|
const LLVMTypeConverter &typeConverter)
|
2020-10-21 08:24:53 +02:00
|
|
|
: ConvertOpToGpuRuntimeCallPattern<gpu::WaitOp>(typeConverter) {}
|
|
|
|
|
|
|
|
|
|
private:
|
|
|
|
|
LogicalResult
|
2021-09-24 17:51:20 +00:00
|
|
|
matchAndRewrite(gpu::WaitOp waitOp, OpAdaptor adaptor,
|
2020-10-21 08:24:53 +02:00
|
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
|
|
|
};
|
|
|
|
|
|
2020-08-01 15:06:25 +02:00
|
|
|
/// A rewrite patter to convert gpu.launch_func operations into a sequence of
|
|
|
|
|
/// GPU runtime calls. Currently it supports CUDA and ROCm (HIP).
|
2019-06-19 06:22:36 -07:00
|
|
|
///
|
|
|
|
|
/// In essence, a gpu.launch_func operations gets compiled into the following
|
|
|
|
|
/// sequence of runtime calls:
|
|
|
|
|
///
|
2020-05-18 13:01:54 -05:00
|
|
|
/// * 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
|
2019-06-19 06:22:36 -07:00
|
|
|
///
|
|
|
|
|
/// Intermediate data structures are allocated on the stack.
|
2020-08-01 15:06:25 +02:00
|
|
|
class ConvertLaunchFuncOpToGpuRuntimeCallPattern
|
|
|
|
|
: public ConvertOpToGpuRuntimeCallPattern<gpu::LaunchFuncOp> {
|
|
|
|
|
public:
|
2023-08-14 08:40:13 +02:00
|
|
|
ConvertLaunchFuncOpToGpuRuntimeCallPattern(
|
|
|
|
|
const LLVMTypeConverter &typeConverter, StringRef gpuBinaryAnnotation,
|
|
|
|
|
bool kernelBarePtrCallConv, SymbolTable *cachedModuleTable)
|
2020-08-01 15:06:25 +02:00
|
|
|
: ConvertOpToGpuRuntimeCallPattern<gpu::LaunchFuncOp>(typeConverter),
|
2022-07-11 18:29:01 +00:00
|
|
|
gpuBinaryAnnotation(gpuBinaryAnnotation),
|
2023-08-12 00:26:13 +00:00
|
|
|
kernelBarePtrCallConv(kernelBarePtrCallConv),
|
|
|
|
|
cachedModuleTable(cachedModuleTable) {}
|
2019-06-19 06:22:36 -07:00
|
|
|
|
2020-08-01 15:06:25 +02:00
|
|
|
private:
|
2021-09-24 17:51:20 +00:00
|
|
|
Value generateParamsArray(gpu::LaunchFuncOp launchOp, OpAdaptor adaptor,
|
|
|
|
|
OpBuilder &builder) const;
|
2020-04-16 13:14:43 +02:00
|
|
|
Value generateKernelNameConstant(StringRef moduleName, StringRef name,
|
2020-08-01 15:06:25 +02:00
|
|
|
Location loc, OpBuilder &builder) const;
|
2019-06-19 06:22:36 -07:00
|
|
|
|
2020-08-01 15:06:25 +02:00
|
|
|
LogicalResult
|
2021-09-24 17:51:20 +00:00
|
|
|
matchAndRewrite(gpu::LaunchFuncOp launchOp, OpAdaptor adaptor,
|
2020-08-01 15:06:25 +02:00
|
|
|
ConversionPatternRewriter &rewriter) const override;
|
[mlir][gpu] Introduce mlir-rocm-runner.
Summary:
`mlir-rocm-runner` is introduced in this commit to execute GPU modules on ROCm
platform. A small wrapper to encapsulate ROCm's HIP runtime API is also inside
the commit.
Due to behavior of ROCm, raw pointers inside memrefs passed to `gpu.launch`
must be modified on the host side to properly capture the pointer values
addressable on the GPU.
LLVM MC is used to assemble AMD GCN ISA coming out from
`ConvertGPUKernelToBlobPass` to binary form, and LLD is used to produce a shared
ELF object which could be loaded by ROCm HIP runtime.
gfx900 is the default target be used right now, although it could be altered via
an option in `mlir-rocm-runner`. Future revisions may consider using ROCm Agent
Enumerator to detect the right target on the system.
Notice AMDGPU Code Object V2 is used in this revision. Future enhancements may
upgrade to AMDGPU Code Object V3.
Bitcode libraries in ROCm-Device-Libs, which implements math routines exposed in
`rocdl` dialect are not yet linked, and is left as a TODO in the logic.
Reviewers: herhut
Subscribers: mgorny, tpr, dexonsmith, mehdi_amini, rriddle, jpienaar, shauheen, antiagainst, nicolasvasilache, csigg, arpith-jacob, mgester, lucyrfox, aartbik, liufengdb, stephenneuendorffer, Joonsoo, grosul1, frgossen, Kayjukh, jurahul, llvm-commits
Tags: #mlir, #llvm
Differential Revision: https://reviews.llvm.org/D80676
2020-05-20 16:07:49 -05:00
|
|
|
|
2020-08-01 15:06:25 +02:00
|
|
|
llvm::SmallString<32> gpuBinaryAnnotation;
|
2022-07-11 18:29:01 +00:00
|
|
|
bool kernelBarePtrCallConv;
|
2023-08-12 00:26:13 +00:00
|
|
|
SymbolTable *cachedModuleTable;
|
2020-08-01 15:06:25 +02:00
|
|
|
};
|
2019-06-19 06:22:36 -07:00
|
|
|
|
2020-08-01 15:06:25 +02:00
|
|
|
class EraseGpuModuleOpPattern : public OpRewritePattern<gpu::GPUModuleOp> {
|
|
|
|
|
using OpRewritePattern<gpu::GPUModuleOp>::OpRewritePattern;
|
2019-10-08 04:29:58 -07:00
|
|
|
|
2020-08-01 15:06:25 +02:00
|
|
|
LogicalResult matchAndRewrite(gpu::GPUModuleOp op,
|
|
|
|
|
PatternRewriter &rewriter) const override {
|
2019-10-08 04:29:58 -07:00
|
|
|
// GPU kernel modules are no longer necessary since we have a global
|
2020-05-18 13:01:54 -05:00
|
|
|
// constant with the CUBIN, or HSACO data.
|
2020-08-01 15:06:25 +02:00
|
|
|
rewriter.eraseOp(op);
|
|
|
|
|
return success();
|
2019-06-19 06:22:36 -07:00
|
|
|
}
|
|
|
|
|
};
|
2020-12-22 17:42:59 +01:00
|
|
|
|
|
|
|
|
/// 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:
|
2023-08-14 08:40:13 +02:00
|
|
|
ConvertMemcpyOpToGpuRuntimeCallPattern(const LLVMTypeConverter &typeConverter)
|
2020-12-22 17:42:59 +01:00
|
|
|
: ConvertOpToGpuRuntimeCallPattern<gpu::MemcpyOp>(typeConverter) {}
|
|
|
|
|
|
|
|
|
|
private:
|
|
|
|
|
LogicalResult
|
2021-09-24 17:51:20 +00:00
|
|
|
matchAndRewrite(gpu::MemcpyOp memcpyOp, OpAdaptor adaptor,
|
2020-12-22 17:42:59 +01:00
|
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
|
|
|
};
|
2021-09-04 08:03:33 +02:00
|
|
|
|
|
|
|
|
/// 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:
|
2023-08-14 08:40:13 +02:00
|
|
|
ConvertMemsetOpToGpuRuntimeCallPattern(const LLVMTypeConverter &typeConverter)
|
2021-09-04 08:03:33 +02:00
|
|
|
: ConvertOpToGpuRuntimeCallPattern<gpu::MemsetOp>(typeConverter) {}
|
|
|
|
|
|
|
|
|
|
private:
|
|
|
|
|
LogicalResult
|
2021-09-24 17:51:20 +00:00
|
|
|
matchAndRewrite(gpu::MemsetOp memsetOp, OpAdaptor adaptor,
|
2021-09-04 08:03:33 +02:00
|
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
|
|
|
};
|
2022-02-15 20:23:44 +00:00
|
|
|
|
|
|
|
|
/// 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(
|
2023-08-14 08:40:13 +02:00
|
|
|
const LLVMTypeConverter &typeConverter)
|
2022-02-15 20:23:44 +00:00
|
|
|
: ConvertOpToGpuRuntimeCallPattern<gpu::SetDefaultDeviceOp>(
|
|
|
|
|
typeConverter) {}
|
|
|
|
|
|
|
|
|
|
LogicalResult
|
|
|
|
|
matchAndRewrite(gpu::SetDefaultDeviceOp op, OpAdaptor adaptor,
|
|
|
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
|
|
|
};
|
2023-05-11 13:56:34 -07:00
|
|
|
|
2023-08-08 10:24:03 -07:00
|
|
|
/// Generic rewriting rule for operation on sparse matrices.
|
|
|
|
|
/// Currently supports CUDA (by means of cuSparse and cuSparseLt).
|
2023-08-01 00:57:47 +00:00
|
|
|
#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( \
|
2023-08-14 08:40:13 +02:00
|
|
|
const LLVMTypeConverter &typeConverter) \
|
2023-08-01 00:57:47 +00:00
|
|
|
: ConvertOpToGpuRuntimeCallPattern<gpu::op_name>(typeConverter) {} \
|
|
|
|
|
\
|
|
|
|
|
private: \
|
|
|
|
|
LogicalResult \
|
|
|
|
|
matchAndRewrite(gpu::op_name op, OpAdaptor adaptor, \
|
|
|
|
|
ConversionPatternRewriter &rewriter) const override; \
|
|
|
|
|
};
|
|
|
|
|
|
2023-08-08 10:24:03 -07:00
|
|
|
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)
|
2023-09-27 09:32:25 -07:00
|
|
|
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(CreateCscOp)
|
|
|
|
|
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(CreateBsrOp)
|
2023-08-08 10:24:03 -07:00
|
|
|
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)
|
2023-08-01 00:57:47 +00:00
|
|
|
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)
|
2023-08-14 13:43:29 -07:00
|
|
|
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(SpMatGetSizeOp)
|
2023-08-10 12:36:15 -07:00
|
|
|
DECLARE_CONVERT_OP_TO_GPU_RUNTIME_CALL_PATTERN(SetCsrPointersOp)
|
2023-08-01 00:57:47 +00:00
|
|
|
|
2020-08-01 15:06:25 +02:00
|
|
|
} // namespace
|
|
|
|
|
|
2022-08-30 22:20:36 +02:00
|
|
|
void GpuToLLVMConversionPass::runOnOperation() {
|
2023-02-21 07:51:44 +01:00
|
|
|
LowerToLLVMOptions options(&getContext());
|
2023-06-17 22:43:43 +05:30
|
|
|
options.useBarePtrCallConv = hostBarePtrCallConv;
|
2023-02-21 07:51:44 +01:00
|
|
|
|
|
|
|
|
LLVMTypeConverter converter(&getContext(), options);
|
2021-03-22 16:58:34 -07:00
|
|
|
RewritePatternSet patterns(&getContext());
|
2021-02-03 20:10:51 +01:00
|
|
|
LLVMConversionTarget target(getContext());
|
|
|
|
|
|
2023-08-12 00:26:13 +00:00
|
|
|
SymbolTable symbolTable = SymbolTable(getOperation());
|
|
|
|
|
// 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() &&
|
2023-08-23 17:51:11 +02:00
|
|
|
!module.getTargetsAttr().empty());
|
2023-08-12 00:26:13 +00:00
|
|
|
});
|
2021-06-14 09:14:30 +02:00
|
|
|
|
2022-09-29 11:14:47 -04:00
|
|
|
mlir::arith::populateArithToLLVMConversionPatterns(converter, patterns);
|
2022-02-06 15:10:03 -08:00
|
|
|
mlir::cf::populateControlFlowToLLVMConversionPatterns(converter, patterns);
|
2021-03-18 12:59:49 -07:00
|
|
|
populateVectorToLLVMConversionPatterns(converter, patterns);
|
2023-01-24 13:05:47 +00:00
|
|
|
populateFinalizeMemRefToLLVMConversionPatterns(converter, patterns);
|
2022-03-01 14:53:41 -08:00
|
|
|
populateFuncToLLVMConversionPatterns(converter, patterns);
|
2021-03-20 16:29:41 -07:00
|
|
|
populateAsyncStructuralTypeConversionsAndLegality(converter, patterns,
|
|
|
|
|
target);
|
2022-07-11 18:29:01 +00:00
|
|
|
populateGpuToLLVMConversionPatterns(converter, patterns, gpuBinaryAnnotation,
|
2023-08-12 00:26:13 +00:00
|
|
|
kernelBarePtrCallConv, &symbolTable);
|
2020-08-01 15:06:25 +02:00
|
|
|
|
2020-10-26 17:25:01 -07:00
|
|
|
if (failed(
|
|
|
|
|
applyPartialConversion(getOperation(), target, std::move(patterns))))
|
2020-08-01 15:06:25 +02:00
|
|
|
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;
|
2021-03-11 23:58:02 +00:00
|
|
|
return OpBuilder::atBlockEnd(module.getBody())
|
2020-08-01 15:06:25 +02:00
|
|
|
.create<LLVM::LLVMFuncOp>(loc, functionName, functionType);
|
|
|
|
|
}();
|
2021-07-28 10:23:06 +02:00
|
|
|
return builder.create<LLVM::CallOp>(loc, function, arguments);
|
2019-06-19 06:22:36 -07:00
|
|
|
}
|
|
|
|
|
|
2023-05-26 21:46:22 +00:00
|
|
|
// Corresponding to cusparseIndexType_t defined in cusparse.h.
|
|
|
|
|
static int32_t getCuSparseIndexTypeFrom(Type type) {
|
2023-09-27 09:32:25 -07:00
|
|
|
if (type.isInteger(16))
|
|
|
|
|
return 1; // CUSPARSE_INDEX_16U
|
|
|
|
|
if (type.isInteger(32))
|
2023-05-26 21:46:22 +00:00
|
|
|
return 2; // CUSPARSE_INDEX_32I
|
2023-09-27 09:32:25 -07:00
|
|
|
return 3; // CUSPARSE_INDEX_64I
|
2023-05-26 21:46:22 +00:00
|
|
|
}
|
|
|
|
|
|
2023-06-06 23:12:45 +00:00
|
|
|
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
|
|
|
|
|
}
|
|
|
|
|
|
2023-05-26 21:46:22 +00:00
|
|
|
// 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");
|
|
|
|
|
}
|
|
|
|
|
|
2023-07-21 02:03:51 +00:00
|
|
|
static gpu::Prune2To4SpMatFlag get2To4PruneFlag(Value spMat) {
|
|
|
|
|
return spMat.getDefiningOp<gpu::Create2To4SpMatOp>().getPruneFlag();
|
|
|
|
|
}
|
2023-09-27 09:32:25 -07:00
|
|
|
|
2023-06-06 23:12:45 +00:00
|
|
|
// 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;
|
2023-09-27 09:32:25 -07:00
|
|
|
if (auto op = spMat.getDefiningOp<gpu::CreateCooAoSOp>())
|
|
|
|
|
return false;
|
2023-06-06 23:12:45 +00:00
|
|
|
if (auto op = spMat.getDefiningOp<gpu::CreateCsrOp>())
|
|
|
|
|
return false;
|
2023-09-27 09:32:25 -07:00
|
|
|
if (auto op = spMat.getDefiningOp<gpu::CreateCscOp>())
|
|
|
|
|
return false;
|
|
|
|
|
if (auto op = spMat.getDefiningOp<gpu::CreateBsrOp>())
|
2023-06-06 23:12:45 +00:00
|
|
|
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;
|
|
|
|
|
}
|
|
|
|
|
|
2020-11-24 22:07:34 +01:00
|
|
|
// Returns whether all operands are of LLVM type.
|
|
|
|
|
static LogicalResult areAllLLVMTypes(Operation *op, ValueRange operands,
|
|
|
|
|
ConversionPatternRewriter &rewriter) {
|
|
|
|
|
if (!llvm::all_of(operands, [](Value value) {
|
2021-01-05 16:22:53 +01:00
|
|
|
return LLVM::isCompatibleType(value.getType());
|
2020-11-24 22:07:34 +01:00
|
|
|
}))
|
|
|
|
|
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();
|
2020-08-10 10:13:57 +02:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
LogicalResult ConvertHostRegisterOpToGpuRuntimeCallPattern::matchAndRewrite(
|
2021-09-24 17:51:20 +00:00
|
|
|
gpu::HostRegisterOp hostRegisterOp, OpAdaptor adaptor,
|
2020-08-10 10:13:57 +02:00
|
|
|
ConversionPatternRewriter &rewriter) const {
|
2020-11-27 21:09:13 +01:00
|
|
|
auto *op = hostRegisterOp.getOperation();
|
2021-09-24 17:51:20 +00:00
|
|
|
if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)))
|
2020-11-24 22:07:34 +01:00
|
|
|
return failure();
|
2020-08-10 10:13:57 +02:00
|
|
|
|
|
|
|
|
Location loc = op->getLoc();
|
|
|
|
|
|
2022-09-30 12:30:41 -07:00
|
|
|
auto memRefType = hostRegisterOp.getValue().getType();
|
2023-05-08 16:33:54 +02:00
|
|
|
auto elementType = cast<UnrankedMemRefType>(memRefType).getElementType();
|
2020-08-10 10:13:57 +02:00
|
|
|
auto elementSize = getSizeInBytes(loc, elementType, rewriter);
|
|
|
|
|
|
2021-09-24 17:51:20 +00:00
|
|
|
auto arguments = getTypeConverter()->promoteOperands(
|
|
|
|
|
loc, op->getOperands(), adaptor.getOperands(), rewriter);
|
2020-08-10 10:13:57 +02:00
|
|
|
arguments.push_back(elementSize);
|
|
|
|
|
hostRegisterCallBuilder.create(loc, rewriter, arguments);
|
|
|
|
|
|
|
|
|
|
rewriter.eraseOp(op);
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2023-04-06 15:07:12 -05:00
|
|
|
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();
|
2023-05-08 16:33:54 +02:00
|
|
|
auto elementType = cast<UnrankedMemRefType>(memRefType).getElementType();
|
2023-04-06 15:07:12 -05:00
|
|
|
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();
|
|
|
|
|
}
|
|
|
|
|
|
2020-11-24 22:07:34 +01:00
|
|
|
LogicalResult ConvertAllocOpToGpuRuntimeCallPattern::matchAndRewrite(
|
2021-09-24 17:51:20 +00:00
|
|
|
gpu::AllocOp allocOp, OpAdaptor adaptor,
|
2020-11-24 22:07:34 +01:00
|
|
|
ConversionPatternRewriter &rewriter) const {
|
2022-09-09 00:04:01 +02:00
|
|
|
|
2020-11-24 22:07:34 +01:00
|
|
|
MemRefType memRefType = allocOp.getType();
|
|
|
|
|
|
2021-09-24 17:51:20 +00:00
|
|
|
if (failed(areAllLLVMTypes(allocOp, adaptor.getOperands(), rewriter)) ||
|
2023-09-26 15:32:11 -07:00
|
|
|
!isConvertibleAndHasIdentityMaps(memRefType))
|
2020-11-24 22:07:34 +01:00
|
|
|
return failure();
|
|
|
|
|
|
2020-11-28 13:46:43 +01:00
|
|
|
auto loc = allocOp.getLoc();
|
2020-11-24 22:07:34 +01:00
|
|
|
|
2023-09-26 15:32:11 -07:00
|
|
|
bool isShared = allocOp.getHostShared();
|
|
|
|
|
|
|
|
|
|
if (isShared && allocOp.getAsyncToken())
|
|
|
|
|
return rewriter.notifyMatchFailure(
|
|
|
|
|
allocOp, "Host Shared allocation cannot be done async");
|
2023-10-20 02:29:50 -07:00
|
|
|
if (!isShared && failed(isAsyncWithOneDependency(rewriter, allocOp)))
|
2023-09-26 15:32:11 -07:00
|
|
|
return failure();
|
|
|
|
|
|
2020-11-24 22:07:34 +01:00
|
|
|
// 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;
|
2022-09-30 12:30:41 -07:00
|
|
|
getMemRefDescriptorSizes(loc, memRefType, adaptor.getDynamicSizes(), rewriter,
|
2021-01-11 14:46:26 +01:00
|
|
|
shape, strides, sizeBytes);
|
2020-11-24 22:07:34 +01:00
|
|
|
|
|
|
|
|
// Allocate the underlying buffer and store a pointer to it in the MemRef
|
|
|
|
|
// descriptor.
|
2023-10-25 12:00:47 -07:00
|
|
|
auto nullPtr = rewriter.create<mlir::LLVM::ZeroOp>(loc, llvmPointerType);
|
|
|
|
|
Value stream = adaptor.getAsyncDependencies().empty()
|
|
|
|
|
? nullPtr
|
|
|
|
|
: adaptor.getAsyncDependencies().front();
|
2023-09-26 15:32:11 -07:00
|
|
|
|
|
|
|
|
auto isHostShared = rewriter.create<mlir::LLVM::ConstantOp>(
|
|
|
|
|
loc, llvmInt8Type, rewriter.getI8IntegerAttr(isShared));
|
|
|
|
|
|
2020-11-24 22:07:34 +01:00
|
|
|
Value allocatedPtr =
|
2023-09-26 15:32:11 -07:00
|
|
|
allocCallBuilder.create(loc, rewriter, {sizeBytes, stream, isHostShared})
|
|
|
|
|
.getResult();
|
2020-11-24 22:07:34 +01:00
|
|
|
|
|
|
|
|
// No alignment.
|
|
|
|
|
Value alignedPtr = allocatedPtr;
|
|
|
|
|
|
|
|
|
|
// Create the MemRef descriptor.
|
|
|
|
|
auto memRefDescriptor = this->createMemRefDescriptor(
|
|
|
|
|
loc, memRefType, allocatedPtr, alignedPtr, shape, strides, rewriter);
|
|
|
|
|
|
2023-10-25 12:00:47 -07:00
|
|
|
if (allocOp.getAsyncToken()) {
|
|
|
|
|
// Async alloc: make dependent ops use the same stream.
|
|
|
|
|
rewriter.replaceOp(allocOp, {memRefDescriptor, stream});
|
|
|
|
|
} else {
|
|
|
|
|
rewriter.replaceOp(allocOp, {memRefDescriptor});
|
|
|
|
|
}
|
2020-11-24 22:07:34 +01:00
|
|
|
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
LogicalResult ConvertDeallocOpToGpuRuntimeCallPattern::matchAndRewrite(
|
2021-09-24 17:51:20 +00:00
|
|
|
gpu::DeallocOp deallocOp, OpAdaptor adaptor,
|
2020-11-24 22:07:34 +01:00
|
|
|
ConversionPatternRewriter &rewriter) const {
|
2021-09-24 17:51:20 +00:00
|
|
|
if (failed(areAllLLVMTypes(deallocOp, adaptor.getOperands(), rewriter)) ||
|
2020-11-28 13:46:43 +01:00
|
|
|
failed(isAsyncWithOneDependency(rewriter, deallocOp)))
|
2020-11-24 22:07:34 +01:00
|
|
|
return failure();
|
|
|
|
|
|
2020-11-28 13:46:43 +01:00
|
|
|
Location loc = deallocOp.getLoc();
|
2020-11-24 22:07:34 +01:00
|
|
|
|
|
|
|
|
Value pointer =
|
2022-09-30 12:30:41 -07:00
|
|
|
MemRefDescriptor(adaptor.getMemref()).allocatedPtr(rewriter, loc);
|
|
|
|
|
Value stream = adaptor.getAsyncDependencies().front();
|
2023-02-21 07:51:44 +01:00
|
|
|
deallocCallBuilder.create(loc, rewriter, {pointer, stream});
|
2020-11-24 22:07:34 +01:00
|
|
|
|
2020-11-28 13:46:43 +01:00
|
|
|
rewriter.replaceOp(deallocOp, {stream});
|
2020-11-24 22:07:34 +01:00
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2021-02-24 15:02:30 +01:00
|
|
|
static bool isGpuAsyncTokenType(Value value) {
|
2023-05-08 16:33:54 +02:00
|
|
|
return isa<gpu::AsyncTokenType>(value.getType());
|
2021-02-24 15:02:30 +01:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// 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(
|
2021-09-24 17:51:20 +00:00
|
|
|
async::YieldOp yieldOp, OpAdaptor adaptor,
|
2021-02-24 15:02:30 +01:00
|
|
|
ConversionPatternRewriter &rewriter) const {
|
2022-10-25 18:29:53 -07:00
|
|
|
if (llvm::none_of(yieldOp.getOperands(), isGpuAsyncTokenType))
|
2021-02-24 15:02:30 +01:00
|
|
|
return rewriter.notifyMatchFailure(yieldOp, "no gpu async token operand");
|
|
|
|
|
|
|
|
|
|
Location loc = yieldOp.getLoc();
|
2021-09-24 17:51:20 +00:00
|
|
|
SmallVector<Value, 4> newOperands(adaptor.getOperands());
|
2021-02-24 15:02:30 +01:00
|
|
|
llvm::SmallDenseSet<Value> streams;
|
|
|
|
|
for (auto &operand : yieldOp->getOpOperands()) {
|
|
|
|
|
if (!isGpuAsyncTokenType(operand.get()))
|
|
|
|
|
continue;
|
|
|
|
|
auto idx = operand.getOperandNumber();
|
2021-09-24 17:51:20 +00:00
|
|
|
auto stream = adaptor.getOperands()[idx];
|
2022-08-11 00:34:02 -04:00
|
|
|
auto event = eventCreateCallBuilder.create(loc, rewriter, {}).getResult();
|
2021-02-24 15:02:30 +01:00
|
|
|
eventRecordCallBuilder.create(loc, rewriter, {event, stream});
|
|
|
|
|
newOperands[idx] = event;
|
|
|
|
|
streams.insert(stream);
|
|
|
|
|
}
|
|
|
|
|
for (auto stream : streams)
|
|
|
|
|
streamDestroyCallBuilder.create(loc, rewriter, {stream});
|
|
|
|
|
|
|
|
|
|
rewriter.updateRootInPlace(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) {
|
2023-05-08 16:33:54 +02:00
|
|
|
assert(isa<LLVM::LLVMPointerType>(value.getType()));
|
2021-02-24 15:02:30 +01:00
|
|
|
if (auto defOp = value.getDefiningOp<LLVM::CallOp>())
|
2021-10-24 18:36:33 -07:00
|
|
|
return defOp.getCallee()->equals(functionName);
|
2021-02-24 15:02:30 +01:00
|
|
|
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.
|
2020-10-21 08:24:53 +02:00
|
|
|
LogicalResult ConvertWaitOpToGpuRuntimeCallPattern::matchAndRewrite(
|
2021-09-24 17:51:20 +00:00
|
|
|
gpu::WaitOp waitOp, OpAdaptor adaptor,
|
2020-10-21 08:24:53 +02:00
|
|
|
ConversionPatternRewriter &rewriter) const {
|
2022-09-30 12:30:41 -07:00
|
|
|
if (waitOp.getAsyncToken())
|
2020-11-27 21:09:13 +01:00
|
|
|
return rewriter.notifyMatchFailure(waitOp, "Cannot convert async op.");
|
2020-10-21 08:24:53 +02:00
|
|
|
|
2020-11-27 21:09:13 +01:00
|
|
|
Location loc = waitOp.getLoc();
|
2020-10-21 08:24:53 +02:00
|
|
|
|
2021-09-24 17:51:20 +00:00
|
|
|
for (auto operand : adaptor.getOperands()) {
|
2021-02-24 15:02:30 +01:00
|
|
|
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});
|
|
|
|
|
}
|
|
|
|
|
}
|
2020-10-21 08:24:53 +02:00
|
|
|
|
2020-11-27 21:09:13 +01:00
|
|
|
rewriter.eraseOp(waitOp);
|
2020-10-21 08:24:53 +02:00
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2021-02-24 15:02:30 +01:00
|
|
|
// 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.
|
2020-10-21 08:24:53 +02:00
|
|
|
LogicalResult ConvertWaitAsyncOpToGpuRuntimeCallPattern::matchAndRewrite(
|
2021-09-24 17:51:20 +00:00
|
|
|
gpu::WaitOp waitOp, OpAdaptor adaptor,
|
2020-10-21 08:24:53 +02:00
|
|
|
ConversionPatternRewriter &rewriter) const {
|
2022-09-30 12:30:41 -07:00
|
|
|
if (!waitOp.getAsyncToken())
|
2020-11-27 21:09:13 +01:00
|
|
|
return rewriter.notifyMatchFailure(waitOp, "Can only convert async op.");
|
2020-10-21 08:24:53 +02:00
|
|
|
|
2020-11-27 21:09:13 +01:00
|
|
|
Location loc = waitOp.getLoc();
|
2020-10-21 08:24:53 +02:00
|
|
|
|
|
|
|
|
auto insertionPoint = rewriter.saveInsertionPoint();
|
|
|
|
|
SmallVector<Value, 1> events;
|
2021-09-24 17:51:20 +00:00
|
|
|
for (auto pair :
|
2022-09-30 12:30:41 -07:00
|
|
|
llvm::zip(waitOp.getAsyncDependencies(), adaptor.getOperands())) {
|
2021-02-24 15:02:30 +01:00
|
|
|
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();
|
2020-10-21 08:24:53 +02:00
|
|
|
rewriter.setInsertionPointAfter(defOp);
|
2022-08-11 00:34:02 -04:00
|
|
|
auto event = eventCreateCallBuilder.create(loc, rewriter, {}).getResult();
|
2021-02-24 15:02:30 +01:00
|
|
|
eventRecordCallBuilder.create(loc, rewriter, {event, operand});
|
|
|
|
|
events.push_back(event);
|
2020-10-21 08:24:53 +02:00
|
|
|
} else {
|
2021-02-24 15:02:30 +01:00
|
|
|
// Otherwise the converted operand is an event. This assumes that we use
|
|
|
|
|
// events in control flow code as well.
|
|
|
|
|
events.push_back(operand);
|
2020-10-21 08:24:53 +02:00
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
rewriter.restoreInsertionPoint(insertionPoint);
|
2022-08-11 00:34:02 -04:00
|
|
|
auto stream = streamCreateCallBuilder.create(loc, rewriter, {}).getResult();
|
2020-10-21 08:24:53 +02:00
|
|
|
for (auto event : events)
|
|
|
|
|
streamWaitEventCallBuilder.create(loc, rewriter, {stream, event});
|
|
|
|
|
for (auto event : events)
|
|
|
|
|
eventDestroyCallBuilder.create(loc, rewriter, {event});
|
2020-11-27 21:09:13 +01:00
|
|
|
rewriter.replaceOp(waitOp, {stream});
|
2020-10-21 08:24:53 +02:00
|
|
|
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2020-08-01 16:40:21 +02:00
|
|
|
// 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.
|
2019-06-19 06:22:36 -07:00
|
|
|
// The generated code is essentially as follows:
|
|
|
|
|
//
|
2020-08-01 16:40:21 +02:00
|
|
|
// %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
|
2019-06-19 06:22:36 -07:00
|
|
|
// return %array
|
2020-08-01 15:06:25 +02:00
|
|
|
Value ConvertLaunchFuncOpToGpuRuntimeCallPattern::generateParamsArray(
|
2021-09-24 17:51:20 +00:00
|
|
|
gpu::LaunchFuncOp launchOp, OpAdaptor adaptor, OpBuilder &builder) const {
|
2020-08-01 16:40:21 +02:00
|
|
|
auto loc = launchOp.getLoc();
|
2019-09-27 09:55:38 -07:00
|
|
|
auto numKernelOperands = launchOp.getNumKernelOperands();
|
2023-08-18 15:23:56 +02:00
|
|
|
// 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);
|
2020-08-01 16:40:21 +02:00
|
|
|
auto numArguments = arguments.size();
|
2021-01-05 16:22:53 +01:00
|
|
|
SmallVector<Type, 4> argumentTypes;
|
2020-08-01 16:40:21 +02:00
|
|
|
argumentTypes.reserve(numArguments);
|
|
|
|
|
for (auto argument : arguments)
|
2021-01-05 16:22:53 +01:00
|
|
|
argumentTypes.push_back(argument.getType());
|
2020-12-22 11:22:56 +01:00
|
|
|
auto structType = LLVM::LLVMStructType::getNewIdentified(context, StringRef(),
|
|
|
|
|
argumentTypes);
|
2022-08-09 14:40:07 -04:00
|
|
|
auto one = builder.create<LLVM::ConstantOp>(loc, llvmInt32Type, 1);
|
2023-10-31 09:22:44 +01:00
|
|
|
auto structPtr =
|
|
|
|
|
builder.create<LLVM::AllocaOp>(loc, llvmPointerType, structType, one,
|
|
|
|
|
/*alignment=*/0);
|
2022-08-09 14:40:07 -04:00
|
|
|
auto arraySize =
|
|
|
|
|
builder.create<LLVM::ConstantOp>(loc, llvmInt32Type, numArguments);
|
2023-02-21 07:51:44 +01:00
|
|
|
auto arrayPtr = builder.create<LLVM::AllocaOp>(
|
2023-10-31 09:22:44 +01:00
|
|
|
loc, llvmPointerType, llvmPointerType, arraySize, /*alignment=*/0);
|
2022-01-02 22:02:14 +00:00
|
|
|
for (const auto &en : llvm::enumerate(arguments)) {
|
2023-10-31 09:22:44 +01:00
|
|
|
Value fieldPtr =
|
|
|
|
|
builder.create<LLVM::GEPOp>(loc, llvmPointerType, structType, structPtr,
|
|
|
|
|
ArrayRef<LLVM::GEPArg>{0, en.index()});
|
2020-08-01 16:40:21 +02:00
|
|
|
builder.create<LLVM::StoreOp>(loc, en.value(), fieldPtr);
|
2023-02-21 07:51:44 +01:00
|
|
|
auto elementPtr = builder.create<LLVM::GEPOp>(
|
2023-10-31 09:22:44 +01:00
|
|
|
loc, llvmPointerType, llvmPointerType, arrayPtr,
|
2023-02-21 07:51:44 +01:00
|
|
|
ArrayRef<LLVM::GEPArg>{en.index()});
|
|
|
|
|
builder.create<LLVM::StoreOp>(loc, fieldPtr, elementPtr);
|
2019-06-19 06:22:36 -07:00
|
|
|
}
|
2020-08-01 16:40:21 +02:00
|
|
|
return arrayPtr;
|
2019-06-19 06:22:36 -07:00
|
|
|
}
|
|
|
|
|
|
2019-08-20 07:51:32 -07:00
|
|
|
// 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:
|
2019-06-19 06:22:36 -07:00
|
|
|
//
|
2019-08-20 07:51:32 -07:00
|
|
|
// 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*">
|
|
|
|
|
// }
|
2020-08-01 15:06:25 +02:00
|
|
|
Value ConvertLaunchFuncOpToGpuRuntimeCallPattern::generateKernelNameConstant(
|
|
|
|
|
StringRef moduleName, StringRef name, Location loc,
|
|
|
|
|
OpBuilder &builder) const {
|
2019-08-20 07:51:32 -07:00
|
|
|
// Make sure the trailing zero is included in the constant.
|
2019-10-08 05:11:00 -07:00
|
|
|
std::vector<char> kernelName(name.begin(), name.end());
|
2019-08-20 07:51:32 -07:00
|
|
|
kernelName.push_back('\0');
|
|
|
|
|
|
2020-04-16 13:14:43 +02:00
|
|
|
std::string globalName =
|
|
|
|
|
std::string(llvm::formatv("{0}_{1}_kernel_name", moduleName, name));
|
2019-08-20 07:51:32 -07:00
|
|
|
return LLVM::createGlobalString(
|
|
|
|
|
loc, builder, globalName, StringRef(kernelName.data(), kernelName.size()),
|
2023-10-31 09:22:44 +01:00
|
|
|
LLVM::Linkage::Internal);
|
2019-06-19 06:22:36 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Emits LLVM IR to launch a kernel function. Expects the module that contains
|
2020-05-18 13:01:54 -05:00
|
|
|
// 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.
|
2019-06-19 06:22:36 -07:00
|
|
|
//
|
2020-05-18 13:01:54 -05:00
|
|
|
// %0 = call %binarygetter
|
2020-07-28 16:29:29 +02:00
|
|
|
// %1 = call %moduleLoad(%0)
|
|
|
|
|
// %2 = <see generateKernelNameConstant>
|
|
|
|
|
// %3 = call %moduleGetFunction(%1, %2)
|
|
|
|
|
// %4 = call %streamCreate()
|
2020-08-01 15:06:25 +02:00
|
|
|
// %5 = <see generateParamsArray>
|
2020-07-28 16:29:29 +02:00
|
|
|
// call %launchKernel(%3, <launchOp operands 0..5>, 0, %4, %5, nullptr)
|
|
|
|
|
// call %streamSynchronize(%4)
|
2020-10-29 17:58:48 +01:00
|
|
|
// 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.
|
2020-08-01 15:06:25 +02:00
|
|
|
LogicalResult ConvertLaunchFuncOpToGpuRuntimeCallPattern::matchAndRewrite(
|
2021-09-24 17:51:20 +00:00
|
|
|
gpu::LaunchFuncOp launchOp, OpAdaptor adaptor,
|
2020-08-01 15:06:25 +02:00
|
|
|
ConversionPatternRewriter &rewriter) const {
|
2021-09-24 17:51:20 +00:00
|
|
|
if (failed(areAllLLVMTypes(launchOp, adaptor.getOperands(), rewriter)))
|
2020-11-24 22:07:34 +01:00
|
|
|
return failure();
|
2020-08-01 16:40:21 +02:00
|
|
|
|
2022-09-30 12:30:41 -07:00
|
|
|
if (launchOp.getAsyncDependencies().size() > 1)
|
2020-10-29 17:58:48 +01:00
|
|
|
return rewriter.notifyMatchFailure(
|
2020-11-27 21:09:13 +01:00
|
|
|
launchOp, "Cannot convert with more than one async dependency.");
|
2020-10-29 17:58:48 +01:00
|
|
|
|
|
|
|
|
// 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.
|
2022-09-30 12:30:41 -07:00
|
|
|
if (!launchOp.getAsyncToken() && !launchOp.getAsyncDependencies().empty())
|
2020-10-29 17:58:48 +01:00
|
|
|
return rewriter.notifyMatchFailure(
|
2020-11-27 21:09:13 +01:00
|
|
|
launchOp, "Cannot convert non-async op with async dependencies.");
|
2020-10-29 17:58:48 +01:00
|
|
|
|
2020-08-01 16:40:21 +02:00
|
|
|
Location loc = launchOp.getLoc();
|
2019-06-19 06:22:36 -07:00
|
|
|
|
2019-10-08 05:03:09 -07:00
|
|
|
// Create an LLVM global with CUBIN extracted from the kernel annotation and
|
|
|
|
|
// obtain a pointer to the first byte in it.
|
2023-08-12 00:26:13 +00:00
|
|
|
gpu::GPUModuleOp kernelModule;
|
|
|
|
|
if (cachedModuleTable)
|
|
|
|
|
kernelModule = cachedModuleTable->lookup<gpu::GPUModuleOp>(
|
|
|
|
|
launchOp.getKernelModuleName());
|
|
|
|
|
else
|
|
|
|
|
kernelModule = SymbolTable::lookupNearestSymbolFrom<gpu::GPUModuleOp>(
|
|
|
|
|
launchOp, launchOp.getKernelModuleName());
|
2019-10-08 04:29:58 -07:00
|
|
|
assert(kernelModule && "expected a kernel module");
|
|
|
|
|
|
2023-08-12 00:26:13 +00:00
|
|
|
// If the module has Targets then just update the op operands.
|
|
|
|
|
if (ArrayAttr targets = kernelModule.getTargetsAttr()) {
|
|
|
|
|
Value stream = Value();
|
2023-08-23 17:51:11 +02:00
|
|
|
if (!adaptor.getAsyncDependencies().empty())
|
2023-08-12 00:26:13 +00:00
|
|
|
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.
|
2023-08-18 15:23:56 +02:00
|
|
|
// 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);
|
2023-08-12 00:26:13 +00:00
|
|
|
|
|
|
|
|
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);
|
|
|
|
|
if (launchOp.getAsyncToken())
|
|
|
|
|
rewriter.replaceOp(launchOp, {stream});
|
|
|
|
|
else
|
|
|
|
|
rewriter.eraseOp(launchOp);
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2020-12-09 11:50:18 +01:00
|
|
|
auto binaryAttr =
|
|
|
|
|
kernelModule->getAttrOfType<StringAttr>(gpuBinaryAnnotation);
|
2020-05-18 13:01:54 -05:00
|
|
|
if (!binaryAttr) {
|
2019-10-08 05:11:00 -07:00
|
|
|
kernelModule.emitOpError()
|
2020-05-18 13:01:54 -05:00
|
|
|
<< "missing " << gpuBinaryAnnotation << " attribute";
|
2020-08-01 15:06:25 +02:00
|
|
|
return failure();
|
2019-06-19 06:22:36 -07:00
|
|
|
}
|
2019-10-10 01:33:33 -07:00
|
|
|
|
Create a gpu.module operation for the GPU Dialect.
Summary:
This is based on the use of code constantly checking for an attribute on
a model and instead represents the distinct operaion with a different
op. Instead, this op can be used to provide better filtering.
Reverts "Revert "[mlir] Create a gpu.module operation for the GPU Dialect.""
This reverts commit ac446302ca4145cdc89f377c0c364c29ee303be5 after
fixing internal Google issues.
This additionally updates ROCDL lowering to use the new gpu.module.
Reviewers: herhut, mravishankar, antiagainst, nicolasvasilache
Subscribers: jholewinski, mgorny, mehdi_amini, jpienaar, burmako, shauheen, csigg, arpith-jacob, mgester, lucyrfox, aartbik, liufengdb, llvm-commits, mravishankar, rriddle, antiagainst, bkramer
Tags: #llvm
Differential Revision: https://reviews.llvm.org/D72921
2020-01-17 15:18:23 +01:00
|
|
|
SmallString<128> nameBuffer(kernelModule.getName());
|
2020-05-18 13:01:54 -05:00
|
|
|
nameBuffer.append(kGpuBinaryStorageSuffix);
|
2023-10-31 09:22:44 +01:00
|
|
|
Value data =
|
|
|
|
|
LLVM::createGlobalString(loc, rewriter, nameBuffer.str(),
|
|
|
|
|
binaryAttr.getValue(), LLVM::Linkage::Internal);
|
2019-10-08 05:03:09 -07:00
|
|
|
|
2023-09-26 11:27:07 -07:00
|
|
|
// 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())));
|
|
|
|
|
|
2019-06-19 06:22:36 -07:00
|
|
|
// Get the function from the module. The name corresponds to the name of
|
|
|
|
|
// the kernel function.
|
2020-04-21 10:16:41 +00:00
|
|
|
auto kernelName = generateKernelNameConstant(
|
2021-08-29 14:22:24 -07:00
|
|
|
launchOp.getKernelModuleName().getValue(),
|
|
|
|
|
launchOp.getKernelName().getValue(), loc, rewriter);
|
2020-08-01 15:06:25 +02:00
|
|
|
auto function = moduleGetFunctionCallBuilder.create(
|
2022-08-11 00:34:02 -04:00
|
|
|
loc, rewriter, {module.getResult(), kernelName});
|
2022-08-16 21:17:03 -07:00
|
|
|
Value zero = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type, 0);
|
2020-10-29 17:58:48 +01:00
|
|
|
Value stream =
|
2022-09-30 12:30:41 -07:00
|
|
|
adaptor.getAsyncDependencies().empty()
|
2022-08-11 00:34:02 -04:00
|
|
|
? streamCreateCallBuilder.create(loc, rewriter, {}).getResult()
|
2022-09-30 12:30:41 -07:00
|
|
|
: adaptor.getAsyncDependencies().front();
|
2020-08-01 16:40:21 +02:00
|
|
|
// Create array of pointers to kernel arguments.
|
2021-09-24 17:51:20 +00:00
|
|
|
auto kernelParams = generateParamsArray(launchOp, adaptor, rewriter);
|
2023-10-31 09:22:44 +01:00
|
|
|
auto nullpointer = rewriter.create<LLVM::ZeroOp>(loc, llvmPointerType);
|
2022-09-30 12:30:41 -07:00
|
|
|
Value dynamicSharedMemorySize = launchOp.getDynamicSharedMemorySize()
|
|
|
|
|
? launchOp.getDynamicSharedMemorySize()
|
2021-09-14 14:13:21 +05:30
|
|
|
: zero;
|
|
|
|
|
launchKernelCallBuilder.create(
|
|
|
|
|
loc, rewriter,
|
2022-09-30 12:30:41 -07:00
|
|
|
{function.getResult(), adaptor.getGridSizeX(), adaptor.getGridSizeY(),
|
|
|
|
|
adaptor.getGridSizeZ(), adaptor.getBlockSizeX(), adaptor.getBlockSizeY(),
|
|
|
|
|
adaptor.getBlockSizeZ(), dynamicSharedMemorySize, stream, kernelParams,
|
2023-09-26 11:27:07 -07:00
|
|
|
/*extra=*/nullpointer, paramsCount});
|
2020-10-29 17:58:48 +01:00
|
|
|
|
2022-09-30 12:30:41 -07:00
|
|
|
if (launchOp.getAsyncToken()) {
|
2020-10-29 17:58:48 +01:00
|
|
|
// Async launch: make dependent ops use the same stream.
|
2020-11-27 21:09:13 +01:00
|
|
|
rewriter.replaceOp(launchOp, {stream});
|
2020-10-29 17:58:48 +01:00
|
|
|
} 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);
|
2020-11-27 21:09:13 +01:00
|
|
|
rewriter.eraseOp(launchOp);
|
2020-10-29 17:58:48 +01:00
|
|
|
}
|
2022-08-11 00:34:02 -04:00
|
|
|
moduleUnloadCallBuilder.create(loc, rewriter, module.getResult());
|
2020-08-01 15:06:25 +02:00
|
|
|
|
|
|
|
|
return success();
|
2019-06-19 06:22:36 -07:00
|
|
|
}
|
|
|
|
|
|
2023-02-12 23:52:16 +01:00
|
|
|
static Value bitAndAddrspaceCast(Location loc,
|
|
|
|
|
ConversionPatternRewriter &rewriter,
|
|
|
|
|
LLVM::LLVMPointerType destinationType,
|
|
|
|
|
Value sourcePtr,
|
2023-08-14 08:40:13 +02:00
|
|
|
const LLVMTypeConverter &typeConverter) {
|
2023-05-08 16:33:54 +02:00
|
|
|
auto sourceTy = cast<LLVM::LLVMPointerType>(sourcePtr.getType());
|
2023-02-12 23:52:16 +01:00
|
|
|
if (destinationType.getAddressSpace() != sourceTy.getAddressSpace())
|
|
|
|
|
sourcePtr = rewriter.create<LLVM::AddrSpaceCastOp>(
|
|
|
|
|
loc,
|
2023-10-31 09:22:44 +01:00
|
|
|
LLVM::LLVMPointerType::get(rewriter.getContext(),
|
|
|
|
|
destinationType.getAddressSpace()),
|
2023-02-12 23:52:16 +01:00
|
|
|
sourcePtr);
|
2023-10-31 09:22:44 +01:00
|
|
|
return sourcePtr;
|
2023-02-12 23:52:16 +01:00
|
|
|
}
|
|
|
|
|
|
2020-12-22 17:42:59 +01:00
|
|
|
LogicalResult ConvertMemcpyOpToGpuRuntimeCallPattern::matchAndRewrite(
|
2021-09-24 17:51:20 +00:00
|
|
|
gpu::MemcpyOp memcpyOp, OpAdaptor adaptor,
|
2020-12-22 17:42:59 +01:00
|
|
|
ConversionPatternRewriter &rewriter) const {
|
2023-05-08 16:33:54 +02:00
|
|
|
auto memRefType = cast<MemRefType>(memcpyOp.getSrc().getType());
|
2020-12-22 17:42:59 +01:00
|
|
|
|
2021-09-24 17:51:20 +00:00
|
|
|
if (failed(areAllLLVMTypes(memcpyOp, adaptor.getOperands(), rewriter)) ||
|
2020-12-23 11:37:49 +01:00
|
|
|
!isConvertibleAndHasIdentityMaps(memRefType) ||
|
2020-12-22 17:42:59 +01:00
|
|
|
failed(isAsyncWithOneDependency(rewriter, memcpyOp)))
|
|
|
|
|
return failure();
|
|
|
|
|
|
|
|
|
|
auto loc = memcpyOp.getLoc();
|
|
|
|
|
|
2022-09-30 12:30:41 -07:00
|
|
|
MemRefDescriptor srcDesc(adaptor.getSrc());
|
2021-09-04 08:03:33 +02:00
|
|
|
Value numElements = getNumElements(rewriter, loc, memRefType, srcDesc);
|
2020-12-22 17:42:59 +01:00
|
|
|
|
|
|
|
|
Type elementPtrType = getElementPtrType(memRefType);
|
2023-09-25 11:11:52 +02:00
|
|
|
Value nullPtr = rewriter.create<LLVM::ZeroOp>(loc, elementPtrType);
|
2023-02-21 07:51:44 +01:00
|
|
|
Value gepPtr = rewriter.create<LLVM::GEPOp>(
|
|
|
|
|
loc, elementPtrType,
|
|
|
|
|
typeConverter->convertType(memRefType.getElementType()), nullPtr,
|
|
|
|
|
numElements);
|
2020-12-22 17:42:59 +01:00
|
|
|
auto sizeBytes =
|
|
|
|
|
rewriter.create<LLVM::PtrToIntOp>(loc, getIndexType(), gepPtr);
|
|
|
|
|
|
2023-02-12 23:52:16 +01:00
|
|
|
auto src = bitAndAddrspaceCast(loc, rewriter, llvmPointerType,
|
|
|
|
|
srcDesc.alignedPtr(rewriter, loc),
|
|
|
|
|
*getTypeConverter());
|
|
|
|
|
auto dst = bitAndAddrspaceCast(
|
|
|
|
|
loc, rewriter, llvmPointerType,
|
|
|
|
|
MemRefDescriptor(adaptor.getDst()).alignedPtr(rewriter, loc),
|
|
|
|
|
*getTypeConverter());
|
2020-12-22 17:42:59 +01:00
|
|
|
|
2022-09-30 12:30:41 -07:00
|
|
|
auto stream = adaptor.getAsyncDependencies().front();
|
2020-12-22 17:42:59 +01:00
|
|
|
memcpyCallBuilder.create(loc, rewriter, {dst, src, sizeBytes, stream});
|
|
|
|
|
|
|
|
|
|
rewriter.replaceOp(memcpyOp, {stream});
|
|
|
|
|
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2021-09-04 08:03:33 +02:00
|
|
|
LogicalResult ConvertMemsetOpToGpuRuntimeCallPattern::matchAndRewrite(
|
2021-09-24 17:51:20 +00:00
|
|
|
gpu::MemsetOp memsetOp, OpAdaptor adaptor,
|
2021-09-04 08:03:33 +02:00
|
|
|
ConversionPatternRewriter &rewriter) const {
|
2023-05-08 16:33:54 +02:00
|
|
|
auto memRefType = cast<MemRefType>(memsetOp.getDst().getType());
|
2021-09-04 08:03:33 +02:00
|
|
|
|
2021-09-24 17:51:20 +00:00
|
|
|
if (failed(areAllLLVMTypes(memsetOp, adaptor.getOperands(), rewriter)) ||
|
2021-09-04 08:03:33 +02:00
|
|
|
!isConvertibleAndHasIdentityMaps(memRefType) ||
|
|
|
|
|
failed(isAsyncWithOneDependency(rewriter, memsetOp)))
|
|
|
|
|
return failure();
|
|
|
|
|
|
|
|
|
|
auto loc = memsetOp.getLoc();
|
|
|
|
|
|
2022-09-30 12:30:41 -07:00
|
|
|
Type valueType = adaptor.getValue().getType();
|
2023-06-05 12:09:35 +05:30
|
|
|
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");
|
2021-09-04 08:03:33 +02:00
|
|
|
}
|
|
|
|
|
|
2023-06-05 12:09:35 +05:30
|
|
|
unsigned valueTypeWidth = valueType.getIntOrFloatBitWidth();
|
|
|
|
|
Type bitCastType = valueTypeWidth == 32 ? llvmInt32Type : llvmInt16Type;
|
|
|
|
|
|
2022-09-30 12:30:41 -07:00
|
|
|
MemRefDescriptor dstDesc(adaptor.getDst());
|
2021-09-04 08:03:33 +02:00
|
|
|
Value numElements = getNumElements(rewriter, loc, memRefType, dstDesc);
|
|
|
|
|
|
|
|
|
|
auto value =
|
2023-06-05 12:09:35 +05:30
|
|
|
rewriter.create<LLVM::BitcastOp>(loc, bitCastType, adaptor.getValue());
|
2023-02-12 23:52:16 +01:00
|
|
|
auto dst = bitAndAddrspaceCast(loc, rewriter, llvmPointerType,
|
|
|
|
|
dstDesc.alignedPtr(rewriter, loc),
|
|
|
|
|
*getTypeConverter());
|
2021-09-04 08:03:33 +02:00
|
|
|
|
2022-09-30 12:30:41 -07:00
|
|
|
auto stream = adaptor.getAsyncDependencies().front();
|
2023-06-05 12:09:35 +05:30
|
|
|
FunctionCallBuilder builder =
|
|
|
|
|
valueTypeWidth == 32 ? memset32CallBuilder : memset16CallBuilder;
|
|
|
|
|
builder.create(loc, rewriter, {dst, value, numElements, stream});
|
2021-09-04 08:03:33 +02:00
|
|
|
|
|
|
|
|
rewriter.replaceOp(memsetOp, {stream});
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2022-02-15 20:23:44 +00:00
|
|
|
LogicalResult ConvertSetDefaultDeviceOpToGpuRuntimeCallPattern::matchAndRewrite(
|
|
|
|
|
gpu::SetDefaultDeviceOp op, OpAdaptor adaptor,
|
|
|
|
|
ConversionPatternRewriter &rewriter) const {
|
|
|
|
|
Location loc = op.getLoc();
|
2022-09-30 12:30:41 -07:00
|
|
|
setDefaultDeviceCallBuilder.create(loc, rewriter, {adaptor.getDevIndex()});
|
2022-02-15 20:23:44 +00:00
|
|
|
rewriter.replaceOp(op, {});
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2023-05-26 21:46:22 +00:00
|
|
|
template <typename T>
|
|
|
|
|
static Value genConstInt32From(OpBuilder &builder, Location loc, T TValue) {
|
2023-05-26 02:29:49 +00:00
|
|
|
Type llvmInt32Type = builder.getIntegerType(32);
|
|
|
|
|
return builder.create<LLVM::ConstantOp>(loc, llvmInt32Type,
|
2023-05-26 21:46:22 +00:00
|
|
|
static_cast<int32_t>(TValue));
|
|
|
|
|
}
|
|
|
|
|
|
2023-08-01 00:57:47 +00:00
|
|
|
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)));
|
|
|
|
|
}
|
|
|
|
|
|
2023-06-09 16:58:44 +00:00
|
|
|
LogicalResult ConvertCreateDnTensorOpToGpuRuntimeCallPattern::matchAndRewrite(
|
|
|
|
|
gpu::CreateDnTensorOp op, OpAdaptor adaptor,
|
2023-05-11 13:56:34 -07:00
|
|
|
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();
|
2023-06-09 16:58:44 +00:00
|
|
|
Value pTensor =
|
2023-05-11 13:56:34 -07:00
|
|
|
MemRefDescriptor(adaptor.getMemref()).allocatedPtr(rewriter, loc);
|
2023-05-26 21:46:22 +00:00
|
|
|
Type dType = op.getMemref().getType().getElementType();
|
|
|
|
|
auto dtp = genConstInt32From(rewriter, loc, getCuSparseDataTypeFrom(dType));
|
2023-05-11 13:56:34 -07:00
|
|
|
|
2023-06-09 16:58:44 +00:00
|
|
|
SmallVector<Value, 4> dims;
|
|
|
|
|
for (Value dim : adaptor.getDims()) {
|
|
|
|
|
dims.push_back(dim);
|
|
|
|
|
}
|
2023-05-11 13:56:34 -07:00
|
|
|
|
2023-06-09 16:58:44 +00:00
|
|
|
Value handle;
|
2023-06-06 23:12:45 +00:00
|
|
|
// 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
|
2023-06-09 16:58:44 +00:00
|
|
|
if (dims.size() == 2) {
|
|
|
|
|
if (isSpMMCusparseLtOp(op.getDnTensor())) {
|
|
|
|
|
auto handleSz = rewriter.create<LLVM::ConstantOp>(
|
|
|
|
|
loc, getIndexType(), rewriter.getIndexAttr(11032));
|
2023-07-12 16:50:42 -07:00
|
|
|
handle = rewriter.create<LLVM::AllocaOp>(
|
2023-10-31 09:22:44 +01:00
|
|
|
loc, llvmPointerType, llvmInt8Type, handleSz, /*alignment=*/16);
|
2023-06-09 16:58:44 +00:00
|
|
|
handle = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, handle);
|
|
|
|
|
|
|
|
|
|
createLtDnMatCallBuilder
|
|
|
|
|
.create(loc, rewriter,
|
2023-06-16 21:38:40 +00:00
|
|
|
{handle, dims[0], dims[1], pTensor, dtp, stream})
|
2023-06-09 16:58:44 +00:00
|
|
|
.getResult();
|
|
|
|
|
} else {
|
|
|
|
|
handle =
|
|
|
|
|
createDnMatCallBuilder
|
|
|
|
|
.create(loc, rewriter, {dims[0], dims[1], pTensor, dtp, stream})
|
|
|
|
|
.getResult();
|
|
|
|
|
}
|
2023-06-06 23:12:45 +00:00
|
|
|
} else {
|
2023-06-09 16:58:44 +00:00
|
|
|
assert(dims.size() == 1 && "Only 1D and 2D tensors are supported");
|
|
|
|
|
handle = createDnVecCallBuilder
|
|
|
|
|
.create(loc, rewriter, {dims[0], pTensor, dtp, stream})
|
|
|
|
|
.getResult();
|
2023-06-06 23:12:45 +00:00
|
|
|
}
|
2023-05-15 14:55:45 -07:00
|
|
|
rewriter.replaceOp(op, {handle, stream});
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2023-06-09 16:58:44 +00:00
|
|
|
LogicalResult ConvertDestroyDnTensorOpToGpuRuntimeCallPattern::matchAndRewrite(
|
|
|
|
|
gpu::DestroyDnTensorOp op, OpAdaptor adaptor,
|
2023-05-15 14:55:45 -07:00
|
|
|
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();
|
2023-06-09 16:58:44 +00:00
|
|
|
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});
|
|
|
|
|
}
|
2023-06-06 23:12:45 +00:00
|
|
|
} else {
|
2023-06-09 16:58:44 +00:00
|
|
|
assert(dims.size() == 1 && "Only 1D and 2D tensors are supported");
|
|
|
|
|
destroyDnVecCallBuilder.create(loc, rewriter,
|
|
|
|
|
{adaptor.getDnTensor(), stream});
|
2023-06-06 23:12:45 +00:00
|
|
|
}
|
2023-05-15 14:55:45 -07:00
|
|
|
rewriter.replaceOp(op, {stream});
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2023-05-11 13:56:34 -07:00
|
|
|
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);
|
2023-05-24 02:21:55 +00:00
|
|
|
Type iType =
|
|
|
|
|
llvm::cast<MemRefType>(op.getColIdxs().getType()).getElementType();
|
|
|
|
|
Type dType =
|
|
|
|
|
llvm::cast<MemRefType>(op.getValues().getType()).getElementType();
|
2023-05-26 21:46:22 +00:00
|
|
|
auto itp = genConstInt32From(rewriter, loc, getCuSparseIndexTypeFrom(iType));
|
|
|
|
|
auto dtp = genConstInt32From(rewriter, loc, getCuSparseDataTypeFrom(dType));
|
2023-05-11 13:56:34 -07:00
|
|
|
auto handle =
|
|
|
|
|
createCooCallBuilder
|
|
|
|
|
.create(loc, rewriter,
|
|
|
|
|
{adaptor.getRows(), adaptor.getCols(), adaptor.getNnz(),
|
2023-05-26 21:46:22 +00:00
|
|
|
pRowIdxs, pColIdxs, pValues, itp, dtp, stream})
|
2023-05-11 13:56:34 -07:00
|
|
|
.getResult();
|
|
|
|
|
rewriter.replaceOp(op, {handle, stream});
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2023-06-06 10:56:58 -07:00
|
|
|
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);
|
2023-06-06 23:12:45 +00:00
|
|
|
Type iType = llvm::cast<MemRefType>(op.getIdxs().getType()).getElementType();
|
2023-06-06 10:56:58 -07:00
|
|
|
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();
|
|
|
|
|
}
|
|
|
|
|
|
2023-05-11 13:56:34 -07:00
|
|
|
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);
|
2023-05-24 02:21:55 +00:00
|
|
|
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();
|
2023-05-26 21:46:22 +00:00
|
|
|
auto ptp = genConstInt32From(rewriter, loc, getCuSparseIndexTypeFrom(pType));
|
|
|
|
|
auto itp = genConstInt32From(rewriter, loc, getCuSparseIndexTypeFrom(iType));
|
|
|
|
|
auto dtp = genConstInt32From(rewriter, loc, getCuSparseDataTypeFrom(dType));
|
2023-05-11 13:56:34 -07:00
|
|
|
auto handle =
|
|
|
|
|
createCsrCallBuilder
|
|
|
|
|
.create(loc, rewriter,
|
|
|
|
|
{adaptor.getRows(), adaptor.getCols(), adaptor.getNnz(),
|
2023-05-26 21:46:22 +00:00
|
|
|
pRowPos, pColIdxs, pValues, ptp, itp, dtp, stream})
|
2023-05-11 13:56:34 -07:00
|
|
|
.getResult();
|
|
|
|
|
rewriter.replaceOp(op, {handle, stream});
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2023-06-06 23:12:45 +00:00
|
|
|
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();
|
2023-06-12 23:24:08 +00:00
|
|
|
auto dtp = genConstInt32From(rewriter, loc, getCuSparseDataTypeFrom(dType));
|
2023-06-06 23:12:45 +00:00
|
|
|
|
2023-06-12 23:24:08 +00:00
|
|
|
// CUDA runner asserts the size is 44104 bytes.
|
2023-06-06 23:12:45 +00:00
|
|
|
auto handleSz = rewriter.create<LLVM::ConstantOp>(
|
|
|
|
|
loc, getIndexType(), rewriter.getIndexAttr(44104));
|
2023-07-12 16:50:42 -07:00
|
|
|
Value handle = rewriter.create<LLVM::AllocaOp>(
|
2023-10-31 09:22:44 +01:00
|
|
|
loc, llvmPointerType, llvmInt8Type, handleSz, /*alignment=*/16);
|
2023-06-06 23:12:45 +00:00
|
|
|
handle = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, handle);
|
|
|
|
|
|
|
|
|
|
create2To4SpMatCallBuilder
|
|
|
|
|
.create(loc, rewriter,
|
2023-06-16 21:38:40 +00:00
|
|
|
{handle, adaptor.getRows(), adaptor.getCols(), pMat, dtp, stream})
|
2023-06-06 23:12:45 +00:00
|
|
|
.getResult();
|
|
|
|
|
rewriter.replaceOp(op, {handle, stream});
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2023-05-11 13:56:34 -07:00
|
|
|
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();
|
2023-06-06 23:12:45 +00:00
|
|
|
// 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});
|
|
|
|
|
}
|
2023-05-11 13:56:34 -07:00
|
|
|
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();
|
2023-05-26 21:46:22 +00:00
|
|
|
auto modeA = genConstInt32From(rewriter, loc, op.getModeA());
|
2023-06-12 23:24:08 +00:00
|
|
|
auto computeType = genConstInt32From(
|
|
|
|
|
rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
|
2023-05-11 13:56:34 -07:00
|
|
|
auto stream = adaptor.getAsyncDependencies().front();
|
2023-06-16 21:38:40 +00:00
|
|
|
auto bufferSize = spMVBufferSizeCallBuilder
|
|
|
|
|
.create(loc, rewriter,
|
|
|
|
|
{modeA, adaptor.getSpmatA(), adaptor.getDnX(),
|
|
|
|
|
adaptor.getDnY(), computeType, stream})
|
|
|
|
|
.getResult();
|
2023-05-11 13:56:34 -07:00
|
|
|
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();
|
2023-05-26 21:46:22 +00:00
|
|
|
auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
|
2023-06-12 23:24:08 +00:00
|
|
|
auto computeType = genConstInt32From(
|
|
|
|
|
rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
|
2023-05-11 13:56:34 -07:00
|
|
|
auto stream = adaptor.getAsyncDependencies().front();
|
|
|
|
|
Value pBuf =
|
|
|
|
|
MemRefDescriptor(adaptor.getBuffer()).allocatedPtr(rewriter, loc);
|
|
|
|
|
spMVCallBuilder.create(loc, rewriter,
|
2023-06-16 21:38:40 +00:00
|
|
|
{modeA, adaptor.getSpmatA(), adaptor.getDnX(),
|
|
|
|
|
adaptor.getDnY(), computeType, pBuf, stream});
|
2023-05-11 13:56:34 -07:00
|
|
|
rewriter.replaceOp(op, {stream});
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2023-05-15 14:55:45 -07:00
|
|
|
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();
|
2023-05-26 21:46:22 +00:00
|
|
|
auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
|
|
|
|
|
auto modeB = genConstInt32From(rewriter, loc, adaptor.getModeB());
|
2023-05-15 14:55:45 -07:00
|
|
|
auto stream = adaptor.getAsyncDependencies().front();
|
2023-06-06 23:12:45 +00:00
|
|
|
Value bufferSize;
|
|
|
|
|
if (is2To4Sparsity(op.getSpmatA())) {
|
2023-07-21 02:03:51 +00:00
|
|
|
auto prune_flag =
|
|
|
|
|
genConstInt32From(rewriter, loc, get2To4PruneFlag(op.getSpmatA()));
|
2023-06-12 23:24:08 +00:00
|
|
|
auto computeType = genConstInt32From(
|
|
|
|
|
rewriter, loc, getCuSparseLtDataTypeFrom(adaptor.getComputeType()));
|
2023-06-06 23:12:45 +00:00
|
|
|
auto three = rewriter.create<LLVM::ConstantOp>(loc, getIndexType(),
|
|
|
|
|
rewriter.getIndexAttr(3));
|
2023-07-12 16:50:42 -07:00
|
|
|
auto bufferSize = rewriter.create<LLVM::AllocaOp>(
|
2023-10-31 09:22:44 +01:00
|
|
|
loc, llvmPointerType, llvmPointerType, three, /*alignment=*/16);
|
2023-06-12 23:24:08 +00:00
|
|
|
createCuSparseLtSpMMBufferSizeBuilder
|
2023-06-06 23:12:45 +00:00
|
|
|
.create(loc, rewriter,
|
2023-06-16 21:38:40 +00:00
|
|
|
{bufferSize, modeA, modeB, adaptor.getSpmatA(),
|
2023-07-21 02:03:51 +00:00
|
|
|
adaptor.getDnmatB(), adaptor.getDnmatC(), computeType,
|
|
|
|
|
prune_flag, stream})
|
2023-06-06 23:12:45 +00:00
|
|
|
.getResult();
|
2023-06-19 15:57:39 +00:00
|
|
|
|
|
|
|
|
auto bufferSizePtr1 = rewriter.create<LLVM::GEPOp>(
|
2023-10-31 09:22:44 +01:00
|
|
|
loc, llvmPointerType, llvmPointerType, bufferSize,
|
2023-06-19 15:57:39 +00:00
|
|
|
ValueRange{rewriter.create<LLVM::ConstantOp>(
|
|
|
|
|
loc, getIndexType(), rewriter.getIndexAttr(1))});
|
|
|
|
|
auto bufferSizePtr2 = rewriter.create<LLVM::GEPOp>(
|
2023-10-31 09:22:44 +01:00
|
|
|
loc, llvmPointerType, llvmPointerType, bufferSize,
|
2023-06-19 15:57:39 +00:00
|
|
|
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});
|
2023-06-06 23:12:45 +00:00
|
|
|
} else {
|
2023-06-12 23:24:08 +00:00
|
|
|
auto computeType = genConstInt32From(
|
|
|
|
|
rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
|
2023-06-16 21:38:40 +00:00
|
|
|
bufferSize =
|
|
|
|
|
createSpMMBufferSizeCallBuilder
|
|
|
|
|
.create(loc, rewriter,
|
|
|
|
|
{modeA, modeB, adaptor.getSpmatA(), adaptor.getDnmatB(),
|
|
|
|
|
adaptor.getDnmatC(), computeType, stream})
|
|
|
|
|
.getResult();
|
2023-06-06 23:12:45 +00:00
|
|
|
rewriter.replaceOp(op, {bufferSize, stream});
|
|
|
|
|
}
|
2023-05-15 14:55:45 -07:00
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2023-05-24 02:21:55 +00:00
|
|
|
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();
|
2023-05-26 21:46:22 +00:00
|
|
|
auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
|
|
|
|
|
auto modeB = genConstInt32From(rewriter, loc, adaptor.getModeB());
|
2023-06-12 23:24:08 +00:00
|
|
|
auto computeType = genConstInt32From(
|
|
|
|
|
rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
|
2023-05-24 02:21:55 +00:00
|
|
|
auto stream = adaptor.getAsyncDependencies().front();
|
2023-06-16 21:38:40 +00:00
|
|
|
auto bufferSize =
|
|
|
|
|
createSDDMMBufferSizeCallBuilder
|
|
|
|
|
.create(loc, rewriter,
|
|
|
|
|
{modeA, modeB, adaptor.getDnmatA(), adaptor.getDnmatB(),
|
|
|
|
|
adaptor.getSpmatC(), computeType, stream})
|
|
|
|
|
.getResult();
|
2023-05-24 02:21:55 +00:00
|
|
|
rewriter.replaceOp(op, {bufferSize, stream});
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2023-05-15 14:55:45 -07:00
|
|
|
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();
|
2023-05-26 21:46:22 +00:00
|
|
|
auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
|
|
|
|
|
auto modeB = genConstInt32From(rewriter, loc, adaptor.getModeB());
|
2023-06-12 23:24:08 +00:00
|
|
|
auto computeType = genConstInt32From(
|
|
|
|
|
rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
|
2023-05-26 21:46:22 +00:00
|
|
|
|
2023-05-15 14:55:45 -07:00
|
|
|
auto stream = adaptor.getAsyncDependencies().front();
|
2023-06-06 23:12:45 +00:00
|
|
|
|
|
|
|
|
// 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);
|
|
|
|
|
}
|
2023-06-12 23:24:08 +00:00
|
|
|
createCuSparseLtSpMMBuilder.create(
|
|
|
|
|
loc, rewriter,
|
2023-06-16 21:38:40 +00:00
|
|
|
{adaptor.getSpmatA(), adaptor.getDnmatB(), adaptor.getDnmatC(),
|
|
|
|
|
pBufs[0], pBufs[1], pBufs[2], stream});
|
2023-06-06 23:12:45 +00:00
|
|
|
} else {
|
|
|
|
|
Value pBuf = MemRefDescriptor(adaptor.getBuffers().front())
|
|
|
|
|
.allocatedPtr(rewriter, loc);
|
2023-06-16 21:38:40 +00:00
|
|
|
createSpMMCallBuilder.create(loc, rewriter,
|
|
|
|
|
{modeA, modeB, adaptor.getSpmatA(),
|
|
|
|
|
adaptor.getDnmatB(), adaptor.getDnmatC(),
|
|
|
|
|
computeType, pBuf, stream});
|
2023-06-06 23:12:45 +00:00
|
|
|
}
|
2023-05-15 14:55:45 -07:00
|
|
|
rewriter.replaceOp(op, {stream});
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2023-05-24 09:43:38 -07:00
|
|
|
template <typename T>
|
|
|
|
|
static void addOpaquePointerConversion(LLVMTypeConverter &converter) {
|
|
|
|
|
converter.addConversion([&converter](T) -> Type {
|
2023-10-31 09:22:44 +01:00
|
|
|
return LLVM::LLVMPointerType::get(&converter.getContext());
|
2023-05-24 09:43:38 -07:00
|
|
|
});
|
|
|
|
|
}
|
|
|
|
|
|
2023-05-24 02:21:55 +00:00
|
|
|
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();
|
2023-06-12 23:24:08 +00:00
|
|
|
auto computeType = genConstInt32From(
|
|
|
|
|
rewriter, loc, getCuSparseDataTypeFrom(adaptor.getComputeType()));
|
2023-05-26 21:46:22 +00:00
|
|
|
auto modeA = genConstInt32From(rewriter, loc, adaptor.getModeA());
|
|
|
|
|
auto modeB = genConstInt32From(rewriter, loc, adaptor.getModeB());
|
2023-05-24 02:21:55 +00:00
|
|
|
auto stream = adaptor.getAsyncDependencies().front();
|
|
|
|
|
Value pBuf =
|
|
|
|
|
MemRefDescriptor(adaptor.getBuffer()).allocatedPtr(rewriter, loc);
|
2023-06-16 21:38:40 +00:00
|
|
|
createSDDMMCallBuilder.create(loc, rewriter,
|
|
|
|
|
{modeA, modeB, adaptor.getDnmatA(),
|
|
|
|
|
adaptor.getDnmatB(), adaptor.getSpmatC(),
|
|
|
|
|
computeType, pBuf, stream});
|
2023-05-24 02:21:55 +00:00
|
|
|
rewriter.replaceOp(op, {stream});
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2023-08-01 00:57:47 +00:00
|
|
|
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();
|
2023-08-10 12:36:15 -07:00
|
|
|
createSpGEMMDestroyDescrBuilder.create(loc, rewriter,
|
|
|
|
|
{adaptor.getDesc(), stream});
|
2023-08-01 00:57:47 +00:00
|
|
|
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(),
|
2023-08-09 12:19:49 -07:00
|
|
|
adaptor.getSpmatB(), adaptor.getSpmatC(), computeType,
|
2023-08-01 00:57:47 +00:00
|
|
|
adaptor.getBufferSz(), pBuf, stream})
|
|
|
|
|
.getResult();
|
|
|
|
|
} else {
|
|
|
|
|
bufferSizeNew =
|
|
|
|
|
createSpGEMMComputeBuilder
|
|
|
|
|
.create(loc, rewriter,
|
|
|
|
|
{adaptor.getDesc(), modeA, modeB, adaptor.getSpmatA(),
|
2023-08-09 12:19:49 -07:00
|
|
|
adaptor.getSpmatB(), adaptor.getSpmatC(), computeType,
|
2023-08-01 00:57:47 +00:00
|
|
|
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();
|
2023-08-09 12:19:49 -07:00
|
|
|
createSpGEMMCopyBuilder.create(loc, rewriter,
|
|
|
|
|
{adaptor.getDesc(), modeA, modeB,
|
|
|
|
|
adaptor.getSpmatA(), adaptor.getSpmatB(),
|
|
|
|
|
adaptor.getSpmatC(), computeType, stream});
|
2023-08-01 00:57:47 +00:00
|
|
|
rewriter.replaceOp(op, {stream});
|
|
|
|
|
return success();
|
|
|
|
|
}
|
|
|
|
|
|
2023-08-14 13:43:29 -07:00
|
|
|
LogicalResult ConvertSpMatGetSizeOpToGpuRuntimeCallPattern::matchAndRewrite(
|
|
|
|
|
gpu::SpMatGetSizeOp op, OpAdaptor adaptor,
|
2023-08-01 00:57:47 +00:00
|
|
|
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>(
|
2023-10-31 09:22:44 +01:00
|
|
|
loc, llvmPointerType, llvmInt64Type, three, /*alignment=*/16);
|
2023-08-01 00:57:47 +00:00
|
|
|
|
|
|
|
|
auto rowsPtr = rewriter.create<LLVM::GEPOp>(
|
2023-10-31 09:22:44 +01:00
|
|
|
loc, llvmPointerType, llvmPointerType, buffer,
|
2023-08-01 00:57:47 +00:00
|
|
|
ValueRange{rewriter.create<LLVM::ConstantOp>(loc, getIndexType(),
|
|
|
|
|
rewriter.getIndexAttr(0))});
|
|
|
|
|
auto colsPtr = rewriter.create<LLVM::GEPOp>(
|
2023-10-31 09:22:44 +01:00
|
|
|
loc, llvmPointerType, llvmPointerType, buffer,
|
2023-08-01 00:57:47 +00:00
|
|
|
ValueRange{rewriter.create<LLVM::ConstantOp>(loc, getIndexType(),
|
|
|
|
|
rewriter.getIndexAttr(1))});
|
|
|
|
|
auto nnzsPtr = rewriter.create<LLVM::GEPOp>(
|
2023-10-31 09:22:44 +01:00
|
|
|
loc, llvmPointerType, llvmPointerType, buffer,
|
2023-08-01 00:57:47 +00:00
|
|
|
ValueRange{rewriter.create<LLVM::ConstantOp>(loc, getIndexType(),
|
|
|
|
|
rewriter.getIndexAttr(2))});
|
2023-08-14 13:43:29 -07:00
|
|
|
createSpMatGetSizeBuilder.create(
|
2023-08-01 00:57:47 +00:00
|
|
|
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();
|
|
|
|
|
}
|
|
|
|
|
|
2023-08-10 12:36:15 -07:00
|
|
|
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();
|
|
|
|
|
}
|
|
|
|
|
|
2023-09-27 09:32:25 -07:00
|
|
|
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();
|
|
|
|
|
}
|
|
|
|
|
|
2022-07-14 13:31:47 -07:00
|
|
|
void mlir::populateGpuToLLVMConversionPatterns(LLVMTypeConverter &converter,
|
|
|
|
|
RewritePatternSet &patterns,
|
2022-07-11 18:29:01 +00:00
|
|
|
StringRef gpuBinaryAnnotation,
|
2023-08-12 00:26:13 +00:00
|
|
|
bool kernelBarePtrCallConv,
|
|
|
|
|
SymbolTable *cachedModuleTable) {
|
2023-05-24 09:43:38 -07:00
|
|
|
addOpaquePointerConversion<gpu::AsyncTokenType>(converter);
|
2023-06-09 16:58:44 +00:00
|
|
|
addOpaquePointerConversion<gpu::SparseDnTensorHandleType>(converter);
|
2023-05-24 09:43:38 -07:00
|
|
|
addOpaquePointerConversion<gpu::SparseSpMatHandleType>(converter);
|
2023-08-01 00:57:47 +00:00
|
|
|
addOpaquePointerConversion<gpu::SparseSpGEMMOpHandleType>(converter);
|
2023-05-11 13:56:34 -07:00
|
|
|
|
2021-07-28 22:31:26 +03:00
|
|
|
patterns.add<ConvertAllocOpToGpuRuntimeCallPattern,
|
|
|
|
|
ConvertDeallocOpToGpuRuntimeCallPattern,
|
|
|
|
|
ConvertHostRegisterOpToGpuRuntimeCallPattern,
|
2023-04-06 15:07:12 -05:00
|
|
|
ConvertHostUnregisterOpToGpuRuntimeCallPattern,
|
2021-07-28 22:31:26 +03:00
|
|
|
ConvertMemcpyOpToGpuRuntimeCallPattern,
|
2021-09-04 08:03:33 +02:00
|
|
|
ConvertMemsetOpToGpuRuntimeCallPattern,
|
2022-02-15 20:23:44 +00:00
|
|
|
ConvertSetDefaultDeviceOpToGpuRuntimeCallPattern,
|
2021-07-28 22:31:26 +03:00
|
|
|
ConvertWaitAsyncOpToGpuRuntimeCallPattern,
|
|
|
|
|
ConvertWaitOpToGpuRuntimeCallPattern,
|
2023-05-11 13:56:34 -07:00
|
|
|
ConvertAsyncYieldToGpuRuntimeCallPattern,
|
2023-06-09 16:58:44 +00:00
|
|
|
ConvertCreateDnTensorOpToGpuRuntimeCallPattern,
|
|
|
|
|
ConvertDestroyDnTensorOpToGpuRuntimeCallPattern,
|
2023-05-11 13:56:34 -07:00
|
|
|
ConvertCreateCooOpToGpuRuntimeCallPattern,
|
2023-06-06 10:56:58 -07:00
|
|
|
ConvertCreateCooAoSOpToGpuRuntimeCallPattern,
|
2023-05-11 13:56:34 -07:00
|
|
|
ConvertCreateCsrOpToGpuRuntimeCallPattern,
|
2023-09-27 09:32:25 -07:00
|
|
|
ConvertCreateCscOpToGpuRuntimeCallPattern,
|
|
|
|
|
ConvertCreateBsrOpToGpuRuntimeCallPattern,
|
2023-06-06 23:12:45 +00:00
|
|
|
ConvertCreate2To4SpMatOpToGpuRuntimeCallPattern,
|
2023-05-11 13:56:34 -07:00
|
|
|
ConvertDestroySpMatOpToGpuRuntimeCallPattern,
|
2023-08-08 10:24:03 -07:00
|
|
|
ConvertSpMVBufferSizeOpToGpuRuntimeCallPattern,
|
|
|
|
|
ConvertSpMVOpToGpuRuntimeCallPattern,
|
|
|
|
|
ConvertSpMMBufferSizeOpToGpuRuntimeCallPattern,
|
|
|
|
|
ConvertSDDMMBufferSizeOpToGpuRuntimeCallPattern,
|
|
|
|
|
ConvertSpMMOpToGpuRuntimeCallPattern,
|
|
|
|
|
ConvertSDDMMOpToGpuRuntimeCallPattern,
|
2023-08-01 00:57:47 +00:00
|
|
|
ConvertSpGEMMCreateDescrOpToGpuRuntimeCallPattern,
|
|
|
|
|
ConvertSpGEMMDestroyDescrOpToGpuRuntimeCallPattern,
|
|
|
|
|
ConvertSpGEMMWorkEstimationOrComputeOpToGpuRuntimeCallPattern,
|
|
|
|
|
ConvertSpGEMMCopyOpToGpuRuntimeCallPattern,
|
2023-08-14 13:43:29 -07:00
|
|
|
ConvertSpMatGetSizeOpToGpuRuntimeCallPattern,
|
2023-08-10 12:36:15 -07:00
|
|
|
ConvertSetCsrPointersOpToGpuRuntimeCallPattern>(converter);
|
2022-07-11 18:29:01 +00:00
|
|
|
patterns.add<ConvertLaunchFuncOpToGpuRuntimeCallPattern>(
|
2023-08-12 00:26:13 +00:00
|
|
|
converter, gpuBinaryAnnotation, kernelBarePtrCallConv, cachedModuleTable);
|
2021-07-28 22:31:26 +03:00
|
|
|
patterns.add<EraseGpuModuleOpPattern>(&converter.getContext());
|
|
|
|
|
}
|