mirror of
https://github.com/sgjzfzzf/triton-tvm-ffi.git
synced 2026-05-02 03:52:11 +08:00
56 lines
2.2 KiB
C++
56 lines
2.2 KiB
C++
#include <ATen/DLConvertor.h>
|
|
#include <ATen/dlpack.h>
|
|
#include <tvm/ffi/extra/cuda/cubin_launcher.h>
|
|
#include <tvm/ffi/tvm_ffi.h>
|
|
|
|
#ifndef MATMUL_KERNEL_STUB
|
|
#define MATMUL_KERNEL_STUB(grid, device, stream, args, kwargs)
|
|
#endif
|
|
|
|
#ifndef MATMUL_NAME
|
|
#define MATMUL_NAME ""
|
|
#endif
|
|
|
|
tvm::ffi::Tensor Matmul(tvm::ffi::Tensor a, tvm::ffi::Tensor b,
|
|
tvm::ffi::String activation) {
|
|
at::Tensor atorch = at::fromDLPack(a.ToDLPack()),
|
|
btorch = at::fromDLPack(b.ToDLPack());
|
|
const int32_t M = atorch.size(0), K = atorch.size(1), N = btorch.size(1);
|
|
at::Tensor ctorch = at::empty({M, N}, atorch.options());
|
|
tvm::ffi::Function grid = tvm::ffi::Function::FromTyped(
|
|
[M, N](const tvm::ffi::Map<tvm::ffi::String, tvm::ffi::Any> &meta)
|
|
-> tvm::ffi::Tuple<int32_t, int32_t, int32_t> {
|
|
const int32_t BLOCK_SIZE_M = meta["BLOCK_SIZE_M"].cast<int32_t>(),
|
|
BLOCK_SIZE_N = meta["BLOCK_SIZE_N"].cast<int32_t>();
|
|
return tvm::ffi::Tuple<int32_t, int32_t, int32_t>{
|
|
(M + BLOCK_SIZE_M - 1) / BLOCK_SIZE_M *
|
|
((N + BLOCK_SIZE_N - 1) / BLOCK_SIZE_N),
|
|
1, 1};
|
|
});
|
|
DLDevice device = a.device();
|
|
void *stream = TVMFFIEnvGetStream(device.device_type, device.device_id);
|
|
tvm::ffi::Tensor c = tvm::ffi::Tensor::FromDLPack(at::toDLPack(ctorch));
|
|
tvm::ffi::Array<tvm::ffi::Any> args = {a,
|
|
b,
|
|
c,
|
|
M,
|
|
N,
|
|
K,
|
|
atorch.stride(0),
|
|
atorch.stride(1),
|
|
btorch.stride(0),
|
|
btorch.stride(1),
|
|
ctorch.stride(0),
|
|
ctorch.stride(1)};
|
|
tvm::ffi::Map<tvm::ffi::String, tvm::ffi::Any> kwargs = {
|
|
{"ACTIVATION", activation},
|
|
};
|
|
MATMUL_KERNEL_STUB(grid, device.device_id, stream, args, kwargs);
|
|
return c;
|
|
}
|
|
|
|
TVM_FFI_STATIC_INIT_BLOCK() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::GlobalDef().def(MATMUL_NAME, Matmul);
|
|
}
|