JinjieLiu/triton-tvm-ffi
1
0
Fork 0
mirror of https://github.com/sgjzfzzf/triton-tvm-ffi.git synced 2026-07-01 08:51:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

put utils in C++ side

Browse Source 
Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>
This commit is contained in:
Jinjie Liu
2026-01-29 18:23:02 +08:00
parent 781a5396cc
commit 1a01c9f2d8
7 changed files with 137 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/utils.cc
+66 -48
View File
@@ -2,8 +2,8 @@
#include "type.h"
#include "value.h"
#include <cuda.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/extra/cuda/cubin_launcher.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/tvm_ffi.h>
#define CUDA_CHECK(code) \
@@ -101,49 +101,6 @@ tvm::ffi::Map<tvm::ffi::String, int32_t> GetDeviceProperties(int device_id) {
{"mem_bus_width", memBusWidth}};
}
tvm::ffi::Tuple<uint64_t, uint64_t, int32_t, int32_t, int32_t>
LoadBinary(const tvm::ffi::String &name, const tvm::ffi::Bytes &data,
int32_t shared, CUdevice device) {
CUcontext pctx;
CUfunction fun;
CUmodule mod;
int32_t nRegs = 0;
int32_t nSpills = 0;
int32_t nMaxThreads = 0;
int32_t sharedOptin = 0;
CUDA_CHECK(cuCtxGetCurrent(&pctx));
if (!pctx) {
CUDA_CHECK(cuDevicePrimaryCtxRetain(&pctx, device));
CUDA_CHECK(cuCtxSetCurrent(pctx));
}
CUDA_CHECK(cuModuleLoadData(&mod, data.data()));
CUDA_CHECK(cuModuleGetFunction(&fun, mod, name.data()));
CUDA_CHECK(cuFuncGetAttribute(&nRegs, CU_FUNC_ATTRIBUTE_NUM_REGS, fun));
CUDA_CHECK(
cuFuncGetAttribute(&nSpills, CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, fun));
CUDA_CHECK(cuFuncGetAttribute(&nMaxThreads,
CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK, fun));
CUDA_CHECK(cuDeviceGetAttribute(
&sharedOptin, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN,
device));
static constexpr int64_t kExpectedMaxDynamicSharedMemory = 49152;
if (shared > kExpectedMaxDynamicSharedMemory &&
sharedOptin > kExpectedMaxDynamicSharedMemory) {
CUDA_CHECK(cuFuncSetCacheConfig(fun, CU_FUNC_CACHE_PREFER_SHARED));
int32_t sharedTotal = 0, sharedStatic = 0;
CUDA_CHECK(cuDeviceGetAttribute(
&sharedTotal, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_MULTIPROCESSOR,
device));
CUDA_CHECK(cuFuncGetAttribute(&sharedStatic,
CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, fun));
CUDA_CHECK(
cuFuncSetAttribute(fun, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
sharedOptin - sharedStatic));
}
return tvm::ffi::Tuple<uint64_t, uint64_t, int32_t, int32_t, int32_t>{
mod, fun, nRegs, nSpills, nMaxThreads};
}
void Launch(tvm::ffi::PackedArgs args, tvm::ffi::Any *ret) {
CUtensorMap x;
int32_t gridX = args[0].cast<int32_t>();
@@ -218,6 +175,7 @@ void Launch(tvm::ffi::PackedArgs args, tvm::ffi::Any *ret) {
}
}
}
// TODO: unwrap PyObject* from scratch pointers and assign to kernel args
params[j] = &globalScratch;
params[j + 1] = &profileScratch;
CUDA_CHECK(cuLaunchKernelEx(&config, function, params, nullptr));
@@ -225,12 +183,72 @@ void Launch(tvm::ffi::PackedArgs args, tvm::ffi::Any *ret) {
// TODO: call `launchExitHook`
}
tvm::ffi::Tuple<uint64_t, uint64_t, int32_t, int32_t, int32_t>
LoadBinary(const tvm::ffi::String &name, const tvm::ffi::Bytes &data,
int32_t shared, CUdevice device) {
CUcontext pctx;
CUfunction fun;
CUmodule mod;
int32_t nRegs = 0;
int32_t nSpills = 0;
int32_t nMaxThreads = 0;
int32_t sharedOptin = 0;
CUDA_CHECK(cuCtxGetCurrent(&pctx));
if (!pctx) {
CUDA_CHECK(cuDevicePrimaryCtxRetain(&pctx, device));
CUDA_CHECK(cuCtxSetCurrent(pctx));
}
CUDA_CHECK(cuModuleLoadData(&mod, data.data()));
CUDA_CHECK(cuModuleGetFunction(&fun, mod, name.data()));
CUDA_CHECK(cuFuncGetAttribute(&nRegs, CU_FUNC_ATTRIBUTE_NUM_REGS, fun));
CUDA_CHECK(
cuFuncGetAttribute(&nSpills, CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, fun));
CUDA_CHECK(cuFuncGetAttribute(&nMaxThreads,
CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK, fun));
CUDA_CHECK(cuDeviceGetAttribute(
&sharedOptin, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN,
device));
static constexpr int64_t kExpectedMaxDynamicSharedMemory = 49152;
if (shared > kExpectedMaxDynamicSharedMemory &&
sharedOptin > kExpectedMaxDynamicSharedMemory) {
CUDA_CHECK(cuFuncSetCacheConfig(fun, CU_FUNC_CACHE_PREFER_SHARED));
int32_t sharedTotal = 0, sharedStatic = 0;
CUDA_CHECK(cuDeviceGetAttribute(
&sharedTotal, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_MULTIPROCESSOR,
device));
CUDA_CHECK(cuFuncGetAttribute(&sharedStatic,
CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, fun));
CUDA_CHECK(
cuFuncSetAttribute(fun, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
sharedOptin - sharedStatic));
}
return tvm::ffi::Tuple<uint64_t, uint64_t, int32_t, int32_t, int32_t>{
mod, fun, nRegs, nSpills, nMaxThreads};
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.def("triton_tvm_ffi.get_device_properties", GetDeviceProperties)
.def_packed("triton_tvm_ffi.launch", Launch)
.def("triton_tvm_ffi.load_binary", LoadBinary);
.def_packed("triton_tvm_ffi.utils.build_signature_metadata",
[](tvm::ffi::PackedArgs args, tvm::ffi::Any *ret) {
throw NotImplementedException("build_signature_metadata");
})
.def_packed("triton_tvm_ffi.utils.cuOccupancyMaxActiveClusters",
[](tvm::ffi::PackedArgs args, tvm::ffi::Any *ret) {
throw NotImplementedException(
"cuOccupancyMaxActiveClusters");
})
.def_packed("triton_tvm_ffi.utils.fill_tma_descriptor",
[](tvm::ffi::PackedArgs args, tvm::ffi::Any *ret) {
throw NotImplementedException("fill_tma_descriptor");
})
.def_packed("triton_tvm_ffi.utils.set_printf_fifo_size",
[](tvm::ffi::PackedArgs args, tvm::ffi::Any *ret) {
throw NotImplementedException("set_printf_fifo_size");
})
.def("triton_tvm_ffi.utils.get_device_properties", GetDeviceProperties)
.def_packed("triton_tvm_ffi.utils.launch", Launch)
.def("triton_tvm_ffi.utils.load_binary", LoadBinary);
}
} // namespace triton_tvm_ffi
} // namespace triton_tvm_ffi