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

implement launch

Browse Source 
Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>
This commit is contained in:
Jinjie Liu
2026-01-29 01:54:16 +08:00
parent b79d880adf
commit fecf48e403
3 changed files with 193 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/utils.cc
+114 -23
View File
@@ -1,7 +1,7 @@
#include "exception.h"
#include <cassert>
#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) \
@@ -14,35 +14,35 @@
tvm::ffi::Map<tvm::ffi::String, int32_t> GetDeviceProperties(int device_id) {
tvm::ffi::cuda_api::DeviceHandle device;
CUDA_CHECK(cuDeviceGet(&device, device_id));
int max_shared_mem = 0;
int max_num_regs = 0;
int multiprocessor_count = 0;
int warp_size = 0;
int sm_clock_rate = 0;
int mem_clock_rate = 0;
int mem_bus_width = 0;
int maxSharedMem = 0;
int maxNumRegs = 0;
int multiprocessorCount = 0;
int warpSize = 0;
int smClockRate = 0;
int memClockRate = 0;
int memBusWidth = 0;
CUDA_CHECK(cuDeviceGetAttribute(
&max_shared_mem, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN,
&maxSharedMem, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN,
device));
CUDA_CHECK(cuDeviceGetAttribute(
&max_num_regs, CU_DEVICE_ATTRIBUTE_MAX_REGISTERS_PER_BLOCK, device));
&maxNumRegs, CU_DEVICE_ATTRIBUTE_MAX_REGISTERS_PER_BLOCK, device));
CUDA_CHECK(cuDeviceGetAttribute(
&multiprocessor_count, CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, device));
&multiprocessorCount, CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, device));
CUDA_CHECK(
cuDeviceGetAttribute(&warp_size, CU_DEVICE_ATTRIBUTE_WARP_SIZE, device));
CUDA_CHECK(cuDeviceGetAttribute(&sm_clock_rate,
CU_DEVICE_ATTRIBUTE_CLOCK_RATE, device));
cuDeviceGetAttribute(&warpSize, CU_DEVICE_ATTRIBUTE_WARP_SIZE, device));
CUDA_CHECK(cuDeviceGetAttribute(&smClockRate, CU_DEVICE_ATTRIBUTE_CLOCK_RATE,
device));
CUDA_CHECK(cuDeviceGetAttribute(
&mem_clock_rate, CU_DEVICE_ATTRIBUTE_MEMORY_CLOCK_RATE, device));
&memClockRate, CU_DEVICE_ATTRIBUTE_MEMORY_CLOCK_RATE, device));
CUDA_CHECK(cuDeviceGetAttribute(
&mem_bus_width, CU_DEVICE_ATTRIBUTE_GLOBAL_MEMORY_BUS_WIDTH, device));
return {{"max_shared_mem", max_shared_mem},
{"max_num_regs", max_num_regs},
{"multiprocessor_count", multiprocessor_count},
{"warp_size", warp_size},
{"sm_clock_rate", sm_clock_rate},
{"mem_clock_rate", mem_clock_rate},
{"mem_bus_width", mem_bus_width}};
&memBusWidth, CU_DEVICE_ATTRIBUTE_GLOBAL_MEMORY_BUS_WIDTH, device));
return {{"max_shared_mem", maxSharedMem},
{"max_num_regs", maxNumRegs},
{"multiprocessor_count", multiprocessorCount},
{"warpSize", warpSize},
{"sm_clock_rate", smClockRate},
{"mem_clock_rate", memClockRate},
{"mem_bus_width", memBusWidth}};
}
tvm::ffi::Tuple<uint64_t, uint64_t, int32_t, int32_t, int32_t>
@@ -88,9 +88,100 @@ LoadBinary(const tvm::ffi::String &name, const tvm::ffi::Bytes &data,
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>();
int32_t gridY = args[1].cast<int32_t>();
int32_t gridZ = args[2].cast<int32_t>();
CUstream stream = reinterpret_cast<CUstream>(args[3].cast<uint64_t>());
CUfunction function = reinterpret_cast<CUfunction>(args[4].cast<uint64_t>());
tvm::ffi::Tuple<int32_t, int32_t, int32_t> kernelMetadata =
args[5].cast<tvm::ffi::Tuple<int32_t, int32_t, int32_t>>();
int32_t numWarps = kernelMetadata.get<0>();
int32_t numCtas = kernelMetadata.get<1>();
int32_t sharedMemory = kernelMetadata.get<2>();
tvm::ffi::ObjectRef launchMetadata = args[6].cast<tvm::ffi::ObjectRef>();
tvm::ffi::ObjectRef launchEnterHook = args[7].cast<tvm::ffi::ObjectRef>();
tvm::ffi::ObjectRef launchExitHook = args[8].cast<tvm::ffi::ObjectRef>();
bool launchCooperativeGrid = args[9].cast<bool>();
bool launchPdl = args[10].cast<bool>();
tvm::ffi::ObjectRef globalScratchObject =
args[11].cast<tvm::ffi::ObjectRef>();
tvm::ffi::ObjectRef profileScratchObject =
args[12].cast<tvm::ffi::ObjectRef>();
tvm::ffi::PackedArgs kernelArgs = args.Slice(13);
// TODO: call `launchEnterHook`
// TODO: check `globalScratchObject`
CUdeviceptr globalScratch = 0;
// TODO: check `profileScratchObject`
CUdeviceptr profileScratch = 0;
if (gridX * gridY * gridZ > 0) {
CUlaunchAttribute launchAttr[4];
CUlaunchConfig config;
config.gridDimX = gridX * numCtas;
config.gridDimY = gridY;
config.gridDimZ = gridZ;
static constexpr int32_t kThreadsPerWarp = 32;
config.blockDimX = kThreadsPerWarp * numWarps;
config.blockDimY = 1;
config.blockDimZ = 1;
config.sharedMemBytes = sharedMemory;
config.hStream = stream;
config.attrs = launchAttr;
int32_t numAttrs = 0;
// TODO: check `launchPdf`
// TODO: check `launchCooperativeGrid`
if (numCtas != 1) {
CUlaunchAttribute clusterAttr;
clusterAttr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
clusterAttr.value.clusterDim.x = numCtas;
clusterAttr.value.clusterDim.y = 1;
clusterAttr.value.clusterDim.z = 1;
launchAttr[numAttrs++] = clusterAttr;
CUlaunchAttribute clusterSchedulingAttr;
clusterSchedulingAttr.id =
CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE;
clusterSchedulingAttr.value.clusterSchedulingPolicyPreference =
CU_CLUSTER_SCHEDULING_POLICY_SPREAD;
launchAttr[numAttrs++] = clusterSchedulingAttr;
}
config.numAttrs = numAttrs;
if (numCtas == 16) {
CUDA_CHECK(cuFuncSetAttribute(
function, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1));
}
const int32_t kernelArgNum = kernelArgs.size();
void **params =
reinterpret_cast<void **>(alloca(sizeof(void *) * (kernelArgNum + 2)));
for (size_t i = 0; i < kernelArgNum; ++i) {
tvm::ffi::AnyView arg = kernelArgs[i];
if (auto val = arg.try_cast<tvm::ffi::TensorView>()) {
void **ptr = reinterpret_cast<void **>(alloca(sizeof(void *)));
*ptr = val->data_ptr();
params[i] = ptr;
} else if (auto val = arg.try_cast<int32_t>()) {
int32_t *ptr = reinterpret_cast<int32_t *>(alloca(sizeof(int32_t)));
*ptr = *val;
params[i] = ptr;
} else if (auto val = arg.try_cast<float>()) {
float *ptr = reinterpret_cast<float *>(alloca(sizeof(float)));
*ptr = *val;
params[i] = ptr;
} else {
assert(false && "unsupported kernel argument type");
}
}
params[kernelArgNum] = &globalScratch;
params[kernelArgNum + 1] = &profileScratch;
CUDA_CHECK(cuLaunchKernelEx(&config, function, params, nullptr));
}
// TODO: call `launchExitHook`
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.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);
}