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

5
examples/mm/mm.cc
View File

@@ -4,7 +4,7 @@
#include <tvm/ffi/tvm_ffi.h> #include <tvm/ffi/tvm_ffi.h>
#ifndef MATMUL_KERNEL_STUB #ifndef MATMUL_KERNEL_STUB
#define MATMUL_KERNEL_STUB(grid, stream, numWarps, numStages, args, kwargs) #define MATMUL_KERNEL_STUB(grid, stream, args, kwargs)
#endif #endif
#ifndef MATMUL_NAME #ifndef MATMUL_NAME
@@ -27,7 +27,6 @@ tvm::ffi::Tensor Matmul(tvm::ffi::Tensor a, tvm::ffi::Tensor b,
((N + BLOCK_SIZE_N - 1) / BLOCK_SIZE_N), ((N + BLOCK_SIZE_N - 1) / BLOCK_SIZE_N),
1, 1}; 1, 1};
}); });
tvm::ffi::Optional<int32_t> numWarps = std::nullopt, numStages = std::nullopt;
DLDevice device = a.device(); DLDevice device = a.device();
void *stream = TVMFFIEnvGetStream(device.device_type, device.device_id); void *stream = TVMFFIEnvGetStream(device.device_type, device.device_id);
tvm::ffi::Tensor c = tvm::ffi::Tensor::FromDLPack(at::toDLPack(ctorch)); tvm::ffi::Tensor c = tvm::ffi::Tensor::FromDLPack(at::toDLPack(ctorch));
@@ -46,7 +45,7 @@ tvm::ffi::Tensor Matmul(tvm::ffi::Tensor a, tvm::ffi::Tensor b,
tvm::ffi::Map<tvm::ffi::String, tvm::ffi::Any> kwargs = { tvm::ffi::Map<tvm::ffi::String, tvm::ffi::Any> kwargs = {
{"ACTIVATION", activation}, {"ACTIVATION", activation},
}; };
MATMUL_KERNEL_STUB(grid, stream, numWarps, numStages, args, kwargs); MATMUL_KERNEL_STUB(grid, stream, args, kwargs);
return c; return c;
} }

12
examples/softmax/softmax.cc
View File

@@ -4,7 +4,7 @@
#include <tvm/ffi/tvm_ffi.h> #include <tvm/ffi/tvm_ffi.h>
#ifndef SOFTMAX_KERNEL_STUB #ifndef SOFTMAX_KERNEL_STUB
#define SOFTMAX_KERNEL_STUB(grid, stream, numWarps, numStages, args, kwargs) #define SOFTMAX_KERNEL_STUB(grid, stream, args, kwargs)
#endif #endif
#ifndef SOFTMAX_NAME #ifndef SOFTMAX_NAME
@@ -14,19 +14,17 @@
tvm::ffi::Tensor Softmax(tvm::ffi::Tensor x) { tvm::ffi::Tensor Softmax(tvm::ffi::Tensor x) {
at::Tensor xtorch = at::fromDLPack(x.ToDLPack()); at::Tensor xtorch = at::fromDLPack(x.ToDLPack());
at::Tensor ytorch = at::empty_like(xtorch); at::Tensor ytorch = at::empty_like(xtorch);
uint32_t nRows = xtorch.size(0), nCols = xtorch.size(1), numWarps = 8, uint32_t nRows = xtorch.size(0), nCols = xtorch.size(1),
numStages = 4, xStride = xtorch.stride(0), xStride = xtorch.stride(0), yStride = ytorch.stride(0),
yStride = ytorch.stride(0),
BLOCK_SIZE = 1u << (32 - __builtin_clz(nCols - 1)); BLOCK_SIZE = 1u << (32 - __builtin_clz(nCols - 1));
tvm::ffi::Tensor y = tvm::ffi::Tensor::FromDLPack(at::toDLPack(ytorch)); tvm::ffi::Tensor y = tvm::ffi::Tensor::FromDLPack(at::toDLPack(ytorch));
tvm::ffi::Tuple<int32_t, int32_t, int32_t> grid{nRows / 1024, 1, 1}; tvm::ffi::Tuple<int32_t, int32_t, int32_t> grid{nRows / 1024, 1, 1};
DLDevice device = x.device(); DLDevice device = x.device();
void* stream = void *stream = TVMFFIEnvGetStream(device.device_type, device.device_id);
TVMFFIEnvGetStream(device.device_type, device.device_id);
tvm::ffi::Array<tvm::ffi::Any> args = {y, x, xStride, yStride, tvm::ffi::Array<tvm::ffi::Any> args = {y, x, xStride, yStride,
nRows, nCols, BLOCK_SIZE}; nRows, nCols, BLOCK_SIZE};
tvm::ffi::Map<tvm::ffi::String, tvm::ffi::Any> kwargs = {}; tvm::ffi::Map<tvm::ffi::String, tvm::ffi::Any> kwargs = {};
SOFTMAX_KERNEL_STUB(grid, stream, numWarps, numStages, args, kwargs); SOFTMAX_KERNEL_STUB(grid, stream, args, kwargs);
return y; return y;
} }

6
python/triton_tvm_ffi/jit.py
View File

@@ -41,8 +41,6 @@ class TVMFFIJITFunction(object):
Callable[[Dict[str, Any]], Tuple[int, int, int]], Tuple[int, int, int] Callable[[Dict[str, Any]], Tuple[int, int, int]], Tuple[int, int, int]
], ],
_: int, _: int,
num_warps: Optional[int],
num_stages: Optional[int],
args: Sequence[Any], args: Sequence[Any],
kwargs: Mapping[str, Any], kwargs: Mapping[str, Any],
): ):
@@ -50,10 +48,6 @@ class TVMFFIJITFunction(object):
kwargs: Dict[str, Any] = { kwargs: Dict[str, Any] = {
k: v for k, v in zip(self.signature, args) if v is not None k: v for k, v in zip(self.signature, args) if v is not None
} | {k: self.canonicalize(v) for k, v in kwargs.items()} } | {k: self.canonicalize(v) for k, v in kwargs.items()}
if num_warps is not None:
kwargs["num_warps"] = num_warps
if num_stages is not None:
kwargs["num_stages"] = num_stages
kernel: CompiledKernel = self.fn[grid](*args, **kwargs) kernel: CompiledKernel = self.fn[grid](*args, **kwargs)
self.num_warps, _, self.shmem = kernel.packed_metadata self.num_warps, _, self.shmem = kernel.packed_metadata
self.ctypes = [type_canonicalize(v) for v in kernel.src.signature.values()] self.ctypes = [type_canonicalize(v) for v in kernel.src.signature.values()]

2
python/triton_tvm_ffi/templates/gendef.cc.j2
View File

@@ -35,7 +35,7 @@ static CUfunction __Get{{ fn.fnname }}Kernel() {
return *function; return *function;
} }
#define {{ fn.fnname | upper }}_STUB(__grid, __stream, __numWarps, __numStages, __args, __kwargs) do { \ #define {{ fn.fnname | upper }}_STUB(__grid, __stream, __args, __kwargs) do { \
const char *__signature[] = { "{{ fn.signature | join("\", \"") }}" }; \ const char *__signature[] = { "{{ fn.signature | join("\", \"") }}" }; \
tvm::ffi::Map<tvm::ffi::String, tvm::ffi::Any> __meta = { \ tvm::ffi::Map<tvm::ffi::String, tvm::ffi::Any> __meta = { \
{% if fn.best_config != none %} {% if fn.best_config != none %}