You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2025/04/14 15:38:35 UTC
(tvm) 02/03: [CODEGEN] TVMFFIAny support in struct get/set/alloca
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch refactor-s1
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit f6f5ba817d6c7366653a52e3fb64d18155303ede
Author: tqchen <ti...@gmail.com>
AuthorDate: Mon Apr 14 11:27:01 2025 -0400
[CODEGEN] TVMFFIAny support in struct get/set/alloca
---
src/target/llvm/codegen_cpu.cc | 41 +++++++++++++++++++++++++++++++++++++
src/target/llvm/codegen_cpu.h | 3 +--
src/target/llvm/codegen_llvm.h | 4 ++++
src/target/source/codegen_c.cc | 22 ++++++++++++++++++++
src/target/source/codegen_c_host.cc | 2 ++
5 files changed, 70 insertions(+), 2 deletions(-)
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index 5e30695517..ded7b91566 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -100,6 +100,8 @@ void CodeGenCPU::Init(const std::string& module_name, LLVMTarget* llvm_target,
// Defined in include/tvm/runtime/c_runtime_api.h:
// typedef union { ... } TVMValue;
t_tvm_value_ = llvm::StructType::create({t_float64_});
+ // Defined in include/tvm/ffi/c_api.h:
+ t_tvm_ffi_any_ = llvm::StructType::create({t_int32_, t_int32_, t_float64_});
// Defined in include/tvm/runtime/c_backend_api.h:
// typedef struct { void* sync_handle; int32_t num_task; } TVMParallelGroupEnv;
t_tvm_parallel_group_env_ = llvm::StructType::create({llvmGetPointerTo(t_int32_, 0), t_int32_});
@@ -345,6 +347,33 @@ CodeGenLLVM::TypedPointer CodeGenCPU::CreateStructRefPtr(DataType t, llvm::Value
builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(1), ConstInt32(0)});
return TypedPointer(member_type, member_addr);
}
+ case builtin::kTVMFFIAnyTypeIndex: {
+ buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(t_tvm_ffi_any_, 0));
+ buf = builder_->CreateInBoundsGEP(t_int32_, buf, {index, ConstInt32(0)});
+ return TypedPointer(t_int32_, buf);
+ }
+ case builtin::kTVMFFIAnyUnionValue: {
+ ICHECK_EQ(t.lanes(), 1);
+ buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(t_tvm_ffi_any_, 0));
+ // field 2 is the union value
+ buf = builder_->CreateInBoundsGEP(t_tvm_ffi_any_, buf, {index, ConstInt32(2)});
+ if (t.is_bool()) {
+ // it should be safe to set the pointer to the first byte of the union value
+ buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(DTypeToLLVMType(t), 0));
+ return TypedPointer(t_int8_, buf);
+ } else if (t.is_int() && t.bits() == 64) {
+ buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(t_int64_, 0));
+ return TypedPointer(t_int64_, buf);
+ } else if (t.is_float() && t.bits() == 64) {
+ buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(t_float64_, 0));
+ return TypedPointer(t_float64_, buf);
+ } else if (t.is_handle()) {
+ builder_->CreatePointerCast(buf, llvmGetPointerTo(t_void_p_, 0));
+ return TypedPointer(t_void_p_, buf);
+ } else {
+ LOG(DEBUG) << "DataType " << t << " cannot be stored into a TVMFFIAny's value field";
+ }
+ }
case builtin::kTVMValueContent: {
ICHECK_EQ(t.lanes(), 1);
if (t.is_bool()) {
@@ -1034,6 +1063,16 @@ llvm::Value* CodeGenCPU::CreateIntrinsic(const CallNode* op) {
if (value->getType()->isPointerTy()) {
value = builder_->CreatePointerCast(value, ref.type);
}
+
+ if (kind == builtin::kTVMFFIAnyUnionValue) {
+ // when we set any union value, we need to be careful to
+ // clear off the union value to zero if the set size is less than 64 bits
+ if (data_layout_->getTypeAllocSize(ref.type) != 8) {
+ llvm::Value* i64_addr = builder_->CreatePointerCast(
+ ref.addr, llvmGetPointerTo(t_int64_, 0));
+ builder_->CreateStore(ConstInt64(0), i64_addr);
+ }
+ }
builder_->CreateStore(value, ref.addr);
return ConstInt32(0);
} else if (op->op.same_as(builtin::tvm_stack_alloca())) {
@@ -1049,6 +1088,8 @@ llvm::Value* CodeGenCPU::CreateIntrinsic(const CallNode* op) {
return builder_->CreateAlloca(t_tvm_value_, num);
} else if (type == "arg_tcode") {
return builder_->CreateAlloca(t_int_, num);
+ } else if (type == "tvm_ffi_any") {
+ return builder_->CreateAlloca(t_tvm_ffi_any_, num);
} else if (type == "array") {
return builder_->CreateAlloca(t_tvm_array_, num);
} else {
diff --git a/src/target/llvm/codegen_cpu.h b/src/target/llvm/codegen_cpu.h
index 182dae81ce..5187b3e3f2 100644
--- a/src/target/llvm/codegen_cpu.h
+++ b/src/target/llvm/codegen_cpu.h
@@ -88,12 +88,11 @@ class CodeGenCPU : public CodeGenLLVM {
llvm::StructType* t_tvm_type_{nullptr};
llvm::StructType* t_tvm_array_{nullptr};
llvm::StructType* t_tvm_value_{nullptr};
+ llvm::StructType* t_tvm_ffi_any_{nullptr};
llvm::StructType* t_tvm_parallel_group_env_{nullptr};
-
llvm::FunctionType* ftype_tvm_backend_packed_c_func_{nullptr};
llvm::StructType* t_tvm_crt_func_registry_{nullptr};
llvm::StructType* t_tvm_crt_module_{nullptr};
-
llvm::FunctionType* ftype_tvm_parallel_lambda_{nullptr};
llvm::FunctionType* ftype_tvm_func_call_{nullptr};
llvm::FunctionType* ftype_tvm_get_func_from_env_{nullptr};
diff --git a/src/target/llvm/codegen_llvm.h b/src/target/llvm/codegen_llvm.h
index 302a0d97b3..c9cbc16b11 100644
--- a/src/target/llvm/codegen_llvm.h
+++ b/src/target/llvm/codegen_llvm.h
@@ -194,6 +194,10 @@ class CodeGenLLVM : public ExprFunctor<llvm::Value*(const PrimExpr&)>,
llvm::Constant* ConstInt32(int64_t value) const {
return llvm::ConstantInt::getSigned(t_int32_, value);
}
+ // Short hande code to get a constant int 64
+ llvm::Constant* ConstInt64(int64_t value) const {
+ return llvm::ConstantInt::getSigned(t_int64_, value);
+ }
// override codegen
llvm::Value* VisitExpr_(const VarNode* op) override;
llvm::Value* VisitExpr_(const CastNode* op) override;
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index a67cb80b91..2d05414f80 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -315,6 +315,28 @@ std::string CodeGenC::GetStructRef(DataType t, const PrimExpr& buffer, const Pri
}
os << ')';
return os.str();
+ } else if (kind == builtin::kTVMFFIAnyTypeIndex) {
+ std::ostringstream os;
+ os << "(((TVMFFIAny*)";
+ this->PrintExpr(buffer, os);
+ os << ")[" << index << "].type_index)";
+ return os.str();
+ } else if (kind == builtin::kTVMFFIAnyUnionValue) {
+ std::ostringstream os;
+ os << "(((TVMFFIAny*)";
+ this->PrintExpr(buffer, os);
+ os << ")[" << index << "].";
+ if (t.is_handle()) {
+ os << "v_ptr";
+ } else if (t.is_float()) {
+ os << "v_float64";
+ } else if (t.is_int()) {
+ os << "v_int64";
+ } else {
+ LOG(FATAL) << "Do not know how to handle type" << t;
+ }
+ os << ")";
+ return os.str();
} else {
ICHECK_LT(kind, builtin::kTVMValueKindBound_);
std::ostringstream os;
diff --git a/src/target/source/codegen_c_host.cc b/src/target/source/codegen_c_host.cc
index b00763a8a3..95c596fbd4 100644
--- a/src/target/source/codegen_c_host.cc
+++ b/src/target/source/codegen_c_host.cc
@@ -335,6 +335,8 @@ void CodeGenCHost::VisitExpr_(const CallNode* op, std::ostream& os) { // NOLINT
size = (num->value * sizeof(tvm_index_t) + unit - 1) / unit;
} else if (type == "arg_value") {
size = (num->value * sizeof(TVMValue) + unit - 1) / unit;
+ } else if (type == "tvm_ffi_any") {
+ size = (num->value * sizeof(TVMFFIAny) + unit - 1) / unit;
} else if (type == "arg_tcode") {
size = (num->value * sizeof(int) + unit - 1) / unit;
} else if (type == "array") {