JinjieLiu/triton-tvm-ffi
1
0
Fork 0
mirror of https://github.com/sgjzfzzf/triton-tvm-ffi.git synced 2026-07-01 08:51:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

unify launch apis

Browse Source 
Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>
This commit is contained in:
Jinjie Liu
2026-01-31 11:29:32 +08:00
parent ac7497b2c8
commit e9576d265e
5 changed files with 59 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/launch.h
+7 -14
View File
@@ -2,6 +2,7 @@
#define TRITON_TVM_FFI_LAUNCH_H_
#include "type.h"
#include <cuda.h>
#include <tvm/ffi/object.h>
namespace triton_tvm_ffi {
@@ -13,13 +14,9 @@ public:
TVMFFILauncherImplObj(const TVMFFILauncherImplObj &other) = default;
TVMFFILauncherImplObj(TVMFFILauncherImplObj &&other) = default;
void Launch(int32_t gridX, int32_t gridY, int32_t gridZ, uint64_t stream,
uint64_t function,
tvm::ffi::Tuple<int32_t, int32_t, int32_t> kernelMetadata,
tvm::ffi::ObjectRef launchMetadata,
tvm::ffi::ObjectRef launchEnterHook,
tvm::ffi::ObjectRef launchExitHook,
tvm::ffi::ObjectRef globalScratchObject,
tvm::ffi::ObjectRef profileScratchObject,
uint64_t function, int32_t numWarps, int32_t numCtas,
int32_t sharedMemory, uint64_t globalScratch,
uint64_t profileScratch,
const tvm::ffi::Array<tvm::ffi::Any> &kernelArgs) const;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("triton_tvm_ffi.TVMFFILauncherImpl",
TVMFFILauncherImplObj, tvm::ffi::Object);
@@ -37,13 +34,9 @@ public:
using tvm::ffi::ObjectRef::ObjectRef;
using tvm::ffi::ObjectRef::operator=;
void Launch(int32_t gridX, int32_t gridY, int32_t gridZ, uint64_t stream,
uint64_t function,
tvm::ffi::Tuple<int32_t, int32_t, int32_t> kernelMetadata,
tvm::ffi::ObjectRef launchMetadata,
tvm::ffi::ObjectRef launchEnterHook,
tvm::ffi::ObjectRef launchExitHook,
tvm::ffi::ObjectRef globalScratchObject,
tvm::ffi::ObjectRef profileScratchObject,
uint64_t function, int32_t numWarps, int32_t numCtas,
int32_t sharedMemory, uint64_t globalScratch,
uint64_t profileScratch,
const tvm::ffi::Array<tvm::ffi::Any> &kernelArgs) const;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TVMFFILauncherImpl,
tvm::ffi::ObjectRef,