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

use templates to substitute parts of macros

Browse Source 
Signed-off-by: jinjieliu <jinjie.liu@usc.edu>
This commit is contained in:
jinjieliu
2026-02-08 22:24:12 +08:00
parent 1c4f13c8f0
commit 213e4fc060
9 changed files with 139 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
examples/add/add.cc
View File

@@ -1,11 +1,11 @@
#include "tvm/ffi/function.h"
#include <ATen/DLConvertor.h>
#include <ATen/dlpack.h>
#include <tvm/ffi/extra/cuda/cubin_launcher.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/tvm_ffi.h>
#ifndef ADD_KERNEL_STUB
#define ADD_KERNEL_STUB(grid, stream, args, kwargs)
#define ADD_KERNEL_STUB(grid, device, stream, args, kwargs)
#endif
#ifndef ADD_NAME
@@ -27,7 +27,7 @@ tvm::ffi::Tensor Add(tvm::ffi::Tensor x, tvm::ffi::Tensor y) {
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, args, kwargs);
ADD_KERNEL_STUB(grid, device.device_id, stream, args, kwargs);
return output;
}

4
examples/mm/mm.cc
View File

@@ -4,7 +4,7 @@
#include <tvm/ffi/tvm_ffi.h>
#ifndef MATMUL_KERNEL_STUB
#define MATMUL_KERNEL_STUB(grid, stream, args, kwargs)
#define MATMUL_KERNEL_STUB(grid, device, stream, args, kwargs)
#endif
#ifndef MATMUL_NAME
@@ -45,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 = {
{"ACTIVATION", activation},
};
MATMUL_KERNEL_STUB(grid, stream, args, kwargs);
MATMUL_KERNEL_STUB(grid, device.device_id, stream, args, kwargs);
return c;
}

4
examples/softmax/softmax.cc
View File

@@ -4,7 +4,7 @@
#include <tvm/ffi/tvm_ffi.h>
#ifndef SOFTMAX_KERNEL_STUB
#define SOFTMAX_KERNEL_STUB(grid, stream, args, kwargs)
#define SOFTMAX_KERNEL_STUB(grid, device, stream, args, kwargs)
#endif
#ifndef SOFTMAX_NAME
@@ -24,7 +24,7 @@ tvm::ffi::Tensor Softmax(tvm::ffi::Tensor x) {
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, args, kwargs);
SOFTMAX_KERNEL_STUB(grid, device.device_id, stream, args, kwargs);
return y;
}