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

put num_warps and num_stages in kwargs

Browse Source 
Signed-off-by: jinjieliu <jinjie.liu@usc.edu>
This commit is contained in:
jinjieliu
2026-02-07 14:25:10 +08:00
parent 2298b6f8c8
commit 6a19a6b06d
5 changed files with 10 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
examples/add/add.cc
View File

@@ -5,7 +5,7 @@
#include <tvm/ffi/tvm_ffi.h>
#ifndef ADD_KERNEL_STUB
#define ADD_KERNEL_STUB(grid, stream, numWarps, numStages, args, kwargs)
#define ADD_KERNEL_STUB(grid, stream, args, kwargs)
#endif
#ifndef ADD_NAME
@@ -23,12 +23,11 @@ tvm::ffi::Tensor Add(tvm::ffi::Tensor x, tvm::ffi::Tensor y) {
const int32_t BLOCK_SIZE = meta["BLOCK_SIZE"].cast<int32_t>();
return tvm::ffi::Tuple((numel + BLOCK_SIZE - 1) / BLOCK_SIZE, 1, 1);
});
tvm::ffi::Optional<int32_t> numWarps = std::nullopt, numStages = std::nullopt;
DLDevice device = x.device();
void *stream = TVMFFIEnvGetStream(device.device_type, device.device_id);
tvm::ffi::Array<tvm::ffi::Any> args = {x, y, output, numel, 1024};
tvm::ffi::Map<tvm::ffi::String, tvm::ffi::Any> kwargs = {};
ADD_KERNEL_STUB(grid, stream, numWarps, numStages, args, kwargs);
ADD_KERNEL_STUB(grid, stream, args, kwargs);
return output;
}