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

support mm and autotune

Browse Source 
Signed-off-by: jinjieliu <jinjie.liu@usc.edu>
This commit is contained in:
jinjieliu
2026-02-07 00:41:23 +08:00
parent f6c7a48c1b
commit 2298b6f8c8
7 changed files with 486 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
examples/softmax/softmax.cc
View File

@@ -4,9 +4,7 @@
#include <tvm/ffi/tvm_ffi.h>
#ifndef SOFTMAX_KERNEL_STUB
#define SOFTMAX_KERNEL_STUB(grid, stream, numWarps, numStages, output, input, \
inputStride, outputStride, nRows, nCols, \
BLOCK_SIZE)
#define SOFTMAX_KERNEL_STUB(grid, stream, numWarps, numStages, args, kwargs)
#endif
#ifndef SOFTMAX_NAME
@@ -23,9 +21,12 @@ tvm::ffi::Tensor Softmax(tvm::ffi::Tensor x) {
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};
DLDevice device = x.device();
void *stream = TVMFFIEnvGetStream(device.device_type, device.device_id);
SOFTMAX_KERNEL_STUB(grid, stream, numWarps, numStages, y, x, xStride, yStride,
nRows, nCols, BLOCK_SIZE);
void* stream =
TVMFFIEnvGetStream(device.device_type, device.device_id);
tvm::ffi::Array<tvm::ffi::Any> args = {y, x, xStride, yStride,
nRows, nCols, BLOCK_SIZE};
tvm::ffi::Map<tvm::ffi::String, tvm::ffi::Any> kwargs = {};
SOFTMAX_KERNEL_STUB(grid, stream, numWarps, numStages, args, kwargs);
return y;
}