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 optional for numwarps and numstages

Browse Source 
Signed-off-by: jinjieliu <jinjie.liu@usc.edu>
This commit is contained in:
jinjieliu
2026-02-05 01:01:49 +08:00
parent b7bf598fde
commit 8b8aa6cb84
3 changed files with 13 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
examples/add/add.cc
View File

@@ -18,7 +18,8 @@ tvm::ffi::Tensor Add(tvm::ffi::Tensor x, tvm::ffi::Tensor y) {
int64_t numel = otorch.numel();
tvm::ffi::Tensor output = tvm::ffi::Tensor::FromDLPack(at::toDLPack(otorch));
tvm::ffi::Tuple<int32_t, int32_t, int32_t> grid{(numel + 1023) / 1024, 1, 1};
size_t numWarps = 4, numStages = 3;
// TODO: check the performance loss after enabling `Optional`
tvm::ffi::Optional<int32_t> numWarps = std::nullopt, numStages = std::nullopt;
DLDevice device = x.device();
void *stream = TVMFFIEnvGetStream(device.device_type, device.device_id);
ADD_KERNEL_STUB(grid, stream, numWarps, numStages, x, y, output, numel, 1024);