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/driver.py
+33 -5
View File
@@ -34,7 +34,35 @@ class TVMLauncher(object):
],
)
launch: tvm_ffi.Function = mod.get_function("launch")
self.launch = launch
self.launch = (
lambda grid_x,
grid_y,
grid_z,
stream,
function,
num_warps,
num_ctas,
shared_memory,
global_scratch,
profile_scratch,
*args: launch(
grid_x,
grid_y,
grid_z,
stream,
function,
num_warps,
num_ctas,
shared_memory,
global_scratch,
profile_scratch,
*(
arg
for arg, type in zip(args, self.signature)
if type != "constexpr"
),
)
)
else:
self.impl: TVMFFILauncherImpl = TVMFFILauncherImpl(
[string_to_type(t) for t in self.signature],
@@ -129,19 +157,19 @@ class TVMLauncher(object):
@cached_property
def codegen(self) -> str:
env: Final[jinja2.Environment] = jinja2.Environment(
env: jinja2.Environment = jinja2.Environment(
loader=jinja2.PackageLoader("triton_tvm_ffi", "templates"),
trim_blocks=True,
lstrip_blocks=True,
)
template = env.get_template("launch.c.j2")
signature = list(
template: jinja2.Template = env.get_template("launch.c.j2")
signature: List[int] = list(
filter(
lambda t: t != "void",
map(lambda t: type_to_ctype(string_to_type(t)), self.signature),
)
)
html = template.render(signature=signature)
html: str = template.render(signature=signature)
return html
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);
}
}