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

remove CudaUtils dependencies

Browse Source 
Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>
This commit is contained in:
Jinjie Liu
2026-01-28 14:30:43 +08:00
parent 15b4b23ec8
commit b79d880adf
3 changed files with 65 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/utils.cc
+47 -2
View File
@@ -1,6 +1,7 @@
#include "exception.h"
#include <cuda.h>
#include <tvm/ffi/extra/cuda/cubin_launcher.h>
#include <tvm/ffi/string.h>
#include <tvm/ffi/tvm_ffi.h>
#define CUDA_CHECK(code) \
@@ -44,8 +45,52 @@ tvm::ffi::Map<tvm::ffi::String, int32_t> GetDeviceProperties(int device_id) {
{"mem_bus_width", mem_bus_width}};
}
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.utils.get_device_properties",
GetDeviceProperties);
refl::GlobalDef()
.def("triton_tvm_ffi.utils.get_device_properties", GetDeviceProperties)
.def("triton_tvm_ffi.utils.load_binary", LoadBinary);
}