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

expand argument extractions with macro

Browse Source 
Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>
This commit is contained in:
Jinjie Liu
2026-01-30 14:03:38 +08:00
parent a953cbe7cc
commit 37a8f4a5be
2 changed files with 34 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/utils.cc
+29 -67
View File
@@ -11,66 +11,11 @@
#define CUDA_CHECK(code) \
do { \
if ((code) != CUDA_SUCCESS) { \
if (__builtin_expect((code) != CUDA_SUCCESS, 0)) { \
throw triton_tvm_ffi::CUDAException(code); \
} \
} while (false)
namespace {
using namespace triton_tvm_ffi;
// --------------- Definitions ---------------
template <Type T> struct ValueCast {
TRITON_TVM_FFI_INLINE static bool apply(void *ptr, const TypedValue &value);
};
template <Type... Ts> struct ValueCastSet {
TRITON_TVM_FFI_INLINE static bool apply(void *ptr, const TypedValue &value);
};
using GenericValueCastSet = ValueCastSet<
#define DEFINE_TYPE_CAST(type, str, ctype) Type::type,
TYPE_TABLE(DEFINE_TYPE_CAST)
#undef DEFINE_TYPE_CAST
Type::PTR>;
// --------------- Implementations ---------------
template <Type T>
TRITON_TVM_FFI_INLINE bool ValueCast<T>::apply(void *ptr,
const TypedValue &value) {
if (value.GetType() == T) {
if constexpr (T == Type::PTR) {
tvm::ffi::TensorView cvalue =
value.GetValue().cast<tvm::ffi::TensorView>();
void **p = reinterpret_cast<void **>(ptr);
*p = cvalue.data_ptr();
} else if constexpr (T == Type::CONSTEXPR) {
#ifdef NDEBUG
__builtin_unreachable();
#else
throw NotImplementedException("CONSTEXPR for value casting");
#endif
} else {
using ctype = type_to_ctype_t<T>;
ctype cvalue = value.GetValue().cast<ctype>();
ctype *p = reinterpret_cast<ctype *>(ptr);
*p = cvalue;
}
return true;
}
return false;
}
template <Type... Ts>
TRITON_TVM_FFI_INLINE bool ValueCastSet<Ts...>::apply(void *ptr,
const TypedValue &value) {
return (ValueCast<Ts>::apply(ptr, value) || ...);
}
} // namespace
namespace triton_tvm_ffi {
tvm::ffi::Map<tvm::ffi::String, int32_t> GetDeviceProperties(int device_id) {
@@ -159,23 +104,40 @@ void Launch(int32_t gridX, int32_t gridY, int32_t gridZ, uint64_t stream,
cFunction, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1));
}
const int32_t kernelArgNum = kernelArgs.size();
uint8_t *buffer =
reinterpret_cast<uint8_t *>(alloca(kMaxOpaqueSize * (kernelArgNum)));
void **params =
reinterpret_cast<void **>(alloca(sizeof(void *) * (kernelArgNum + 2)));
size_t j = 0;
for (size_t i = 0; i < kernelArgNum; ++i) {
TypedValue value = kernelArgs[i].cast<TypedValue>();
if (value.GetType() != Type::CONSTEXPR) {
#ifndef NDEBUG
assert(j < kernelArgNum);
#endif
void *ptr = buffer + j * kMaxOpaqueSize;
params[j] = ptr;
if (!GenericValueCastSet::apply(ptr, value)) {
throw UnknownTypeException(value.GetType());
}
switch (value.GetType()) {
#define CASE_STMT(type, str, ctype) \
case Type::type: { \
using cpptype = type_to_ctype_t<Type::type>; \
params[j] = reinterpret_cast<void *>(alloca(sizeof(cpptype))); \
*reinterpret_cast<cpptype *>(params[j]) = \
value.GetValue().cast<cpptype>(); \
++j; \
break; \
}
TYPE_TABLE_NATIVE(CASE_STMT)
#undef CASE_STMT
case Type::PTR: {
params[j] = reinterpret_cast<void *>(alloca(sizeof(void *)));
*reinterpret_cast<void **>(params[j]) =
value.GetValue().cast<tvm::ffi::TensorView>().data_ptr();
++j;
break;
}
case Type::CONSTEXPR: {
break;
}
default: {
#ifdef NDEBUG
__builtin_unreachable();
#else
throw NotImplementedException("CONSTEXPR for value casting");
#endif
}
}
}
// TODO: unwrap PyObject* from scratch pointers and assign to kernel args