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 <[email protected]> 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") {
