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

fix bugs on failure of mm example

Browse Source 
Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>
This commit is contained in:
Jinjie Liu
2026-01-31 17:48:59 +08:00
parent e9576d265e
commit 875145b333
2 changed files with 39 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/triton_tvm_ffi/templates/launch.c.j2
+6 -6
View File
@@ -1,11 +1,13 @@
#include <cassert>
#include <assert.h>
#include <cuda.h>
#include <tvm/ffi/tvm_ffi.h>
#ifdef __cplusplus
extern "C"
#endif
TVM_FFI_DLL_EXPORT void __tvm_ffi_launch(void *handle, const TVMFFIAny *args, int32_t num_args, TVMFFIAny *result) {
TVM_FFI_DLL_EXPORT void
__tvm_ffi_launch(void *handle, const TVMFFIAny *args, int32_t num_args,
TVMFFIAny *result) {
int32_t gridX = args[0].v_int64;
int32_t gridY = args[1].v_int64;
int32_t gridZ = args[2].v_int64;
@@ -22,8 +24,7 @@ TVM_FFI_DLL_EXPORT void __tvm_ffi_launch(void *handle, const TVMFFIAny *args, in
config.gridDimX = gridX * numCtas;
config.gridDimY = gridY;
config.gridDimZ = gridZ;
static constexpr int32_t kThreadsPerWarp = 32;
config.blockDimX = kThreadsPerWarp * numWarps;
config.blockDimX = 32 * numWarps;
config.blockDimY = 1;
config.blockDimZ = 1;
config.sharedMemBytes = sharedMemory;
@@ -56,8 +57,7 @@ TVM_FFI_DLL_EXPORT void __tvm_ffi_launch(void *handle, const TVMFFIAny *args, in
assert(false, "unsupported type yet {{ type }}");
{% endif %}
{% endfor %}
void *foo = NULL, *bar = NULL;
void *params[] = { {% for type in signature %} &arg{{ loop.index0 }}, {% endfor %}&foo, &bar };
void *params[] = { {% for type in signature %} &arg{{ loop.index0 }}, {% endfor %}&globalScratch, &profileScratch };
cuLaunchKernelEx(&config, function, params, NULL);
}
}