JinjieLiu/triton-tvm-ffi
1
0
Fork 0
mirror of https://github.com/sgjzfzzf/triton-tvm-ffi.git synced 2026-05-02 03:52:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

enable lambda function for grid descriptor

Browse Source 
Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>
This commit is contained in:
Jinjie Liu
2026-02-05 15:59:22 +08:00
parent 8b8aa6cb84
commit f6c7a48c1b
6 changed files with 104 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
examples/add/add.py
View File

@@ -33,8 +33,9 @@ def add_triton(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE
n_elements: int = output.numel()
BLOCK_SIZE: int = 1024
grid = (triton.cdiv(n_elements, BLOCK_SIZE), 1, 1)
add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE)
add_kernel[lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), 1, 1)](
x, y, output, n_elements, BLOCK_SIZE
)
return output