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 void branch in template

Browse Source 
Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>
This commit is contained in:
Jinjie Liu
2026-02-01 00:27:49 +08:00
parent 875145b333
commit 9a237d40e2
2 changed files with 3 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/triton_tvm_ffi/driver.py
+1 -8
View File
@@ -56,11 +56,7 @@ class TVMLauncher(object):
shared_memory, shared_memory,
global_scratch, global_scratch,
profile_scratch, profile_scratch,
*( *args,
arg
for arg, type in zip(args, self.signature)
if type != "constexpr"
),
) )
) )
else: else:
@@ -164,10 +160,7 @@ class TVMLauncher(object):
) )
template: jinja2.Template = env.get_template("launch.c.j2") template: jinja2.Template = env.get_template("launch.c.j2")
signature: List[int] = list( signature: List[int] = list(
filter(
lambda t: t != "void",
map(lambda t: type_to_ctype(string_to_type(t)), self.signature), map(lambda t: type_to_ctype(string_to_type(t)), self.signature),
)
) )
html: str = template.render(signature=signature) html: str = template.render(signature=signature)
return html return html
python/triton_tvm_ffi/templates/launch.c.j2
+2 -1
View File
@@ -53,11 +53,12 @@ extern "C"
{{ type }} arg{{ loop.index0 }} = ((DLTensor*)(args[{{ loop.index0 + 10 }}].v_c_str + sizeof(TVMFFIObject)))->data; {{ type }} arg{{ loop.index0 }} = ((DLTensor*)(args[{{ loop.index0 + 10 }}].v_c_str + sizeof(TVMFFIObject)))->data;
{% elif type == "int32_t" %} {% elif type == "int32_t" %}
{{ type }} arg{{ loop.index0 }} = args[{{ loop.index0 + 10 }}].v_int64; {{ type }} arg{{ loop.index0 }} = args[{{ loop.index0 + 10 }}].v_int64;
{% elif type == "void" %}
{% else %} {% else %}
assert(false, "unsupported type yet {{ type }}"); assert(false, "unsupported type yet {{ type }}");
{% endif %} {% endif %}
{% endfor %} {% endfor %}
void *params[] = { {% for type in signature %} &arg{{ loop.index0 }}, {% endfor %}&globalScratch, &profileScratch }; void *params[] = { {% for type in signature %} {% if type != "void" %} &arg{{ loop.index0 }}, {% endif %} {% endfor %}&globalScratch, &profileScratch };
cuLaunchKernelEx(&config, function, params, NULL); cuLaunchKernelEx(&config, function, params, NULL);
} }
} }