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/05/02 13:37:32 UTC
(tvm) 03/03: Update to enable explicit Any cast
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch refactor-s2
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 7126c6569522fb0bb69fbcb4d9d1579386a34d20
Author: tqchen <ti...@gmail.com>
AuthorDate: Thu May 1 06:56:38 2025 -0400
Update to enable explicit Any cast
---
apps/cpp_rpc/rpc_server.cc | 4 +-
apps/ios_rpc/tvmrpc/TVMRuntime.mm | 4 +-
docs/arch/runtime.rst | 4 +-
include/tvm/ir/attrs.h | 16 +--
include/tvm/node/attr_registry_map.h | 14 +-
include/tvm/node/structural_equal.h | 4 +-
include/tvm/node/structural_hash.h | 4 +-
include/tvm/relax/expr_functor.h | 4 +-
include/tvm/relax/tuning_api.h | 4 +-
include/tvm/runtime/packed_func.h | 59 +++-----
include/tvm/runtime/relax_vm/vm.h | 2 +-
include/tvm/script/printer/ir_docsifier.h | 1 +
include/tvm/script/printer/ir_docsifier_functor.h | 6 +-
include/tvm/topi/utils.h | 2 +-
jvm/README.md | 2 +-
python/tvm/contrib/cutlass/attention_operation.py | 8 +-
python/tvm/contrib/cutlass/conv2d_operation.py | 2 +-
python/tvm/contrib/cutlass/gemm_operation.py | 4 +-
python/tvm/contrib/cutlass/layer_norm_operation.py | 2 +-
python/tvm/contrib/cutlass/rms_norm_operation.py | 2 +-
src/arith/analyzer.cc | 43 +++---
src/arith/int_constraints.cc | 4 +-
src/arith/solve_linear_equation.cc | 8 +-
src/arith/solve_linear_inequality.cc | 19 +--
src/contrib/msc/core/codegen/cpp_codegen.h | 18 +--
src/contrib/msc/core/printer/prototxt_printer.cc | 4 +-
src/contrib/msc/core/transform/set_expr_layout.cc | 2 +-
src/contrib/msc/core/utils.cc | 2 +-
src/contrib/msc/framework/tensorrt/codegen.cc | 9 +-
src/ir/attrs.cc | 2 +-
src/ir/env_func.cc | 2 +-
src/ir/function.cc | 2 +-
src/ir/op.cc | 2 +-
src/ir/transform.cc | 14 +-
src/meta_schedule/arg_info.cc | 6 +-
src/meta_schedule/database/database.cc | 16 +--
src/meta_schedule/database/json_database.cc | 8 +-
.../mutator/mutate_compute_location.cc | 8 +-
src/meta_schedule/mutator/mutate_parallel.cc | 6 +-
src/meta_schedule/mutator/mutate_tile_size.cc | 4 +-
.../schedule_rule/apply_custom_rule.cc | 3 +-
.../schedule_rule/multi_level_tiling.cc | 2 +-
src/meta_schedule/space_generator/schedule_fn.cc | 2 +-
.../space_generator/space_generator.cc | 8 +-
src/meta_schedule/utils.h | 11 +-
src/node/reflection.cc | 40 +++---
src/node/repr_printer.cc | 19 +--
src/node/script_printer.cc | 16 +--
src/node/serialization.cc | 15 +-
src/node/structural_equal.cc | 12 +-
src/node/structural_hash.cc | 6 +-
src/relax/backend/contrib/cublas/codegen.cc | 2 +-
src/relax/backend/contrib/tensorrt/codegen.cc | 2 +-
src/relax/backend/task_extraction.cc | 2 +-
src/relax/backend/vm/codegen_vm.cc | 4 +-
src/relax/backend/vm/exec_builder.cc | 2 +-
src/relax/ir/py_expr_functor.cc | 2 +-
src/relax/transform/few_shot_tuning.cc | 6 +-
src/relax/transform/fold_constant.cc | 4 +-
src/relax/transform/gradient.cc | 3 +-
src/relax/transform/meta_schedule.cc | 4 +-
src/relax/transform/run_codegen.cc | 7 +-
src/relax/transform/tuning_api/database.cc | 8 +-
src/relax/transform/tuning_api/primitives.cc | 8 +-
src/runtime/c_runtime_api.cc | 24 ++--
src/runtime/const_loader_module.cc | 2 +-
src/runtime/container.cc | 2 +-
src/runtime/contrib/amx/amx_config.cc | 4 +-
src/runtime/contrib/cblas/cblas.cc | 6 +-
src/runtime/contrib/cblas/dnnl_blas.cc | 2 +-
src/runtime/contrib/cblas/gemm_common.h | 42 +++---
src/runtime/contrib/cblas/mkl.cc | 12 +-
src/runtime/contrib/cublas/cublas.cc | 54 ++++----
src/runtime/contrib/cublas/cublas_json_runtime.cc | 4 +-
src/runtime/contrib/cudnn/conv_backward.cc | 56 ++++----
src/runtime/contrib/cudnn/conv_forward.cc | 66 ++++-----
src/runtime/contrib/cudnn/softmax.cc | 6 +-
src/runtime/contrib/cutlass/fp16_group_gemm.cu | 2 +-
.../contrib/cutlass/fp8_blockwise_scaled_gemm.cu | 4 +-
src/runtime/contrib/cutlass/fp8_gemm.cu | 2 +-
src/runtime/contrib/cutlass/fp8_group_gemm.cu | 2 +-
src/runtime/contrib/dnnl/dnnl.cc | 12 +-
src/runtime/contrib/dnnl/dnnl_json_runtime.cc | 2 +-
src/runtime/contrib/hipblas/hipblas.cc | 28 ++--
.../contrib/hipblas/hipblas_json_runtime.cc | 4 +-
src/runtime/contrib/json/json_runtime.h | 6 +-
src/runtime/contrib/miopen/conv_forward.cc | 60 ++++----
src/runtime/contrib/miopen/softmax.cc | 6 +-
src/runtime/contrib/mps/conv.mm | 18 +--
src/runtime/contrib/mps/gemm.mm | 10 +-
src/runtime/contrib/mrvl/mrvl_sw_runtime_lib.cc | 2 +-
src/runtime/contrib/random/random.cc | 22 +--
src/runtime/contrib/rocblas/rocblas.cc | 20 +--
src/runtime/contrib/sort/sort.cc | 44 +++---
src/runtime/contrib/tflite/tflite_runtime.cc | 4 +-
src/runtime/contrib/thrust/thrust.cu | 22 +--
src/runtime/disco/bcast_session.cc | 6 +-
src/runtime/disco/builtin.cc | 2 +-
src/runtime/disco/disco_worker.cc | 22 +--
src/runtime/disco/distributed/socket_session.cc | 20 +--
src/runtime/disco/loader.cc | 2 +-
src/runtime/disco/nccl/nccl_context.h | 2 +-
src/runtime/disco/process_session.cc | 8 +-
src/runtime/disco/protocol.h | 4 +-
src/runtime/disco/session.cc | 2 +-
src/runtime/dso_library.cc | 9 +-
src/runtime/hexagon/hexagon_device_api.cc | 36 ++---
src/runtime/hexagon/rpc/android/session.cc | 6 +-
src/runtime/hexagon/rpc/hexagon/rpc_server.cc | 10 +-
src/runtime/hexagon/rpc/simulator/rpc_server.cc | 10 +-
src/runtime/hexagon/rpc/simulator/session.cc | 6 +-
src/runtime/library_module.cc | 2 +-
src/runtime/memory/memory_manager.cc | 2 +-
src/runtime/module.cc | 4 +-
src/runtime/opencl/opencl_device_api.cc | 22 +--
src/runtime/opencl/opencl_module.cc | 2 +-
src/runtime/profiling.cc | 16 +--
src/runtime/regex.cc | 2 +-
src/runtime/relax_vm/attn_backend.h | 30 ++--
src/runtime/relax_vm/builtin.cc | 50 +++----
src/runtime/relax_vm/cuda/cuda_graph_builtin.cc | 16 +--
src/runtime/relax_vm/kv_state.cc | 8 +-
src/runtime/relax_vm/lm_support.cc | 4 +-
src/runtime/relax_vm/ndarray_cache_support.cc | 8 +-
src/runtime/relax_vm/paged_kv_cache.cc | 58 ++++----
src/runtime/relax_vm/vm.cc | 49 +++----
src/runtime/rpc/rpc_channel.cc | 2 +-
src/runtime/rpc/rpc_endpoint.cc | 85 ++++++------
src/runtime/rpc/rpc_endpoint.h | 2 +-
src/runtime/rpc/rpc_event_impl.cc | 2 +-
src/runtime/rpc/rpc_module.cc | 26 ++--
src/runtime/rpc/rpc_pipe_impl.cc | 2 +-
src/runtime/rpc/rpc_server_env.cc | 10 +-
src/runtime/rpc/rpc_socket_impl.cc | 12 +-
src/runtime/system_library.cc | 2 +-
src/runtime/thread_pool.cc | 6 +-
src/runtime/vulkan/vulkan_wrapped_func.cc | 2 +-
src/script/ir_builder/ir/ir.cc | 2 +-
src/script/ir_builder/tir/ir.cc | 2 +-
src/script/printer/legacy_repr.cc | 6 +-
src/script/printer/tir/expr.cc | 3 +-
src/support/array.h | 12 +-
src/support/ffi_testing.cc | 18 +--
src/target/codegen.cc | 5 +-
src/target/datatype/registry.cc | 11 +-
src/target/llvm/codegen_amdgpu.cc | 8 +-
src/target/llvm/codegen_hexagon.cc | 2 +-
src/target/llvm/codegen_llvm.cc | 4 +-
src/target/llvm/codegen_nvptx.cc | 2 +-
src/target/llvm/llvm_instance.cc | 4 +-
src/target/opt/build_cuda_on.cc | 4 +-
src/target/parsers/cpu.cc | 2 +-
src/target/source/codegen_metal.cc | 2 +-
src/target/source/codegen_opencl.cc | 2 +-
src/target/spirv/spirv_utils.cc | 2 +-
src/target/target.cc | 17 +--
src/target/target_info.cc | 2 +-
src/target/target_kind.cc | 10 +-
src/te/operation/create_primfunc.cc | 8 +-
src/tir/analysis/var_use_def_analysis.cc | 4 +-
src/tir/ir/buffer.cc | 16 ++-
src/tir/ir/data_type_rewriter.cc | 8 +-
src/tir/ir/expr.cc | 4 +-
src/tir/ir/stmt_functor.cc | 6 +-
src/tir/ir/utils.cc | 4 +-
src/tir/op/op.cc | 28 ++--
src/tir/schedule/instruction_traits.h | 5 +-
src/tir/schedule/primitive/blockize_tensorize.cc | 2 +-
src/tir/schedule/trace.cc | 25 ++--
src/tir/schedule/traced_schedule.cc | 2 +-
src/tir/transforms/lower_custom_datatypes.cc | 8 +-
src/tir/transforms/unsupported_dtype_legalize.cc | 4 +-
src/topi/broadcast.cc | 28 ++--
src/topi/einsum.cc | 2 +-
src/topi/elemwise.cc | 68 +++++-----
src/topi/nn.cc | 108 +++++++++------
src/topi/reduction.cc | 20 +--
src/topi/transform.cc | 151 +++++++++++----------
src/topi/utils.cc | 10 +-
src/topi/vision.cc | 2 +-
tests/python/contrib/test_hexagon/README_RPC.md | 14 +-
web/emcc/tvmjs_support.cc | 4 +-
web/emcc/wasm_runtime.cc | 8 +-
web/emcc/webgpu_runtime.cc | 4 +-
184 files changed, 1188 insertions(+), 1109 deletions(-)
diff --git a/apps/cpp_rpc/rpc_server.cc b/apps/cpp_rpc/rpc_server.cc
index 00fa489d01..c4ee4d3545 100644
--- a/apps/cpp_rpc/rpc_server.cc
+++ b/apps/cpp_rpc/rpc_server.cc
@@ -398,8 +398,6 @@ void RPCServerCreate(std::string host, int port, int port_end, std::string track
rpc.Start();
}
-TVM_REGISTER_GLOBAL("rpc.ServerCreate").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- RPCServerCreate(args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7]);
-});
+TVM_REGISTER_GLOBAL("rpc.ServerCreate").set_body_typed(RPCServerCreate);
} // namespace runtime
} // namespace tvm
diff --git a/apps/ios_rpc/tvmrpc/TVMRuntime.mm b/apps/ios_rpc/tvmrpc/TVMRuntime.mm
index f33ae12d4b..a14bb50e0b 100644
--- a/apps/ios_rpc/tvmrpc/TVMRuntime.mm
+++ b/apps/ios_rpc/tvmrpc/TVMRuntime.mm
@@ -53,13 +53,13 @@ void LogMessageImpl(const std::string& file, int lineno, int level, const std::s
TVM_REGISTER_GLOBAL("tvm.rpc.server.workpath").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
static const std::string base_ = NSTemporaryDirectory().UTF8String;
- const std::string path = args[0];
+ const auto path = args[0].cast<std::string>();
*rv = base_ + "/" + path;
});
TVM_REGISTER_GLOBAL("tvm.rpc.server.load_module")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- std::string name = args[0];
+ auto name = args[0].cast<std::string>();
std::string fmt = GetFileFormat(name, "");
NSString* base;
if (fmt == "dylib") {
diff --git a/docs/arch/runtime.rst b/docs/arch/runtime.rst
index 9f640654ae..90e02a4f0c 100644
--- a/docs/arch/runtime.rst
+++ b/docs/arch/runtime.rst
@@ -56,8 +56,8 @@ The following code block provides an example in C++
void MyAdd(TVMArgs args, TVMRetValue* rv) {
// automatically convert arguments to desired type.
- int a = args[0];
- int b = args[1];
+ int a = args[0].cast<int>();
+ int b = args[1].cast<int>();
// automatically assign value return to rv
*rv = a + b;
}
diff --git a/include/tvm/ir/attrs.h b/include/tvm/ir/attrs.h
index ed2d38d539..eef32b4773 100644
--- a/include/tvm/ir/attrs.h
+++ b/include/tvm/ir/attrs.h
@@ -263,7 +263,7 @@ class DictAttrs : public Attrs {
const DictAttrsNode* node = this->as<DictAttrsNode>();
auto it = node->dict.find(attr_key);
if (it != node->dict.end()) {
- return (*it).second;
+ return (*it).second.cast<TObjectRef>();
} else {
return default_value;
}
@@ -599,7 +599,7 @@ struct AttrInitEntry {
// from Expr types into the constants.
template <typename T>
inline void SetValue(T* ptr, const ffi::AnyView& val) {
- *ptr = val.operator T();
+ *ptr = val.cast<T>();
}
template <typename T>
@@ -607,7 +607,7 @@ inline void SetIntValue(T* ptr, const ffi::AnyView& val) {
if (auto opt_int = val.as<int64_t>()) {
*ptr = static_cast<T>(opt_int.value());
} else {
- IntImm expr = val;
+ IntImm expr = val.cast<IntImm>();
*ptr = static_cast<T>(expr->value);
}
}
@@ -615,7 +615,7 @@ inline void SetIntValue(T* ptr, const ffi::AnyView& val) {
// Workaround for GCC8.1 / GCC8.2
template <>
inline void SetValue<DataType>(DataType* ptr, const ffi::AnyView& val) {
- *ptr = DataType(val.operator DLDataType());
+ *ptr = DataType(val.cast<DLDataType>());
}
template <>
@@ -632,7 +632,7 @@ inline void SetValue<double>(double* ptr, const ffi::AnyView& val) {
if (auto opt_double = val.as<double>()) {
*ptr = opt_double.value();
} else {
- ObjectRef expr = val;
+ ObjectRef expr = val.cast<ObjectRef>();
ICHECK(expr.defined());
if (const IntImmNode* op = expr.as<IntImmNode>()) {
*ptr = static_cast<double>(op->value);
@@ -887,7 +887,7 @@ class AttrsNode : public BaseAttrsNode {
// linear search.
auto ffind = [&args](const char* key, ffi::AnyView* val) {
for (int i = 0; i < args.size(); i += 2) {
- if (!std::strcmp(key, args[i].operator const char*())) {
+ if (!std::strcmp(key, args[i].cast<const char*>())) {
*val = args[i + 1];
return true;
}
@@ -901,7 +901,7 @@ class AttrsNode : public BaseAttrsNode {
// construct a map then do lookup.
std::unordered_map<std::string, runtime::TVMArgValue> kwargs;
for (int i = 0; i < args.size(); i += 2) {
- kwargs[args[i].operator std::string()] = args[i + 1];
+ kwargs[args[i].cast<std::string>()] = args[i + 1];
}
auto ffind = [&kwargs](const char* key, ffi::AnyView* val) {
auto it = kwargs.find(key);
@@ -919,7 +919,7 @@ class AttrsNode : public BaseAttrsNode {
if (hit_count * 2 != args.size() && !allow_unknown) {
for (int i = 0; i < args.size(); i += 2) {
::tvm::detail::AttrExistVisitor visitor;
- visitor.key_ = args[i].operator std::string();
+ visitor.key_ = args[i].cast<std::string>();
self()->_tvm_VisitAttrs(visitor);
if (!visitor.exist_) {
std::ostringstream os;
diff --git a/include/tvm/node/attr_registry_map.h b/include/tvm/node/attr_registry_map.h
index c4b54ef0f2..909d0ad0d4 100644
--- a/include/tvm/node/attr_registry_map.h
+++ b/include/tvm/node/attr_registry_map.h
@@ -74,7 +74,11 @@ class AttrRegistryMapContainerMap {
ICHECK(key.defined());
const uint32_t idx = key->AttrRegistryIndex();
if (idx < data_.size() && data_[idx].second != 0) {
- return data_[idx].first;
+ if constexpr (std::is_same_v<ValueType, ffi::Any>) {
+ return data_[idx].first;
+ } else {
+ return data_[idx].first.template cast<ValueType>();
+ }
} else {
return def_value;
}
@@ -116,7 +120,13 @@ class AttrRegistryMap {
* \param key The key to the map
* \return the const reference to the content value.
*/
- ValueType operator[](const KeyType& key) const { return map_[key]; }
+ ValueType operator[](const KeyType& key) const {
+ if constexpr (std::is_same_v<ValueType, ffi::Any>) {
+ return map_[key];
+ } else {
+ return map_[key].template cast<ValueType>();
+ }
+ }
/*!
* \brief get the corresponding value element at key with default value.
* \param key The key to the map
diff --git a/include/tvm/node/structural_equal.h b/include/tvm/node/structural_equal.h
index 7ab6949366..7f56fd6ca9 100644
--- a/include/tvm/node/structural_equal.h
+++ b/include/tvm/node/structural_equal.h
@@ -140,7 +140,7 @@ class StructuralEqual : public BaseValueEqual {
bool map_free_params = false) const {
if (lhs.type_index() != rhs.type_index()) return false;
if (lhs.type_index() >= ffi::TypeIndex::kTVMFFIStaticObjectBegin) {
- return operator()(lhs.operator ObjectRef(), rhs.operator ObjectRef(), map_free_params);
+ return operator()(lhs.cast<ObjectRef>(), rhs.cast<ObjectRef>(), map_free_params);
}
// POD value can always use v_int64 to get the hash value
return (ffi::details::AnyUnsafe::TVMFFIAnyPtrFromAny(lhs)->v_uint64 ==
@@ -216,7 +216,7 @@ class SEqualReducer {
*/
TVM_FFI_INLINE ffi::Any MapLhsToRhs(const ffi::Any& lhs) {
if (lhs.type_index() >= ffi::TypeIndex::kTVMFFIStaticObjectBegin) {
- return MapLhsToRhs(lhs.operator ObjectRef());
+ return MapLhsToRhs(lhs.cast<ObjectRef>());
} else {
return lhs;
}
diff --git a/include/tvm/node/structural_hash.h b/include/tvm/node/structural_hash.h
index bbeceb8cbe..ca547b1714 100644
--- a/include/tvm/node/structural_hash.h
+++ b/include/tvm/node/structural_hash.h
@@ -131,7 +131,7 @@ class StructuralHash : public BaseValueHash {
*/
TVM_FFI_INLINE uint64_t operator()(const ffi::Any& key) const {
if (key.type_index() >= ffi::TypeIndex::kTVMFFIStaticObjectBegin) {
- return operator()(key.operator ObjectRef());
+ return operator()(key.cast<ObjectRef>());
}
return HashPODValueInAny(key);
}
@@ -225,7 +225,7 @@ class SHashReducer {
*/
void operator()(const ffi::Any& key) const {
if (key.type_index() >= ffi::TypeIndex::kTVMFFIStaticObjectBegin) {
- return operator()(key.operator ObjectRef());
+ return operator()(key.cast<ObjectRef>());
}
// POD value can always use v_int64 to get the hash value
handler_->SHashReduceHashedValue(BaseValueHash().HashPODValueInAny(key));
diff --git a/include/tvm/relax/expr_functor.h b/include/tvm/relax/expr_functor.h
index 4904b02960..96b5b20d1e 100644
--- a/include/tvm/relax/expr_functor.h
+++ b/include/tvm/relax/expr_functor.h
@@ -71,7 +71,7 @@ class ExprFunctor;
#define PY_EXPR_MUTATOR_DEFAULT(N, PY_FUNC, DEFAULT_FUNC, RET_TYPE) \
{ \
if (PY_FUNC != nullptr) { \
- RET_TYPE ret = PY_FUNC(N); \
+ RET_TYPE ret = PY_FUNC(N).cast<RET_TYPE>(); \
return ret; \
} else { \
return DEFAULT_FUNC; \
@@ -89,7 +89,7 @@ class ExprFunctor;
#define PY_EXPR_MUTATOR_DISPATCH(OP, PY_FUNC) \
vtable.template set_dispatch<OP>([](const ObjectRef& n, TSelf* self) { \
if (self->PY_FUNC != nullptr) { \
- Expr expr = self->PY_FUNC(n); \
+ Expr expr = self->PY_FUNC(n).cast<Expr>(); \
return expr; \
} else { \
return self->VisitExpr_(static_cast<const OP*>(n.get())); \
diff --git a/include/tvm/relax/tuning_api.h b/include/tvm/relax/tuning_api.h
index 6d72c56025..bcbfad2c7a 100644
--- a/include/tvm/relax/tuning_api.h
+++ b/include/tvm/relax/tuning_api.h
@@ -86,7 +86,7 @@ class ChoiceNode : public runtime::Object {
bool CheckConstr(IRModule mod) {
Array<Any> args(constr_func_args);
args.insert(args.begin(), GetRef<IRModule>(mod.CopyOnWrite()));
- return CallPackedWithArgsInArray(GetConstrFunc(), args);
+ return CallPackedWithArgsInArray(GetConstrFunc(), args).cast<bool>();
}
/*! \brief Perform transform_func. */
@@ -95,7 +95,7 @@ class ChoiceNode : public runtime::Object {
if (CheckConstr(mod)) {
Array<Any> args(transform_func_args);
args.insert(args.begin(), GetRef<IRModule>(mod.CopyOnWrite()));
- return CallPackedWithArgsInArray(GetTransformFunc(), args);
+ return CallPackedWithArgsInArray(GetTransformFunc(), args).cast<IRModule>();
}
return mod;
}
diff --git a/include/tvm/runtime/packed_func.h b/include/tvm/runtime/packed_func.h
index 268aec8429..192c9a883c 100644
--- a/include/tvm/runtime/packed_func.h
+++ b/include/tvm/runtime/packed_func.h
@@ -341,52 +341,38 @@ struct ModuleVTableEntryHelper {};
template <typename T, typename R, typename... Args>
struct ModuleVTableEntryHelper<R (T::*)(Args...) const> {
using MemFnType = R (T::*)(Args...) const;
- using IndexSeq = std::index_sequence_for<Args...>;
- static constexpr const std::size_t LenArgs = sizeof...(Args);
-
- template <std::size_t... Is>
- static TVM_ALWAYS_INLINE void Call(TVMRetValue* rv, T* self, MemFnType f, TVMArgs args,
- std::index_sequence<Is...>) {
- *rv = (self->*f)(args[Is]...);
+ static TVM_ALWAYS_INLINE void Call(TVMRetValue* rv, T* self, MemFnType f, TVMArgs args) {
+ auto wrapped = [self, f](Args... args) -> R { return (self->*f)(std::forward<Args>(args)...); };
+ ffi::details::unpack_call<R, sizeof...(Args)>(nullptr, wrapped, args.data(), args.size(), rv);
}
};
template <typename T, typename R, typename... Args>
struct ModuleVTableEntryHelper<R (T::*)(Args...)> {
using MemFnType = R (T::*)(Args...);
- using IndexSeq = std::index_sequence_for<Args...>;
- static constexpr const std::size_t LenArgs = sizeof...(Args);
-
- template <std::size_t... Is>
- static TVM_ALWAYS_INLINE void Call(TVMRetValue* rv, T* self, MemFnType f, TVMArgs args,
- std::index_sequence<Is...>) {
- *rv = (self->*f)(args[Is]...);
+ static TVM_ALWAYS_INLINE void Call(TVMRetValue* rv, T* self, MemFnType f, TVMArgs args) {
+ auto wrapped = [self, f](Args... args) -> R { return (self->*f)(std::forward<Args>(args)...); };
+ ffi::details::unpack_call<R, sizeof...(Args)>(nullptr, wrapped, args.data(), args.size(), rv);
}
};
template <typename T, typename... Args>
struct ModuleVTableEntryHelper<void (T::*)(Args...) const> {
using MemFnType = void (T::*)(Args...) const;
- using IndexSeq = std::index_sequence_for<Args...>;
- static constexpr const std::size_t LenArgs = sizeof...(Args);
-
- template <std::size_t... Is>
- static TVM_ALWAYS_INLINE void Call(TVMRetValue* rv, T* self, MemFnType f, TVMArgs args,
- std::index_sequence<Is...>) {
- (self->*f)(args[Is]...);
+ static TVM_ALWAYS_INLINE void Call(TVMRetValue* rv, T* self, MemFnType f, TVMArgs args) {
+ auto wrapped = [self, f](Args... args) -> void { (self->*f)(std::forward<Args>(args)...); };
+ ffi::details::unpack_call<void, sizeof...(Args)>(nullptr, wrapped, args.data(), args.size(),
+ rv);
}
};
template <typename T, typename... Args>
struct ModuleVTableEntryHelper<void (T::*)(Args...)> {
using MemFnType = void (T::*)(Args...);
- using IndexSeq = std::index_sequence_for<Args...>;
- static constexpr const std::size_t LenArgs = sizeof...(Args);
-
- template <std::size_t... Is>
- static TVM_ALWAYS_INLINE void Call(TVMRetValue* rv, T* self, MemFnType f, TVMArgs args,
- std::index_sequence<Is...>) {
- (self->*f)(args[Is]...);
+ static TVM_ALWAYS_INLINE void Call(TVMRetValue* rv, T* self, MemFnType f, TVMArgs args) {
+ auto wrapped = [self, f](Args... args) -> void { (self->*f)(std::forward<Args>(args)...); };
+ ffi::details::unpack_call<void, sizeof...(Args)>(nullptr, wrapped, args.data(), args.size(),
+ rv);
}
};
} // namespace details
@@ -404,16 +390,13 @@ struct ModuleVTableEntryHelper<void (T::*)(Args...)> {
return (this->*f)(_name); \
} \
} // NOLINT(*)
-#define TVM_MODULE_VTABLE_ENTRY(Name, MemFunc) \
- if (_name == Name) { \
- return ffi::Function::FromPacked([_self](ffi::PackedArgs args, Any* rv) -> void { \
- using Helper = ::tvm::runtime::details::ModuleVTableEntryHelper<decltype(MemFunc)>; \
- SelfPtr self = static_cast<SelfPtr>(_self.get()); \
- CHECK_EQ(args.size(), Helper::LenArgs) \
- << "Function `" << self->type_key() << "::" << Name << "` requires " << Helper::LenArgs \
- << " arguments, but got " << args.size(); \
- Helper::Call(rv, self, MemFunc, args, Helper::IndexSeq{}); \
- }); \
+#define TVM_MODULE_VTABLE_ENTRY(Name, MemFunc) \
+ if (_name == Name) { \
+ return ffi::Function::FromPacked([_self](ffi::PackedArgs args, Any* rv) -> void { \
+ using Helper = ::tvm::runtime::details::ModuleVTableEntryHelper<decltype(MemFunc)>; \
+ SelfPtr self = static_cast<SelfPtr>(_self.get()); \
+ Helper::Call(rv, self, MemFunc, args); \
+ }); \
}
#define TVM_MODULE_VTABLE_ENTRY_PACKED(Name, MemFunc) \
if (_name == Name) { \
diff --git a/include/tvm/runtime/relax_vm/vm.h b/include/tvm/runtime/relax_vm/vm.h
index 7c21b12c27..67439c102a 100644
--- a/include/tvm/runtime/relax_vm/vm.h
+++ b/include/tvm/runtime/relax_vm/vm.h
@@ -210,7 +210,7 @@ class VirtualMachine : public runtime::ModuleNode {
* \param arg The argument value.
*/
static VirtualMachine* GetContextPtr(TVMArgValue arg) {
- return static_cast<VirtualMachine*>(arg.operator void*());
+ return static_cast<VirtualMachine*>(arg.cast<void*>());
}
~VirtualMachine() {}
diff --git a/include/tvm/script/printer/ir_docsifier.h b/include/tvm/script/printer/ir_docsifier.h
index 0c47865055..e316088c06 100644
--- a/include/tvm/script/printer/ir_docsifier.h
+++ b/include/tvm/script/printer/ir_docsifier.h
@@ -335,6 +335,7 @@ inline TDoc IRDocsifierNode::AsDoc(const Any& value, const ObjectPath& path) con
return Downcast<TDoc>(d);
} else {
LOG(FATAL) << "TypeError: Cannot handle Any type: `" << value.GetTypeKey() << "`";
+ TVM_FFI_UNREACHABLE();
}
}
}
diff --git a/include/tvm/script/printer/ir_docsifier_functor.h b/include/tvm/script/printer/ir_docsifier_functor.h
index e63c00f68a..5ce105ea09 100644
--- a/include/tvm/script/printer/ir_docsifier_functor.h
+++ b/include/tvm/script/printer/ir_docsifier_functor.h
@@ -65,13 +65,13 @@ class IRDocsifierFunctor {
uint32_t type_index = obj.defined() ? obj->type_index() : 0;
const runtime::PackedFunc* pf = nullptr;
if ((pf = LookupDispatchTable(token, type_index)) != nullptr) {
- return (*pf)(obj, args...);
+ return (*pf)(obj, args...).template cast<R>();
}
if ((pf = LookupDispatchTable("", type_index)) != nullptr) {
- return (*pf)(obj, args...);
+ return (*pf)(obj, args...).template cast<R>();
}
if ((pf = LookupFallback()) != nullptr) {
- return (*pf)(obj, args...);
+ return (*pf)(obj, args...).template cast<R>();
}
LOG(WARNING) << "ObjectFunctor calls un-registered function on type: "
diff --git a/include/tvm/topi/utils.h b/include/tvm/topi/utils.h
index 3f8575ee9b..368674c0b6 100644
--- a/include/tvm/topi/utils.h
+++ b/include/tvm/topi/utils.h
@@ -42,7 +42,7 @@ inline Optional<Array<Integer>> ArrayOrInt(AnyView arg) {
result.push_back(opt_int.value());
return result;
} else {
- return arg;
+ return arg.cast<Array<Integer>>();
}
}
} // namespace topi
diff --git a/jvm/README.md b/jvm/README.md
index 62b685010c..0f53f4e561 100644
--- a/jvm/README.md
+++ b/jvm/README.md
@@ -103,7 +103,7 @@ import java.util.Arrays;
public class LoadAddFunc {
public static void main(String[] args) {
- String loadingDir = args[0];
+ String loadingDir = args[0].cast<String>();
Module fadd = Module.load(loadingDir + File.separator + "add_cpu.so");
Device dev = Device.cpu();
diff --git a/python/tvm/contrib/cutlass/attention_operation.py b/python/tvm/contrib/cutlass/attention_operation.py
index e89e2f747e..4c876142d3 100644
--- a/python/tvm/contrib/cutlass/attention_operation.py
+++ b/python/tvm/contrib/cutlass/attention_operation.py
@@ -148,7 +148,7 @@ def instantiate_attention_template(attrs):
CHECK(Attention::check_supported(p));
auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
- cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>(func().cast<void*>());
kernel_fn<<<p.getBlocksGrid(), p.getThreadsGrid(), smem_bytes, stream>>>(p);
@@ -187,7 +187,7 @@ def instantiate_flash_attention_template(attrs):
int o_batch_stride = o_row_stride * ${num_queries};
auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
- cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>(func().cast<void*>());
flash_attn::flash_attention_forward(
static_cast<const cutlass::half_t*>(${query}->data),
@@ -238,7 +238,7 @@ def instantiate_flash_attention_template(attrs):
int o_batch_stride = o_row_stride * ${num_queries};
auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
- cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>(func().cast<void*>());
flash_attn::flash_attention_forward(
static_cast<const cutlass::half_t*>(${qkv}->data),
@@ -295,7 +295,7 @@ def instantiate_flash_attention_var_len_template(attrs):
int o_row_stride = o_head_stride * ${num_q_heads};
auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
- cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>(func().cast<void*>());
flash_attn::flash_attention_var_len_forward(
static_cast<const cutlass::half_t*>(${query}->data),
diff --git a/python/tvm/contrib/cutlass/conv2d_operation.py b/python/tvm/contrib/cutlass/conv2d_operation.py
index 4585933c04..a37e46f404 100644
--- a/python/tvm/contrib/cutlass/conv2d_operation.py
+++ b/python/tvm/contrib/cutlass/conv2d_operation.py
@@ -425,7 +425,7 @@ def instantiate_conv2d_template(attrs):
${split_k_update}
auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
- cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>(func().cast<void*>());
status = conv2d_op(stream);
CHECK(status == cutlass::Status::kSuccess);
diff --git a/python/tvm/contrib/cutlass/gemm_operation.py b/python/tvm/contrib/cutlass/gemm_operation.py
index dce354b711..46b68c29ee 100644
--- a/python/tvm/contrib/cutlass/gemm_operation.py
+++ b/python/tvm/contrib/cutlass/gemm_operation.py
@@ -346,7 +346,7 @@ def instantiate_gemm_template(attrs):
CHECK(status == cutlass::Status::kSuccess);
auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
- cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>(func().cast<void*>());
status = gemm_op(stream);
CHECK(status == cutlass::Status::kSuccess);
@@ -429,7 +429,7 @@ def emit_fp16A_intB_matmul(attrs):
int k = ${B_arg}->shape[0];
auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
- cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>(func().cast<void*>());
""",
attrs,
)
diff --git a/python/tvm/contrib/cutlass/layer_norm_operation.py b/python/tvm/contrib/cutlass/layer_norm_operation.py
index 0afe4cba19..74f397b39a 100644
--- a/python/tvm/contrib/cutlass/layer_norm_operation.py
+++ b/python/tvm/contrib/cutlass/layer_norm_operation.py
@@ -40,7 +40,7 @@ def instantiate_layer_norm_template(attrs):
cutlass::TensorRef<data_type, RowMajor> _output((data_type*)out0->data, layout_2D);
auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
- cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>(func().cast<void*>());
cutlass::layernorm(size, _output, _input, _gamma, _beta, stream);
"""
diff --git a/python/tvm/contrib/cutlass/rms_norm_operation.py b/python/tvm/contrib/cutlass/rms_norm_operation.py
index ce0843c88f..27e98fb251 100644
--- a/python/tvm/contrib/cutlass/rms_norm_operation.py
+++ b/python/tvm/contrib/cutlass/rms_norm_operation.py
@@ -39,7 +39,7 @@ def instantiate_rms_norm_template(attrs):
cutlass::TensorRef<data_type, RowMajor> _output((data_type*)out0->data, layout_2D);
auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
- cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>(func().cast<void*>());
cutlass::rmsnorm(size, _output, _input, _weight, stream, ${rms_eps});
"""
diff --git a/src/arith/analyzer.cc b/src/arith/analyzer.cc
index cab660096f..fa1f100a5f 100644
--- a/src/arith/analyzer.cc
+++ b/src/arith/analyzer.cc
@@ -274,28 +274,33 @@ TVM_REGISTER_GLOBAL("arith.CreateAnalyzer").set_body_packed([](TVMArgs args, TVM
auto self = std::make_shared<Analyzer>();
auto f = [self](std::string name) -> PackedFunc {
if (name == "const_int_bound") {
- return PackedFunc(
- [self](TVMArgs args, TVMRetValue* ret) { *ret = self->const_int_bound(args[0]); });
+ return PackedFunc([self](TVMArgs args, TVMRetValue* ret) {
+ *ret = self->const_int_bound(args[0].cast<PrimExpr>());
+ });
} else if (name == "modular_set") {
- return PackedFunc(
- [self](TVMArgs args, TVMRetValue* ret) { *ret = self->modular_set(args[0]); });
+ return PackedFunc([self](TVMArgs args, TVMRetValue* ret) {
+ *ret = self->modular_set(args[0].cast<PrimExpr>());
+ });
} else if (name == "const_int_bound_update") {
return PackedFunc([self](TVMArgs args, TVMRetValue* ret) {
- self->const_int_bound.Update(args[0], args[1], args[2]);
+ self->const_int_bound.Update(args[0].cast<Var>(), args[1].cast<ConstIntBound>(),
+ args[2].cast<bool>());
});
} else if (name == "Simplify") {
return PackedFunc([self](TVMArgs args, TVMRetValue* ret) {
if (args.size() == 1) {
- *ret = self->Simplify(args[0]);
+ *ret = self->Simplify(args[0].cast<PrimExpr>());
} else if (args.size() == 2) {
- *ret = self->Simplify(args[0], args[1]);
+ *ret = self->Simplify(args[0].cast<PrimExpr>(), args[1].cast<int>());
} else {
LOG(FATAL) << "Invalid size of argument (" << args.size() << ")";
}
});
} else if (name == "rewrite_simplify") {
return PackedFunc(
- [self](TVMArgs args, TVMRetValue* ret) { *ret = self->rewrite_simplify(args[0]); });
+ [self](TVMArgs args, TVMRetValue* ret) {
+ *ret = self->rewrite_simplify(args[0].cast<PrimExpr>());
+ });
} else if (name == "get_rewrite_simplify_stats") {
return PackedFunc([self](TVMArgs args, TVMRetValue* ret) {
*ret = self->rewrite_simplify.GetStatsCounters();
@@ -305,42 +310,46 @@ TVM_REGISTER_GLOBAL("arith.CreateAnalyzer").set_body_packed([](TVMArgs args, TVM
[self](TVMArgs args, TVMRetValue* ret) { self->rewrite_simplify.ResetStatsCounters(); });
} else if (name == "canonical_simplify") {
return PackedFunc(
- [self](TVMArgs args, TVMRetValue* ret) { *ret = self->canonical_simplify(args[0]); });
+ [self](TVMArgs args, TVMRetValue* ret) {
+ *ret = self->canonical_simplify(args[0].cast<PrimExpr>());
+ });
} else if (name == "int_set") {
return PackedFunc(
- [self](TVMArgs args, TVMRetValue* ret) { *ret = self->int_set(args[0], args[1]); });
+ [self](TVMArgs args, TVMRetValue* ret) { *ret = self->int_set(args[0].cast<Var>()); });
} else if (name == "bind") {
return PackedFunc([self](TVMArgs args, TVMRetValue* ret) {
if (auto opt_range = args[1].as<Range>()) {
- self->Bind(args[0], opt_range.value());
+ self->Bind(args[0].cast<Var>(), opt_range.value());
} else {
- self->Bind(args[0], args[1].operator PrimExpr());
+ self->Bind(args[0].cast<Var>(), args[1].cast<PrimExpr>());
}
});
} else if (name == "can_prove") {
return PackedFunc([self](TVMArgs args, TVMRetValue* ret) {
- int strength = args[1];
- *ret = self->CanProve(args[0], static_cast<ProofStrength>(strength));
+ int strength = args[1].cast<int>();
+ *ret = self->CanProve(args[0].cast<PrimExpr>(), static_cast<ProofStrength>(strength));
});
} else if (name == "enter_constraint_context") {
return PackedFunc([self](TVMArgs args, TVMRetValue* ret) {
// can't use make_shared due to noexcept(false) decl in destructor,
// see https://stackoverflow.com/a/43907314
auto ctx = std::shared_ptr<With<ConstraintContext>>(
- new With<ConstraintContext>(self.get(), args[0]));
+ new With<ConstraintContext>(self.get(), args[0].cast<PrimExpr>()));
auto fexit = [ctx](TVMArgs, TVMRetValue*) mutable { ctx.reset(); };
*ret = ffi::Function::FromPacked(fexit);
});
} else if (name == "can_prove_equal") {
return PackedFunc(
- [self](TVMArgs args, TVMRetValue* ret) { *ret = self->CanProveEqual(args[0], args[1]); });
+ [self](TVMArgs args, TVMRetValue* ret) {
+ *ret = self->CanProveEqual(args[0].cast<PrimExpr>(), args[1].cast<PrimExpr>());
+ });
} else if (name == "get_enabled_extensions") {
return PackedFunc([self](TVMArgs args, TVMRetValue* ret) {
*ret = static_cast<std::int64_t>(self->rewrite_simplify.GetEnabledExtensions());
});
} else if (name == "set_enabled_extensions") {
return PackedFunc([self](TVMArgs args, TVMRetValue* ret) {
- std::int64_t flags = args[0];
+ int64_t flags = args[0].cast<int64_t>();
self->rewrite_simplify.SetEnabledExtensions(
static_cast<RewriteSimplifier::Extension>(flags));
});
diff --git a/src/arith/int_constraints.cc b/src/arith/int_constraints.cc
index 923f6f2241..5d3bc074bc 100644
--- a/src/arith/int_constraints.cc
+++ b/src/arith/int_constraints.cc
@@ -206,11 +206,11 @@ TVM_REGISTER_GLOBAL("arith.IntGroupBounds_from_range").set_body_typed(IntGroupBo
TVM_REGISTER_GLOBAL("arith.IntGroupBounds_FindBestRange")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
ICHECK(args.size() == 1 || args.size() == 2);
- IntGroupBounds bounds = args[0];
+ auto bounds = args[0].cast<IntGroupBounds>();
if (args.size() == 1) {
*ret = bounds.FindBestRange();
} else if (args.size() == 2) {
- *ret = bounds.FindBestRange(args[1]);
+ *ret = bounds.FindBestRange(args[1].cast<Map<Var, Range>>());
}
});
diff --git a/src/arith/solve_linear_equation.cc b/src/arith/solve_linear_equation.cc
index bff1ddea71..3ad1cdbf5a 100644
--- a/src/arith/solve_linear_equation.cc
+++ b/src/arith/solve_linear_equation.cc
@@ -457,11 +457,11 @@ IntConstraintsTransform SolveLinearEquations(const IntConstraints& system_to_sol
TVM_REGISTER_GLOBAL("arith.SolveLinearEquations")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
if (args.size() == 1) {
- *ret = SolveLinearEquations(args[0]);
+ *ret = SolveLinearEquations(args[0].cast<IntConstraints>());
} else if (args.size() == 3) {
- Optional<Array<Var>> opt_vars = args[0];
- Optional<Map<Var, Range>> opt_map = args[1];
- Optional<Array<PrimExpr>> opt_relations = args[2];
+ auto opt_vars = args[0].cast<Optional<Array<Var>>>();
+ auto opt_map = args[1].cast<Optional<Map<Var, Range>>>();
+ auto opt_relations = args[2].cast<Optional<Array<PrimExpr>>>();
IntConstraints problem(opt_vars.value_or({}), opt_map.value_or({}),
opt_relations.value_or({}));
*ret = SolveLinearEquations(problem);
diff --git a/src/arith/solve_linear_inequality.cc b/src/arith/solve_linear_inequality.cc
index cd0bb3409d..949ee97cc5 100644
--- a/src/arith/solve_linear_inequality.cc
+++ b/src/arith/solve_linear_inequality.cc
@@ -540,10 +540,11 @@ TVM_REGISTER_GLOBAL("arith.SolveInequalitiesAsCondition")
IntConstraints problem;
PartialSolvedInequalities ret_ineq;
if (args.size() == 1) {
- problem = args[0];
+ problem = args[0].cast<IntConstraints>();
ret_ineq = SolveLinearInequalities(problem);
} else if (args.size() == 3) {
- problem = IntConstraints(args[0], args[1], args[2]);
+ problem = IntConstraints(args[0].cast<Array<Var>>(), args[1].cast<Map<Var, Range>>(),
+ args[2].cast<Array<PrimExpr>>());
ret_ineq = SolveLinearInequalities(problem);
} else {
LOG(FATAL) << "arith.SolveInequalitiesAsCondition expects 1 or 3 arguments, gets "
@@ -555,10 +556,11 @@ TVM_REGISTER_GLOBAL("arith.SolveInequalitiesAsCondition")
TVM_REGISTER_GLOBAL("arith.SolveInequalitiesToRange")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
if (args.size() == 1) {
- *ret = SolveInequalitiesToRange(args[0]);
+ *ret = SolveInequalitiesToRange(args[0].cast<IntConstraints>());
} else if (args.size() == 3) {
- Optional<Map<Var, Range>> opt_map = args[1];
- IntConstraints problem(args[0], opt_map.value_or({}), args[2]);
+ auto opt_map = args[1].cast<Optional<Map<Var, Range>>>();
+ IntConstraints problem(args[0].cast<Array<Var>>(), opt_map.value_or({}),
+ args[2].cast<Array<PrimExpr>>());
*ret = SolveInequalitiesToRange(problem);
} else {
LOG(FATAL) << "arith.SolveInequalitiesToRange expects 1 or 3 arguments, gets "
@@ -569,10 +571,11 @@ TVM_REGISTER_GLOBAL("arith.SolveInequalitiesToRange")
TVM_REGISTER_GLOBAL("arith.SolveInequalitiesDeskewRange")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
if (args.size() == 1) {
- *ret = SolveInequalitiesDeskewRange(args[0]);
+ *ret = SolveInequalitiesDeskewRange(args[0].cast<IntConstraints>());
} else if (args.size() == 3) {
- Optional<Map<Var, Range>> opt_map = args[1];
- IntConstraints problem(args[0], opt_map.value_or({}), args[2]);
+ auto opt_map = args[1].cast<Optional<Map<Var, Range>>>();
+ IntConstraints problem(args[0].cast<Array<Var>>(), opt_map.value_or({}),
+ args[2].cast<Array<PrimExpr>>());
*ret = SolveInequalitiesDeskewRange(problem);
} else {
LOG(FATAL) << "arith.SolveInequalitiesDeskewRange expects 1 or 3 arguments, gets "
diff --git a/src/contrib/msc/core/codegen/cpp_codegen.h b/src/contrib/msc/core/codegen/cpp_codegen.h
index 397bf50bc7..260bd27ca3 100644
--- a/src/contrib/msc/core/codegen/cpp_codegen.h
+++ b/src/contrib/msc/core/codegen/cpp_codegen.h
@@ -117,16 +117,16 @@ class CppCodeGen : public BaseCodeGen<ConfigType, HelperType> {
const auto pf = tvm::ffi::Function::GetGlobalRequired("msc_tool.codegen_tensor");
for (size_t i = 0; i < node->inputs.size(); i++) {
const auto& input = node->InputAt(i);
- const Array<String>& lines = pf(GetTensorCtx(input), input->name, node->name,
- this->config()->tools_scope, this->config()->tools_tag);
- for (const auto& l : lines) {
+ ffi::Any lines = pf(GetTensorCtx(input), input->name, node->name,
+ this->config()->tools_scope, this->config()->tools_tag);
+ for (const auto& l : lines.cast<ffi::Array<ffi::String>>()) {
this->stack_.line(l);
}
}
for (const auto& pair : node->weights) {
- const Array<String>& lines = pf(GetTensorCtx(pair.second), pair.second->name, node->name,
- this->config()->tools_scope, this->config()->tools_tag);
- for (const auto& l : lines) {
+ ffi::Any lines = pf(GetTensorCtx(pair.second), pair.second->name, node->name,
+ this->config()->tools_scope, this->config()->tools_tag);
+ for (const auto& l : lines.cast<ffi::Array<ffi::String>>()) {
this->stack_.line(l);
}
}
@@ -141,9 +141,9 @@ class CppCodeGen : public BaseCodeGen<ConfigType, HelperType> {
int index = static_cast<int>(i);
if (graph_outputs_.count(node->OutputAt(index))) {
const auto& output = node->OutputAt(index);
- const Array<String>& lines = pf(GetTensorCtx(output), output->name, node->name,
- this->config()->tools_scope, this->config()->tools_tag);
- for (const auto& l : lines) {
+ ffi::Any lines = pf(GetTensorCtx(output), output->name, node->name,
+ this->config()->tools_scope, this->config()->tools_tag);
+ for (const auto& l : lines.cast<ffi::Array<ffi::String>>()) {
this->stack_.line(l);
}
}
diff --git a/src/contrib/msc/core/printer/prototxt_printer.cc b/src/contrib/msc/core/printer/prototxt_printer.cc
index c541dc3094..06a70a5498 100644
--- a/src/contrib/msc/core/printer/prototxt_printer.cc
+++ b/src/contrib/msc/core/printer/prototxt_printer.cc
@@ -51,7 +51,7 @@ DictDoc PrototxtPrinter::ToDictDoc(const Map<String, ffi::Any>& dict) {
if (pair.second.as<DictDocNode>()) {
values.push_back(Downcast<DictDoc>(pair.second));
} else {
- values.push_back(ToLiteralDoc(pair.second));
+ values.push_back(ToLiteralDoc(pair.second.cast<ObjectRef>()));
}
}
return DictDoc(keys, values);
@@ -65,7 +65,7 @@ DictDoc PrototxtPrinter::ToDictDoc(const std::vector<std::pair<String, Any>>& di
if (pair.second.as<DictDocNode>()) {
values.push_back(Downcast<DictDoc>(pair.second));
} else {
- values.push_back(ToLiteralDoc(pair.second));
+ values.push_back(ToLiteralDoc(pair.second.cast<ObjectRef>()));
}
}
return DictDoc(keys, values);
diff --git a/src/contrib/msc/core/transform/set_expr_layout.cc b/src/contrib/msc/core/transform/set_expr_layout.cc
index 22da70c7f0..80416bafd0 100644
--- a/src/contrib/msc/core/transform/set_expr_layout.cc
+++ b/src/contrib/msc/core/transform/set_expr_layout.cc
@@ -519,7 +519,7 @@ InferLayoutOutput ForwardInferLayoutPlugin(const Call& call,
return InferLayoutOutput();
}
const auto& args = Downcast<Tuple>(call->args[1]);
- return (*pf)(args->fields, var_layout_map);
+ return (*pf)(args->fields, var_layout_map).cast<InferLayoutOutput>();
}
// nn ops
diff --git a/src/contrib/msc/core/utils.cc b/src/contrib/msc/core/utils.cc
index dcef830452..ff2cb30c3b 100644
--- a/src/contrib/msc/core/utils.cc
+++ b/src/contrib/msc/core/utils.cc
@@ -271,7 +271,7 @@ const String StringUtils::ToString(const runtime::ObjectRef& obj) {
obj_string = std::to_string(n->value);
} else if (const auto* n = obj.as<ArrayObj>()) {
for (size_t i = 0; i < n->size(); i++) {
- obj_string = obj_string + ToString((*n)[i]);
+ obj_string = obj_string + ToString((*n)[i].cast<ObjectRef>());
if (n->size() == 1 || i < n->size() - 1) {
obj_string = obj_string + ",";
}
diff --git a/src/contrib/msc/framework/tensorrt/codegen.cc b/src/contrib/msc/framework/tensorrt/codegen.cc
index 32eabb771c..8992d3713a 100644
--- a/src/contrib/msc/framework/tensorrt/codegen.cc
+++ b/src/contrib/msc/framework/tensorrt/codegen.cc
@@ -120,7 +120,8 @@ void TensorRTCodeGen::CodeGenClassDefine() {
// save codegen before build
if (config()->use_tools) {
const auto pf = tvm::ffi::Function::GetGlobalRequired("msc_tool.codegen_step");
- before_build_codes_ = pf(GetStepCtx(), "before_build", graph()->name, config()->tools_tag);
+ before_build_codes_ =
+ pf(GetStepCtx(), "before_build", graph()->name, config()->tools_tag).cast<Array<String>>();
}
if (graph()->weight_holders.size() > 0) {
stack_.func_call("TRTUtils::LoadWeights", "mWeights")
@@ -205,7 +206,8 @@ void TensorRTCodeGen::CodeGenClassDefine() {
// save codegen after build
if (config()->use_tools) {
const auto pf = tvm::ffi::Function::GetGlobalRequired("msc_tool.codegen_step");
- after_build_codes_ = pf(GetStepCtx(), "after_build", graph()->name, config()->tools_tag);
+ after_build_codes_ =
+ pf(GetStepCtx(), "after_build", graph()->name, config()->tools_tag).cast<Array<String>>();
}
// end define build method
stack_.func_end("true");
@@ -610,7 +612,8 @@ Array<runtime::Module> MSCTensorRTCompiler(Array<Function> functions,
std::string graph_json = serializer.GetJSON();
const auto pf = tvm::ffi::Function::GetGlobalRequired("runtime.msc_tensorrt_runtime_create");
VLOG(1) << "Creating msc_tensorrt runtime::Module for '" << func_name << "'";
- compiled_functions.push_back(pf(func_name, graph_json, serializer.GetConstantNames()));
+ compiled_functions.push_back(
+ pf(func_name, graph_json, serializer.GetConstantNames()).cast<runtime::Module>());
}
return compiled_functions;
}
diff --git a/src/ir/attrs.cc b/src/ir/attrs.cc
index 922e5328a2..2ebca45c88 100644
--- a/src/ir/attrs.cc
+++ b/src/ir/attrs.cc
@@ -55,7 +55,7 @@ DictAttrs WithoutAttr(DictAttrs attrs, const std::string& key) {
void DictAttrsNode::InitByPackedArgs(const runtime::TVMArgs& args, bool allow_unknown) {
for (int i = 0; i < args.size(); i += 2) {
- String key = args[i];
+ String key = args[i].cast<String>();
ffi::AnyView val = args[i + 1];
dict.Set(key, val);
}
diff --git a/src/ir/env_func.cc b/src/ir/env_func.cc
index 6cd245bc99..a94223668b 100644
--- a/src/ir/env_func.cc
+++ b/src/ir/env_func.cc
@@ -50,7 +50,7 @@ EnvFunc EnvFunc::Get(const String& name) { return EnvFunc(CreateEnvNode(name));
TVM_REGISTER_GLOBAL("ir.EnvFuncGet").set_body_typed(EnvFunc::Get);
TVM_REGISTER_GLOBAL("ir.EnvFuncCall").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- EnvFunc env = args[0];
+ EnvFunc env = args[0].cast<EnvFunc>();
ICHECK_GE(args.size(), 1);
env->func.CallPacked(args.Slice(1), rv);
});
diff --git a/src/ir/function.cc b/src/ir/function.cc
index bb5a01551e..b8de83612b 100644
--- a/src/ir/function.cc
+++ b/src/ir/function.cc
@@ -53,7 +53,7 @@ TVM_REGISTER_GLOBAL("ir.BaseFuncWithAttrs")
return WithAttrs(Downcast<tir::PrimFunc>(std::move(func)), attr_map);
}
if (const auto f = tvm::ffi::Function::GetGlobal("relax.FuncWithAttrs")) {
- if (Optional<BaseFunc> ret = (*f)(func, attr_map)) {
+ if (auto ret = (*f)(func, attr_map).cast<Optional<BaseFunc>>()) {
return ret.value();
}
}
diff --git a/src/ir/op.cc b/src/ir/op.cc
index e5c98b5501..d69b558a11 100644
--- a/src/ir/op.cc
+++ b/src/ir/op.cc
@@ -138,7 +138,7 @@ TVM_REGISTER_GLOBAL("ir.RegisterOpAttr")
auto& reg = OpRegistry::Global()->RegisterOrGet(op_name).set_name();
// enable resgiteration and override of certain properties
if (attr_key == "num_inputs" && plevel > 128) {
- reg.set_num_inputs(value);
+ reg.set_num_inputs(value.cast<int>());
} else if (attr_key == "attrs_type_key" && plevel > 128) {
LOG(FATAL) << "attrs type key no longer supported";
} else {
diff --git a/src/ir/transform.cc b/src/ir/transform.cc
index a7383bf2e2..0756a86f52 100644
--- a/src/ir/transform.cc
+++ b/src/ir/transform.cc
@@ -466,7 +466,7 @@ Pass GetPass(const String& pass_name) {
f = tvm::ffi::Function::GetGlobal("transform." + pass_name);
}
ICHECK(f.has_value()) << "Cannot use " << pass_name << " to create the pass";
- return (*f)();
+ return (*f)().cast<Pass>();
}
// TODO(zhiics): we currently only sequentially execute each pass in
@@ -539,7 +539,7 @@ TVM_REGISTER_GLOBAL("transform.PassInfo")
});
TVM_REGISTER_GLOBAL("transform.Info").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- Pass pass = args[0];
+ Pass pass = args[0].cast<Pass>();
*ret = pass->Info();
});
@@ -587,11 +587,11 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
TVM_REGISTER_NODE_TYPE(SequentialNode);
TVM_REGISTER_GLOBAL("transform.Sequential").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- tvm::Array<Pass> passes = args[0];
- int opt_level = args[1];
- std::string name = args[2];
- tvm::Array<runtime::String> required = args[3];
- bool traceable = args[4];
+ auto passes = args[0].cast<tvm::Array<Pass>>();
+ int opt_level = args[1].cast<int>();
+ std::string name = args[2].cast<std::string>();
+ auto required = args[3].cast<tvm::Array<runtime::String>>();
+ bool traceable = args[4].cast<bool>();
PassInfo pass_info = PassInfo(opt_level, name, required, /* traceable */ traceable);
*ret = Sequential(passes, pass_info);
});
diff --git a/src/meta_schedule/arg_info.cc b/src/meta_schedule/arg_info.cc
index f3b8345141..ad3966ed1e 100644
--- a/src/meta_schedule/arg_info.cc
+++ b/src/meta_schedule/arg_info.cc
@@ -71,7 +71,7 @@ ArgInfo ArgInfo::FromJSON(const ObjectRef& json_obj) {
try {
const ArrayObj* json_array = json_obj.as<ArrayObj>();
CHECK(json_array && json_array->size() >= 1);
- tag = json_array->at(0);
+ tag = json_array->at(0).cast<String>();
} catch (const std::runtime_error& e) { // includes tvm::Error and dmlc::Error
LOG(FATAL) << "ValueError: Unable to parse the JSON object: " << json_obj
<< "\nThe error is: " << e.what();
@@ -133,11 +133,11 @@ TensorInfo TensorInfo::FromJSON(const ObjectRef& json_obj) {
CHECK(json_array && json_array->size() == 3);
// Load json[1] => dtype
{
- String dtype_str = json_array->at(1);
+ String dtype_str = json_array->at(1).cast<String>();
dtype = runtime::StringToDLDataType(dtype_str);
}
// Load json[2] => shape
- shape = AsIntArray(json_array->at(2));
+ shape = AsIntArray(json_array->at(2).cast<ObjectRef>());
} catch (const std::runtime_error& e) { // includes tvm::Error and dmlc::Error
LOG(FATAL) << "ValueError: Unable to parse the JSON object: " << json_obj
<< "\nThe error is: " << e.what();
diff --git a/src/meta_schedule/database/database.cc b/src/meta_schedule/database/database.cc
index 8e800bdb9b..f859648464 100644
--- a/src/meta_schedule/database/database.cc
+++ b/src/meta_schedule/database/database.cc
@@ -54,12 +54,12 @@ Workload Workload::FromJSON(const ObjectRef& json_obj) {
const ArrayObj* json_array = json_obj.as<ArrayObj>();
CHECK(json_array && json_array->size() == 2);
// Load json[0] => shash
- String str_shash = json_array->at(0);
+ String str_shash = json_array->at(0).cast<String>();
// Load json[1] => mod
{
- String b64_mod = json_array->at(1);
+ String b64_mod = json_array->at(1).cast<String>();
std::string json_mod = Base64Decode(b64_mod);
- mod = LoadJSON(json_mod).operator IRModule();
+ mod = LoadJSON(json_mod).cast<IRModule>();
std::stringstream(str_shash) >> shash;
}
} catch (const std::runtime_error& e) { // includes tvm::Error and dmlc::Error
@@ -138,25 +138,25 @@ TuningRecord TuningRecord::FromJSON(const ObjectRef& json_obj, const Workload& w
CHECK(json_array && json_array->size() == 4);
// Load json[1] => run_secs
if (json_array->at(1) != nullptr) {
- run_secs = AsFloatArray(json_array->at(1));
+ run_secs = AsFloatArray(json_array->at(1).cast<ObjectRef>());
}
// Load json[2] => target
if (json_array->at(2) != nullptr) {
- target = Target(json_array->at(2).operator Map<String, ffi::Any>());
+ target = Target(json_array->at(2).cast<Map<String, ffi::Any>>());
}
// Load json[3] => args_info
if (json_array->at(3) != nullptr) {
- const ArrayObj* json_args_info = json_array->at(3).operator const ArrayObj*();
+ const ArrayObj* json_args_info = json_array->at(3).cast<const ArrayObj*>();
Array<ArgInfo> info;
info.reserve(json_args_info->size());
for (Any json_arg_info : *json_args_info) {
- info.push_back(ArgInfo::FromJSON(json_arg_info));
+ info.push_back(ArgInfo::FromJSON(json_arg_info.cast<ObjectRef>()));
}
args_info = info;
}
// Load json[0] => trace
{
- const ObjectRef& json_trace = json_array->at(0);
+ auto json_trace = json_array->at(0).cast<ObjectRef>();
tir::Schedule sch =
tir::Schedule::Traced(workload->mod, /*seed=*/-1, /*debug_mask=*/0,
/*error_render_level=*/tir::ScheduleErrorRenderLevel::kNone);
diff --git a/src/meta_schedule/database/json_database.cc b/src/meta_schedule/database/json_database.cc
index ff6b363bb7..facfc9b880 100644
--- a/src/meta_schedule/database/json_database.cc
+++ b/src/meta_schedule/database/json_database.cc
@@ -166,7 +166,7 @@ Database Database::JSONDatabase(String path_workload, String path_tuning_record,
n->workloads2idx_.reserve(n_objs);
workloads.reserve(n_objs);
for (int i = 0; i < n_objs; ++i) {
- Workload workload = Workload::FromJSON(json_objs[i]);
+ Workload workload = Workload::FromJSON(json_objs[i].cast<ObjectRef>());
auto recalc_hash = n->GetModuleEquality().Hash(workload->mod);
// Todo(tvm-team): re-enable the shash check when we get environment
// independent structural hash values.
@@ -186,15 +186,15 @@ Database Database::JSONDatabase(String path_workload, String path_tuning_record,
records.resize(json_objs.size(), TuningRecord{nullptr});
support::parallel_for_dynamic(
0, json_objs.size(), num_threads, [&](int thread_id, int task_id) {
- const ObjectRef& json_obj = json_objs[task_id];
+ auto json_obj = json_objs[task_id].cast<ObjectRef>();
Workload workload{nullptr};
try {
const ArrayObj* arr = json_obj.as<ArrayObj>();
ICHECK_EQ(arr->size(), 2);
- int64_t workload_index = arr->at(0).operator IntImm()->value;
+ int64_t workload_index = arr->at(0).cast<IntImm>()->value;
ICHECK(workload_index >= 0 && static_cast<size_t>(workload_index) < workloads.size());
workload = workloads[workload_index];
- records[task_id] = TuningRecord::FromJSON(arr->at(1), workload);
+ records[task_id] = TuningRecord::FromJSON(arr->at(1).cast<ObjectRef>(), workload);
} catch (std::runtime_error& e) {
LOG(FATAL) << "ValueError: Unable to parse TuningRecord, on line " << (task_id + 1)
<< " of file " << path_tuning_record << ". The workload is:\n"
diff --git a/src/meta_schedule/mutator/mutate_compute_location.cc b/src/meta_schedule/mutator/mutate_compute_location.cc
index 4490d7a408..98926c3405 100644
--- a/src/meta_schedule/mutator/mutate_compute_location.cc
+++ b/src/meta_schedule/mutator/mutate_compute_location.cc
@@ -70,10 +70,10 @@ class MutateComputeLocationNode : public MutatorNode {
*/
std::vector<MutateComputeLocationNode::Candidate> MutateComputeLocationNode::FindCandidates(
const Trace& trace, TRandState* rand_state) {
- tir::Schedule sch = tir::Schedule::Traced( //
- /*mod=*/LoadJSON(this->json_mod_), //
- /*rand_state=*/ForkSeed(rand_state), //
- /*debug_mode=*/0, //
+ tir::Schedule sch = tir::Schedule::Traced( //
+ /*mod=*/LoadJSON(this->json_mod_).cast<IRModule>(), //
+ /*rand_state=*/ForkSeed(rand_state), //
+ /*debug_mode=*/0, //
/*error_render_level=*/tir::ScheduleErrorRenderLevel::kNone);
static InstructionKind inst_sample_compute_location =
diff --git a/src/meta_schedule/mutator/mutate_parallel.cc b/src/meta_schedule/mutator/mutate_parallel.cc
index 4d24f698d6..e50e7410f6 100644
--- a/src/meta_schedule/mutator/mutate_parallel.cc
+++ b/src/meta_schedule/mutator/mutate_parallel.cc
@@ -253,9 +253,9 @@ Optional<Trace> MutateParallelNode::Apply(const Trace& trace, TRandState* rand_s
return NullOpt;
}
// Step 2. Replay the instructions to recover loop extents
- tir::Schedule sch = tir::Schedule::Traced( //
- /*mod=*/LoadJSON(this->json_mod_), //
- /*rand_state=*/ForkSeed(rand_state), //
+ tir::Schedule sch = tir::Schedule::Traced( //
+ /*mod=*/LoadJSON(this->json_mod_).cast<IRModule>(), //
+ /*rand_state=*/ForkSeed(rand_state), //
/*debug_mode=*/0,
/*error_render_level=*/tir::ScheduleErrorRenderLevel::kNone);
trace->ApplyToSchedule(sch, /*remove_postproc=*/true);
diff --git a/src/meta_schedule/mutator/mutate_tile_size.cc b/src/meta_schedule/mutator/mutate_tile_size.cc
index d7b3e707fe..f06736abfb 100644
--- a/src/meta_schedule/mutator/mutate_tile_size.cc
+++ b/src/meta_schedule/mutator/mutate_tile_size.cc
@@ -88,7 +88,7 @@ void FindSamplePerfectTile(const Trace& trace, std::vector<Instruction>* inst,
decisions.reserve(trace->decisions.size());
for (const auto& kv : trace->decisions) {
const Instruction& inst = kv.first;
- const ObjectRef& decision = kv.second;
+ const ObjectRef& decision = kv.second.cast<ObjectRef>();
if (inst->kind.same_as(inst_sample_perfect_tile)) {
std::vector<int64_t> tiles = DowncastTilingDecision(decision);
if (tiles.size() >= 2 && Product(tiles) >= 2) {
@@ -124,7 +124,7 @@ void FindSampleVectorize(const Trace& trace, std::vector<Instruction>* inst,
// Find sampling instruction that generates the annotation
for (const auto& kv : trace->decisions) {
const Instruction& inst = kv.first;
- const ObjectRef& decision = kv.second;
+ const ObjectRef& decision = kv.second.cast<ObjectRef>();
if (inst->kind.same_as(inst_sample_categorical)) {
ICHECK_EQ(inst->outputs.size(), 1);
if (annotated.count(inst->outputs[0].as<Object>())) {
diff --git a/src/meta_schedule/schedule_rule/apply_custom_rule.cc b/src/meta_schedule/schedule_rule/apply_custom_rule.cc
index e972be32e1..b89dc824c5 100644
--- a/src/meta_schedule/schedule_rule/apply_custom_rule.cc
+++ b/src/meta_schedule/schedule_rule/apply_custom_rule.cc
@@ -43,7 +43,8 @@ class ApplyCustomRuleNode : public ScheduleRuleNode {
for (const String& key : keys) {
if (const auto custom_schedule_fn =
tvm::ffi::Function::GetGlobal(GetCustomRuleName(ann.value(), key))) {
- Array<tir::Schedule> result = (*custom_schedule_fn)(sch, block_rv);
+ Array<tir::Schedule> result =
+ (*custom_schedule_fn)(sch, block_rv).cast<Array<tir::Schedule>>();
return result;
}
}
diff --git a/src/meta_schedule/schedule_rule/multi_level_tiling.cc b/src/meta_schedule/schedule_rule/multi_level_tiling.cc
index b1709cf9a6..b7e7cd0964 100644
--- a/src/meta_schedule/schedule_rule/multi_level_tiling.cc
+++ b/src/meta_schedule/schedule_rule/multi_level_tiling.cc
@@ -102,7 +102,7 @@ void MultiLevelTilingNode::InitializeWithTuneContext(const TuneContext& context)
// Entry of the mega rule; Inherited from ScheduleRuleNode
Array<Schedule> MultiLevelTilingNode::Apply(const Schedule& sch, const BlockRV& block_rv) {
- if ((filter_fn_ && filter_fn_.value()(sch, sch->GetSRef(block_rv))) ||
+ if ((filter_fn_ && filter_fn_.value()(sch, sch->GetSRef(block_rv)).cast<bool>()) ||
NeedsMultiLevelTiling(sch->state(), sch->GetSRef(block_rv))) {
sch->Annotate(block_rv, tir::attr::meta_schedule_tiling_structure, structure);
diff --git a/src/meta_schedule/space_generator/schedule_fn.cc b/src/meta_schedule/space_generator/schedule_fn.cc
index 4d583dcb5f..1c1b5078a5 100644
--- a/src/meta_schedule/space_generator/schedule_fn.cc
+++ b/src/meta_schedule/space_generator/schedule_fn.cc
@@ -50,7 +50,7 @@ class ScheduleFnNode : public SpaceGeneratorNode {
if (rv == nullptr) {
return {sch};
}
- ObjectRef obj = rv;
+ ObjectRef obj = rv.cast<ObjectRef>();
if (auto sch = obj.as<tir::Schedule>()) {
return {sch.value()};
}
diff --git a/src/meta_schedule/space_generator/space_generator.cc b/src/meta_schedule/space_generator/space_generator.cc
index 036705e43b..8712f5ad48 100644
--- a/src/meta_schedule/space_generator/space_generator.cc
+++ b/src/meta_schedule/space_generator/space_generator.cc
@@ -26,13 +26,13 @@ String GetRuleKindFromTarget(const Target& target) {
if (target->kind->name == "llvm") {
static auto target_has_feature_fn_ptr =
tvm::ffi::Function::GetGlobalRequired("target.target_has_feature");
- bool have_avx512vnni = target_has_feature_fn_ptr("avx512vnni", target);
- bool have_avxvnni = target_has_feature_fn_ptr("avxvnni", target);
+ bool have_avx512vnni = target_has_feature_fn_ptr("avx512vnni", target).cast<bool>();
+ bool have_avxvnni = target_has_feature_fn_ptr("avxvnni", target).cast<bool>();
if (have_avx512vnni || have_avxvnni) {
return "vnni";
} else {
- bool have_avx512f = target_has_feature_fn_ptr("avx512f", target);
- bool have_avx512bw = target_has_feature_fn_ptr("avx512bw", target);
+ bool have_avx512f = target_has_feature_fn_ptr("avx512f", target).cast<bool>();
+ bool have_avx512bw = target_has_feature_fn_ptr("avx512bw", target).cast<bool>();
if (have_avx512bw && have_avx512f) {
return "avx512";
}
diff --git a/src/meta_schedule/utils.h b/src/meta_schedule/utils.h
index decd8731b8..8b36405baf 100644
--- a/src/meta_schedule/utils.h
+++ b/src/meta_schedule/utils.h
@@ -126,7 +126,7 @@ inline bool using_ipython() {
bool flag = false;
const auto f_using_ipython = tvm::ffi::Function::GetGlobal("meta_schedule.using_ipython");
if (f_using_ipython) {
- flag = (*f_using_ipython)();
+ flag = (*f_using_ipython)().cast<bool>();
}
return flag;
}
@@ -263,9 +263,7 @@ inline std::vector<support::LinearCongruentialEngine::TRandState> ForkSeed(
* \param mod The IRModule to make a deep copy.
* \return The deep copy of the IRModule.
*/
-inline IRModule DeepCopyIRModule(IRModule mod) {
- return LoadJSON(SaveJSON(mod)).operator IRModule();
-}
+inline IRModule DeepCopyIRModule(IRModule mod) { return LoadJSON(SaveJSON(mod)).cast<IRModule>(); }
/*!
* \brief Concatenate strings
@@ -383,7 +381,7 @@ inline int GetTargetNumCores(const Target& target) {
static const auto f_cpu_count = tvm::ffi::Function::GetGlobal("meta_schedule.cpu_count");
ICHECK(f_cpu_count.has_value())
<< "ValueError: Cannot find the packed function \"meta_schedule._cpu_count\"";
- num_cores = (*f_cpu_count)(false);
+ num_cores = (*f_cpu_count)(false).cast<int>();
LOG(FATAL)
<< "Target does not have attribute \"num-cores\", physical core number must be "
"defined! For example, on the local machine, the target must be \"llvm -num-cores "
@@ -456,6 +454,7 @@ inline Array<Integer> AsIntArray(const ObjectRef& obj) {
return (*opt_int_imm)->value;
} else {
LOG(FATAL) << "TypeError: Expect an array of integers, but gets: " << val.GetTypeKey();
+ TVM_FFI_UNREACHABLE();
}
}();
results.push_back(Integer(int_value));
@@ -618,7 +617,7 @@ class BlockCollector : public tir::StmtVisitor {
// Otherwise collect all blocks.
Bool collect_block = Bool(true);
if (f_block_filter_ != nullptr) {
- collect_block = f_block_filter_(GetRef<tir::Block>(block));
+ collect_block = f_block_filter_(GetRef<tir::Block>(block)).cast<Bool>();
}
if (collect_block) {
blocks_to_collect_.push_back(block->name_hint);
diff --git a/src/node/reflection.cc b/src/node/reflection.cc
index 92687a3bb6..b49440dce3 100644
--- a/src/node/reflection.cc
+++ b/src/node/reflection.cc
@@ -180,28 +180,26 @@ class NodeAttrSetter : public AttrVisitor {
std::string type_key;
std::unordered_map<std::string, runtime::TVMArgValue> attrs;
- void Visit(const char* key, double* value) final { *value = GetAttr(key).operator double(); }
- void Visit(const char* key, int64_t* value) final { *value = GetAttr(key).operator int64_t(); }
- void Visit(const char* key, uint64_t* value) final { *value = GetAttr(key).operator uint64_t(); }
- void Visit(const char* key, int* value) final { *value = GetAttr(key).operator int(); }
- void Visit(const char* key, bool* value) final { *value = GetAttr(key).operator bool(); }
+ void Visit(const char* key, double* value) final { *value = GetAttr(key).cast<double>(); }
+ void Visit(const char* key, int64_t* value) final { *value = GetAttr(key).cast<int64_t>(); }
+ void Visit(const char* key, uint64_t* value) final { *value = GetAttr(key).cast<uint64_t>(); }
+ void Visit(const char* key, int* value) final { *value = GetAttr(key).cast<int>(); }
+ void Visit(const char* key, bool* value) final { *value = GetAttr(key).cast<bool>(); }
void Visit(const char* key, std::string* value) final {
- *value = GetAttr(key).operator std::string();
+ *value = GetAttr(key).cast<std::string>();
}
- void Visit(const char* key, void** value) final { *value = GetAttr(key).operator void*(); }
- void Visit(const char* key, DataType* value) final { *value = GetAttr(key).operator DataType(); }
+ void Visit(const char* key, void** value) final { *value = GetAttr(key).cast<void*>(); }
+ void Visit(const char* key, DataType* value) final { *value = GetAttr(key).cast<DataType>(); }
void Visit(const char* key, runtime::NDArray* value) final {
- *value = GetAttr(key).operator runtime::NDArray();
- }
- void Visit(const char* key, ObjectRef* value) final {
- *value = GetAttr(key).operator ObjectRef();
+ *value = GetAttr(key).cast<runtime::NDArray>();
}
+ void Visit(const char* key, ObjectRef* value) final { *value = GetAttr(key).cast<ObjectRef>(); }
void Visit(const char* key, Optional<double>* value) final {
- *value = GetAttr(key).operator Optional<double>();
+ *value = GetAttr(key).cast<Optional<double>>();
}
void Visit(const char* key, Optional<int64_t>* value) final {
- *value = GetAttr(key).operator Optional<int64_t>();
+ *value = GetAttr(key).cast<Optional<int64_t>>();
}
private:
@@ -221,7 +219,7 @@ void InitNodeByPackedArgs(ReflectionVTable* reflection, Object* n, const TVMArgs
setter.type_key = n->GetTypeKey();
ICHECK_EQ(args.size() % 2, 0);
for (int i = 0; i < args.size(); i += 2) {
- setter.attrs.emplace(args[i].operator std::string(), args[i + 1]);
+ setter.attrs.emplace(args[i].cast<std::string>(), args[i + 1]);
}
reflection->VisitAttrs(n, &setter);
@@ -254,7 +252,7 @@ ObjectRef ReflectionVTable::CreateObject(const std::string& type_key,
int index = 0;
for (const auto& kv : *static_cast<const MapObj*>(kwargs.get())) {
- packed_args[index] = kv.first.operator String().c_str();
+ packed_args[index] = kv.first.cast<String>().c_str();
packed_args[index + 1] = kv.second;
index += 2;
}
@@ -264,18 +262,18 @@ ObjectRef ReflectionVTable::CreateObject(const std::string& type_key,
// Expose to FFI APIs.
void NodeGetAttr(TVMArgs args, TVMRetValue* ret) {
- Object* self = const_cast<Object*>(args[0].operator const Object*());
- *ret = ReflectionVTable::Global()->GetAttr(self, args[1]);
+ Object* self = const_cast<Object*>(args[0].cast<const Object*>());
+ *ret = ReflectionVTable::Global()->GetAttr(self, args[1].cast<std::string>());
}
void NodeListAttrNames(TVMArgs args, TVMRetValue* ret) {
- Object* self = const_cast<Object*>(args[0].operator const Object*());
+ Object* self = const_cast<Object*>(args[0].cast<const Object*>());
auto names =
std::make_shared<std::vector<std::string>>(ReflectionVTable::Global()->ListAttrNames(self));
*ret = PackedFunc([names](TVMArgs args, TVMRetValue* rv) {
- int64_t i = args[0];
+ int64_t i = args[0].cast<int64_t>();
if (i == -1) {
*rv = static_cast<int64_t>(names->size());
} else {
@@ -288,7 +286,7 @@ void NodeListAttrNames(TVMArgs args, TVMRetValue* ret) {
// args format:
// key1, value1, ..., key_n, value_n
void MakeNode(const TVMArgs& args, TVMRetValue* rv) {
- std::string type_key = args[0];
+ auto type_key = args[0].cast<std::string>();
*rv = ReflectionVTable::Global()->CreateObject(type_key, args.Slice(1));
}
diff --git a/src/node/repr_printer.cc b/src/node/repr_printer.cc
index 22d223fe53..3e80751e66 100644
--- a/src/node/repr_printer.cc
+++ b/src/node/repr_printer.cc
@@ -48,30 +48,31 @@ void ReprPrinter::Print(const ffi::Any& node) {
break;
}
case ffi::TypeIndex::kTVMFFIInt: {
- stream << node.operator int64_t();
+ stream << node.cast<int64_t>();
break;
}
case ffi::TypeIndex::kTVMFFIBool: {
- stream << node.operator bool();
+ stream << node.cast<bool>();
break;
}
case ffi::TypeIndex::kTVMFFIFloat: {
- stream << node.operator double();
+ stream << node.cast<double>();
break;
}
case ffi::TypeIndex::kTVMFFIOpaquePtr: {
- stream << node.operator void*();
+ stream << node.cast<void*>();
+ break;
+ }
+ case ffi::TypeIndex::kTVMFFIDataType: {
+ stream << node.cast<DataType>();
break;
- case ffi::TypeIndex::kTVMFFIDataType:
- stream << node.operator DataType();
- break;
}
case ffi::TypeIndex::kTVMFFIDevice: {
- runtime::operator<<(stream, node.operator Device());
+ runtime::operator<<(stream, node.cast<Device>());
break;
}
case ffi::TypeIndex::kTVMFFIObject: {
- Print(node.operator ObjectRef());
+ Print(node.cast<ObjectRef>());
break;
}
default: {
diff --git a/src/node/script_printer.cc b/src/node/script_printer.cc
index 2f4a9fb9e2..ab1be75568 100644
--- a/src/node/script_printer.cc
+++ b/src/node/script_printer.cc
@@ -57,7 +57,7 @@ PrinterConfig::PrinterConfig(Map<String, Any> config_dict) {
n->binding_names.push_back(Downcast<String>(v.value()));
}
if (auto v = config_dict.Get("show_meta")) {
- n->show_meta = v.value();
+ n->show_meta = v.value().cast<bool>();
}
if (auto v = config_dict.Get("ir_prefix")) {
n->ir_prefix = Downcast<String>(v.value());
@@ -81,16 +81,16 @@ PrinterConfig::PrinterConfig(Map<String, Any> config_dict) {
n->float_dtype = DataType(runtime::StringToDLDataType(Downcast<String>(v.value())));
}
if (auto v = config_dict.Get("verbose_expr")) {
- n->verbose_expr = v.value();
+ n->verbose_expr = v.value().cast<bool>();
}
if (auto v = config_dict.Get("indent_spaces")) {
- n->indent_spaces = v.value();
+ n->indent_spaces = v.value().cast<int>();
}
if (auto v = config_dict.Get("print_line_numbers")) {
- n->print_line_numbers = v.value();
+ n->print_line_numbers = v.value().cast<bool>();
}
if (auto v = config_dict.Get("num_context_lines")) {
- n->num_context_lines = v.value();
+ n->num_context_lines = v.value().cast<int>();
}
if (auto v = config_dict.Get("path_to_underline")) {
n->path_to_underline = Downcast<Optional<Array<ObjectPath>>>(v).value_or(Array<ObjectPath>());
@@ -107,13 +107,13 @@ PrinterConfig::PrinterConfig(Map<String, Any> config_dict) {
Downcast<Optional<Map<ObjectRef, String>>>(v).value_or(Map<ObjectRef, String>());
}
if (auto v = config_dict.Get("syntax_sugar")) {
- n->syntax_sugar = v.value();
+ n->syntax_sugar = v.value().cast<bool>();
}
if (auto v = config_dict.Get("show_object_address")) {
- n->show_object_address = v.value();
+ n->show_object_address = v.value().cast<bool>();
}
if (auto v = config_dict.Get("show_all_struct_info")) {
- n->show_all_struct_info = v.value();
+ n->show_all_struct_info = v.value().cast<bool>();
}
// Checking prefixes if they are valid Python identifiers.
diff --git a/src/node/serialization.cc b/src/node/serialization.cc
index 86f1fe1703..96c6696b33 100644
--- a/src/node/serialization.cc
+++ b/src/node/serialization.cc
@@ -284,7 +284,7 @@ class JSONAttrGetter : public AttrVisitor {
});
if (is_str_map) {
for (const auto& kv : *n) {
- node_->keys.push_back(kv.first.operator String());
+ node_->keys.push_back(kv.first.cast<String>());
node_->data.push_back(node_index_->at(kv.second));
}
} else {
@@ -306,19 +306,19 @@ class JSONAttrGetter : public AttrVisitor {
switch (node.type_index()) {
case ffi::TypeIndex::kTVMFFIBool:
case ffi::TypeIndex::kTVMFFIInt: {
- node_->attrs["v_int64"] = std::to_string(node.operator int64_t());
+ node_->attrs["v_int64"] = std::to_string(node.cast<int64_t>());
break;
}
case ffi::TypeIndex::kTVMFFIFloat: {
- node_->attrs["v_float64"] = std::to_string(node.operator double());
+ node_->attrs["v_float64"] = std::to_string(node.cast<double>());
break;
}
case ffi::TypeIndex::kTVMFFIDataType: {
- node_->attrs["v_type"] = Type2String(DataType(node.operator DLDataType()));
+ node_->attrs["v_type"] = Type2String(DataType(node.cast<DLDataType>()));
break;
}
case ffi::TypeIndex::kTVMFFIDevice: {
- DLDevice dev = node.operator DLDevice();
+ DLDevice dev = node.cast<DLDevice>();
node_->attrs["v_device_type"] = std::to_string(dev.device_type);
node_->attrs["v_device_id"] = std::to_string(dev.device_id);
break;
@@ -376,7 +376,8 @@ class FieldDependencyFinder : public AttrVisitor {
return;
}
// Skip the objects that have their own string repr
- if (jnode->repr_bytes.length() > 0 || reflection_->GetReprBytes(node, nullptr)) {
+ if (jnode->repr_bytes.length() > 0 ||
+ reflection_->GetReprBytes(node.cast<const Object*>(), nullptr)) {
return;
}
// Skip containers
@@ -476,7 +477,7 @@ class JSONAttrSetter : public AttrVisitor {
size_t index;
ParseValue(key, &index);
ICHECK_LE(index, node_list_->size());
- *value = node_list_->at(index);
+ *value = node_list_->at(index).cast<ObjectRef>();
}
static Any CreateInitAny(ReflectionVTable* reflection, JSONNode* jnode) {
diff --git a/src/node/structural_equal.cc b/src/node/structural_equal.cc
index 01331d2546..1a20810e79 100644
--- a/src/node/structural_equal.cc
+++ b/src/node/structural_equal.cc
@@ -191,9 +191,9 @@ bool SEqualReducer::AnyEqual(const ffi::Any& lhs, const ffi::Any& rhs,
}
if (lhs.type_index() >= ffi::TypeIndex::kTVMFFIStaticObjectBegin) {
if (paths) {
- return operator()(lhs.operator ObjectRef(), rhs.operator ObjectRef(), paths.value());
+ return operator()(lhs.cast<ObjectRef>(), rhs.cast<ObjectRef>(), paths.value());
} else {
- return operator()(lhs.operator ObjectRef(), rhs.operator ObjectRef());
+ return operator()(lhs.cast<ObjectRef>(), rhs.cast<ObjectRef>());
}
}
if (ffi::details::AnyUnsafe::TVMFFIAnyPtrFromAny(lhs)->v_uint64 ==
@@ -368,7 +368,7 @@ class SEqualHandlerDefault::Impl {
}
// normal object ref path
- if (!SEqualReduce(lhs, rhs, map_free_vars, current_paths)) {
+ if (!SEqualReduce(lhs.cast<ObjectRef>(), rhs.cast<ObjectRef>(), map_free_vars, current_paths)) {
return false;
}
@@ -422,7 +422,8 @@ class SEqualHandlerDefault::Impl {
// The TVMScriptPrinter::Script will fallback to Repr printer,
// if the root node to print is not supported yet,
// e.g. Relax nodes, ArrayObj, MapObj, etc.
- oss << ":" << std::endl << TVMScriptPrinter::Script(root_lhs_.value(), cfg);
+ oss << ":" << std::endl
+ << TVMScriptPrinter::Script(root_lhs_.value().cast<ObjectRef>(), cfg);
}
} else {
oss << ":" << std::endl << lhs;
@@ -437,7 +438,8 @@ class SEqualHandlerDefault::Impl {
// The TVMScriptPrinter::Script will fallback to Repr printer,
// if the root node to print is not supported yet,
// e.g. Relax nodes, ArrayObj, MapObj, etc.
- oss << ":" << std::endl << TVMScriptPrinter::Script(root_rhs_.value(), cfg);
+ oss << ":" << std::endl
+ << TVMScriptPrinter::Script(root_rhs_.value().cast<ObjectRef>(), cfg);
}
} else {
oss << ":" << std::endl << rhs;
diff --git a/src/node/structural_hash.cc b/src/node/structural_hash.cc
index 982c915759..bba8a59647 100644
--- a/src/node/structural_hash.cc
+++ b/src/node/structural_hash.cc
@@ -137,7 +137,7 @@ class SHashHandlerDefault::Impl {
if (value.type_index() < ffi::TypeIndex::kTVMFFIStaticObjectBegin) {
return BaseValueHash().HashPODValueInAny(value);
}
- ObjectRef object = value.operator ObjectRef();
+ ObjectRef object = value.cast<ObjectRef>();
ICHECK_EQ(task_stack_.size(), 0U);
ICHECK_EQ(pending_tasks_.size(), 0U);
ICHECK_EQ(result_stack_.size(), 0U);
@@ -551,7 +551,7 @@ struct MapObjTrait {
uint64_t hashed_value;
// skip non-object keys
if (kv.first.type_index() >= ffi::TypeIndex::kTVMFFIStaticObjectBegin) {
- if (hash_reduce->LookupHashedValue(kv.first, &hashed_value)) {
+ if (hash_reduce->LookupHashedValue(kv.first.cast<ObjectRef>(), &hashed_value)) {
temp.emplace_back(hashed_value, kv.second);
}
}
@@ -583,7 +583,7 @@ struct MapObjTrait {
using KV = std::pair<String, Any>;
std::vector<KV> temp;
for (const auto& kv : *key) {
- temp.push_back(std::make_pair(kv.first.operator String(), kv.second));
+ temp.push_back(std::make_pair(kv.first.cast<String>(), kv.second));
}
// sort by the hash key of the keys.
std::sort(temp.begin(), temp.end(),
diff --git a/src/relax/backend/contrib/cublas/codegen.cc b/src/relax/backend/contrib/cublas/codegen.cc
index bea9ed7baa..431265d1a7 100644
--- a/src/relax/backend/contrib/cublas/codegen.cc
+++ b/src/relax/backend/contrib/cublas/codegen.cc
@@ -119,7 +119,7 @@ Array<runtime::Module> CublasCompiler(Array<Function> functions, Map<String, ffi
auto constant_names = serializer.GetConstantNames();
const auto pf = tvm::ffi::Function::GetGlobalRequired("runtime.CublasJSONRuntimeCreate");
auto func_name = GetExtSymbol(func);
- compiled_functions.push_back(pf(func_name, graph_json, constant_names));
+ compiled_functions.push_back(pf(func_name, graph_json, constant_names).cast<runtime::Module>());
}
return compiled_functions;
diff --git a/src/relax/backend/contrib/tensorrt/codegen.cc b/src/relax/backend/contrib/tensorrt/codegen.cc
index f458ff85ef..c9cf80695f 100644
--- a/src/relax/backend/contrib/tensorrt/codegen.cc
+++ b/src/relax/backend/contrib/tensorrt/codegen.cc
@@ -226,7 +226,7 @@ Array<runtime::Module> TensorRTCompiler(Array<Function> functions, Map<String, f
const auto pf = tvm::ffi::Function::GetGlobalRequired("runtime.tensorrt_runtime_create");
std::string func_name = GetExtSymbol(func);
VLOG(1) << "Creating tensorrt runtime::Module for '" << func_name << "'";
- compiled_functions.push_back(pf(func_name, graph_json, constant_names));
+ compiled_functions.push_back(pf(func_name, graph_json, constant_names).cast<runtime::Module>());
}
return compiled_functions;
}
diff --git a/src/relax/backend/task_extraction.cc b/src/relax/backend/task_extraction.cc
index aeee9d2285..af6e83cf5f 100644
--- a/src/relax/backend/task_extraction.cc
+++ b/src/relax/backend/task_extraction.cc
@@ -104,7 +104,7 @@ class TaskExtractor : public ExprVisitor {
const GlobalVar& global_var = Downcast<GlobalVar>(call->args[0]);
const tir::PrimFunc& func = Downcast<tir::PrimFunc>(mod_->Lookup(global_var));
- IRModule mod = (*normalize_mod_func_)(func);
+ IRModule mod = (*normalize_mod_func_)(func).cast<IRModule>();
size_t weight = 1;
auto it = func2task_.find(mod);
if (it != func2task_.end()) {
diff --git a/src/relax/backend/vm/codegen_vm.cc b/src/relax/backend/vm/codegen_vm.cc
index bd56b0fd7b..b096028945 100644
--- a/src/relax/backend/vm/codegen_vm.cc
+++ b/src/relax/backend/vm/codegen_vm.cc
@@ -446,8 +446,8 @@ void LinkModules(ObjectPtr<VMExecutable> exec, const Map<String, runtime::NDArra
auto pf_var = mod.GetFunction("get_const_vars");
std::vector<std::string> symbol_const_vars;
if (pf_sym != nullptr && pf_var != nullptr) {
- String symbol = pf_sym();
- Array<String> variables = pf_var();
+ String symbol = pf_sym().cast<String>();
+ Array<String> variables = pf_var().cast<Array<String>>();
for (size_t i = 0; i < variables.size(); i++) {
symbol_const_vars.push_back(variables[i].operator std::string());
}
diff --git a/src/relax/backend/vm/exec_builder.cc b/src/relax/backend/vm/exec_builder.cc
index e32c043b06..85f013d986 100644
--- a/src/relax/backend/vm/exec_builder.cc
+++ b/src/relax/backend/vm/exec_builder.cc
@@ -331,7 +331,7 @@ TVM_REGISTER_GLOBAL("relax.ExecBuilderCreate").set_body_typed(ExecBuilderNode::C
TVM_REGISTER_GLOBAL("relax.ExecBuilderConvertConstant")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- ExecBuilder builder = args[0];
+ ExecBuilder builder = args[0].cast<ExecBuilder>();
TVMRetValue rt;
rt = args[1];
*ret = builder->ConvertConstant(rt).data();
diff --git a/src/relax/ir/py_expr_functor.cc b/src/relax/ir/py_expr_functor.cc
index eb286b4ef6..1da1c29b20 100644
--- a/src/relax/ir/py_expr_functor.cc
+++ b/src/relax/ir/py_expr_functor.cc
@@ -327,7 +327,7 @@ class PyExprMutatorNode : public Object, public ExprMutator {
Expr VisitExpr(const Expr& expr) {
if (f_visit_expr != nullptr) {
- return builder_->Normalize(f_visit_expr(expr));
+ return builder_->Normalize(f_visit_expr(expr).cast<Expr>());
} else {
static FType vtable = InitVTable();
return builder_->Normalize(vtable(expr, this));
diff --git a/src/relax/transform/few_shot_tuning.cc b/src/relax/transform/few_shot_tuning.cc
index d7e0ead986..3df818cf3c 100644
--- a/src/relax/transform/few_shot_tuning.cc
+++ b/src/relax/transform/few_shot_tuning.cc
@@ -30,14 +30,14 @@ tir::PrimFunc FewShotTunePrimFunc(const tir::PrimFunc& prim_func, const Target&
// fetch a local builder
static const auto f_get_local_builder =
tvm::ffi::Function::GetGlobalRequired("meta_schedule.builder.get_local_builder");
- meta_schedule::Builder builder = f_get_local_builder();
+ meta_schedule::Builder builder = f_get_local_builder().cast<meta_schedule::Builder>();
ICHECK(builder.defined()) << "ValueError: The local builder is not defined!";
// fetch a local runner
meta_schedule::Runner runner{nullptr};
if (benchmark) {
static const auto f_get_local_runner =
tvm::ffi::Function::GetGlobalRequired("meta_schedule.runner.get_local_runner");
- runner = f_get_local_runner();
+ runner = f_get_local_runner().cast<meta_schedule::Runner>();
ICHECK(runner.defined()) << "ValueError: The local runner is not defined!";
}
// create an IRModule
@@ -45,7 +45,7 @@ tir::PrimFunc FewShotTunePrimFunc(const tir::PrimFunc& prim_func, const Target&
{{GlobalVar("main"), WithAttr(prim_func, tvm::attr::kGlobalSymbol, String("main"))}}));
// fetch the number of physical cores
static const auto f_cpu_count = tvm::ffi::Function::GetGlobalRequired("meta_schedule.cpu_count");
- int num_threads = f_cpu_count(false);
+ int num_threads = f_cpu_count(false).cast<int>();
// store the results
Array<IRModule> results;
std::vector<double> costs;
diff --git a/src/relax/transform/fold_constant.cc b/src/relax/transform/fold_constant.cc
index 352092c433..af38645139 100644
--- a/src/relax/transform/fold_constant.cc
+++ b/src/relax/transform/fold_constant.cc
@@ -117,7 +117,7 @@ class ConstantFolder : public ExprMutator {
// TODO(Hongyi): further check and narrow the scope of foldable function
const auto pf = tvm::ffi::Function::GetGlobalRequired("tir.build");
func = WithAttr(func, tvm::attr::kGlobalSymbol, String("tir_function"));
- runtime::Module rt_module = pf(func, eval_cpu_target);
+ runtime::Module rt_module = pf(func, eval_cpu_target).cast<runtime::Module>();
build_func = rt_module.GetFunction("tir_function");
} catch (const tvm::Error& err) {
// build failure may happen in which case we skip
@@ -294,7 +294,7 @@ class ConstantFolder : public ExprMutator {
}
if (is_known) {
const auto func = tvm::ffi::Function::GetGlobalRequired("relax.run.shape_to_tensor");
- runtime::NDArray vals = func(arr);
+ runtime::NDArray vals = func(arr).cast<runtime::NDArray>();
return Constant(vals);
}
}
diff --git a/src/relax/transform/gradient.cc b/src/relax/transform/gradient.cc
index 1100e2d457..fa96f31cf5 100644
--- a/src/relax/transform/gradient.cc
+++ b/src/relax/transform/gradient.cc
@@ -368,7 +368,8 @@ class BackwardBindingGenerator : private ExprVisitor {
const auto grad_func =
tvm::ffi::Function::GetGlobalRequired(te_grad_func_prefix + te_grad_name);
Var partials =
- grad_func(checkpoint_var, Downcast<Call>(checkpoint_call), adjoint_var, builder_);
+ grad_func(checkpoint_var, Downcast<Call>(checkpoint_call), adjoint_var, builder_)
+ .cast<Var>();
Tuple args = Downcast<Tuple>(call->args[1]);
auto* tuple_sinfo = GetStructInfoAs<TupleStructInfoNode>(partials);
if (!tuple_sinfo) {
diff --git a/src/relax/transform/meta_schedule.cc b/src/relax/transform/meta_schedule.cc
index a1fc299ec3..c0123a43ab 100644
--- a/src/relax/transform/meta_schedule.cc
+++ b/src/relax/transform/meta_schedule.cc
@@ -76,7 +76,7 @@ class MetaScheduleTuner {
Choice choice("tvm.meta_schedule.tune_tir", {target_, work_dir_, max_trials_global_},
"relax.tuning_api.Choice.default_constr_func", {});
Knob knob("meta_schedule.tune_primfunc", {{"0", choice}});
- knob->Apply((*normalize_mod_func_)(f), "0");
+ knob->Apply((*normalize_mod_func_)(f).cast<IRModule>(), "0");
/*
// TODO(@sunggg): revisit when we have a solution for large params
Trace trace = Trace((*normalize_mod_func_)(f), {}, {});
@@ -125,7 +125,7 @@ Pass MetaScheduleApplyDatabase(Optional<String> work_dir, bool enable_warning =
if (const auto* prim_func_node = base_func.as<tir::PrimFuncNode>()) {
tir::PrimFunc prim_func = GetRef<tir::PrimFunc>(prim_func_node);
- IRModule tir_mod = (*normalize_mod_func_)(prim_func);
+ IRModule tir_mod = (*normalize_mod_func_)(prim_func).cast<IRModule>();
if (Optional<meta_schedule::TuningRecord> opt_record =
database->QueryTuningRecord(tir_mod, target, gv->name_hint)) {
meta_schedule::TuningRecord record = opt_record.value();
diff --git a/src/relax/transform/run_codegen.cc b/src/relax/transform/run_codegen.cc
index 73949a19f8..3940ecd70b 100644
--- a/src/relax/transform/run_codegen.cc
+++ b/src/relax/transform/run_codegen.cc
@@ -130,8 +130,8 @@ class CodeGenRunner : ExprMutator {
// removed the module.
const auto RemoveFuncAttrFunc = tvm::ffi::Function::GetGlobal("ir.BaseFuncWithoutAttr");
ICHECK(RemoveFuncAttrFunc.has_value());
- func = (*RemoveFuncAttrFunc)(func, tvm::attr::kGlobalSymbol);
- func = (*RemoveFuncAttrFunc)(func, attr::kCodegen);
+ func = (*RemoveFuncAttrFunc)(func, tvm::attr::kGlobalSymbol).cast<Function>();
+ func = (*RemoveFuncAttrFunc)(func, attr::kCodegen).cast<Function>();
builder_->UpdateFunction(gvar, func);
return create_call_dps_packed(new_func, ret_sinfo);
}
@@ -194,7 +194,8 @@ class CodeGenRunner : ExprMutator {
const auto codegen = tvm::ffi::Function::GetGlobal(codegen_name);
ICHECK(codegen.has_value()) << "Codegen is not found: " << codegen_name << "\n";
- Array<runtime::Module> compiled_functions = (*codegen)(functions, options, constant_names);
+ Array<runtime::Module> compiled_functions =
+ (*codegen)(functions, options, constant_names).cast<Array<runtime::Module>>();
ext_mods.insert(ext_mods.end(), compiled_functions.begin(), compiled_functions.end());
}
diff --git a/src/relax/transform/tuning_api/database.cc b/src/relax/transform/tuning_api/database.cc
index cccd8e7f87..c4e706ebe4 100644
--- a/src/relax/transform/tuning_api/database.cc
+++ b/src/relax/transform/tuning_api/database.cc
@@ -62,13 +62,13 @@ TuningRecord TuningRecord::FromJSON(const ObjectRef& json_obj) {
CHECK(json_array && json_array->size() == 2);
// Load json[0] => trace
{
- const ObjectRef& json_trace = json_array->at(0);
+ const ObjectRef& json_trace = json_array->at(0).cast<ObjectRef>();
trace = Trace::FromJSON(json_trace);
}
// Load json[1] => run_secs
if (json_array->at(1) != nullptr) {
- run_secs = meta_schedule::AsFloatArray(json_array->at(1));
+ run_secs = meta_schedule::AsFloatArray(json_array->at(1).cast<ObjectRef>());
}
} catch (const std::runtime_error& e) { // includes tvm::Error and dmlc::Error
LOG(FATAL) << "ValueError: Unable to parse the JSON object: " << json_obj
@@ -268,7 +268,7 @@ Database Database::JSONDatabase(String path_workload, String path_tuning_record,
ICHECK_EQ(arr->size(), 3);
workload_idxs[task_id] = Downcast<Integer>(arr->at(0)).IntValue();
targets[task_id] = Target(Downcast<Map<String, ffi::Any>>(arr->at(1)));
- records[task_id] = TuningRecord::FromJSON(arr->at(2));
+ records[task_id] = TuningRecord::FromJSON(arr->at(2).cast<ObjectRef>());
} catch (std::runtime_error& e) {
LOG(FATAL) << "ValueError: Unable to parse the JSON object: " << json_obj
<< "\nThe error is: " << e.what();
@@ -300,7 +300,7 @@ Database Database::JSONDatabase(String path_workload, String path_tuning_record,
ICHECK_EQ(arr->size(), 3);
workload_idxs[task_id] = Downcast<Integer>(arr->at(0)).IntValue();
targets[task_id] = Target(Downcast<Map<String, ffi::Any>>(arr->at(1)));
- measurements[task_id] = meta_schedule::AsFloatArray(arr->at(2));
+ measurements[task_id] = meta_schedule::AsFloatArray(arr->at(2).cast<ObjectRef>());
} catch (std::runtime_error& e) {
LOG(FATAL) << "ValueError: Unable to parse the JSON object: " << json_obj
<< "\nThe error is: " << e.what();
diff --git a/src/relax/transform/tuning_api/primitives.cc b/src/relax/transform/tuning_api/primitives.cc
index 9d60a3e185..cdad8e58c3 100644
--- a/src/relax/transform/tuning_api/primitives.cc
+++ b/src/relax/transform/tuning_api/primitives.cc
@@ -41,7 +41,7 @@ Choice::Choice(String transform_func_key, Array<Any> transform_func_args, String
// TODO(sunggg): Currently, it only supports an array of primitive data types.
ObjectRef ChoiceNode::AsJSON() const {
Array<ObjectRef> json_transfrom_args, json_constr_args;
- for (ObjectRef arg : this->transform_func_args) {
+ for (Any arg : this->transform_func_args) {
std::string json_arg = tvm::SaveJSON(arg);
std::string b64_arg = meta_schedule::Base64Encode(json_arg);
json_transfrom_args.push_back(String(b64_arg));
@@ -131,7 +131,7 @@ Knob Knob::FromJSON(const ObjectRef& json) {
name = GetRef<String>(arr0);
for (auto const& x : GetRef<Map<String, ffi::Any>>(arr1)) {
String decision = x.first;
- Choice choice = Choice::FromJSON(x.second);
+ Choice choice = Choice::FromJSON(x.second.cast<ObjectRef>());
choices.Set(decision, choice);
}
} catch (const tvm::Error& e) {
@@ -150,7 +150,7 @@ Trace::Trace(IRModule in_mod, Array<Knob> knobs, Array<String> decisions) {
// Deep-copy IRModule
const auto func_deepcopy =
tvm::ffi::Function::GetGlobalRequired("relax.tuning_api.deepcopy_irmodule");
- IRModule out_mod = func_deepcopy(in_mod);
+ IRModule out_mod = func_deepcopy(in_mod).cast<IRModule>();
// Apply the decision history if provided
int size = knobs.size();
for (int i = 0; i < size; i++) {
@@ -207,7 +207,7 @@ Trace Trace::FromJSON(const ObjectRef& json) {
ICHECK(arr0 && arr1);
for (const Any& elem : *arr0) {
- knobs.push_back(Knob::FromJSON(elem));
+ knobs.push_back(Knob::FromJSON(elem.cast<ObjectRef>()));
}
for (const Any& elem : *arr1) {
diff --git a/src/runtime/c_runtime_api.cc b/src/runtime/c_runtime_api.cc
index 9f17a03216..c69fcfae41 100644
--- a/src/runtime/c_runtime_api.cc
+++ b/src/runtime/c_runtime_api.cc
@@ -47,17 +47,17 @@ namespace runtime {
std::string GetCustomTypeName(uint8_t type_code) {
const auto f = tvm::ffi::Function::GetGlobalRequired("runtime._datatype_get_type_name");
- return f(type_code).operator String();
+ return f(type_code).cast<std::string>();
}
uint8_t GetCustomTypeCode(const std::string& type_name) {
const auto f = tvm::ffi::Function::GetGlobalRequired("runtime._datatype_get_type_code");
- return f(type_name).operator int();
+ return f(type_name).cast<uint8_t>();
}
bool GetCustomTypeRegistered(uint8_t type_code) {
const auto f = tvm::ffi::Function::GetGlobalRequired("runtime._datatype_get_type_registered");
- return f(type_code).operator bool();
+ return f(type_code).cast<bool>();
}
uint8_t ParseCustomDatatype(const std::string& s, const char** scan) {
@@ -131,7 +131,7 @@ class DeviceAPIManager {
ICHECK(allow_missing) << "Device API " << name << " is not enabled.";
return nullptr;
}
- void* ptr = (*f)();
+ void* ptr = (*f)().cast<void*>();
return static_cast<DeviceAPI*>(ptr);
}
};
@@ -580,17 +580,17 @@ int TVMFuncCall(TVMFunctionHandle func, TVMValue* args, int* arg_type_codes, int
// TODO(tvm-team): handle bytes return type here
TVMRuntimeEntry* e = TVMAPIRuntimeStore::Get();
if (rv.type_index() == tvm::ffi::TypeIndex::kTVMFFIDataType) {
- e->ret_str = DLDataTypeToString(rv.operator DLDataType());
+ e->ret_str = DLDataTypeToString(rv.cast<DLDataType>());
*ret_type_code = kTVMStr;
ret_val->v_str = e->ret_str.c_str();
} else if (rv.type_index() == tvm::ffi::TypeIndex::kTVMFFIBytes) {
- e->ret_str = rv.operator std::string();
+ e->ret_str = rv.cast<std::string>();
e->ret_bytes.data = e->ret_str.c_str();
e->ret_bytes.size = e->ret_str.length();
*ret_type_code = kTVMBytes;
ret_val->v_handle = &(e->ret_bytes);
} else if (rv.type_index() == tvm::ffi::TypeIndex::kTVMFFIStr) {
- e->ret_str = rv.operator std::string();
+ e->ret_str = rv.cast<std::string>();
*ret_type_code = kTVMStr;
ret_val->v_str = e->ret_str.c_str();
}
@@ -748,18 +748,18 @@ int TVMDeviceCopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHandle stream
TVM_REGISTER_GLOBAL(tvm::runtime::symbol::tvm_set_device)
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLDevice dev;
- dev.device_type = static_cast<DLDeviceType>(args[0].operator int());
- dev.device_id = args[1];
+ dev.device_type = static_cast<DLDeviceType>(args[0].cast<int>());
+ dev.device_id = args[1].cast<int>();
DeviceAPIManager::Get(dev)->SetDevice(dev);
});
// set device api
TVM_REGISTER_GLOBAL("runtime.GetDeviceAttr").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLDevice dev;
- dev.device_type = static_cast<DLDeviceType>(args[0].operator int());
- dev.device_id = args[1];
+ dev.device_type = static_cast<DLDeviceType>(args[0].cast<int>());
+ dev.device_id = args[1].cast<int>();
- DeviceAttrKind kind = static_cast<DeviceAttrKind>(args[2].operator int());
+ DeviceAttrKind kind = static_cast<DeviceAttrKind>(args[2].cast<int>());
if (kind == kExist) {
DeviceAPI* api = DeviceAPIManager::Get(dev.device_type, true);
if (api != nullptr) {
diff --git a/src/runtime/const_loader_module.cc b/src/runtime/const_loader_module.cc
index 05833f5859..e202129f6a 100644
--- a/src/runtime/const_loader_module.cc
+++ b/src/runtime/const_loader_module.cc
@@ -141,7 +141,7 @@ class ConstLoaderModuleNode : public ModuleNode {
if (init != nullptr) {
auto md = GetRequiredConstants(symbol);
// Initialize the module with constants.
- int ret = init(md);
+ int ret = init(md).cast<int>();
// Report the error if initialization is failed.
ICHECK_EQ(ret, 0) << TVMGetLastError();
break;
diff --git a/src/runtime/container.cc b/src/runtime/container.cc
index 838e1c0f03..004feb3afb 100644
--- a/src/runtime/container.cc
+++ b/src/runtime/container.cc
@@ -94,7 +94,7 @@ TVM_REGISTER_OBJECT_TYPE(ShapeTupleObj);
TVM_REGISTER_GLOBAL("runtime.ShapeTuple").set_body_packed([](ffi::PackedArgs args, Any* ret) {
std::vector<ShapeTuple::index_type> shape;
for (int i = 0; i < args.size(); ++i) {
- shape.push_back(args[i]);
+ shape.push_back(args[i].cast<ShapeTuple::index_type>());
}
*ret = ShapeTuple(shape);
});
diff --git a/src/runtime/contrib/amx/amx_config.cc b/src/runtime/contrib/amx/amx_config.cc
index a3a216b993..80da3e71eb 100644
--- a/src/runtime/contrib/amx/amx_config.cc
+++ b/src/runtime/contrib/amx/amx_config.cc
@@ -77,8 +77,8 @@ void init_tile_config(__tilecfg_u* dst, uint16_t cols, uint8_t rows) {
}
TVM_REGISTER_GLOBAL("runtime.amx_tileconfig").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- int rows = args[0];
- int cols = args[1];
+ int rows = args[0].cast<int>();
+ int cols = args[1].cast<int>();
LOG(INFO) << "rows: " << rows << ", cols:" << cols;
// -----------Config for AMX tile resgister----------------------
__tilecfg_u cfg;
diff --git a/src/runtime/contrib/cblas/cblas.cc b/src/runtime/contrib/cblas/cblas.cc
index 9d7a86e09d..edb4d7e0af 100644
--- a/src/runtime/contrib/cblas/cblas.cc
+++ b/src/runtime/contrib/cblas/cblas.cc
@@ -124,7 +124,7 @@ struct CblasDgemmBatchIterativeOp {
// matrix multiplication for row major
TVM_REGISTER_GLOBAL("tvm.contrib.cblas.matmul").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
+ auto A = args[0].cast<DLTensor*>();
ICHECK(TypeMatch(A->dtype, kDLFloat, 32) || TypeMatch(A->dtype, kDLFloat, 64));
if (TypeMatch(A->dtype, kDLFloat, 32))
@@ -135,7 +135,7 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cblas.matmul").set_body_packed([](TVMArgs args,
TVM_REGISTER_GLOBAL("tvm.contrib.cblas.batch_matmul")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
+ auto A = args[0].cast<DLTensor*>();
ICHECK(TypeMatch(A->dtype, kDLFloat, 32) || TypeMatch(A->dtype, kDLFloat, 64));
if (TypeMatch(A->dtype, kDLFloat, 32)) {
CallBatchGemm(args, ret, CblasSgemmBatchOp());
@@ -146,7 +146,7 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cblas.batch_matmul")
TVM_REGISTER_GLOBAL("tvm.contrib.cblas.batch_matmul_iterative")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
+ auto A = args[0].cast<DLTensor*>();
ICHECK(TypeMatch(A->dtype, kDLFloat, 32) || TypeMatch(A->dtype, kDLFloat, 64));
if (TypeMatch(A->dtype, kDLFloat, 32)) {
CallBatchGemm(args, ret, CblasSgemmBatchIterativeOp());
diff --git a/src/runtime/contrib/cblas/dnnl_blas.cc b/src/runtime/contrib/cblas/dnnl_blas.cc
index 6b9391f9c4..a742d94933 100644
--- a/src/runtime/contrib/cblas/dnnl_blas.cc
+++ b/src/runtime/contrib/cblas/dnnl_blas.cc
@@ -47,7 +47,7 @@ struct DNNLSgemmOp {
// matrix multiplication for row major
TVM_REGISTER_GLOBAL("tvm.contrib.dnnl.matmul").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
+ auto A = args[0].cast<DLTensor*>();
ICHECK(TypeMatch(A->dtype, kDLFloat, 32));
CallGemm(args, ret, DNNLSgemmOp());
});
diff --git a/src/runtime/contrib/cblas/gemm_common.h b/src/runtime/contrib/cblas/gemm_common.h
index 0c33d7dd21..d211d7c7ff 100644
--- a/src/runtime/contrib/cblas/gemm_common.h
+++ b/src/runtime/contrib/cblas/gemm_common.h
@@ -74,11 +74,11 @@ inline int ColumnCount(const DLTensor* tensor, bool trans, int batch_offset = 0)
// major, so this we switch the arguments.
template <typename TGemmOp>
inline void CallGemm(TVMArgs args, TVMRetValue* ret, TGemmOp op) {
- DLTensor* A = args[0];
- DLTensor* B = args[1];
- DLTensor* C = args[2];
- bool transa = args[3];
- bool transb = args[4];
+ auto A = args[0].cast<DLTensor*>();
+ auto B = args[1].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
+ bool transa = args[3].cast<bool>();
+ bool transb = args[4].cast<bool>();
int bit_depth = sizeof(typename TGemmOp::TDatatype) * 8;
ICHECK_EQ(A->ndim, 2);
ICHECK_EQ(B->ndim, 2);
@@ -97,8 +97,8 @@ inline void CallGemm(TVMArgs args, TVMRetValue* ret, TGemmOp op) {
ICHECK(TypeMatch(B->dtype, kDLFloat, bit_depth));
ICHECK(TypeMatch(C->dtype, kDLFloat, bit_depth));
- double alpha = args.size() > 5 ? args[5].operator double() : 1.0;
- double beta = args.size() > 6 ? args[6].operator double() : 0.0;
+ double alpha = args.size() > 5 ? args[5].cast<double>() : 1.0;
+ double beta = args.size() > 6 ? args[6].cast<double>() : 0.0;
op(transb, transa, ColumnCount(B, transb), RowCount(A, transa), ColumnCount(A, transa),
static_cast<typename TGemmOp::TDatatype>(alpha),
reinterpret_cast<typename TGemmOp::TDatatype*>(static_cast<char*>(B->data) + B->byte_offset),
@@ -113,11 +113,11 @@ inline void CallGemm(TVMArgs args, TVMRetValue* ret, TGemmOp op) {
// major, so this we switch the arguments.
template <typename TGemmOp>
inline void CallU8S8S32Gemm(TVMArgs args, TVMRetValue* ret, TGemmOp op) {
- DLTensor* A = args[0];
- DLTensor* B = args[1];
- DLTensor* C = args[2];
- bool transa = args[3];
- bool transb = args[4];
+ auto A = args[0].cast<DLTensor*>();
+ auto B = args[1].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
+ bool transa = args[3].cast<bool>();
+ bool transb = args[4].cast<bool>();
// Set the sgemm attributes. Currently, support is limited to CblasFixOffset with all offsets
// equal to 0. This is sufficient for dense.
@@ -145,8 +145,8 @@ inline void CallU8S8S32Gemm(TVMArgs args, TVMRetValue* ret, TGemmOp op) {
ICHECK(TypeMatch(A->dtype, kDLUInt, 8));
ICHECK(TypeMatch(B->dtype, kDLInt, 8));
ICHECK(TypeMatch(C->dtype, kDLInt, 32));
- double alpha = args.size() > 5 ? args[5].operator double() : 1.0;
- double beta = args.size() > 6 ? args[6].operator double() : 0.0;
+ double alpha = args.size() > 5 ? args[5].cast<double>() : 1.0;
+ double beta = args.size() > 6 ? args[6].cast<double>() : 0.0;
op(transb, transa, ColumnCount(B, transb), RowCount(A, transa), ColumnCount(A, transa),
static_cast<float>(alpha),
reinterpret_cast<void*>(static_cast<char*>(B->data) + B->byte_offset), ColumnStride(B),
@@ -183,11 +183,11 @@ inline int ColumnCount3D(DLTensor* tensor, bool trans) { return tensor->shape[tr
template <typename TBatchGemmOp>
inline void CallBatchGemm(TVMArgs args, TVMRetValue* ret, TBatchGemmOp op) {
using DType = typename TBatchGemmOp::TDatatype;
- DLTensor* A = args[0];
- DLTensor* B = args[1];
- DLTensor* C = args[2];
- bool transa = args[3];
- bool transb = args[4];
+ auto A = args[0].cast<DLTensor*>();
+ auto B = args[1].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
+ bool transa = args[3].cast<bool>();
+ bool transb = args[4].cast<bool>();
int bit_depth = sizeof(DType) * 8;
@@ -209,8 +209,8 @@ inline void CallBatchGemm(TVMArgs args, TVMRetValue* ret, TBatchGemmOp op) {
ICHECK(TypeMatch(B->dtype, kDLFloat, bit_depth));
ICHECK(TypeMatch(C->dtype, kDLFloat, bit_depth));
- double alpha = args.size() > 5 ? args[5].operator double() : 1.0;
- double beta = args.size() > 6 ? args[6].operator double() : 0.0;
+ double alpha = args.size() > 5 ? args[5].cast<double>() : 1.0;
+ double beta = args.size() > 6 ? args[6].cast<double>() : 0.0;
int A_stride = A->shape[1] * A->shape[2];
int B_stride = B->shape[1] * B->shape[2];
diff --git a/src/runtime/contrib/cblas/mkl.cc b/src/runtime/contrib/cblas/mkl.cc
index 6eac35418d..56a45b3f81 100644
--- a/src/runtime/contrib/cblas/mkl.cc
+++ b/src/runtime/contrib/cblas/mkl.cc
@@ -155,7 +155,7 @@ struct MKLDgemmBatchIterativeOp {
// matrix multiplication for row major
TVM_REGISTER_GLOBAL("tvm.contrib.mkl.matmul").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
+ auto A = args[0].cast<DLTensor*>();
ICHECK(TypeMatch(A->dtype, kDLFloat, 32) || TypeMatch(A->dtype, kDLFloat, 64));
if (TypeMatch(A->dtype, kDLFloat, 32))
@@ -167,9 +167,9 @@ TVM_REGISTER_GLOBAL("tvm.contrib.mkl.matmul").set_body_packed([](TVMArgs args, T
// integer matrix multiplication for row major
TVM_REGISTER_GLOBAL("tvm.contrib.mkl.matmul_u8s8s32")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
- DLTensor* B = args[1];
- DLTensor* C = args[2];
+ auto A = args[0].cast<DLTensor*>();
+ auto B = args[1].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
ICHECK(TypeMatch(A->dtype, kDLUInt, 8) && TypeMatch(B->dtype, kDLInt, 8) &&
TypeMatch(C->dtype, kDLInt, 32));
@@ -178,7 +178,7 @@ TVM_REGISTER_GLOBAL("tvm.contrib.mkl.matmul_u8s8s32")
TVM_REGISTER_GLOBAL("tvm.contrib.mkl.batch_matmul")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
+ auto A = args[0].cast<DLTensor*>();
ICHECK(TypeMatch(A->dtype, kDLFloat, 32) || TypeMatch(A->dtype, kDLFloat, 64));
if (TypeMatch(A->dtype, kDLFloat, 32)) {
CallBatchGemm(args, ret, MKLSgemmBatchOp());
@@ -189,7 +189,7 @@ TVM_REGISTER_GLOBAL("tvm.contrib.mkl.batch_matmul")
TVM_REGISTER_GLOBAL("tvm.contrib.mkl.batch_matmul_iterative")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
+ auto A = args[0].cast<DLTensor*>();
ICHECK(TypeMatch(A->dtype, kDLFloat, 32) || TypeMatch(A->dtype, kDLFloat, 64));
if (TypeMatch(A->dtype, kDLFloat, 32)) {
CallBatchGemm(args, ret, MKLSgemmBatchIterativeOp());
diff --git a/src/runtime/contrib/cublas/cublas.cc b/src/runtime/contrib/cublas/cublas.cc
index f5d72db859..4907bc9a30 100644
--- a/src/runtime/contrib/cublas/cublas.cc
+++ b/src/runtime/contrib/cublas/cublas.cc
@@ -306,11 +306,11 @@ void CallCublasLt(cublasLtHandle_t hdl, cudaStream_t stream,
}
inline void CallLtIgemm(TVMArgs args, TVMRetValue* ret, cublasLtHandle_t hdl, cudaStream_t stream) {
- DLTensor* A = args[0];
- DLTensor* B = args[1];
- DLTensor* C = args[2];
- bool transa = args[3];
- bool transb = args[4];
+ auto A = args[0].cast<DLTensor*>();
+ auto B = args[1].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
+ bool transa = args[3].cast<bool>();
+ bool transb = args[4].cast<bool>();
// Reversed strides indicates an in-place transpose operation.
transa = IsInPlaceTransposed(A) ? !transa : transa;
transb = IsInPlaceTransposed(B) ? !transb : transb;
@@ -337,8 +337,8 @@ inline void CallLtIgemm(TVMArgs args, TVMRetValue* ret, cublasLtHandle_t hdl, cu
ICHECK(TypeMatch(C->dtype, kDLInt, 32));
ICHECK(CheckMixPrecisionType(A->dtype, C->dtype)) << "Unsupported data type";
- int32_t alpha = args.size() > 5 ? args[5].operator int32_t() : 1;
- int32_t beta = args.size() > 6 ? args[6].operator int32_t() : 0;
+ int32_t alpha = args.size() > 5 ? args[5].cast<int32_t>() : 1;
+ int32_t beta = args.size() > 6 ? args[6].cast<int32_t>() : 0;
cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr;
auto A_data = reinterpret_cast<void*>(static_cast<char*>(A->data) + A->byte_offset);
auto B_data = reinterpret_cast<void*>(static_cast<char*>(B->data) + B->byte_offset);
@@ -378,11 +378,11 @@ inline void CallLtIgemm(TVMArgs args, TVMRetValue* ret, cublasLtHandle_t hdl, cu
#endif
inline void CallGemmEx(TVMArgs args, TVMRetValue* ret, cublasHandle_t hdl) {
- DLTensor* A = args[0];
- DLTensor* B = args[1];
- DLTensor* C = args[2];
- bool transa = args[3];
- bool transb = args[4];
+ auto A = args[0].cast<DLTensor*>();
+ auto B = args[1].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
+ bool transa = args[3].cast<bool>();
+ bool transb = args[4].cast<bool>();
ICHECK_EQ(A->ndim, 2);
ICHECK_EQ(B->ndim, 2);
ICHECK_EQ(C->ndim, 2);
@@ -405,8 +405,8 @@ inline void CallGemmEx(TVMArgs args, TVMRetValue* ret, cublasHandle_t hdl) {
<< "leading dimension must divide 4 for int8 gemm";
ICHECK(!TypeMatch(B->dtype, kDLInt, 8) || ColumnStride(B) % 4 == 0)
<< "leading dimension must divide 4 for int8 gemm";
- double alpha = args.size() > 5 ? args[5].operator double() : 1.0;
- double beta = args.size() > 6 ? args[6].operator double() : 0.0;
+ double alpha = args.size() > 5 ? args[5].cast<double>() : 1.0;
+ double beta = args.size() > 6 ? args[6].cast<double>() : 0.0;
cudaDataType_t cuda_in_type = GetCudaDataType(A->dtype);
cudaDataType_t cuda_out_type = GetCudaDataType(C->dtype);
@@ -436,11 +436,11 @@ inline void CallGemmEx(TVMArgs args, TVMRetValue* ret, cublasHandle_t hdl) {
}
inline void CallBatchGemmEx(TVMArgs args, TVMRetValue* ret, cublasHandle_t hdl) {
- DLTensor* A = args[0];
- DLTensor* B = args[1];
- DLTensor* C = args[2];
- bool transa = args[3];
- bool transb = args[4];
+ auto A = args[0].cast<DLTensor*>();
+ auto B = args[1].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
+ bool transa = args[3].cast<bool>();
+ bool transb = args[4].cast<bool>();
ICHECK_EQ(A->ndim, 3);
ICHECK_EQ(B->ndim, 3);
ICHECK_EQ(C->ndim, 3);
@@ -464,8 +464,8 @@ inline void CallBatchGemmEx(TVMArgs args, TVMRetValue* ret, cublasHandle_t hdl)
<< "leading dimension must divide 4 for int8 gemm";
ICHECK(!TypeMatch(B->dtype, kDLInt, 8) || ColumnStride3D(B) % 4 == 0)
<< "leading dimension must divide 4 for int8 gemm";
- double alpha = args.size() > 5 ? args[5].operator double() : 1.0;
- double beta = args.size() > 6 ? args[6].operator double() : 0.0;
+ double alpha = args.size() > 5 ? args[5].cast<double>() : 1.0;
+ double beta = args.size() > 6 ? args[6].cast<double>() : 0.0;
int A_stride = A->shape[1] * A->shape[2];
int B_stride = B->shape[1] * B->shape[2];
@@ -515,8 +515,8 @@ inline void CallBatchGemmEx(TVMArgs args, TVMRetValue* ret, cublasHandle_t hdl)
// matrix multiplication for row major
TVM_REGISTER_GLOBAL("tvm.contrib.cublas.matmul")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
- DLTensor* C = args[2];
+ auto A = args[0].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
CuBlasThreadEntry* entry_ptr = CuBlasThreadEntry::ThreadLocal();
@@ -540,7 +540,7 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cublas.matmul")
#if CUDART_VERSION >= 10010
TVM_REGISTER_GLOBAL("tvm.contrib.cublaslt.matmul")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
+ auto A = args[0].cast<DLTensor*>();
CuBlasThreadEntry* entry_ptr = CuBlasThreadEntry::ThreadLocal();
@@ -550,7 +550,7 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cublaslt.matmul")
cublasLtHandle_t ltHandle;
CHECK_CUBLAS_ERROR(cublasLtCreate(<Handle));
auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
- cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>(func().cast<void*>());
CallLtIgemm(args, ret, ltHandle, stream);
CHECK_CUBLAS_ERROR(cublasLtDestroy(ltHandle));
});
@@ -558,8 +558,8 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cublaslt.matmul")
TVM_REGISTER_GLOBAL("tvm.contrib.cublas.batch_matmul")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
- DLTensor* C = args[2];
+ auto A = args[0].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
CuBlasThreadEntry* entry_ptr = CuBlasThreadEntry::ThreadLocal();
diff --git a/src/runtime/contrib/cublas/cublas_json_runtime.cc b/src/runtime/contrib/cublas/cublas_json_runtime.cc
index 1d1593ae97..79ed696a87 100644
--- a/src/runtime/contrib/cublas/cublas_json_runtime.cc
+++ b/src/runtime/contrib/cublas/cublas_json_runtime.cc
@@ -68,7 +68,7 @@ class CublasJSONRuntime : public JSONRuntimeBase {
auto* entry_ptr = tvm::contrib::CuBlasLtThreadEntry::ThreadLocal();
auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
- cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>(func().cast<void*>());
std::vector<const DLTensor*> dl_tensors(NumEntries());
@@ -81,7 +81,7 @@ class CublasJSONRuntime : public JSONRuntimeBase {
NDArray arr = opt_nd.value();
arg = arr.operator->();
} else {
- arg = args[i].operator DLTensor*();
+ arg = args[i].cast<DLTensor*>();
}
dl_tensors[eid] = arg;
diff --git a/src/runtime/contrib/cudnn/conv_backward.cc b/src/runtime/contrib/cudnn/conv_backward.cc
index d27bdccfcb..f387183c33 100644
--- a/src/runtime/contrib/cudnn/conv_backward.cc
+++ b/src/runtime/contrib/cudnn/conv_backward.cc
@@ -187,20 +187,20 @@ void BackwardFilterFindAlgo(int format, int dims, int groups, const int pad[], c
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d.backward_data")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- int mode = args[0];
- int format = args[1];
- int algo = args[2];
+ int mode = args[0].cast<int>();
+ int format = args[1].cast<int>();
+ int algo = args[2].cast<int>();
int pad_v[2], stride_v[2], dilation_v[2];
for (int i = 0; i < 2; i++) {
pad_v[i] = args[3 + i];
stride_v[i] = args[5 + i];
dilation_v[i] = args[7 + i];
}
- DLTensor* dy = args[9];
- DLTensor* w = args[10];
- DLTensor* dx = args[11];
- std::string conv_dtype = args[12];
- int groups = args[13];
+ auto dy = args[9].cast<DLTensor*>();
+ auto w = args[10].cast<DLTensor*>();
+ auto dx = args[11].cast<DLTensor*>();
+ auto conv_dtype = args[12].cast<std::string>();
+ int groups = args[13].cast<int>();
ConvolutionBackwardData(mode, format, algo, 2, groups, pad_v, stride_v, dilation_v, dy, w, dx,
conv_dtype);
@@ -208,18 +208,18 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d.backward_data")
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv.backward_data_find_algo")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- int format = args[0];
- int dims = args[1];
+ int format = args[0].cast<int>();
+ int dims = args[1].cast<int>();
int* pad = static_cast<int*>(static_cast<void*>(args[2]));
int* stride = static_cast<int*>(static_cast<void*>(args[3]));
int* dilation = static_cast<int*>(static_cast<void*>(args[4]));
int* dy_dim = static_cast<int*>(static_cast<void*>(args[5]));
int* w_dim = static_cast<int*>(static_cast<void*>(args[6]));
int* dx_dim = static_cast<int*>(static_cast<void*>(args[7]));
- std::string data_dtype = args[8];
- std::string conv_dtype = args[9];
- int groups = args[10];
- bool verbose = args[11];
+ auto data_dtype = args[8].cast<std::string>();
+ auto conv_dtype = args[9].cast<std::string>();
+ int groups = args[10].cast<int>();
+ bool verbose = args[11].cast<bool>();
BackwardDataFindAlgo(format, dims, groups, pad, stride, dilation, dy_dim, w_dim, dx_dim,
data_dtype, conv_dtype, verbose, ret);
@@ -227,20 +227,20 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv.backward_data_find_algo")
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d.backward_filter")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- int mode = args[0];
- int format = args[1];
- int algo = args[2];
+ int mode = args[0].cast<int>();
+ int format = args[1].cast<int>();
+ int algo = args[2].cast<int>();
int pad_v[2], stride_v[2], dilation_v[2];
for (int i = 0; i < 2; i++) {
pad_v[i] = args[3 + i];
stride_v[i] = args[5 + i];
dilation_v[i] = args[7 + i];
}
- DLTensor* dy = args[9];
- DLTensor* x = args[10];
- DLTensor* dw = args[11];
- std::string conv_dtype = args[12];
- int groups = args[13];
+ auto dy = args[9].cast<DLTensor*>();
+ auto x = args[10].cast<DLTensor*>();
+ auto dw = args[11].cast<DLTensor*>();
+ auto conv_dtype = args[12].cast<std::string>();
+ int groups = args[13].cast<int>();
ConvolutionBackwardFilter(mode, format, algo, 2, groups, pad_v, stride_v, dilation_v, dy, x,
dw, conv_dtype);
@@ -248,18 +248,18 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d.backward_filter")
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv.backward_filter_find_algo")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- int format = args[0];
- int dims = args[1];
+ int format = args[0].cast<int>();
+ int dims = args[1].cast<int>();
int* pad = static_cast<int*>(static_cast<void*>(args[2]));
int* stride = static_cast<int*>(static_cast<void*>(args[3]));
int* dilation = static_cast<int*>(static_cast<void*>(args[4]));
int* dy_dim = static_cast<int*>(static_cast<void*>(args[5]));
int* x_dim = static_cast<int*>(static_cast<void*>(args[6]));
int* dw_dim = static_cast<int*>(static_cast<void*>(args[7]));
- std::string data_dtype = args[8];
- std::string conv_dtype = args[9];
- int groups = args[10];
- bool verbose = args[11];
+ auto data_dtype = args[8].cast<std::string>();
+ auto conv_dtype = args[9].cast<std::string>();
+ int groups = args[10].cast<int>();
+ bool verbose = args[11].cast<bool>();
BackwardFilterFindAlgo(format, dims, groups, pad, stride, dilation, dy_dim, x_dim, dw_dim,
data_dtype, conv_dtype, verbose, ret);
diff --git a/src/runtime/contrib/cudnn/conv_forward.cc b/src/runtime/contrib/cudnn/conv_forward.cc
index 5f21f31d61..c330162729 100644
--- a/src/runtime/contrib/cudnn/conv_forward.cc
+++ b/src/runtime/contrib/cudnn/conv_forward.cc
@@ -155,20 +155,20 @@ void FindAlgo(int format, int dims, int groups, const int pad[], const int strid
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d.forward")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- int mode = args[0];
- int format = args[1];
- int algo = args[2];
+ int mode = args[0].cast<int>();
+ int format = args[1].cast<int>();
+ int algo = args[2].cast<int>();
int pad_v[2], stride_v[2], dilation_v[2];
for (int i = 0; i < 2; i++) {
pad_v[i] = args[3 + i];
stride_v[i] = args[5 + i];
dilation_v[i] = args[7 + i];
}
- DLTensor* x = args[9];
- DLTensor* w = args[10];
- DLTensor* y = args[11];
- std::string conv_dtype = args[12];
- int groups = args[13];
+ auto x = args[9].cast<DLTensor*>();
+ auto w = args[10].cast<DLTensor*>();
+ auto y = args[11].cast<DLTensor*>();
+ auto conv_dtype = args[12].cast<std::string>();
+ int groups = args[13].cast<int>();
ConvolutionForward(mode, format, algo, 2, groups, pad_v, stride_v, dilation_v, x, w, y,
conv_dtype);
@@ -176,23 +176,23 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d.forward")
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d+bias+act.forward")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- int mode = args[0];
- int format = args[1];
- int algo = args[2];
+ int mode = args[0].cast<int>();
+ int format = args[1].cast<int>();
+ int algo = args[2].cast<int>();
int pad_v[2], stride_v[2], dilation_v[2];
for (int i = 0; i < 2; i++) {
pad_v[i] = args[3 + i];
stride_v[i] = args[5 + i];
dilation_v[i] = args[7 + i];
}
- int act = args[9];
- double coef = args[10];
- DLTensor* x = args[11];
- DLTensor* w = args[12];
- DLTensor* bias = args[13];
- DLTensor* y = args[14];
- std::string conv_dtype = args[15];
- int groups = args[16];
+ int act = args[9].cast<int>();
+ double coef = args[10].cast<double>();
+ auto x = args[11].cast<DLTensor*>();
+ auto w = args[12].cast<DLTensor*>();
+ auto bias = args[13].cast<DLTensor*>();
+ auto y = args[14].cast<DLTensor*>();
+ auto conv_dtype = args[15].cast<std::string>();
+ int groups = args[16].cast<int>();
ConvolutionBiasActivationForward(mode, format, algo, 2, groups, act, coef, pad_v, stride_v,
dilation_v, x, w, y, bias, conv_dtype);
@@ -200,20 +200,20 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d+bias+act.forward")
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv3d.forward")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- int mode = args[0];
- int format = args[1];
- int algo = args[2];
+ int mode = args[0].cast<int>();
+ int format = args[1].cast<int>();
+ int algo = args[2].cast<int>();
int pad_v[3], stride_v[3], dilation_v[3];
for (int i = 0; i < 3; i++) {
pad_v[i] = args[3 + i];
stride_v[i] = args[6 + i];
dilation_v[i] = args[9 + i];
}
- DLTensor* x = args[12];
- DLTensor* w = args[13];
- DLTensor* y = args[14];
- std::string conv_dtype = args[15];
- int groups = args[16];
+ auto x = args[12].cast<DLTensor*>();
+ auto w = args[13].cast<DLTensor*>();
+ auto y = args[14].cast<DLTensor*>();
+ auto conv_dtype = args[15].cast<std::string>();
+ int groups = args[16].cast<int>();
ConvolutionForward(mode, format, algo, 3, groups, pad_v, stride_v, dilation_v, x, w, y,
conv_dtype);
@@ -221,18 +221,18 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv3d.forward")
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv.forward_find_algo")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- int format = args[0];
- int dims = args[1];
+ int format = args[0].cast<int>();
+ int dims = args[1].cast<int>();
int* pad = static_cast<int*>(static_cast<void*>(args[2]));
int* stride = static_cast<int*>(static_cast<void*>(args[3]));
int* dilation = static_cast<int*>(static_cast<void*>(args[4]));
int* x_dim = static_cast<int*>(static_cast<void*>(args[5]));
int* w_dim = static_cast<int*>(static_cast<void*>(args[6]));
int* y_dim = static_cast<int*>(static_cast<void*>(args[7]));
- std::string data_dtype = args[8];
- std::string conv_dtype = args[9];
- int groups = args[10];
- bool verbose = args[11];
+ auto data_dtype = args[8].cast<std::string>();
+ auto conv_dtype = args[9].cast<std::string>();
+ int groups = args[10].cast<int>();
+ bool verbose = args[11].cast<bool>();
FindAlgo(format, dims, groups, pad, stride, dilation, x_dim, w_dim, y_dim, data_dtype,
conv_dtype, verbose, ret);
});
diff --git a/src/runtime/contrib/cudnn/softmax.cc b/src/runtime/contrib/cudnn/softmax.cc
index e178df812a..141a84f670 100644
--- a/src/runtime/contrib/cudnn/softmax.cc
+++ b/src/runtime/contrib/cudnn/softmax.cc
@@ -32,9 +32,9 @@ namespace contrib {
using namespace runtime;
void softmax_impl(cudnnSoftmaxAlgorithm_t alg, TVMArgs args, TVMRetValue* ret) {
- DLTensor* x = args[0];
- DLTensor* y = args[1];
- int axis = args[2];
+ auto x = args[0].cast<DLTensor*>();
+ auto y = args[1].cast<DLTensor*>();
+ int axis = args[2].cast<int>();
int ndim = x->ndim;
int64_t* shape = x->shape;
if (axis < 0) axis += ndim;
diff --git a/src/runtime/contrib/cutlass/fp16_group_gemm.cu b/src/runtime/contrib/cutlass/fp16_group_gemm.cu
index 123bf1b8d6..f09925ceec 100644
--- a/src/runtime/contrib/cutlass/fp16_group_gemm.cu
+++ b/src/runtime/contrib/cutlass/fp16_group_gemm.cu
@@ -43,7 +43,7 @@ void tvm_cutlass_group_gemm_sm90(NDArray x, NDArray weight, NDArray indptr, NDAr
// Workspace is used for storing device-side group gemm arguments and cutlass internal workspace.
// Recommened size is 4MB.
static auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
- cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>(func().cast<void*>());
CHECK_EQ(x->ndim, 2);
CHECK_EQ(weight->ndim, 3);
CHECK_EQ(indptr->ndim, 1);
diff --git a/src/runtime/contrib/cutlass/fp8_blockwise_scaled_gemm.cu b/src/runtime/contrib/cutlass/fp8_blockwise_scaled_gemm.cu
index 6de7201856..bf257622f1 100644
--- a/src/runtime/contrib/cutlass/fp8_blockwise_scaled_gemm.cu
+++ b/src/runtime/contrib/cutlass/fp8_blockwise_scaled_gemm.cu
@@ -41,7 +41,7 @@ void tvm_cutlass_fp8_blockwise_scaled_gemm(NDArray a, NDArray b, NDArray scales_
// Recommened size is 4MB.
const auto get_stream_func = tvm::ffi::Function::GetGlobal("runtime.get_cuda_stream");
ICHECK(get_stream_func.has_value());
- cudaStream_t stream = static_cast<cudaStream_t>((*get_stream_func)().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>((*get_stream_func)().cast<void*>());
CHECK_GE(a->ndim, 2);
CHECK_EQ(scales_a->ndim, a->ndim);
@@ -102,7 +102,7 @@ void tvm_cutlass_fp8_blockwise_scaled_bmm(NDArray a, NDArray b, NDArray scales_a
// Recommened size is 4MB.
const auto get_stream_func = tvm::ffi::Function::GetGlobal("runtime.get_cuda_stream");
ICHECK(get_stream_func.has_value());
- cudaStream_t stream = static_cast<cudaStream_t>((*get_stream_func)().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>((*get_stream_func)().cast<void*>());
CHECK_EQ(a->ndim, 3);
CHECK_EQ(scales_a->ndim, 3);
diff --git a/src/runtime/contrib/cutlass/fp8_gemm.cu b/src/runtime/contrib/cutlass/fp8_gemm.cu
index dfffc34e92..4859295705 100644
--- a/src/runtime/contrib/cutlass/fp8_gemm.cu
+++ b/src/runtime/contrib/cutlass/fp8_gemm.cu
@@ -43,7 +43,7 @@ void tvm_cutlass_fp8_gemm(NDArray x, NDArray weight, NDArray workspace, NDArray
// Workspace is used for storing device-side gemm arguments and cutlass internal workspace.
// Recommened size is 4MB.
static auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
- cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>(func().cast<void*>());
CHECK_GE(x->ndim, 2);
CHECK_EQ(weight->ndim, 2);
diff --git a/src/runtime/contrib/cutlass/fp8_group_gemm.cu b/src/runtime/contrib/cutlass/fp8_group_gemm.cu
index 4c1689ca80..fd528a22cc 100644
--- a/src/runtime/contrib/cutlass/fp8_group_gemm.cu
+++ b/src/runtime/contrib/cutlass/fp8_group_gemm.cu
@@ -46,7 +46,7 @@ void tvm_cutlass_fp8_group_gemm(NDArray x, NDArray weight, NDArray indptr, NDArr
// Workspace is used for storing device-side group gemm arguments and cutlass internal workspace.
// Recommened size is 4MB.
static auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
- cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+ cudaStream_t stream = static_cast<cudaStream_t>(func().cast<void*>());
CHECK_EQ(x->ndim, 2);
CHECK_EQ(weight->ndim, 3);
CHECK_EQ(indptr->ndim, 1);
diff --git a/src/runtime/contrib/dnnl/dnnl.cc b/src/runtime/contrib/dnnl/dnnl.cc
index c3afebcc47..822e8ac376 100644
--- a/src/runtime/contrib/dnnl/dnnl.cc
+++ b/src/runtime/contrib/dnnl/dnnl.cc
@@ -349,14 +349,14 @@ extern "C" void dnnl_binary_op(float* data, float* weight, float* out, int algo_
// DNNL Conv2d single OP
TVM_REGISTER_GLOBAL("tvm.contrib.dnnl.conv2d").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* input = args[0];
- DLTensor* weights = args[1];
- DLTensor* output = args[2];
+ auto input = args[0].cast<DLTensor*>();
+ auto weights = args[1].cast<DLTensor*>();
+ auto output = args[2].cast<DLTensor*>();
int p_Ph0_ = args[3], p_Pw0_ = args[4], p_Ph1_ = args[5], p_Pw1_ = args[6], p_Sh_ = args[7],
p_Sw_ = args[8], p_G_ = args[9];
- bool channel_last = args[10];
- bool pre_cast = args[11];
- bool post_cast = args[12];
+ bool channel_last = args[10].cast<bool>();
+ bool pre_cast = args[11].cast<bool>();
+ bool post_cast = args[12].cast<bool>();
int p_N_ = input->shape[0], p_C_ = input->shape[1], p_H_ = input->shape[2],
p_W_ = input->shape[3], p_O_ = output->shape[1], p_Kh_ = weights->shape[2],
diff --git a/src/runtime/contrib/dnnl/dnnl_json_runtime.cc b/src/runtime/contrib/dnnl/dnnl_json_runtime.cc
index 33feadbb8b..323d00f4a6 100644
--- a/src/runtime/contrib/dnnl/dnnl_json_runtime.cc
+++ b/src/runtime/contrib/dnnl/dnnl_json_runtime.cc
@@ -118,7 +118,7 @@ class DNNLJSONRuntime : public JSONRuntimeBase {
TensorRegistry::DLTensorProvider makeIODataProvider(const TVMArgs& args) const {
std::map<uint32_t, const DLTensor*> io_map; // eid to dl tensor map
for (size_t i = 0; i < run_arg_eid_.size(); i++) {
- io_map[run_arg_eid_[i]] = args[i].operator DLTensor*();
+ io_map[run_arg_eid_[i]] = args[i].cast<DLTensor*>();
}
// lambda with captured IO data handlers
diff --git a/src/runtime/contrib/hipblas/hipblas.cc b/src/runtime/contrib/hipblas/hipblas.cc
index e402121d04..a6d4b04482 100644
--- a/src/runtime/contrib/hipblas/hipblas.cc
+++ b/src/runtime/contrib/hipblas/hipblas.cc
@@ -273,11 +273,11 @@ void CallHipblasLt(hipblasLtHandle_t hdl, hipStream_t stream,
}
inline void CallGemmEx(TVMArgs args, TVMRetValue* ret, hipblasHandle_t hdl) {
- DLTensor* A = args[0];
- DLTensor* B = args[1];
- DLTensor* C = args[2];
- bool transa = args[3];
- bool transb = args[4];
+ auto A = args[0].cast<DLTensor*>();
+ auto B = args[1].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
+ bool transa = args[3].cast<bool>();
+ bool transb = args[4].cast<bool>();
ICHECK_EQ(A->ndim, 2);
ICHECK_EQ(B->ndim, 2);
ICHECK_EQ(C->ndim, 2);
@@ -331,11 +331,11 @@ inline void CallGemmEx(TVMArgs args, TVMRetValue* ret, hipblasHandle_t hdl) {
}
inline void CallBatchGemmEx(TVMArgs args, TVMRetValue* ret, hipblasHandle_t hdl) {
- DLTensor* A = args[0];
- DLTensor* B = args[1];
- DLTensor* C = args[2];
- bool transa = args[3];
- bool transb = args[4];
+ auto A = args[0].cast<DLTensor*>();
+ auto B = args[1].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
+ bool transa = args[3].cast<bool>();
+ bool transb = args[4].cast<bool>();
ICHECK_EQ(A->ndim, 3);
ICHECK_EQ(B->ndim, 3);
ICHECK_EQ(C->ndim, 3);
@@ -409,8 +409,8 @@ inline void CallBatchGemmEx(TVMArgs args, TVMRetValue* ret, hipblasHandle_t hdl)
// matrix multiplication for row major
TVM_REGISTER_GLOBAL("tvm.contrib.hipblas.matmul")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
- DLTensor* C = args[2];
+ auto A = args[0].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
HipBlasThreadEntry* entry_ptr = HipBlasThreadEntry::ThreadLocal();
@@ -432,8 +432,8 @@ TVM_REGISTER_GLOBAL("tvm.contrib.hipblas.matmul")
TVM_REGISTER_GLOBAL("tvm.contrib.hipblas.batch_matmul")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
- DLTensor* C = args[2];
+ auto A = args[0].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
HipBlasThreadEntry* entry_ptr = HipBlasThreadEntry::ThreadLocal();
diff --git a/src/runtime/contrib/hipblas/hipblas_json_runtime.cc b/src/runtime/contrib/hipblas/hipblas_json_runtime.cc
index 78f0673e91..ac3e8b87d0 100644
--- a/src/runtime/contrib/hipblas/hipblas_json_runtime.cc
+++ b/src/runtime/contrib/hipblas/hipblas_json_runtime.cc
@@ -65,7 +65,7 @@ class HipblasJSONRuntime : public JSONRuntimeBase {
void Run(TVMArgs args) {
auto* entry_ptr = tvm::contrib::HipBlasLtThreadEntry::ThreadLocal();
static auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_rocm_stream");
- hipStream_t stream = static_cast<hipStream_t>(func().operator void*());
+ hipStream_t stream = static_cast<hipStream_t>(func().cast<void*>());
std::vector<const DLTensor*> dl_tensors(NumEntries());
@@ -80,7 +80,7 @@ class HipblasJSONRuntime : public JSONRuntimeBase {
NDArray arr = args[i];
arg = arr.operator->();
} else {
- arg = args[i].operator DLTensor*();
+ arg = args[i].cast<DLTensor*>();
}
dl_tensors[eid] = arg;
diff --git a/src/runtime/contrib/json/json_runtime.h b/src/runtime/contrib/json/json_runtime.h
index 32a9b754b1..e60132affd 100644
--- a/src/runtime/contrib/json/json_runtime.h
+++ b/src/runtime/contrib/json/json_runtime.h
@@ -131,7 +131,7 @@ class JSONRuntimeBase : public ModuleNode {
}
} else {
// Profile the subgraph.
- profiling::Profiler* prof = static_cast<profiling::Profiler*>(rv->operator void*());
+ profiling::Profiler* prof = static_cast<profiling::Profiler*>(rv->cast<void*>());
this->RunProfile(prof);
}
// String vendor_prof = this->RunProfile(prof);
@@ -142,7 +142,7 @@ class JSONRuntimeBase : public ModuleNode {
ICHECK_EQ(args.size(), 1U);
std::lock_guard<std::mutex> guard(this->initialize_mutex_);
if (!this->initialized_) {
- this->Init(args[0]);
+ this->Init(args[0].cast<Array<NDArray>>());
this->initialized_ = true;
}
*rv = 0;
@@ -212,7 +212,7 @@ class JSONRuntimeBase : public ModuleNode {
NDArray arr = opt_nd.value();
arg = arr.operator->();
} else {
- arg = args[i].operator DLTensor*();
+ arg = args[i].cast<DLTensor*>();
}
// Assign input/output the NDArray pointers to data entry so that we can directly
diff --git a/src/runtime/contrib/miopen/conv_forward.cc b/src/runtime/contrib/miopen/conv_forward.cc
index 70772209c3..bd309a948f 100644
--- a/src/runtime/contrib/miopen/conv_forward.cc
+++ b/src/runtime/contrib/miopen/conv_forward.cc
@@ -36,24 +36,24 @@ using namespace runtime;
TVM_REGISTER_GLOBAL("tvm.contrib.miopen.conv2d.setup")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- const int mode = args[0];
- const int dtype = args[1];
- const int pad_h = args[2];
- const int pad_w = args[3];
- const int stride_h = args[4];
- const int stride_w = args[5];
- const int dilation_h = args[6];
- const int dilation_w = args[7];
- const int x_dim0 = args[8];
- const int x_dim1 = args[9];
- const int x_dim2 = args[10];
- const int x_dim3 = args[11];
- const int w_dim0 = args[12];
- const int w_dim1 = args[13];
- const int w_dim2 = args[14];
- const int w_dim3 = args[15];
- const int n_group = args[16];
- void* out_shape = args[17];
+ const int mode = args[0].cast<int>();
+ const int dtype = args[1].cast<int>();
+ const int pad_h = args[2].cast<int>();
+ const int pad_w = args[3].cast<int>();
+ const int stride_h = args[4].cast<int>();
+ const int stride_w = args[5].cast<int>();
+ const int dilation_h = args[6].cast<int>();
+ const int dilation_w = args[7].cast<int>();
+ const int x_dim0 = args[8].cast<int>();
+ const int x_dim1 = args[9].cast<int>();
+ const int x_dim2 = args[10].cast<int>();
+ const int x_dim3 = args[11].cast<int>();
+ const int w_dim0 = args[12].cast<int>();
+ const int w_dim1 = args[13].cast<int>();
+ const int w_dim2 = args[14].cast<int>();
+ const int w_dim3 = args[15].cast<int>();
+ const int n_group = args[16].cast<int>();
+ void* out_shape = args[17].cast<void*>();
MIOpenThreadEntry* entry_ptr = MIOpenThreadEntry::ThreadLocal();
assert(n_group > 0 && "Group Size > 0 is expected");
@@ -150,18 +150,18 @@ TVM_REGISTER_GLOBAL("tvm.contrib.miopen.conv2d.setup")
TVM_REGISTER_GLOBAL("tvm.contrib.miopen.conv2d.forward")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- const int mode = args[0];
- const int dtype = args[1];
- const int pad_h = args[2];
- const int pad_w = args[3];
- const int stride_h = args[4];
- const int stride_w = args[5];
- const int dilation_h = args[6];
- const int dilation_w = args[7];
- const int algo = args[8];
- const DLTensor* x = args[9];
- const DLTensor* w = args[10];
- const DLTensor* y = args[11];
+ const int mode = args[0].cast<int>();
+ const int dtype = args[1].cast<int>();
+ const int pad_h = args[2].cast<int>();
+ const int pad_w = args[3].cast<int>();
+ const int stride_h = args[4].cast<int>();
+ const int stride_w = args[5].cast<int>();
+ const int dilation_h = args[6].cast<int>();
+ const int dilation_w = args[7].cast<int>();
+ const int algo = args[8].cast<int>();
+ const auto x = args[9].cast<DLTensor*>();
+ const auto w = args[10].cast<DLTensor*>();
+ const auto y = args[11].cast<DLTensor*>();
MIOpenThreadEntry* entry_ptr = MIOpenThreadEntry::ThreadLocal();
entry_ptr->conv_entry.fwd_algo = static_cast<miopenConvFwdAlgorithm_t>(algo);
diff --git a/src/runtime/contrib/miopen/softmax.cc b/src/runtime/contrib/miopen/softmax.cc
index 32c648edee..897372524d 100644
--- a/src/runtime/contrib/miopen/softmax.cc
+++ b/src/runtime/contrib/miopen/softmax.cc
@@ -33,9 +33,9 @@ namespace miopen {
using namespace runtime;
void softmax_impl(TVMArgs args, TVMRetValue* ret, miopenSoftmaxAlgorithm_t alg) {
- DLTensor* x = args[0];
- DLTensor* y = args[1];
- int axis = args[2];
+ auto x = args[0].cast<DLTensor*>();
+ auto y = args[1].cast<DLTensor*>();
+ int axis = args[2].cast<int>();
int ndim = x->ndim;
int64_t* shape = x->shape;
if (axis < 0) axis += ndim;
diff --git a/src/runtime/contrib/mps/conv.mm b/src/runtime/contrib/mps/conv.mm
index d268eae4c3..c8788950a2 100644
--- a/src/runtime/contrib/mps/conv.mm
+++ b/src/runtime/contrib/mps/conv.mm
@@ -26,8 +26,8 @@ using namespace runtime;
TVM_REGISTER_GLOBAL("tvm.contrib.mps.buffer2img")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* buf = args[0];
- DLTensor* img = args[1];
+ auto buf = args[0].cast<DLTensor*>();
+ auto img = args[1].cast<DLTensor*>();
// copy to temp
id<MTLBuffer> mtlbuf = (__bridge id<MTLBuffer>)(buf->data);
MetalThreadEntry* entry_ptr = MetalThreadEntry::ThreadLocal();
@@ -59,8 +59,8 @@ TVM_REGISTER_GLOBAL("tvm.contrib.mps.buffer2img")
TVM_REGISTER_GLOBAL("tvm.contrib.mps.img2buffer")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* img = args[0];
- DLTensor* buf = args[1];
+ auto img = args[0].cast<DLTensor*>();
+ auto buf = args[1].cast<DLTensor*>();
id<MTLBuffer> mtlbuf = (__bridge id<MTLBuffer>)(buf->data);
MPSImage* mpsimg = (__bridge MPSImage*)(img->data);
MetalThreadEntry* entry_ptr = MetalThreadEntry::ThreadLocal();
@@ -78,11 +78,11 @@ TVM_REGISTER_GLOBAL("tvm.contrib.mps.img2buffer")
TVM_REGISTER_GLOBAL("tvm.contrib.mps.conv2d").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
// MPS-NHWC
- DLTensor* data = args[0];
- DLTensor* weight = args[1];
- DLTensor* output = args[2];
- int pad = args[3];
- int stride = args[4];
+ auto data = args[0].cast<DLTensor*>();
+ auto weight = args[1].cast<DLTensor*>();
+ auto output = args[2].cast<DLTensor*>();
+ int pad = args[3].cast<int>();
+ int stride = args[4].cast<int>();
ICHECK_EQ(data->ndim, 4);
ICHECK_EQ(weight->ndim, 4);
diff --git a/src/runtime/contrib/mps/gemm.mm b/src/runtime/contrib/mps/gemm.mm
index 288277ab5b..a87f437563 100644
--- a/src/runtime/contrib/mps/gemm.mm
+++ b/src/runtime/contrib/mps/gemm.mm
@@ -25,11 +25,11 @@ namespace contrib {
using namespace runtime;
TVM_REGISTER_GLOBAL("tvm.contrib.mps.matmul").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
- DLTensor* B = args[1];
- DLTensor* C = args[2];
- bool transa = args[3];
- bool transb = args[4];
+ auto A = args[0].cast<DLTensor*>();
+ auto B = args[1].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
+ bool transa = args[3].cast<bool>();
+ bool transb = args[4].cast<bool>();
// call gemm for simple compact code.
ICHECK_EQ(A->ndim, 2);
ICHECK_EQ(B->ndim, 2);
diff --git a/src/runtime/contrib/mrvl/mrvl_sw_runtime_lib.cc b/src/runtime/contrib/mrvl/mrvl_sw_runtime_lib.cc
index 38b6ba760b..d993072aa2 100644
--- a/src/runtime/contrib/mrvl/mrvl_sw_runtime_lib.cc
+++ b/src/runtime/contrib/mrvl/mrvl_sw_runtime_lib.cc
@@ -74,7 +74,7 @@ static void ReadInputsAndGenerateInputBin(TVMArgs args, const std::string& input
NDArray arr = args[i];
tensor = arr.operator->();
} else {
- tensor = args[i].operator DLTensor*();
+ tensor = args[i].cast<DLTensor*>();
}
std::vector<int64_t> shape;
for (int64_t i = 0; i < tensor->ndim; i++) {
diff --git a/src/runtime/contrib/random/random.cc b/src/runtime/contrib/random/random.cc
index 70a3e59f15..ebc60762c0 100644
--- a/src/runtime/contrib/random/random.cc
+++ b/src/runtime/contrib/random/random.cc
@@ -72,9 +72,9 @@ RandomThreadLocalEntry* RandomThreadLocalEntry::ThreadLocal() {
TVM_REGISTER_GLOBAL("tvm.contrib.random.randint")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
RandomThreadLocalEntry* entry = RandomThreadLocalEntry::ThreadLocal();
- int64_t low = args[0];
- int64_t high = args[1];
- DLTensor* out = args[2];
+ int64_t low = args[0].cast<int64_t>();
+ int64_t high = args[1].cast<int64_t>();
+ auto out = args[2].cast<DLTensor*>();
ICHECK_GT(high, low) << "high must be bigger than low";
ICHECK(out->strides == nullptr);
@@ -106,32 +106,32 @@ TVM_REGISTER_GLOBAL("tvm.contrib.random.randint")
TVM_REGISTER_GLOBAL("tvm.contrib.random.uniform")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
RandomThreadLocalEntry* entry = RandomThreadLocalEntry::ThreadLocal();
- double low = args[0];
- double high = args[1];
- DLTensor* out = args[2];
+ double low = args[0].cast<double>();
+ double high = args[1].cast<double>();
+ auto out = args[2].cast<DLTensor*>();
entry->random_engine.SampleUniform(out, low, high);
});
TVM_REGISTER_GLOBAL("tvm.contrib.random.normal")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
RandomThreadLocalEntry* entry = RandomThreadLocalEntry::ThreadLocal();
- double loc = args[0];
- double scale = args[1];
- DLTensor* out = args[2];
+ double loc = args[0].cast<double>();
+ double scale = args[1].cast<double>();
+ auto out = args[2].cast<DLTensor*>();
entry->random_engine.SampleNormal(out, loc, scale);
});
TVM_REGISTER_GLOBAL("tvm.contrib.random.random_fill")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
RandomThreadLocalEntry* entry = RandomThreadLocalEntry::ThreadLocal();
- DLTensor* out = args[0];
+ auto out = args[0].cast<DLTensor*>();
entry->random_engine.RandomFill(out);
});
TVM_REGISTER_GLOBAL("tvm.contrib.random.random_fill_for_measure")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) -> void {
const auto curand = tvm::ffi::Function::GetGlobal("runtime.contrib.curand.RandomFill");
- DLTensor* out = args[0];
+ auto out = args[0].cast<DLTensor*>();
if (curand.has_value() && out->device.device_type == DLDeviceType::kDLCUDA) {
if (out->dtype.code == DLDataTypeCode::kDLFloat) {
(*curand)(out);
diff --git a/src/runtime/contrib/rocblas/rocblas.cc b/src/runtime/contrib/rocblas/rocblas.cc
index 97951eda6e..bf3f5b998a 100644
--- a/src/runtime/contrib/rocblas/rocblas.cc
+++ b/src/runtime/contrib/rocblas/rocblas.cc
@@ -67,11 +67,11 @@ typedef dmlc::ThreadLocalStore<RocBlasThreadEntry> RocBlasThreadStore;
// matrix multiplication for row major
TVM_REGISTER_GLOBAL("tvm.contrib.rocblas.matmul")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
- DLTensor* B = args[1];
- DLTensor* C = args[2];
- bool transa = args[3];
- bool transb = args[4];
+ auto A = args[0].cast<DLTensor*>();
+ auto B = args[1].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
+ bool transa = args[3].cast<bool>();
+ bool transb = args[4].cast<bool>();
// call gemm for simple compact code.
ICHECK_EQ(A->ndim, 2);
ICHECK_EQ(B->ndim, 2);
@@ -105,11 +105,11 @@ TVM_REGISTER_GLOBAL("tvm.contrib.rocblas.matmul")
TVM_REGISTER_GLOBAL("tvm.contrib.rocblas.batch_matmul")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
- DLTensor* B = args[1];
- DLTensor* C = args[2];
- bool transa = args[3];
- bool transb = args[4];
+ auto A = args[0].cast<DLTensor*>();
+ auto B = args[1].cast<DLTensor*>();
+ auto C = args[2].cast<DLTensor*>();
+ bool transa = args[3].cast<bool>();
+ bool transb = args[4].cast<bool>();
// call gemm for simple compact code.
ICHECK_EQ(A->ndim, 3);
ICHECK_EQ(B->ndim, 3);
diff --git a/src/runtime/contrib/sort/sort.cc b/src/runtime/contrib/sort/sort.cc
index 81f498fc6b..1efaca220f 100644
--- a/src/runtime/contrib/sort/sort.cc
+++ b/src/runtime/contrib/sort/sort.cc
@@ -79,11 +79,11 @@ struct float16 {
// (d1, d2, ..., d(k-1), d(k+1), ..., dn).
TVM_REGISTER_GLOBAL("tvm.contrib.sort.argsort_nms")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* input = args[0];
- DLTensor* sort_num = args[1];
- DLTensor* output = args[2];
- int32_t axis = args[3];
- bool is_ascend = args[4];
+ auto input = args[0].cast<DLTensor*>();
+ auto sort_num = args[1].cast<DLTensor*>();
+ auto output = args[2].cast<DLTensor*>();
+ int32_t axis = args[3].cast<int>();
+ bool is_ascend = args[4].cast<bool>();
auto dtype = input->dtype;
auto data_ptr = static_cast<float*>(input->data);
@@ -217,10 +217,10 @@ void sort(DLTensor* input, DLTensor* output, int32_t axis, bool is_ascend) {
// and sort axis is dk. sort_num should have dimension of
// (d1, d2, ..., d(k-1), d(k+1), ..., dn).
TVM_REGISTER_GLOBAL("tvm.contrib.sort.argsort").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* input = args[0];
- DLTensor* output = args[1];
- int32_t axis = args[2];
- bool is_ascend = args[3];
+ auto input = args[0].cast<DLTensor*>();
+ auto output = args[1].cast<DLTensor*>();
+ int32_t axis = args[2].cast<int32_t>();
+ bool is_ascend = args[3].cast<bool>();
if (axis < 0) {
axis = input->ndim + axis;
}
@@ -312,10 +312,10 @@ TVM_REGISTER_GLOBAL("tvm.contrib.sort.argsort").set_body_packed([](TVMArgs args,
// and sort axis is dk. sort_num should have dimension of
// (d1, d2, ..., d(k-1), d(k+1), ..., dn).
TVM_REGISTER_GLOBAL("tvm.contrib.sort.sort").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* input = args[0];
- DLTensor* output = args[1];
- int32_t axis = args[2];
- bool is_ascend = args[3];
+ auto input = args[0].cast<DLTensor*>();
+ auto output = args[1].cast<DLTensor*>();
+ int32_t axis = args[2].cast<int32_t>();
+ bool is_ascend = args[3].cast<bool>();
if (axis < 0) {
axis = input->ndim + axis;
}
@@ -441,20 +441,20 @@ void topk(DLTensor* input, DLTensor* out_values, DLTensor* out_indices, int k, i
// and sort axis is dk. sort_num should have dimension of
// (d1, d2, ..., d(k-1), d(k+1), ..., dn).
TVM_REGISTER_GLOBAL("tvm.contrib.sort.topk").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* input = args[0];
+ auto input = args[0].cast<DLTensor*>();
DLTensor* values_out = nullptr;
DLTensor* indices_out = nullptr;
- int k = args[args.size() - 4];
- int axis = args[args.size() - 3];
- std::string ret_type = args[args.size() - 2];
- bool is_ascend = args[args.size() - 1];
+ int k = args[args.size() - 4].cast<int>();
+ int axis = args[args.size() - 3].cast<int>();
+ std::string ret_type = args[args.size() - 2].cast<std::string>();
+ bool is_ascend = args[args.size() - 1].cast<bool>();
if (ret_type == "both") {
- values_out = args[1];
- indices_out = args[2];
+ values_out = args[1].cast<DLTensor*>();
+ indices_out = args[2].cast<DLTensor*>();
} else if (ret_type == "values") {
- values_out = args[1];
+ values_out = args[1].cast<DLTensor*>();
} else if (ret_type == "indices") {
- indices_out = args[1];
+ indices_out = args[1].cast<DLTensor*>();
} else {
LOG(FATAL) << "Unsupported ret type: " << ret_type;
}
diff --git a/src/runtime/contrib/tflite/tflite_runtime.cc b/src/runtime/contrib/tflite/tflite_runtime.cc
index 70231eac01..82257e2a82 100644
--- a/src/runtime/contrib/tflite/tflite_runtime.cc
+++ b/src/runtime/contrib/tflite/tflite_runtime.cc
@@ -154,7 +154,7 @@ PackedFunc TFLiteRuntime::GetFunction(const String& name, const ObjectPtr<Object
// Return member functions during query.
if (name == "set_input") {
return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
- int in_idx = args[0];
+ int in_idx = args[0].cast<int>();
ICHECK_GE(in_idx, 0);
this->SetInput(in_idx, args[1]);
});
@@ -165,7 +165,7 @@ PackedFunc TFLiteRuntime::GetFunction(const String& name, const ObjectPtr<Object
return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { this->Invoke(); });
} else if (name == "set_num_threads") {
return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
- int num_threads = args[0];
+ int num_threads = args[0].cast<int>();
CHECK_GE(num_threads, 1);
this->SetNumThreads(num_threads);
});
diff --git a/src/runtime/contrib/thrust/thrust.cu b/src/runtime/contrib/thrust/thrust.cu
index 03517c366d..22a6cf425d 100644
--- a/src/runtime/contrib/thrust/thrust.cu
+++ b/src/runtime/contrib/thrust/thrust.cu
@@ -234,10 +234,10 @@ void thrust_sort_common(DLTensor* input, DLTensor* values_out, DLTensor* indices
TVM_REGISTER_GLOBAL("tvm.contrib.thrust.sort").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
ICHECK_GE(args.num_args, 4);
- DLTensor* input = args[0];
- DLTensor* values_out = args[1];
- DLTensor* indices_out = args[2];
- bool is_ascend = args[3];
+ auto input = args[0].cast<DLTensor*>();
+ auto values_out = args[1].cast<DLTensor*>();
+ auto indices_out = args[2].cast<DLTensor*>();
+ bool is_ascend = args[3].cast<bool>();
DLTensor* workspace = nullptr;
if (args.num_args == 5) {
workspace = args[4];
@@ -281,11 +281,11 @@ void thrust_stable_sort_by_key(DLTensor* keys_in, DLTensor* values_in, DLTensor*
TVM_REGISTER_GLOBAL("tvm.contrib.thrust.stable_sort_by_key")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
ICHECK_GE(args.num_args, 5);
- DLTensor* keys_in = args[0];
- DLTensor* values_in = args[1];
- DLTensor* keys_out = args[2];
- DLTensor* values_out = args[3];
- bool for_scatter = args[4];
+ auto keys_in = args[0].cast<DLTensor*>();
+ auto values_in = args[1].cast<DLTensor*>();
+ auto keys_out = args[2].cast<DLTensor*>();
+ auto values_out = args[3].cast<DLTensor*>();
+ bool for_scatter = args[4].cast<bool>();
DLTensor* workspace = nullptr;
if (args.num_args == 6) {
workspace = args[5];
@@ -396,8 +396,8 @@ void thrust_scan(DLTensor* data, DLTensor* output, bool exclusive, DLTensor* wor
TVM_REGISTER_GLOBAL("tvm.contrib.thrust.sum_scan")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
ICHECK(args.num_args == 2 || args.num_args == 3 || args.num_args == 4);
- DLTensor* data = args[0];
- DLTensor* output = args[1];
+ auto data = args[0].cast<DLTensor*>();
+ auto output = args[1].cast<DLTensor*>();
bool exclusive = false;
DLTensor* workspace = nullptr;
diff --git a/src/runtime/disco/bcast_session.cc b/src/runtime/disco/bcast_session.cc
index 5effbe37a4..2346f26bd0 100644
--- a/src/runtime/disco/bcast_session.cc
+++ b/src/runtime/disco/bcast_session.cc
@@ -79,8 +79,8 @@ void BcastSessionObj::SyncWorker(int worker_id) {
BcastSessionObj::Internal::BroadcastUnpacked(this, DiscoAction::kSyncWorker, worker_id);
TVMArgs args = this->RecvReplyPacked(worker_id);
ICHECK_EQ(args.size(), 2);
- DiscoAction action = static_cast<DiscoAction>(args[0].operator int());
- int ret_worker_id = args[1];
+ DiscoAction action = static_cast<DiscoAction>(args[0].cast<int>());
+ int ret_worker_id = args[1].cast<int>();
ICHECK(action == DiscoAction::kSyncWorker);
ICHECK_EQ(ret_worker_id, worker_id);
}
@@ -92,7 +92,7 @@ DRef BcastSessionObj::CallWithPacked(const TVMArgs& args) {
// tranlsate args into remote calling convention
int reg_id = AllocateReg();
{
- DRef func = args[2];
+ DRef func = args[2].cast<DRef>();
args_vec[0] = static_cast<int>(DiscoAction::kCallPacked);
args_vec[1] = reg_id;
args_vec[2] = func->reg_id;
diff --git a/src/runtime/disco/builtin.cc b/src/runtime/disco/builtin.cc
index 57707656c3..889a152f61 100644
--- a/src/runtime/disco/builtin.cc
+++ b/src/runtime/disco/builtin.cc
@@ -55,7 +55,7 @@ Module LoadVMModule(std::string path, Device device) {
CHECK(vm_load_executable != nullptr)
<< "ValueError: File `" << path
<< "` is not built by RelaxVM, because `vm_load_executable` does not exist";
- Module mod = vm_load_executable();
+ auto mod = vm_load_executable().cast<Module>();
PackedFunc vm_initialization = mod.GetFunction("vm_initialization");
CHECK(vm_initialization != nullptr)
<< "ValueError: File `" << path
diff --git a/src/runtime/disco/disco_worker.cc b/src/runtime/disco/disco_worker.cc
index 541d293a5a..486b9d2ae1 100644
--- a/src/runtime/disco/disco_worker.cc
+++ b/src/runtime/disco/disco_worker.cc
@@ -39,8 +39,8 @@ void DiscoWorker::SetRegister(int reg_id, AnyView value) {
TVMRetValue& rv = register_file.at(reg_id);
if (rv.type_index() == ffi::TypeIndex::kTVMFFINDArray &&
value.type_index() == ffi::TypeIndex::kTVMFFINDArray) {
- NDArray dst = rv;
- NDArray src = value;
+ NDArray dst = rv.cast<NDArray>();
+ NDArray src = value.cast<NDArray>();
dst.CopyFrom(src);
} else {
rv = value;
@@ -53,8 +53,8 @@ struct DiscoWorker::Impl {
using namespace tvm;
while (true) {
TVMArgs args = self->channel->Recv();
- DiscoAction action = static_cast<DiscoAction>(args[0].operator int());
- int64_t reg_id = args[1];
+ DiscoAction action = static_cast<DiscoAction>(args[0].cast<int>());
+ int64_t reg_id = args[1].cast<int64_t>();
switch (action) {
case DiscoAction::kShutDown: {
Shutdown(self);
@@ -65,13 +65,13 @@ struct DiscoWorker::Impl {
break;
}
case DiscoAction::kGetGlobalFunc: {
- GetGlobalFunc(self, reg_id, args[2]);
+ GetGlobalFunc(self, reg_id, args[2].cast<std::string>());
break;
}
case DiscoAction::kCallPacked: {
- int func_reg_id = args[2];
+ int func_reg_id = args[2].cast<int>();
CHECK_LT(func_reg_id, self->register_file.size());
- PackedFunc func = GetReg(self, func_reg_id);
+ PackedFunc func = GetReg(self, func_reg_id).cast<PackedFunc>();
CHECK(func.defined());
CallPacked(self, reg_id, func, args.Slice(3));
break;
@@ -89,12 +89,12 @@ struct DiscoWorker::Impl {
break;
}
case DiscoAction::kDebugGetFromRemote: {
- int worker_id = args[2];
+ int worker_id = args[2].cast<int>();
DebugGetFromRemote(self, reg_id, worker_id);
break;
}
case DiscoAction::kDebugSetRegister: {
- int worker_id = args[2];
+ int worker_id = args[2].cast<int>();
AnyView value = args[3];
DebugSetRegister(self, reg_id, worker_id, value);
break;
@@ -123,7 +123,7 @@ struct DiscoWorker::Impl {
static void CopyFromWorker0(DiscoWorker* self, int reg_id) {
if (self->worker_id == 0) {
NDArray tgt = GetNDArrayFromHost(self);
- NDArray src = GetReg(self, reg_id);
+ NDArray src = GetReg(self, reg_id).cast<NDArray>();
tgt.CopyFrom(src);
}
}
@@ -131,7 +131,7 @@ struct DiscoWorker::Impl {
static void CopyToWorker0(DiscoWorker* self, int reg_id) {
if (self->worker_id == 0) {
NDArray src = GetNDArrayFromHost(self);
- NDArray tgt = GetReg(self, reg_id);
+ NDArray tgt = GetReg(self, reg_id).cast<NDArray>();
tgt.CopyFrom(src);
}
}
diff --git a/src/runtime/disco/distributed/socket_session.cc b/src/runtime/disco/distributed/socket_session.cc
index f56c6d35d4..4a0f038094 100644
--- a/src/runtime/disco/distributed/socket_session.cc
+++ b/src/runtime/disco/distributed/socket_session.cc
@@ -61,7 +61,7 @@ class SocketSessionObj : public BcastSessionObj {
tvm::ffi::Function::GetGlobal("runtime.disco.create_socket_session_local_workers");
ICHECK(f_create_local_session.has_value())
<< "Cannot find function runtime.disco.create_socket_session_local_workers";
- local_session_ = ((*f_create_local_session)(num_workers_per_node)).operator BcastSession();
+ local_session_ = ((*f_create_local_session)(num_workers_per_node)).cast<BcastSession>();
DRef f_init_workers =
local_session_->GetGlobalFunc("runtime.disco.socket_session_init_workers");
local_session_->CallPacked(f_init_workers, num_nodes_, /*node_id=*/0, num_groups,
@@ -107,7 +107,7 @@ class SocketSessionObj : public BcastSessionObj {
remote_channels_[node_id - 1]->Send(ffi::PackedArgs(packed_args, 5));
TVMArgs args = this->RecvReplyPacked(worker_id);
ICHECK_EQ(args.size(), 2);
- ICHECK(static_cast<DiscoAction>(args[0].operator int()) == DiscoAction::kDebugGetFromRemote);
+ ICHECK(static_cast<DiscoAction>(args[0].cast<int>()) == DiscoAction::kDebugGetFromRemote);
TVMRetValue result;
result = args[1];
return result;
@@ -134,7 +134,7 @@ class SocketSessionObj : public BcastSessionObj {
TVMRetValue result;
TVMArgs args = this->RecvReplyPacked(worker_id);
ICHECK_EQ(args.size(), 1);
- ICHECK(static_cast<DiscoAction>(args[0].operator int()) == DiscoAction::kDebugSetRegister);
+ ICHECK(static_cast<DiscoAction>(args[0].cast<int>()) == DiscoAction::kDebugSetRegister);
}
}
@@ -222,10 +222,10 @@ class RemoteSocketSession {
channel_ = std::make_unique<DiscoSocketChannel>(socket_);
TVMArgs metadata = channel_->Recv();
ICHECK_EQ(metadata.size(), 4);
- num_nodes_ = metadata[0].operator int();
- num_workers_per_node_ = metadata[1].operator int();
- num_groups_ = metadata[2].operator int();
- node_id_ = metadata[3].operator int();
+ num_nodes_ = metadata[0].cast<int>();
+ num_workers_per_node_ = metadata[1].cast<int>();
+ num_groups_ = metadata[2].cast<int>();
+ node_id_ = metadata[3].cast<int>();
CHECK_GE(num_local_workers, num_workers_per_node_);
InitLocalSession();
}
@@ -233,8 +233,8 @@ class RemoteSocketSession {
void MainLoop() {
while (true) {
TVMArgs args = channel_->Recv();
- DiscoSocketAction action = static_cast<DiscoSocketAction>(args[0].operator int());
- int worker_id = args[1].operator int();
+ DiscoSocketAction action = static_cast<DiscoSocketAction>(args[0].cast<int>());
+ int worker_id = args[1].cast<int>();
int local_worker_id = worker_id - node_id_ * num_workers_per_node_;
switch (action) {
case DiscoSocketAction::kSend: {
@@ -271,7 +271,7 @@ class RemoteSocketSession {
void InitLocalSession() {
const auto f_create_local_session =
tvm::ffi::Function::GetGlobal("runtime.disco.create_socket_session_local_workers");
- local_session_ = ((*f_create_local_session)(num_workers_per_node_)).operator BcastSession();
+ local_session_ = ((*f_create_local_session)(num_workers_per_node_)).cast<BcastSession>();
DRef f_init_workers =
local_session_->GetGlobalFunc("runtime.disco.socket_session_init_workers");
diff --git a/src/runtime/disco/loader.cc b/src/runtime/disco/loader.cc
index 79618fcd0f..19d2b80f40 100644
--- a/src/runtime/disco/loader.cc
+++ b/src/runtime/disco/loader.cc
@@ -180,7 +180,7 @@ ObjectRef ShardLoaderObj::Create(const std::string& path_to_metadata, const std:
std::string shard_info, Module mod) {
if (shard_info.empty() && mod.defined()) {
if (PackedFunc get_shard_info = mod->GetFunction("get_shard_info"); get_shard_info != nullptr) {
- shard_info = get_shard_info().operator String();
+ shard_info = get_shard_info().cast<String>();
}
}
ObjectPtr<ShardLoaderObj> n = make_object<ShardLoaderObj>();
diff --git a/src/runtime/disco/nccl/nccl_context.h b/src/runtime/disco/nccl/nccl_context.h
index aa466c5263..d70efebdc8 100644
--- a/src/runtime/disco/nccl/nccl_context.h
+++ b/src/runtime/disco/nccl/nccl_context.h
@@ -151,7 +151,7 @@ struct CCLThreadLocalContext {
deviceStream_t GetDefaultStream() {
const auto func = tvm::ffi::Function::GetGlobal("runtime.get_" TVM_DISCO_DEVICE_NAME "_stream");
ICHECK(func.has_value());
- deviceStream_t stream = static_cast<deviceStream_t>((*func)().operator void*());
+ deviceStream_t stream = static_cast<deviceStream_t>((*func)().cast<void*>());
return stream == nullptr ? default_stream : stream;
}
diff --git a/src/runtime/disco/process_session.cc b/src/runtime/disco/process_session.cc
index eb8c3c7355..e0eb0215e2 100644
--- a/src/runtime/disco/process_session.cc
+++ b/src/runtime/disco/process_session.cc
@@ -70,7 +70,7 @@ class ProcessSessionObj final : public BcastSessionObj {
read_fds.reserve(num_workers - 1);
write_fds.reserve(num_workers - 1);
for (int i = 1; i < num_workers; ++i) {
- IntTuple fds = process_pool(i);
+ IntTuple fds = process_pool(i).cast<IntTuple>();
CHECK_EQ(fds.size(), 2) << "ValueError: process_pool(" << i << ") should return a tuple of "
<< "size 2, but got a tuple of size " << fds.size() << ".";
read_fds.push_back(fds[0]);
@@ -107,7 +107,7 @@ class ProcessSessionObj final : public BcastSessionObj {
}
TVMArgs args = this->RecvReplyPacked(worker_id);
ICHECK_EQ(args.size(), 2);
- ICHECK(static_cast<DiscoAction>(args[0].operator int()) == DiscoAction::kDebugGetFromRemote);
+ ICHECK(static_cast<DiscoAction>(args[0].cast<int>()) == DiscoAction::kDebugGetFromRemote);
TVMRetValue result;
result = args[1];
return result;
@@ -133,7 +133,7 @@ class ProcessSessionObj final : public BcastSessionObj {
TVMRetValue result;
TVMArgs args = this->RecvReplyPacked(worker_id);
ICHECK_EQ(args.size(), 1);
- ICHECK(static_cast<DiscoAction>(args[0].operator int()) == DiscoAction::kDebugSetRegister);
+ ICHECK(static_cast<DiscoAction>(args[0].cast<int>()) == DiscoAction::kDebugSetRegister);
}
void BroadcastPacked(const TVMArgs& args) final {
@@ -183,7 +183,7 @@ Session Session::ProcessSession(int num_workers, int num_group, String process_p
const auto pf = tvm::ffi::Function::GetGlobal(process_pool_creator);
CHECK(pf) << "ValueError: Cannot find function " << process_pool_creator
<< " in the registry. Please check if it is registered.";
- PackedFunc process_pool = (*pf)(num_workers, num_group, entrypoint);
+ auto process_pool = (*pf)(num_workers, num_group, entrypoint).cast<PackedFunc>();
auto n = make_object<ProcessSessionObj>(num_workers, num_group, process_pool);
return Session(n);
}
diff --git a/src/runtime/disco/protocol.h b/src/runtime/disco/protocol.h
index c5964cc6df..6023deaf97 100644
--- a/src/runtime/disco/protocol.h
+++ b/src/runtime/disco/protocol.h
@@ -207,7 +207,7 @@ inline void DiscoProtocol<SubClassType>::ReadObject(int* tcode, TVMValue* value)
self->template Read<uint64_t>(&size);
std::string data(size, '\0');
self->template ReadArray<char>(data.data(), size);
- result = ffi::ObjectRef(DiscoDebugObject::LoadFromStr(std::move(data))->data);
+ result = DiscoDebugObject::LoadFromStr(std::move(data))->data.cast<ObjectRef>();
} else {
LOG(FATAL) << "ValueError: Object type is not supported in Disco calling convention: "
<< Object::TypeIndex2Key(type_index) << " (type_index = " << type_index << ")";
@@ -235,7 +235,7 @@ inline std::string DiscoDebugObject::SaveToStr() const {
const auto f = tvm::ffi::Function::GetGlobal("node.SaveJSON");
CHECK(f.has_value()) << "ValueError: Cannot serialize object in non-debugging mode: "
<< obj->GetTypeKey();
- std::string result = (*f)(obj);
+ std::string result = (*f)(obj).cast<std::string>();
result.push_back('0');
return result;
}
diff --git a/src/runtime/disco/session.cc b/src/runtime/disco/session.cc
index bea6f6eef8..ae94f24fbe 100644
--- a/src/runtime/disco/session.cc
+++ b/src/runtime/disco/session.cc
@@ -49,7 +49,7 @@ TVM_REGISTER_GLOBAL("runtime.disco.SessionInitCCL") //
.set_body_method(&SessionObj::InitCCL);
TVM_REGISTER_GLOBAL("runtime.disco.SessionCallPacked")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- Session self = args[0];
+ Session self = args[0].cast<Session>();
*rv = SessionObj::FFI::CallWithPacked(self, args.Slice(1));
});
TVM_REGISTER_GLOBAL("runtime.disco.SessionShutdown").set_body_method(&SessionObj::Shutdown);
diff --git a/src/runtime/dso_library.cc b/src/runtime/dso_library.cc
index f162499787..2600f3fdb9 100644
--- a/src/runtime/dso_library.cc
+++ b/src/runtime/dso_library.cc
@@ -149,10 +149,9 @@ ObjectPtr<Library> CreateDSOLibraryObject(std::string library_path) {
return n;
}
-TVM_REGISTER_GLOBAL("runtime.module.loadfile_so")
- .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- ObjectPtr<Library> n = CreateDSOLibraryObject(args[0]);
- *rv = CreateModuleFromLibrary(n);
- });
+TVM_REGISTER_GLOBAL("runtime.module.loadfile_so").set_body_typed([](std::string library_path) {
+ ObjectPtr<Library> n = CreateDSOLibraryObject(library_path);
+ return CreateModuleFromLibrary(n);
+});
} // namespace runtime
} // namespace tvm
diff --git a/src/runtime/hexagon/hexagon_device_api.cc b/src/runtime/hexagon/hexagon_device_api.cc
index 9eda329cdd..e17779c58f 100644
--- a/src/runtime/hexagon/hexagon_device_api.cc
+++ b/src/runtime/hexagon/hexagon_device_api.cc
@@ -192,11 +192,11 @@ void HexagonDeviceAPI::CopyDataFromTo(const void* from, size_t from_offset, void
TVM_REGISTER_GLOBAL("device_api.hexagon.dma_copy_dltensor")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- DLTensor* dst = args[0];
- DLTensor* src = args[1];
- int size = args[2];
+ auto dst = args[0].cast<DLTensor*>();
+ auto src = args[1].cast<DLTensor*>();
+ int size = args[2].cast<int>();
ICHECK(size > 0);
- bool bypass_cache = args[3];
+ bool bypass_cache = args[3].cast<bool>();
int ret = DMA_RETRY;
do {
@@ -212,11 +212,11 @@ TVM_REGISTER_GLOBAL("device_api.hexagon.dma_copy_dltensor")
TVM_REGISTER_GLOBAL("device_api.hexagon.dma_copy")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
uint32_t queue_id = static_cast<int>(args[0]);
- void* dst = args[1];
- void* src = args[2];
+ void* dst = args[1].cast<void*>();
+ void* src = args[2].cast<void*>();
uint32_t size = static_cast<int>(args[3]);
ICHECK(size > 0);
- bool bypass_cache = args[4];
+ bool bypass_cache = args[4].cast<bool>();
int ret = DMA_RETRY;
do {
@@ -229,7 +229,7 @@ TVM_REGISTER_GLOBAL("device_api.hexagon.dma_copy")
TVM_REGISTER_GLOBAL("device_api.hexagon.dma_wait")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
uint32_t queue_id = static_cast<int>(args[0]);
- int inflight = args[1];
+ int inflight = args[1].cast<int>();
ICHECK(inflight >= 0);
HexagonDeviceAPI::Global()->UserDMA()->Wait(queue_id, inflight);
*rv = static_cast<int32_t>(0);
@@ -251,13 +251,13 @@ TVM_REGISTER_GLOBAL("device_api.hexagon.dma_end_group")
TVM_REGISTER_GLOBAL("device_api.hexagon.alloc_nd")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- int32_t device_type = args[0];
- int32_t device_id = args[1];
- int32_t dtype_code_hint = args[2];
- int32_t dtype_bits_hint = args[3];
- std::string scope = args[4];
+ int32_t device_type = args[0].cast<int32_t>();
+ int32_t device_id = args[1].cast<int32_t>();
+ int32_t dtype_code_hint = args[2].cast<int32_t>();
+ int32_t dtype_bits_hint = args[3].cast<int32_t>();
+ auto scope = args[4].cast<std::string>();
CHECK(scope.find("global.vtcm") != std::string::npos);
- int64_t ndim = args[5];
+ int64_t ndim = args[5].cast<int64_t>();
CHECK((ndim == 1 || ndim == 2) && "Hexagon Device API supports only 1d and 2d allocations");
int64_t* shape = static_cast<int64_t*>(static_cast<void*>(args[6]));
@@ -276,11 +276,11 @@ TVM_REGISTER_GLOBAL("device_api.hexagon.alloc_nd")
TVM_REGISTER_GLOBAL("device_api.hexagon.free_nd")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- int32_t device_type = args[0];
- int32_t device_id = args[1];
- std::string scope = args[2];
+ int32_t device_type = args[0].cast<int32_t>();
+ int32_t device_id = args[1].cast<int32_t>();
+ auto scope = args[2].cast<std::string>();
CHECK(scope.find("global.vtcm") != std::string::npos);
- void* ptr = args[3];
+ void* ptr = args[3].cast<void*>();
Device dev;
dev.device_type = static_cast<DLDeviceType>(device_type);
diff --git a/src/runtime/hexagon/rpc/android/session.cc b/src/runtime/hexagon/rpc/android/session.cc
index aafd5e82df..dc1fa90909 100644
--- a/src/runtime/hexagon/rpc/android/session.cc
+++ b/src/runtime/hexagon/rpc/android/session.cc
@@ -113,10 +113,10 @@ TVM_REGISTER_GLOBAL("tvm.contrib.hexagon.create_hexagon_session")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
ICHECK(args.size() >= 4) << args.size() << " is less than 4";
- std::string session_name = args[0];
- int remote_stack_size_bytes = args[1];
+ auto session_name = args[0].cast<std::string>();
+ int remote_stack_size_bytes = args[1].cast<int>();
// For simulator, the third parameter is sim_args, ignore it.
- int hexagon_rpc_receive_buf_size_bytes = args[3];
+ int hexagon_rpc_receive_buf_size_bytes = args[3].cast<int>();
HexagonTransportChannel* hexagon_channel =
new HexagonTransportChannel(hexagon_rpc_URI CDSP_DOMAIN, remote_stack_size_bytes,
static_cast<uint32_t>(hexagon_rpc_receive_buf_size_bytes));
diff --git a/src/runtime/hexagon/rpc/hexagon/rpc_server.cc b/src/runtime/hexagon/rpc/hexagon/rpc_server.cc
index 690fb7e0c5..97b648bfa3 100644
--- a/src/runtime/hexagon/rpc/hexagon/rpc_server.cc
+++ b/src/runtime/hexagon/rpc/hexagon/rpc_server.cc
@@ -331,15 +331,15 @@ __attribute__((weak)) void _Parse_fde_instr() {}
TVM_REGISTER_GLOBAL("tvm.hexagon.load_module")
.set_body_packed([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
- std::string soname = args[0];
+ auto soname = args[0].cast<std::string>();
tvm::ObjectPtr<tvm::runtime::Library> n = tvm::runtime::CreateDSOLibraryObject(soname);
*rv = CreateModuleFromLibrary(n);
});
TVM_REGISTER_GLOBAL("tvm.hexagon.get_profile_output")
.set_body_packed([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
- std::string profiling_mode = args[0];
- std::string out_file = args[1];
+ auto profiling_mode = args[0].cast<std::string>();
+ auto out_file = args[1].cast<std::string>();
if (profiling_mode.compare("lwp") == 0) {
*rv = WriteLWPOutput(out_file);
} else {
@@ -356,7 +356,7 @@ void SaveBinaryToFile(const std::string& file_name, const std::string& data) {
TVM_REGISTER_GLOBAL("tvm.rpc.server.upload")
.set_body_packed([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
- std::string file_name = args[0];
- std::string data = args[1];
+ auto file_name = args[0].cast<std::string>();
+ auto data = args[1].cast<std::string>();
SaveBinaryToFile(file_name, data);
});
diff --git a/src/runtime/hexagon/rpc/simulator/rpc_server.cc b/src/runtime/hexagon/rpc/simulator/rpc_server.cc
index 17b380e25a..b0d31bf59b 100644
--- a/src/runtime/hexagon/rpc/simulator/rpc_server.cc
+++ b/src/runtime/hexagon/rpc/simulator/rpc_server.cc
@@ -334,15 +334,15 @@ __attribute__((weak)) void _Parse_fde_instr() {}
TVM_REGISTER_GLOBAL("tvm.hexagon.load_module")
.set_body_packed([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
- std::string soname = args[0];
+ auto soname = args[0].cast<std::string>();
tvm::ObjectPtr<tvm::runtime::Library> n = tvm::runtime::CreateDSOLibraryObject(soname);
*rv = CreateModuleFromLibrary(n);
});
TVM_REGISTER_GLOBAL("tvm.hexagon.get_profile_output")
.set_body_packed([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
- std::string profiling_mode = args[0];
- std::string out_file = args[1];
+ auto profiling_mode = args[0].cast<std::string>();
+ auto out_file = args[1].cast<std::string>();
if (profiling_mode.compare("lwp") == 0) {
*rv = WriteLWPOutput(out_file);
} else {
@@ -359,7 +359,7 @@ void SaveBinaryToFile(const std::string& file_name, const std::string& data) {
TVM_REGISTER_GLOBAL("tvm.rpc.server.upload")
.set_body_packed([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
- std::string file_name = args[0];
- std::string data = args[1];
+ auto file_name = args[0].cast<std::string>();
+ auto data = args[1].cast<std::string>();
SaveBinaryToFile(file_name, data);
});
diff --git a/src/runtime/hexagon/rpc/simulator/session.cc b/src/runtime/hexagon/rpc/simulator/session.cc
index 5d7a550629..8c02f55466 100644
--- a/src/runtime/hexagon/rpc/simulator/session.cc
+++ b/src/runtime/hexagon/rpc/simulator/session.cc
@@ -1374,9 +1374,9 @@ TVM_REGISTER_GLOBAL("tvm.contrib.hexagon.create_hexagon_session")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
ICHECK(args.size() >= 4) << args.size() << " is less than 4";
- std::string session_name = args[0];
- int stack_size = args[1];
- std::string sim_args = args[2];
+ auto session_name = args[0].cast<std::string>();
+ int stack_size = args[1].cast<int>();
+ auto sim_args = args[2].cast<std::string>();
auto channel = std::make_unique<SimulatorRPCChannel>(stack_size, sim_args);
std::shared_ptr<RPCEndpoint> endpoint =
RPCEndpoint::Create(std::move(channel), session_name, "", nullptr);
diff --git a/src/runtime/library_module.cc b/src/runtime/library_module.cc
index c4a73f2d29..13aca3df45 100644
--- a/src/runtime/library_module.cc
+++ b/src/runtime/library_module.cc
@@ -103,7 +103,7 @@ Module LoadModuleFromBinary(const std::string& type_key, dmlc::Stream* stream) {
<< "Perhaps you need to recompile with this runtime enabled.";
}
- return (*f)(static_cast<void*>(stream));
+ return (*f)(static_cast<void*>(stream)).cast<Module>();
}
/*!
diff --git a/src/runtime/memory/memory_manager.cc b/src/runtime/memory/memory_manager.cc
index ed53b04840..9b6eae8039 100644
--- a/src/runtime/memory/memory_manager.cc
+++ b/src/runtime/memory/memory_manager.cc
@@ -149,7 +149,7 @@ Allocator* GetDeviceSpecificAllocator(Device dev, AllocatorType type) {
void* valloc;
Allocator* allocator = nullptr;
if (device_alloc_helper) {
- valloc = (*device_alloc_helper)(dev, static_cast<int>(type));
+ valloc = (*device_alloc_helper)(dev, static_cast<int>(type)).cast<void*>();
allocator = static_cast<Allocator*>(valloc);
}
if (nullptr == allocator) {
diff --git a/src/runtime/module.cc b/src/runtime/module.cc
index 39a48a8e76..a52adbd57f 100644
--- a/src/runtime/module.cc
+++ b/src/runtime/module.cc
@@ -85,7 +85,7 @@ Module Module::LoadFromFile(const String& file_name, const String& format) {
<< " resolved to (" << load_f_name << ") in the global registry."
<< "Ensure that you have loaded the correct runtime code, and"
<< "that you are on the correct hardware architecture.";
- Module m = (*f)(file_name, format);
+ Module m = (*f)(file_name, format).cast<Module>();
return m;
}
@@ -159,7 +159,7 @@ bool RuntimeEnabled(const String& target_str) {
} else if (target.length() >= 4 && target.substr(0, 4) == "llvm") {
const auto pf = tvm::ffi::Function::GetGlobal("codegen.llvm_target_enabled");
if (!pf.has_value()) return false;
- return (*pf)(target);
+ return (*pf)(target).cast<bool>();
} else {
LOG(FATAL) << "Unknown optional runtime " << target;
}
diff --git a/src/runtime/opencl/opencl_device_api.cc b/src/runtime/opencl/opencl_device_api.cc
index 0df96b0b62..18c265697b 100644
--- a/src/runtime/opencl/opencl_device_api.cc
+++ b/src/runtime/opencl/opencl_device_api.cc
@@ -762,15 +762,15 @@ void OpenCLWorkspace::Init(const std::string& type_key, const std::string& devic
TVM_REGISTER_GLOBAL("device_api.opencl.alloc_nd")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- int32_t device_type = args[0];
- int32_t device_id = args[1];
- int32_t dtype_code_hint = args[2];
- int32_t dtype_bits_hint = args[3];
- std::string scope = args[4];
+ int32_t device_type = args[0].cast<int32_t>();
+ int32_t device_id = args[1].cast<int32_t>();
+ int32_t dtype_code_hint = args[2].cast<int32_t>();
+ int32_t dtype_bits_hint = args[3].cast<int32_t>();
+ auto scope = args[4].cast<std::string>();
CHECK(scope.find("texture") != std::string::npos);
- int64_t ndim = args[5];
+ int64_t ndim = args[5].cast<int64_t>();
CHECK_EQ(ndim, 2);
- int64_t* shape = static_cast<int64_t*>(static_cast<void*>(args[6]));
+ int64_t* shape = static_cast<int64_t*>(args[6].cast<void*>());
int64_t width = shape[0];
int64_t height = shape[1];
@@ -789,11 +789,11 @@ TVM_REGISTER_GLOBAL("device_api.opencl.alloc_nd")
});
TVM_REGISTER_GLOBAL("device_api.opencl.free_nd").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- int32_t device_type = args[0];
- int32_t device_id = args[1];
- std::string scope = args[2];
+ int32_t device_type = args[0].cast<int32_t>();
+ int32_t device_id = args[1].cast<int32_t>();
+ auto scope = args[2].cast<std::string>();
CHECK(scope.find("texture") != std::string::npos);
- void* data = args[3];
+ void* data = args[3].cast<void*>();
OpenCLWorkspace* ptr = OpenCLWorkspace::Global();
Device dev;
dev.device_type = static_cast<DLDeviceType>(device_type);
diff --git a/src/runtime/opencl/opencl_module.cc b/src/runtime/opencl/opencl_module.cc
index a8508bfb38..3b0dafc243 100644
--- a/src/runtime/opencl/opencl_module.cc
+++ b/src/runtime/opencl/opencl_module.cc
@@ -354,7 +354,7 @@ PackedFunc OpenCLModuleNode::GetFunction(const String& name,
});
} else if (name == "opencl.SetPreCompiledPrograms") {
return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
- this->SetPreCompiledPrograms(args[0]);
+ this->SetPreCompiledPrograms(args[0].cast<std::string>());
});
}
return OpenCLModuleNodeBase::GetFunction(name, sptr_to_self);
diff --git a/src/runtime/profiling.cc b/src/runtime/profiling.cc
index 97a03ba9ae..08fd1fb11a 100644
--- a/src/runtime/profiling.cc
+++ b/src/runtime/profiling.cc
@@ -109,7 +109,7 @@ Timer Timer::Start(Device dev) {
t->Start();
return t;
} else {
- Timer t = f->operator()(dev);
+ Timer t = f->operator()(dev).cast<Timer>();
t->Start();
return t;
}
@@ -166,7 +166,7 @@ void Profiler::StopCall(std::unordered_map<std::string, ObjectRef> extra_metrics
for (const auto& obj : cf.extra_collectors) {
auto collector_metrics = obj.first->Stop(obj.second);
for (auto& p : collector_metrics) {
- cf.extra_metrics[p.first] = p.second;
+ cf.extra_metrics[p.first] = p.second.cast<ObjectRef>();
}
}
in_flight_.pop();
@@ -334,7 +334,7 @@ String ReportNode::AsJSON() const {
s << "{";
for (const auto& kv : calls[i]) {
s << "\"" << kv.first << "\":";
- metric_as_json(s, kv.second);
+ metric_as_json(s, kv.second.cast<ObjectRef>());
if (j < calls[i].size() - 1) {
s << ",";
}
@@ -354,7 +354,7 @@ String ReportNode::AsJSON() const {
s << "\"" << dev_kv.first << "\":{";
for (const auto& metric_kv : dev_kv.second) {
s << "\"" << metric_kv.first << "\":";
- metric_as_json(s, metric_kv.second);
+ metric_as_json(s, metric_kv.second.cast<ObjectRef>());
if (j < dev_kv.second.size() - 1) {
s << ",";
}
@@ -372,7 +372,7 @@ String ReportNode::AsJSON() const {
size_t k = 0;
for (const auto& kv : configuration) {
s << "\"" << kv.first << "\":";
- metric_as_json(s, kv.second);
+ metric_as_json(s, kv.second.cast<ObjectRef>());
if (k < configuration.size() - 1) {
s << ",";
}
@@ -511,7 +511,7 @@ String ReportNode::AsTable(bool sort, bool aggregate, bool compute_col_sums) con
return std::string(call_metric.first) == metric;
});
if (it != call.end()) {
- per_call.push_back((*it).second);
+ per_call.push_back((*it).second.cast<ObjectRef>());
}
}
if (per_call.size() > 0) {
@@ -607,7 +607,7 @@ String ReportNode::AsTable(bool sort, bool aggregate, bool compute_col_sums) con
// fill empty data with empty strings
cols[i].push_back("");
} else {
- cols[i].push_back(print_metric((*it).second));
+ cols[i].push_back(print_metric((*it).second.cast<ObjectRef>()));
}
}
}
@@ -647,7 +647,7 @@ String ReportNode::AsTable(bool sort, bool aggregate, bool compute_col_sums) con
// Add configuration information. It will not be aligned with the columns.
s << std::endl << "Configuration" << std::endl << "-------------" << std::endl;
for (auto kv : configuration) {
- s << kv.first << ": " << print_metric(kv.second) << std::endl;
+ s << kv.first << ": " << print_metric(kv.second.cast<ObjectRef>()) << std::endl;
}
return s.str();
}
diff --git a/src/runtime/regex.cc b/src/runtime/regex.cc
index c99ec88a93..0a3215832e 100644
--- a/src/runtime/regex.cc
+++ b/src/runtime/regex.cc
@@ -35,7 +35,7 @@ bool regex_match(const std::string& match_against, const std::string& regex_patt
<< "RuntimeError: "
<< "The PackedFunc 'tvm.runtime.regex_match' has not been registered. "
<< "This can occur if the TVM Python library has not yet been imported.";
- return (*regex_match_func)(regex_pattern, match_against);
+ return (*regex_match_func)(regex_pattern, match_against).cast<bool>();
}
} // namespace runtime
diff --git a/src/runtime/relax_vm/attn_backend.h b/src/runtime/relax_vm/attn_backend.h
index fe35bc5952..10142460a6 100644
--- a/src/runtime/relax_vm/attn_backend.h
+++ b/src/runtime/relax_vm/attn_backend.h
@@ -169,16 +169,19 @@ class FlashInferPagedPrefillFunc : public PagedPrefillFunc {
IntTuple plan_info_vec;
if (attn_kind == AttnKind::kMHA) {
// Todo(tvm-team): enable cuda graph
- plan_info_vec = plan_func_(
- float_workspace_buffer, int_workspace_buffer, page_locked_int_workspace_buffer,
- qo_indptr->as_ndarray(), page_indptr->as_ndarray(), IntTuple(std::move(kv_len)),
- total_qo_len, batch_size, num_qo_heads, num_kv_heads, page_size,
- /*enable_cuda_graph=*/false, qk_head_dim, v_head_dim, causal, copy_stream);
+ plan_info_vec =
+ plan_func_(float_workspace_buffer, int_workspace_buffer, page_locked_int_workspace_buffer,
+ qo_indptr->as_ndarray(), page_indptr->as_ndarray(),
+ IntTuple(std::move(kv_len)), total_qo_len, batch_size, num_qo_heads,
+ num_kv_heads, page_size,
+ /*enable_cuda_graph=*/false, qk_head_dim, v_head_dim, causal, copy_stream)
+ .cast<IntTuple>();
} else if (attn_kind == AttnKind::kMLA) {
plan_info_vec =
plan_func_(float_workspace_buffer, int_workspace_buffer, page_locked_int_workspace_buffer,
qo_indptr->as_ndarray(), page_indptr->as_ndarray(),
- IntTuple(std::move(kv_len)), num_qo_heads, v_head_dim, causal, copy_stream);
+ IntTuple(std::move(kv_len)), num_qo_heads, v_head_dim, causal, copy_stream)
+ .cast<IntTuple>();
}
if (cached_buffers_.size() <= static_cast<size_t>(depth)) {
@@ -274,7 +277,8 @@ class FlashInferRaggedPrefillFunc : public RaggedPrefillFunc {
plan_func_(float_workspace_buffer, int_workspace_buffer, page_locked_int_workspace_buffer,
qo_indptr->as_ndarray(), kv_indptr->as_ndarray(), IntTuple(std::move(kv_len)),
total_qo_len, batch_size, num_qo_heads, num_kv_heads, /*page_size=*/1,
- /*enable_cuda_graph=*/false, qk_head_dim, v_head_dim, causal, copy_stream);
+ /*enable_cuda_graph=*/false, qk_head_dim, v_head_dim, causal, copy_stream)
+ .cast<IntTuple>();
}
private:
@@ -366,11 +370,13 @@ class FlashInferPagedDecodeFunc : public PagedDecodeFunc {
RoPEMode rope_mode, DataType q_dtype, DataType kv_dtype,
TVMStreamHandle copy_stream) final {
// Todo(tvm-team): enable cuda graph
- IntTuple plan_info_vec = plan_func_(
- float_workspace_buffer, int_workspace_buffer, page_locked_int_workspace_buffer,
- page_indptr->as_ndarray(), batch_size, num_qo_heads, num_kv_heads, page_size,
- /*enable_cuda_graph=*/false, static_cast<int64_t>(rope_mode == RoPEMode::kInline),
- /*window_left=*/-1, qk_head_dim, v_head_dim, q_dtype, kv_dtype, copy_stream);
+ IntTuple plan_info_vec =
+ plan_func_(float_workspace_buffer, int_workspace_buffer, page_locked_int_workspace_buffer,
+ page_indptr->as_ndarray(), batch_size, num_qo_heads, num_kv_heads, page_size,
+ /*enable_cuda_graph=*/false,
+ static_cast<int64_t>(rope_mode == RoPEMode::kInline),
+ /*window_left=*/-1, qk_head_dim, v_head_dim, q_dtype, kv_dtype, copy_stream)
+ .cast<IntTuple>();
if (cached_buffers_.size() <= static_cast<size_t>(depth)) {
cached_buffers_.resize(depth + 1);
diff --git a/src/runtime/relax_vm/builtin.cc b/src/runtime/relax_vm/builtin.cc
index 002134c209..340792e63b 100644
--- a/src/runtime/relax_vm/builtin.cc
+++ b/src/runtime/relax_vm/builtin.cc
@@ -121,23 +121,23 @@ void MatchShape(ffi::PackedArgs args, Any* rv) {
if (auto opt_nd = args[0].as<NDArray>()) {
input_shape = opt_nd.value().Shape();
} else {
- input_shape = args[0];
+ input_shape = args[0].cast<ShapeTuple>();
}
auto heap = args[1].as<DLTensor*>();
int64_t* heap_data = heap.has_value() ? static_cast<int64_t*>((*heap)->data) : nullptr;
- int64_t size = args[2];
+ int64_t size = args[2].cast<int64_t>();
const int64_t kBeginCode = 3;
ICHECK_LE(kBeginCode + size * 2, args.size());
// a function that lazily get context for error reporting
const int64_t kErrorContextOffset = kBeginCode + size * 2;
- Optional<String> err_ctx = args[kErrorContextOffset];
+ Optional<String> err_ctx = args[kErrorContextOffset].cast<String>();
CHECK_EQ(input_shape.size(), size)
<< "RuntimeError: " << err_ctx.value_or("") << " match_cast shape size mismatch.";
for (int64_t i = 0; i < size; ++i) {
- MatchShapeCode code = static_cast<MatchShapeCode>(args[kBeginCode + i * 2].operator int());
- int64_t reg = args[kBeginCode + i * 2 + 1];
+ MatchShapeCode code = static_cast<MatchShapeCode>(args[kBeginCode + i * 2].cast<int>());
+ int64_t reg = args[kBeginCode + i * 2 + 1].cast<int64_t>();
if (code == MatchShapeCode::kAssertEqualToImm) {
CHECK_EQ(input_shape[i], reg)
@@ -194,14 +194,14 @@ void MakeShape(ffi::PackedArgs args, Any* rv) {
// NOTE: heap can be nullptr
auto heap = args[0].as<DLTensor*>();
int64_t* heap_data = heap.has_value() ? static_cast<int64_t*>((*heap)->data) : nullptr;
- int64_t size = args[1];
+ int64_t size = args[1].cast<int64_t>();
const int64_t kBeginCode = 2;
std::vector<int64_t> shape(size);
for (int64_t i = 0; i < size; ++i) {
- MakeShapeCode code = static_cast<MakeShapeCode>(args[kBeginCode + i * 2].operator int());
- int64_t reg = args[kBeginCode + i * 2 + 1];
+ MakeShapeCode code = static_cast<MakeShapeCode>(args[kBeginCode + i * 2].cast<int>());
+ int64_t reg = args[kBeginCode + i * 2 + 1].cast<int64_t>();
if (code == MakeShapeCode::kUseImm) {
shape[i] = reg;
} else {
@@ -223,16 +223,16 @@ TVM_REGISTER_GLOBAL("vm.builtin.make_shape").set_body_packed(MakeShape);
*/
void CheckTensorInfo(ffi::PackedArgs args, Any* rv) {
AnyView arg = args[0];
- int ndim = args[1];
+ int ndim = args[1].cast<int>();
DataType dtype;
Optional<String> err_ctx;
if (args.size() == 3) {
dtype = DataType::Void();
- err_ctx = args[2].operator Optional<String>();
+ err_ctx = args[2].cast<Optional<String>>();
} else {
- dtype = args[2];
- err_ctx = args[3].operator Optional<String>();
+ dtype = args[2].cast<DataType>();
+ err_ctx = args[3].cast<Optional<String>>();
}
auto opt_ptr = arg.as<DLTensor*>();
@@ -285,15 +285,15 @@ void CheckPrimValueInfo(AnyView arg, DataType dtype, Optional<String> err_ctx) {
LOG(FATAL) << "TypeError: " << err_ctx.value_or("") << ", expected dtype " << dtype
<< ", but received ObjectRef of type " << opt_obj.value()->GetTypeKey();
} else if (dtype.is_bool()) {
- arg.operator bool();
+ arg.cast<bool>();
} else if (dtype.is_int()) {
- arg.operator int64_t();
+ arg.cast<int64_t>();
} else if (dtype.is_uint()) {
- arg.operator uint64_t();
+ arg.cast<uint64_t>();
} else if (dtype.is_float()) {
- arg.operator double();
+ arg.cast<double>();
} else if (dtype.is_handle()) {
- arg.operator void*();
+ arg.cast<void*>();
} else {
LOG(FATAL) << "TypeError: " << err_ctx.value_or("") << ", unsupported dtype " << dtype;
}
@@ -364,7 +364,7 @@ TVM_REGISTER_GLOBAL("vm.builtin.alloc_tensor").set_body_method(&StorageObj::Allo
// Closure function handling, calling convention
//-------------------------------------------------
TVM_REGISTER_GLOBAL("vm.builtin.make_closure").set_body_packed([](ffi::PackedArgs args, Any* rv) {
- VMClosure clo = args[0];
+ VMClosure clo = args[0].cast<VMClosure>();
std::vector<Any> saved_args;
saved_args.resize(args.size() - 1);
for (size_t i = 0; i < saved_args.size(); ++i) {
@@ -377,13 +377,13 @@ TVM_REGISTER_GLOBAL("vm.builtin.make_closure").set_body_packed([](ffi::PackedArg
TVM_REGISTER_GLOBAL("vm.builtin.invoke_closure").set_body_packed([](ffi::PackedArgs args, Any* rv) {
// args[0]: vm; args[1]: closure; args[2, 3, ...]: function arguments
VirtualMachine* vm = VirtualMachine::GetContextPtr(args[0]);
- ObjectRef vm_closure = args[1];
+ ObjectRef vm_closure = args[1].cast<ObjectRef>();
vm->InvokeClosurePacked(vm_closure, args.Slice(2), rv);
});
TVM_REGISTER_GLOBAL("vm.builtin.call_tir_dyn").set_body_packed([](ffi::PackedArgs args, Any* rv) {
- ffi::Function func = args[0];
- ShapeTuple to_unpack = args[args.size() - 1];
+ ffi::Function func = args[0].cast<ffi::Function>();
+ ShapeTuple to_unpack = args[args.size() - 1].cast<ShapeTuple>();
size_t num_tensor_args = args.size() - 2;
std::vector<AnyView> packed_args(num_tensor_args + to_unpack.size());
@@ -425,7 +425,7 @@ bool ReadIfCond(AnyView cond) {
if (auto opt_int = cond.as<bool>()) {
return opt_int.value();
}
- NDArray arr = cond.operator tvm::runtime::NDArray();
+ NDArray arr = cond.cast<tvm::runtime::NDArray>();
if (arr->device.device_type != kDLCPU) {
arr = arr.CopyTo(DLDevice{kDLCPU, 0});
}
@@ -469,14 +469,14 @@ TVM_REGISTER_GLOBAL("vm.builtin.invoke_debug_func")
.set_body_packed([](ffi::PackedArgs args, Any* rv) -> void {
ICHECK_GE(args.size(), 3);
int num_args = args.size() - 3;
- ObjectRef io_effect = args[0];
+ ObjectRef io_effect = args[0].cast<ObjectRef>();
ICHECK(!io_effect.defined()) << "ValueError: IOEffect is expected to be lowered to None.";
- String debug_func_name = args[1];
+ String debug_func_name = args[1].cast<String>();
const auto debug_func = tvm::ffi::Function::GetGlobal(debug_func_name);
CHECK(debug_func.has_value()) << "ValueError: " << debug_func_name << " is not found. "
<< "Use the decorator `@tvm.register_func(\"" << debug_func_name
<< "\")` to register it.";
- String line_info = args[2];
+ String line_info = args[2].cast<String>();
std::vector<AnyView> call_args(num_args + 1);
{
call_args[0] = line_info;
diff --git a/src/runtime/relax_vm/cuda/cuda_graph_builtin.cc b/src/runtime/relax_vm/cuda/cuda_graph_builtin.cc
index 91bf54bdbf..60905468f6 100644
--- a/src/runtime/relax_vm/cuda/cuda_graph_builtin.cc
+++ b/src/runtime/relax_vm/cuda/cuda_graph_builtin.cc
@@ -182,7 +182,7 @@ class CUDAGraphExtensionNode : public VMExtensionNode {
}
CUDAGraphCapturedState entry;
- entry.states = capture_func_rv.operator ObjectRef();
+ entry.states = capture_func_rv.cast<ObjectRef>();
CUDA_CALL(cudaGraphInstantiate(&entry.exec, graph, NULL, NULL, 0));
CUDA_CALL(cudaGraphDestroy(graph));
@@ -207,7 +207,7 @@ class CUDAGraphExtensionNode : public VMExtensionNode {
}
TVMRetValue alloc_func_rv;
vm->InvokeClosurePacked(alloc_func, TVMArgs(nullptr, 0), &alloc_func_rv);
- ObjectRef alloc_result = alloc_func_rv.operator ObjectRef();
+ ObjectRef alloc_result = alloc_func_rv.cast<ObjectRef>();
alloc_cache_[entry_index] = alloc_result;
return alloc_result;
}
@@ -244,12 +244,12 @@ TVM_REGISTER_GLOBAL("vm.builtin.cuda_graph.run_or_capture")
ICHECK(args.size() == 5 || args.size() == 4);
VirtualMachine* vm = VirtualMachine::GetContextPtr(args[0]);
auto extension = vm->GetOrCreateExtension<CUDAGraphExtension>();
- ObjectRef capture_func = args[1];
- ObjectRef func_args = args[2];
- int64_t entry_index = args[3];
+ auto capture_func = args[1].cast<ObjectRef>();
+ auto func_args = args[2].cast<ObjectRef>();
+ int64_t entry_index = args[3].cast<int64_t>();
Optional<ShapeTuple> shape_expr = NullOpt;
if (args.size() == 5) {
- shape_expr = args[4].operator ShapeTuple();
+ shape_expr = args[4].cast<ShapeTuple>();
}
*rv = extension->RunOrCapture(vm, capture_func, func_args, entry_index, shape_expr);
});
@@ -259,8 +259,8 @@ TVM_REGISTER_GLOBAL("vm.builtin.cuda_graph.get_cached_alloc")
ICHECK_EQ(args.size(), 3);
VirtualMachine* vm = VirtualMachine::GetContextPtr(args[0]);
auto extension = vm->GetOrCreateExtension<CUDAGraphExtension>();
- ObjectRef alloc_func = args[1];
- int64_t entry_index = args[2];
+ auto alloc_func = args[1].cast<ObjectRef>();
+ int64_t entry_index = args[2].cast<int64_t>();
*rv = extension->GetCachedAllocation(vm, alloc_func, entry_index);
});
diff --git a/src/runtime/relax_vm/kv_state.cc b/src/runtime/relax_vm/kv_state.cc
index 7a3d97fe7c..0a4ce1b41f 100644
--- a/src/runtime/relax_vm/kv_state.cc
+++ b/src/runtime/relax_vm/kv_state.cc
@@ -41,12 +41,12 @@ TVM_REGISTER_GLOBAL("vm.builtin.kv_state_begin_forward")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
CHECK(args.size() == 3 || args.size() == 4)
<< "KVState BeginForward only accepts 3 or 4 arguments";
- KVState kv_state = args[0];
- IntTuple seq_ids = args[1];
- IntTuple append_lengths = args[2];
+ KVState kv_state = args[0].cast<KVState>();
+ IntTuple seq_ids = args[1].cast<IntTuple>();
+ IntTuple append_lengths = args[2].cast<IntTuple>();
Optional<IntTuple> token_tree_parent_ptr;
if (args.size() == 4) {
- token_tree_parent_ptr = args[3].operator Optional<IntTuple>();
+ token_tree_parent_ptr = args[3].cast<Optional<IntTuple>>();
}
kv_state->BeginForward(seq_ids, append_lengths, token_tree_parent_ptr);
});
diff --git a/src/runtime/relax_vm/lm_support.cc b/src/runtime/relax_vm/lm_support.cc
index e3f5cd3c3d..797f19815f 100644
--- a/src/runtime/relax_vm/lm_support.cc
+++ b/src/runtime/relax_vm/lm_support.cc
@@ -305,9 +305,9 @@ TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_view")
CHECK(args.size() == 1 || args.size() == 2)
<< "ValueError: `vm.builtin.attention_kv_cache_view` expects 1 or 2 arguments, but got "
<< args.size() << ".";
- AttentionKVCacheLegacy cache = args[0];
+ AttentionKVCacheLegacy cache = args[0].cast<AttentionKVCacheLegacy>();
if (args.size() == 2) {
- ShapeTuple shape = args[1];
+ ShapeTuple shape = args[1].cast<ShapeTuple>();
*rv = cache->View(shape);
} else {
std::vector<ShapeTuple::index_type> shape;
diff --git a/src/runtime/relax_vm/ndarray_cache_support.cc b/src/runtime/relax_vm/ndarray_cache_support.cc
index 40e73df0a9..5d554528a6 100644
--- a/src/runtime/relax_vm/ndarray_cache_support.cc
+++ b/src/runtime/relax_vm/ndarray_cache_support.cc
@@ -270,15 +270,15 @@ TVM_REGISTER_GLOBAL("vm.builtin.ndarray_cache.get").set_body_typed(NDArrayCache:
TVM_REGISTER_GLOBAL("vm.builtin.ndarray_cache.update")
.set_body_packed([](ffi::PackedArgs args, Any* rv) {
CHECK(args.size() == 2 || args.size() == 3);
- String name = args[0];
- bool is_override = args.size() == 2 ? false : args[2].operator bool();
+ String name = args[0].cast<String>();
+ bool is_override = args.size() == 2 ? false : args[2].cast<bool>();
NDArray arr;
if (auto opt_nd = args[1].as<NDArray>()) {
arr = opt_nd.value();
} else {
// We support converting DLTensors to NDArrays as RPC references are always DLTensors
- DLTensor* tensor = args[1];
+ auto tensor = args[1].cast<DLTensor*>();
std::vector<int64_t> shape;
for (int64_t i = 0; i < tensor->ndim; i++) {
shape.push_back(tensor->shape[i]);
@@ -368,7 +368,7 @@ TVM_REGISTER_GLOBAL("vm.builtin.param_array_from_cache_by_name_unpacked")
LOG(FATAL) << "ValueError: Expect string as input, but get " << args[i].GetTypeKey()
<< " at " << i;
}
- names.push_back(args[i]);
+ names.push_back(args[i].cast<String>());
}
*rv = ParamModuleNode::GetParamByName(names);
});
diff --git a/src/runtime/relax_vm/paged_kv_cache.cc b/src/runtime/relax_vm/paged_kv_cache.cc
index 0d12def981..73b1fc9227 100644
--- a/src/runtime/relax_vm/paged_kv_cache.cc
+++ b/src/runtime/relax_vm/paged_kv_cache.cc
@@ -340,9 +340,11 @@ class PagedAttentionKVCacheObj : public AttentionKVCacheObj {
<< "NVSHMEM is not enabled. Please make sure NVSHMEM is enabled when compiling TVM.";
const auto f_nvshmem_empty = tvm::ffi::Function::GetGlobal("runtime.disco.nvshmem.empty");
ICHECK(f_nvshmem_empty.has_value());
- nvshmem_pages_ = (*f_nvshmem_empty)(
- ShapeTuple({num_layers, num_total_pages, 2, num_kv_heads, page_size, qk_head_dim}), dtype,
- device);
+ nvshmem_pages_ =
+ (*f_nvshmem_empty)(
+ ShapeTuple({num_layers, num_total_pages, 2, num_kv_heads, page_size, qk_head_dim}),
+ dtype, device)
+ .cast<NDArray>();
for (int i = 0; i < num_layers; ++i) {
pages_.push_back(nvshmem_pages_.CreateView(
{num_total_pages_, 2, num_kv_heads_, page_size_, qk_head_dim_}, nvshmem_pages_->dtype,
@@ -2285,8 +2287,8 @@ TVM_REGISTER_GLOBAL("vm.builtin.paged_attention_kv_cache_create")
// Todo: cuda graph arg
CHECK(args.size() == 28 || args.size() == 29)
<< "Invalid number of KV cache constructor args: " << args.size();
- ShapeTuple cache_config = args[0];
- ShapeTuple layer_indptr_tuple = args[1];
+ ShapeTuple cache_config = args[0].cast<ShapeTuple>();
+ ShapeTuple layer_indptr_tuple = args[1].cast<ShapeTuple>();
int num_groups = 1;
int group_id = 0;
if (DiscoWorker* disco_worker = ThreadLocalDiscoWorker::Get()->worker) {
@@ -2298,40 +2300,40 @@ TVM_REGISTER_GLOBAL("vm.builtin.paged_attention_kv_cache_create")
int64_t num_layers = layer_indptr_tuple[group_id + 1] - layer_indptr_tuple[group_id];
int64_t layer_id_begin_offset = layer_indptr_tuple[group_id];
int64_t layer_id_end_offset = layer_indptr_tuple[group_id + 1];
- int64_t num_qo_heads = args[2];
- int64_t num_kv_heads = args[3];
- int64_t qk_head_dim = args[4];
- int64_t v_head_dim = args[5];
- IntTuple attn_kinds = args[6];
- bool enable_kv_transfer = args[7];
- int rope_mode = args[8];
- double rotary_scale = args[9];
- double rotary_theta = args[10];
+ int64_t num_qo_heads = args[2].cast<int64_t>();
+ int64_t num_kv_heads = args[3].cast<int64_t>();
+ int64_t qk_head_dim = args[4].cast<int64_t>();
+ int64_t v_head_dim = args[5].cast<int64_t>();
+ IntTuple attn_kinds = args[6].cast<IntTuple>();
+ bool enable_kv_transfer = args[7].cast<bool>();
+ int rope_mode = args[8].cast<int>();
+ double rotary_scale = args[9].cast<double>();
+ double rotary_theta = args[10].cast<double>();
Optional<NDArray> rope_ext_factors = NullOpt; // args[11]
- NDArray init = args[12];
+ NDArray init = args[12].cast<NDArray>();
Optional<PackedFunc> f_transpose_append_mha = NullOpt; // args[13]
Optional<PackedFunc> f_transpose_append_mla = NullOpt; // args[14]
std::unique_ptr<RaggedPrefillFunc> f_attention_prefill_ragged =
- ConvertRaggedPrefillFunc(args[15], AttnKind::kMHA);
+ ConvertRaggedPrefillFunc(args[15].cast<Array<ObjectRef>>(), AttnKind::kMHA);
std::unique_ptr<PagedPrefillFunc> f_attention_prefill =
- ConvertPagedPrefillFunc(args[16], AttnKind::kMHA);
+ ConvertPagedPrefillFunc(args[16].cast<Array<ObjectRef>>(), AttnKind::kMHA);
std::unique_ptr<PagedDecodeFunc> f_attention_decode =
- ConvertPagedDecodeFunc(args[17], AttnKind::kMHA);
+ ConvertPagedDecodeFunc(args[17].cast<Array<ObjectRef>>(), AttnKind::kMHA);
std::unique_ptr<PagedPrefillFunc> f_attention_prefill_sliding_window =
- ConvertPagedPrefillFunc(args[18], AttnKind::kMHA);
+ ConvertPagedPrefillFunc(args[18].cast<Array<ObjectRef>>(), AttnKind::kMHA);
std::unique_ptr<PagedDecodeFunc> f_attention_decode_sliding_window =
- ConvertPagedDecodeFunc(args[19], AttnKind::kMHA);
+ ConvertPagedDecodeFunc(args[19].cast<Array<ObjectRef>>(), AttnKind::kMHA);
std::unique_ptr<PagedPrefillTreeMaskFunc> f_attention_prefill_with_tree_mask_paged_kv =
- ConvertPagedPrefillTreeMaskFunc(args[20], AttnKind::kMHA);
+ ConvertPagedPrefillTreeMaskFunc(args[20].cast<Array<ObjectRef>>(), AttnKind::kMHA);
std::unique_ptr<RaggedPrefillTreeMaskFunc> f_attention_prefill_with_tree_mask =
- ConvertRaggedPrefillTreeMaskFunc(args[21], AttnKind::kMHA);
+ ConvertRaggedPrefillTreeMaskFunc(args[21].cast<Array<ObjectRef>>(), AttnKind::kMHA);
std::unique_ptr<PagedPrefillFunc> f_mla_prefill =
- ConvertPagedPrefillFunc(args[22], AttnKind::kMLA);
- Array<PackedFunc> f_merge_inplace = args[23];
- PackedFunc f_split_rotary = args[24];
- PackedFunc f_copy_single_page = args[25];
- PackedFunc f_debug_get_kv = args[26];
- PackedFunc f_compact_copy = args[27];
+ ConvertPagedPrefillFunc(args[22].cast<Array<ObjectRef>>(), AttnKind::kMLA);
+ Array<PackedFunc> f_merge_inplace = args[23].cast<Array<PackedFunc>>();
+ PackedFunc f_split_rotary = args[24].cast<PackedFunc>();
+ PackedFunc f_copy_single_page = args[25].cast<PackedFunc>();
+ PackedFunc f_debug_get_kv = args[26].cast<PackedFunc>();
+ PackedFunc f_compact_copy = args[27].cast<PackedFunc>();
if (auto opt_nd = args[11].as<NDArray>()) {
rope_ext_factors = opt_nd.value();
diff --git a/src/runtime/relax_vm/vm.cc b/src/runtime/relax_vm/vm.cc
index e53c50aa1f..5f9440d7f7 100644
--- a/src/runtime/relax_vm/vm.cc
+++ b/src/runtime/relax_vm/vm.cc
@@ -74,7 +74,7 @@ Any IndexIntoNestedObject(Any obj, TVMArgs args, int starting_arg_idx) {
if (!obj.as<ffi::ArrayObj>()) {
LOG(FATAL) << "ValueError: Attempted to index into an object that is not an Array.";
}
- int index = args[i];
+ int index = args[i].cast<int>();
auto arr = Downcast<ffi::Array<Any>>(obj);
// make sure the index is in bounds
if (index >= static_cast<int>(arr.size())) {
@@ -515,7 +515,7 @@ void VirtualMachineImpl::SetInput(std::string func_name, bool with_param_module,
for (int i = 0; i < args.size(); ++i) {
if (with_param_module && i == args.size() - 1) {
// call param func to get the arguments(usually corresponds to param pack.)
- func_args[i] = (args[i].operator Module()).GetFunction("get_params")();
+ func_args[i] = (args[i].cast<Module>()).GetFunction("get_params")();
} else {
func_args[i] = ConvertArgToDevice(args[i], devices[0], allocators[0]);
}
@@ -615,7 +615,7 @@ Optional<VMClosure> VirtualMachineImpl::GetClosureInternal(const String& func_na
// NOTE: should not capture strong ref to self and avoid cyclic ref.
auto impl = PackedFunc([gf_idx](TVMArgs args, TVMRetValue* rv) {
// Per convention, ctx ptr is a VirtualMachine*
- VirtualMachine* ctx_ptr = static_cast<VirtualMachine*>(args[0].operator void*());
+ VirtualMachine* ctx_ptr = static_cast<VirtualMachine*>(args[0].cast<void*>());
std::vector<RegType> inputs(args.size() - 1);
for (size_t i = 0; i < inputs.size(); ++i) {
@@ -632,7 +632,7 @@ Optional<VMClosure> VirtualMachineImpl::GetClosureInternal(const String& func_na
<< finfo.name;
auto impl = PackedFunc([this, finfo, tir_func](TVMArgs args, TVMRetValue* rv) {
// Per convention, ctx ptr is a VirtualMachine*
- VirtualMachine* ctx_ptr = static_cast<VirtualMachine*>(args[0].operator void*());
+ VirtualMachine* ctx_ptr = static_cast<VirtualMachine*>(args[0].cast<void*>());
ICHECK(ctx_ptr == this);
ICHECK_EQ(args.size() - 1, finfo.num_args)
<< "Function " << finfo.name << " expects " << finfo.num_args << " arguments";
@@ -766,7 +766,7 @@ void VirtualMachineImpl::RunInstrCall(VMFrame* curr_frame, Instruction instr) {
ICHECK_LT(static_cast<size_t>(instr.func_idx), this->func_pool_.size());
if (instrument_ == nullptr) {
- this->InvokeClosurePacked(func_pool_[instr.func_idx], args, &ret);
+ this->InvokeClosurePacked(func_pool_[instr.func_idx].cast<ObjectRef>(), args, &ret);
} else {
// insert light-weight instrument callback
call_args[0] = func_pool_[instr.func_idx];
@@ -779,7 +779,8 @@ void VirtualMachineImpl::RunInstrCall(VMFrame* curr_frame, Instruction instr) {
std::vector<std::unique_ptr<std::string>> temp_dtype;
for (int i = 0; i < instr.num_args; ++i) {
if (call_args[i + args_begin_offset].type_index() == ffi::TypeIndex::kTVMFFIDataType) {
- std::string str_dtype = DLDataTypeToString(call_args[i + args_begin_offset]);
+ std::string str_dtype =
+ DLDataTypeToString(call_args[i + args_begin_offset].cast<DLDataType>());
temp_dtype.emplace_back(std::make_unique<std::string>(str_dtype));
call_args[i + args_begin_offset] = *temp_dtype.back();
}
@@ -790,7 +791,7 @@ void VirtualMachineImpl::RunInstrCall(VMFrame* curr_frame, Instruction instr) {
ret_kind = opt_int.value();
}
if (ret_kind != static_cast<int>(VMInstrumentReturnKind::kSkipRun)) {
- this->InvokeClosurePacked(func_pool_[instr.func_idx], args, &ret);
+ this->InvokeClosurePacked(func_pool_[instr.func_idx].cast<ObjectRef>(), args, &ret);
call_args[2] = false;
call_args[3] = ret;
instrument_.CallPacked(call_args.data(), call_args.size(), &rv);
@@ -838,7 +839,7 @@ void VirtualMachineImpl::RunLoop() {
break;
}
case Opcode::If: {
- int64_t cond_val = ReadRegister(curr_frame, instr.cond);
+ int64_t cond_val = ReadRegister(curr_frame, instr.cond).cast<int64_t>();
if (cond_val != 0) {
pc_++;
} else {
@@ -862,9 +863,9 @@ void VirtualMachineImpl::_Init(TVMArgs args, TVMRetValue* rv) {
std::vector<Device> devices;
std::vector<AllocatorType> alloc_types;
for (int i = 0; i < args.size(); i += 3) {
- int device_type = args[i];
- int device_id = args[i + 1];
- int alloc_type = args[i + 2];
+ int device_type = args[i].cast<int>();
+ int device_id = args[i + 1].cast<int>();
+ int alloc_type = args[i + 2].cast<int>();
devices.push_back(Device{DLDeviceType(device_type), device_id});
alloc_types.push_back(AllocatorType(alloc_type));
}
@@ -873,12 +874,12 @@ void VirtualMachineImpl::_Init(TVMArgs args, TVMRetValue* rv) {
void VirtualMachineImpl::_SaveClosure(TVMArgs args, TVMRetValue* rv) {
ICHECK_GE(args.size(), 3);
- std::string func_name = args[0];
- this->SaveClosure(func_name, args[1], args[2], args.Slice(3));
+ std::string func_name = args[0].cast<std::string>();
+ this->SaveClosure(func_name, args[1].cast<String>(), args[2].cast<bool>(), args.Slice(3));
}
void VirtualMachineImpl::_InvokeClosure(TVMArgs args, TVMRetValue* rv) {
- this->InvokeClosurePacked(args[0], args.Slice(1), rv);
+ this->InvokeClosurePacked(args[0].cast<ObjectRef>(), args.Slice(1), rv);
}
void VirtualMachineImpl::_InvokeClosureStateful(std::string func_name) {
@@ -891,25 +892,25 @@ void VirtualMachineImpl::_InvokeClosureStateful(std::string func_name) {
<< "; use `set_input` first.";
return;
}
- outputs_[func_name] =
- this->InvokeClosureInternal(func_pool_[m.at(func_name)], inputs_[func_name]);
+ outputs_[func_name] = this->InvokeClosureInternal(func_pool_[m.at(func_name)].cast<ObjectRef>(),
+ inputs_[func_name]);
}
void VirtualMachineImpl::_SetInstrument(TVMArgs args, TVMRetValue* rv) {
if (args[0].as<ffi::Function>()) {
- this->SetInstrument(args[0]);
+ this->SetInstrument(args[0].cast<PackedFunc>());
} else {
- String func_name = args[0];
+ String func_name = args[0].cast<String>();
const auto factory = tvm::ffi::Function::GetGlobal(func_name);
CHECK(factory.has_value()) << "Cannot find factory " << func_name;
TVMRetValue rv;
factory->CallPacked(args.Slice(1), &rv);
- this->SetInstrument(rv);
+ this->SetInstrument(rv.cast<PackedFunc>());
}
}
void VirtualMachineImpl::_GetOutputArity(TVMArgs args, TVMRetValue* rv) {
- std::string func_name = args[0];
+ std::string func_name = args[0].cast<std::string>();
RegType out = LookupVMOutput(func_name);
Any obj = IndexIntoNestedObject(out, args, 1);
if (const auto* arr = obj.as<ffi::ArrayObj>()) {
@@ -920,7 +921,7 @@ void VirtualMachineImpl::_GetOutputArity(TVMArgs args, TVMRetValue* rv) {
}
void VirtualMachineImpl::_GetOutput(TVMArgs args, TVMRetValue* rv) {
- std::string func_name = args[0];
+ std::string func_name = args[0].cast<std::string>();
RegType out = LookupVMOutput(func_name);
Any obj = IndexIntoNestedObject(out, args, 1);
if (obj.as<ffi::ArrayObj>()) {
@@ -932,12 +933,12 @@ void VirtualMachineImpl::_GetOutput(TVMArgs args, TVMRetValue* rv) {
}
void VirtualMachineImpl::_SetInputWithoutParamModule(TVMArgs args, TVMRetValue* rv) {
- std::string func_name = args[0];
+ std::string func_name = args[0].cast<std::string>();
this->SetInput(func_name, false, args.Slice(1));
}
void VirtualMachineImpl::_SetInputWithParamModule(TVMArgs args, TVMRetValue* rv) {
- std::string func_name = args[0];
+ std::string func_name = args[0].cast<std::string>();
this->SetInput(func_name, true, args.Slice(1));
}
@@ -981,7 +982,7 @@ class VirtualMachineProfiler : public VirtualMachineImpl {
PackedFunc GetFunction(const String& name, const ObjectPtr<Object>& sptr_to_self) override {
if (name == "profile") {
return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
- std::string f_name = args[0];
+ std::string f_name = args[0].cast<std::string>();
VMClosure clo = this->GetClosure(f_name);
std::vector<Device> devices;
diff --git a/src/runtime/rpc/rpc_channel.cc b/src/runtime/rpc/rpc_channel.cc
index 69c6950856..a9e897f1a3 100644
--- a/src/runtime/rpc/rpc_channel.cc
+++ b/src/runtime/rpc/rpc_channel.cc
@@ -31,7 +31,7 @@ size_t CallbackChannel::Send(const void* data, size_t size) {
TVMFFIByteArray bytes;
bytes.data = static_cast<const char*>(data);
bytes.size = size;
- int64_t n = fsend_(&bytes);
+ int64_t n = fsend_(&bytes).cast<int64_t>();
if (n == -1) {
LOG(FATAL) << "CallbackChannel::Send";
}
diff --git a/src/runtime/rpc/rpc_endpoint.cc b/src/runtime/rpc/rpc_endpoint.cc
index 9a7657fce5..47fa2e3b44 100644
--- a/src/runtime/rpc/rpc_endpoint.cc
+++ b/src/runtime/rpc/rpc_endpoint.cc
@@ -242,6 +242,7 @@ class RPCEndpoint::EventHandler : public dmlc::Stream {
} else {
LOG(FATAL) << "ValueError: Object type is not supported in RPC calling convention: "
<< obj->GetTypeKey() << " (type_index = " << obj->type_index() << ")";
+ TVM_FFI_UNREACHABLE();
}
}
@@ -443,7 +444,7 @@ class RPCEndpoint::EventHandler : public dmlc::Stream {
if (code == RPCCode::kException) {
// switch to the state before sending exception.
this->SwitchToState(kRecvPacketNumBytes);
- String msg = args[0];
+ String msg = args[0].cast<String>();
if (!support::StartsWith(msg, "RPCSessionTimeoutError: ")) {
msg = "RPCError: Error caught from RPC call:\n" + msg;
}
@@ -485,7 +486,7 @@ class RPCEndpoint::EventHandler : public dmlc::Stream {
auto on_copy_complete = [this, elem_bytes, data_bytes, temp_data, fcopyack](RPCCode status,
TVMArgs args) {
if (status == RPCCode::kException) {
- this->ReturnException(args[0].operator const char*());
+ this->ReturnException(args[0].cast<const char*>());
this->SwitchToState(kRecvPacketNumBytes);
} else {
// endian aware handling
@@ -529,7 +530,7 @@ class RPCEndpoint::EventHandler : public dmlc::Stream {
auto on_copy_complete = [this](RPCCode status, TVMArgs args) {
if (status == RPCCode::kException) {
- this->ReturnException(args[0].operator const char*());
+ this->ReturnException(args[0].cast<const char*>());
this->SwitchToState(kRecvPacketNumBytes);
} else {
this->ReturnVoid();
@@ -553,7 +554,7 @@ class RPCEndpoint::EventHandler : public dmlc::Stream {
GetServingSession()->AsyncCallFunc(reinterpret_cast<void*>(call_handle), args,
[this](RPCCode status, TVMArgs args) {
if (status == RPCCode::kException) {
- this->ReturnException(args[0].operator const char*());
+ this->ReturnException(args[0].cast<const char*>());
} else {
ValidateArguments(args);
this->ReturnPackedSeq(args);
@@ -588,7 +589,7 @@ class RPCEndpoint::EventHandler : public dmlc::Stream {
constructor_name = "rpc.LocalSession";
serving_session_ = std::make_shared<LocalSession>();
} else {
- constructor_name = args[0].operator std::string();
+ constructor_name = args[0].cast<std::string>();
constructor_args = args.Slice(1);
}
@@ -623,13 +624,13 @@ class RPCEndpoint::EventHandler : public dmlc::Stream {
void HandleSyscallStreamSync() {
TVMArgs args = RecvPackedSeq();
try {
- Device dev = args[0];
- TVMStreamHandle handle = args[1];
+ auto dev = args[0].cast<Device>();
+ TVMStreamHandle handle = args[1].cast<void*>();
this->SwitchToState(kWaitForAsyncCallback);
GetServingSession()->AsyncStreamWait(dev, handle, [this](RPCCode status, TVMArgs args) {
if (status == RPCCode::kException) {
- this->ReturnException(args[0]);
+ this->ReturnException(args[0].cast<const char*>());
} else {
this->ReturnVoid();
}
@@ -741,7 +742,7 @@ void RPCEndpoint::Init() {
// Quick function to for syscall remote.
syscall_remote_ = PackedFunc([this](TVMArgs all_args, TVMRetValue* rv) {
std::lock_guard<std::mutex> lock(mutex_);
- RPCCode code = static_cast<RPCCode>(all_args[0].operator int());
+ RPCCode code = static_cast<RPCCode>(all_args[0].cast<int>());
ffi::PackedArgs args = all_args.Slice(1);
// Legacy ABI translation
// TODO(tqchen): remove this once we have upgraded to new ABI
@@ -940,23 +941,23 @@ void RPCEndpoint::CopyFromRemote(DLTensor* from, void* to_bytes, uint64_t nbytes
// SysCallEventHandler functions
void RPCGetGlobalFunc(RPCSession* handler, TVMArgs args, TVMRetValue* rv) {
- std::string name = args[0];
+ auto name = args[0].cast<std::string>();
*rv = handler->GetFunction(name);
}
void RPCFreeHandle(RPCSession* handler, TVMArgs args, TVMRetValue* rv) {
- void* handle = args[0];
+ void* handle = args[0].cast<void*>();
handler->FreeHandle(handle);
}
void RPCDevSetDevice(RPCSession* handler, TVMArgs args, TVMRetValue* rv) {
- Device dev = args[0];
+ auto dev = args[0].cast<Device>();
handler->GetDeviceAPI(dev)->SetDevice(dev);
}
void RPCDevGetAttr(RPCSession* handler, TVMArgs args, TVMRetValue* rv) {
- Device dev = args[0];
- DeviceAttrKind kind = static_cast<DeviceAttrKind>(args[1].operator int());
+ auto dev = args[0].cast<Device>();
+ DeviceAttrKind kind = static_cast<DeviceAttrKind>(args[1].cast<int>());
if (kind == kExist) {
DeviceAPI* api = handler->GetDeviceAPI(dev, true);
if (api != nullptr) {
@@ -970,35 +971,35 @@ void RPCDevGetAttr(RPCSession* handler, TVMArgs args, TVMRetValue* rv) {
}
void RPCDevAllocData(RPCSession* handler, TVMArgs args, TVMRetValue* rv) {
- Device dev = args[0];
- uint64_t nbytes = args[1];
- uint64_t alignment = args[2];
- DLDataType type_hint = args[3];
+ auto dev = args[0].cast<Device>();
+ uint64_t nbytes = args[1].cast<uint64_t>();
+ uint64_t alignment = args[2].cast<uint64_t>();
+ DLDataType type_hint = args[3].cast<DLDataType>();
void* data = handler->GetDeviceAPI(dev)->AllocDataSpace(dev, nbytes, alignment, type_hint);
*rv = data;
}
void RPCDevAllocDataWithScope(RPCSession* handler, TVMArgs args, TVMRetValue* rv) {
- DLTensor* arr = args[0];
+ auto arr = args[0].cast<DLTensor*>();
Device dev = arr->device;
int ndim = arr->ndim;
int64_t* shape = arr->shape;
DLDataType dtype = arr->dtype;
- Optional<String> mem_scope = args[1];
+ auto mem_scope = args[1].cast<Optional<String>>();
void* data = handler->GetDeviceAPI(dev)->AllocDataSpace(dev, ndim, shape, dtype, mem_scope);
*rv = data;
}
void RPCDevFreeData(RPCSession* handler, TVMArgs args, TVMRetValue* rv) {
- Device dev = args[0];
- void* ptr = args[1];
+ auto dev = args[0].cast<Device>();
+ void* ptr = args[1].cast<void*>();
handler->GetDeviceAPI(dev)->FreeDataSpace(dev, ptr);
}
void RPCCopyAmongRemote(RPCSession* handler, TVMArgs args, TVMRetValue* rv) {
- DLTensor* from = args[0];
- DLTensor* to = args[1];
- TVMStreamHandle stream = args[2];
+ auto from = args[0].cast<DLTensor*>();
+ auto to = args[1].cast<DLTensor*>();
+ TVMStreamHandle stream = args[2].cast<void*>();
Device dev = from->device;
if (dev.device_type == kDLCPU) {
@@ -1011,25 +1012,25 @@ void RPCCopyAmongRemote(RPCSession* handler, TVMArgs args, TVMRetValue* rv) {
}
void RPCDevCreateStream(RPCSession* handler, TVMArgs args, TVMRetValue* rv) {
- Device dev = args[0];
+ auto dev = args[0].cast<Device>();
void* data = handler->GetDeviceAPI(dev)->CreateStream(dev);
*rv = data;
}
void RPCDevFreeStream(RPCSession* handler, TVMArgs args, TVMRetValue* rv) {
- Device dev = args[0];
- TVMStreamHandle stream = args[1];
+ auto dev = args[0].cast<Device>();
+ TVMStreamHandle stream = args[1].cast<void*>();
handler->GetDeviceAPI(dev)->FreeStream(dev, stream);
}
void RPCDevSetStream(RPCSession* handler, TVMArgs args, TVMRetValue* rv) {
- Device dev = args[0];
- TVMStreamHandle stream = args[1];
+ auto dev = args[0].cast<Device>();
+ TVMStreamHandle stream = args[1].cast<void*>();
handler->GetDeviceAPI(dev)->SetStream(dev, stream);
}
void RPCDevGetCurrentStream(RPCSession* handler, TVMArgs args, TVMRetValue* rv) {
- Device dev = args[0];
+ auto dev = args[0].cast<Device>();
*rv = handler->GetDeviceAPI(dev)->GetCurrentStream(dev);
}
@@ -1097,7 +1098,7 @@ class RPCClientSession : public RPCSession, public DeviceAPI {
// function overrides
PackedFuncHandle GetFunction(const std::string& name) final {
- return endpoint_->SysCallRemote(RPCCode::kGetGlobalFunc, name);
+ return endpoint_->SysCallRemote(RPCCode::kGetGlobalFunc, name).cast<void*>();
}
void CallFunc(PackedFuncHandle func, ffi::PackedArgs args,
@@ -1171,7 +1172,8 @@ class RPCClientSession : public RPCSession, public DeviceAPI {
}
void* AllocDataSpace(Device dev, size_t nbytes, size_t alignment, DLDataType type_hint) final {
- return endpoint_->SysCallRemote(RPCCode::kDevAllocData, dev, nbytes, alignment, type_hint);
+ return endpoint_->SysCallRemote(RPCCode::kDevAllocData, dev, nbytes, alignment, type_hint)
+ .cast<void*>();
}
void* AllocDataSpace(Device dev, int ndim, const int64_t* shape, DLDataType dtype,
@@ -1185,10 +1187,13 @@ class RPCClientSession : public RPCSession, public DeviceAPI {
temp.strides = nullptr;
temp.byte_offset = 0;
if (mem_scope.defined()) {
- return endpoint_->SysCallRemote(RPCCode::kDevAllocDataWithScope, &temp,
- static_cast<std::string>(mem_scope.value()));
+ return endpoint_
+ ->SysCallRemote(RPCCode::kDevAllocDataWithScope, &temp,
+ static_cast<std::string>(mem_scope.value()))
+ .cast<void*>();
} else {
- return endpoint_->SysCallRemote(RPCCode::kDevAllocDataWithScope, &temp, nullptr);
+ return endpoint_->SysCallRemote(RPCCode::kDevAllocDataWithScope, &temp, nullptr)
+ .cast<void*>();
}
}
@@ -1197,11 +1202,11 @@ class RPCClientSession : public RPCSession, public DeviceAPI {
}
void CopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHandle stream) final {
- endpoint_->SysCallRemote(RPCCode::kCopyAmongRemote, from, to, stream);
+ endpoint_->SysCallRemote(RPCCode::kCopyAmongRemote, from, to, stream).cast<void*>();
}
TVMStreamHandle CreateStream(Device dev) final {
- return endpoint_->SysCallRemote(RPCCode::kDevCreateStream, dev);
+ return endpoint_->SysCallRemote(RPCCode::kDevCreateStream, dev).cast<void*>();
}
void FreeStream(Device dev, TVMStreamHandle stream) final {
@@ -1217,7 +1222,7 @@ class RPCClientSession : public RPCSession, public DeviceAPI {
}
TVMStreamHandle GetCurrentStream(Device dev) final {
- return endpoint_->SysCallRemote(RPCCode::kDevGetCurrentStream, dev);
+ return endpoint_->SysCallRemote(RPCCode::kDevGetCurrentStream, dev).cast<void*>();
}
DeviceAPI* GetDeviceAPI(Device dev, bool allow_missing) final { return this; }
@@ -1239,7 +1244,7 @@ class RPCClientSession : public RPCSession, public DeviceAPI {
CallFunc(rpc_func, ffi::PackedArgs(nullptr, 0), [this](TVMArgs args) {
// Use args[1] as return value, args[0] is tcode
// Look at RPCWrappedFunc in src/runtime/rpc/rpc_module.cc
- rpc_chunk_max_size_bytes_ = (int64_t)args[1];
+ rpc_chunk_max_size_bytes_ = args[1].cast<int64_t>();
ICHECK_GT(rpc_chunk_max_size_bytes_, 0)
<< "RPC max transfer size is <= 0! (remote value = " << rpc_chunk_max_size_bytes_
<< ")";
diff --git a/src/runtime/rpc/rpc_endpoint.h b/src/runtime/rpc/rpc_endpoint.h
index 2ce0855430..7da1340a7e 100644
--- a/src/runtime/rpc/rpc_endpoint.h
+++ b/src/runtime/rpc/rpc_endpoint.h
@@ -118,7 +118,7 @@ class RPCEndpoint {
* auto args = session_constructor_args;
* int n = args.size();
* if (n != 0) {
- * std::string constructor = args[0];
+ * auto constructor = args[0].cast<std::string>();
* server.serving_session_ = GetGlobalFunc(constructor)(
* args[1], args[2] ... args[n - 1])
* } else {
diff --git a/src/runtime/rpc/rpc_event_impl.cc b/src/runtime/rpc/rpc_event_impl.cc
index 64cc025f00..a47d5cd16c 100644
--- a/src/runtime/rpc/rpc_event_impl.cc
+++ b/src/runtime/rpc/rpc_event_impl.cc
@@ -38,7 +38,7 @@ PackedFunc CreateEventDrivenServer(PackedFunc fsend, std::string name, std::stri
auto ch = std::make_unique<CallbackChannel>(fsend, frecv);
std::shared_ptr<RPCEndpoint> sess = RPCEndpoint::Create(std::move(ch), name, remote_key);
return PackedFunc([sess](TVMArgs args, TVMRetValue* rv) {
- int ret = sess->ServerAsyncIOEventHandler(args[0], args[1]);
+ int ret = sess->ServerAsyncIOEventHandler(args[0].cast<std::string>(), args[1].cast<int>());
*rv = ret;
});
}
diff --git a/src/runtime/rpc/rpc_module.cc b/src/runtime/rpc/rpc_module.cc
index b7be8ef098..23b982127c 100644
--- a/src/runtime/rpc/rpc_module.cc
+++ b/src/runtime/rpc/rpc_module.cc
@@ -97,7 +97,7 @@ class RPCWrappedFunc : public Object {
switch (args[i].type_index()) {
case ffi::TypeIndex::kTVMFFINDArray: {
// Pass NDArray as DLTensor
- auto dptr = std::make_unique<DLTensor>(*args[i].operator NDArray().operator->());
+ auto dptr = std::make_unique<DLTensor>(*args[i].cast<NDArray>().operator->());
dptr->device = RemoveSessMask(dptr->device);
dptr->data = static_cast<RemoteSpace*>(dptr->data)->data;
packed_args[i] = dptr.get();
@@ -106,7 +106,7 @@ class RPCWrappedFunc : public Object {
}
case ffi::TypeIndex::kTVMFFIDLTensorPtr: {
// translate to a remote view of DLTensor
- auto dptr = std::make_unique<DLTensor>(*args[i].operator DLTensor*());
+ auto dptr = std::make_unique<DLTensor>(*args[i].cast<DLTensor*>());
dptr->device = RemoveSessMask(dptr->device);
dptr->data = static_cast<RemoteSpace*>(dptr->data)->data;
packed_args[i] = dptr.get();
@@ -114,7 +114,7 @@ class RPCWrappedFunc : public Object {
break;
}
case ffi::TypeIndex::kTVMFFIDevice: {
- packed_args[i] = RemoveSessMask(args[i].operator DLDevice());
+ packed_args[i] = RemoveSessMask(args[i].cast<DLDevice>());
break;
}
case ffi::TypeIndex::kTVMFFIFunc:
@@ -274,7 +274,7 @@ class RPCModuleNode final : public ModuleNode {
void* RPCWrappedFunc::UnwrapRemoteValueToHandle(const AnyView& arg) const {
// TODO(tqchen): only support Module unwrapping for now.
if (arg.type_index() == ffi::TypeIndex::kTVMFFIModule) {
- Module mod = arg;
+ Module mod = arg.cast<Module>();
std::string tkey = mod->type_key();
ICHECK_EQ(tkey, "rpc") << "ValueError: Cannot pass a non-RPC module to remote";
auto* rmod = static_cast<RPCModuleNode*>(mod.operator->());
@@ -289,31 +289,31 @@ void* RPCWrappedFunc::UnwrapRemoteValueToHandle(const AnyView& arg) const {
}
void RPCWrappedFunc::WrapRemoteReturnToValue(TVMArgs args, TVMRetValue* rv) const {
- int tcode = args[0];
+ int tcode = args[0].cast<int>();
// TODO(tqchen): move to RPC to new ABI
if (tcode == kTVMNullptr) {
*rv = nullptr;
return;
} else if (tcode == kTVMPackedFuncHandle) {
ICHECK_EQ(args.size(), 2);
- void* handle = args[1];
+ void* handle = args[1].cast<void*>();
auto wf = std::make_shared<RPCWrappedFunc>(handle, sess_);
*rv = PackedFunc([wf](TVMArgs args, TVMRetValue* rv) { return wf->operator()(args, rv); });
} else if (tcode == kTVMModuleHandle) {
ICHECK_EQ(args.size(), 2);
- void* handle = args[1];
+ void* handle = args[1].cast<void*>();
auto n = make_object<RPCModuleNode>(handle, sess_);
*rv = Module(n);
} else if (tcode == kTVMNDArrayHandle || tcode == kTVMDLTensorHandle) {
ICHECK_EQ(args.size(), 3);
- DLTensor* tensor = args[1];
- void* nd_handle = args[2];
+ auto tensor = args[1].cast<DLTensor*>();
+ void* nd_handle = args[2].cast<void*>();
*rv = NDArrayFromRemoteOpaqueHandle(sess_, tensor->data, tensor,
AddRPCSessionMask(tensor->device, sess_->table_index()),
nd_handle);
} else if (tcode == kTVMObjectHandle) {
ICHECK_EQ(args.size(), 2);
- void* handle = args[1];
+ void* handle = args[1].cast<void*>();
auto n = make_object<RPCObjectRefObj>(handle, sess_);
*rv = ObjectRef(n);
} else {
@@ -380,8 +380,8 @@ inline void CPUCacheFlushImpl(const char* addr, unsigned int len) {
inline void CPUCacheFlush(int begin_index, const TVMArgs& args) {
for (int i = begin_index; i < args.size(); i++) {
- CPUCacheFlushImpl(static_cast<char*>((args[i].operator DLTensor*()->data)),
- GetDataSize(*(args[i].operator DLTensor*())));
+ CPUCacheFlushImpl(static_cast<char*>((args[i].cast<DLTensor*>()->data)),
+ GetDataSize(*(args[i].cast<DLTensor*>())));
}
}
@@ -458,7 +458,7 @@ TVM_REGISTER_GLOBAL("rpc.ImportRemoteModule").set_body_typed([](Module parent, M
});
TVM_REGISTER_GLOBAL("rpc.SessTableIndex").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- Module m = args[0];
+ Module m = args[0].cast<Module>();
std::string tkey = m->type_key();
ICHECK_EQ(tkey, "rpc");
*rv = static_cast<RPCModuleNode*>(m.operator->())->sess()->table_index();
diff --git a/src/runtime/rpc/rpc_pipe_impl.cc b/src/runtime/rpc/rpc_pipe_impl.cc
index bfcea856ca..f36b70f641 100644
--- a/src/runtime/rpc/rpc_pipe_impl.cc
+++ b/src/runtime/rpc/rpc_pipe_impl.cc
@@ -115,7 +115,7 @@ Module CreatePipeClient(std::vector<std::string> cmd) {
TVM_REGISTER_GLOBAL("rpc.CreatePipeClient").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
std::vector<std::string> cmd;
for (int i = 0; i < args.size(); ++i) {
- cmd.push_back(args[i].operator std::string());
+ cmd.push_back(args[i].cast<std::string>());
}
*rv = CreatePipeClient(cmd);
});
diff --git a/src/runtime/rpc/rpc_server_env.cc b/src/runtime/rpc/rpc_server_env.cc
index 605b2bc0c6..d7529115f1 100644
--- a/src/runtime/rpc/rpc_server_env.cc
+++ b/src/runtime/rpc/rpc_server_env.cc
@@ -32,17 +32,17 @@ std::string RPCGetPath(const std::string& name) {
// do live lookup everytime as workpath can change.
const auto f = tvm::ffi::Function::GetGlobal("tvm.rpc.server.workpath");
ICHECK(f.has_value()) << "require tvm.rpc.server.workpath";
- return (*f)(name);
+ return (*f)(name).cast<std::string>();
}
TVM_REGISTER_GLOBAL("tvm.rpc.server.upload").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- std::string file_name = RPCGetPath(args[0]);
- std::string data = args[1];
+ std::string file_name = RPCGetPath(args[0].cast<std::string>());
+ auto data = args[1].cast<std::string>();
SaveBinaryToFile(file_name, data);
});
TVM_REGISTER_GLOBAL("tvm.rpc.server.download").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- std::string file_name = RPCGetPath(args[0]);
+ std::string file_name = RPCGetPath(args[0].cast<std::string>());
std::string data;
LoadBinaryFromFile(file_name, &data);
LOG(INFO) << "Download " << file_name << "... nbytes=" << data.size();
@@ -50,7 +50,7 @@ TVM_REGISTER_GLOBAL("tvm.rpc.server.download").set_body_packed([](TVMArgs args,
});
TVM_REGISTER_GLOBAL("tvm.rpc.server.remove").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- std::string file_name = RPCGetPath(args[0]);
+ std::string file_name = RPCGetPath(args[0].cast<std::string>());
RemoveFile(file_name);
});
diff --git a/src/runtime/rpc/rpc_socket_impl.cc b/src/runtime/rpc/rpc_socket_impl.cc
index 11d198b227..7b66725ce5 100644
--- a/src/runtime/rpc/rpc_socket_impl.cc
+++ b/src/runtime/rpc/rpc_socket_impl.cc
@@ -125,10 +125,10 @@ void RPCServerLoop(PackedFunc fsend, PackedFunc frecv) {
}
TVM_REGISTER_GLOBAL("rpc.Connect").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- std::string url = args[0];
- int port = args[1];
- std::string key = args[2];
- bool enable_logging = args[3];
+ auto url = args[0].cast<std::string>();
+ int port = args[1].cast<int>();
+ auto key = args[2].cast<std::string>();
+ bool enable_logging = args[3].cast<bool>();
*rv = RPCClientConnect(url, port, key, enable_logging, args.Slice(4));
});
@@ -136,8 +136,8 @@ TVM_REGISTER_GLOBAL("rpc.ServerLoop").set_body_packed([](TVMArgs args, TVMRetVal
if (auto opt_int = args[0].as<int64_t>()) {
RPCServerLoop(opt_int.value());
} else {
- RPCServerLoop(args[0].operator tvm::runtime::PackedFunc(),
- args[1].operator tvm::runtime::PackedFunc());
+ RPCServerLoop(args[0].cast<tvm::runtime::PackedFunc>(),
+ args[1].cast<tvm::runtime::PackedFunc>());
}
});
diff --git a/src/runtime/system_library.cc b/src/runtime/system_library.cc
index 3d8fa86abe..d27c2d4c03 100644
--- a/src/runtime/system_library.cc
+++ b/src/runtime/system_library.cc
@@ -115,7 +115,7 @@ class SystemLibModuleRegistry {
TVM_REGISTER_GLOBAL("runtime.SystemLib").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
std::string symbol_prefix = "";
if (args.size() != 0) {
- symbol_prefix = args[0].operator std::string();
+ symbol_prefix = args[0].cast<std::string>();
}
*rv = SystemLibModuleRegistry::Global()->GetOrCreateModule(symbol_prefix);
});
diff --git a/src/runtime/thread_pool.cc b/src/runtime/thread_pool.cc
index 56fa1e7e0b..7537df718f 100644
--- a/src/runtime/thread_pool.cc
+++ b/src/runtime/thread_pool.cc
@@ -381,11 +381,11 @@ class ThreadPool {
*/
TVM_REGISTER_GLOBAL("runtime.config_threadpool").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
threading::ThreadGroup::AffinityMode mode =
- static_cast<threading::ThreadGroup::AffinityMode>(static_cast<int>(args[0]));
- int nthreads = args[1];
+ static_cast<threading::ThreadGroup::AffinityMode>(args[0].cast<int>());
+ int nthreads = args[1].cast<int>();
std::vector<unsigned int> cpus;
if (args.size() >= 3) {
- Array<String> cpu_array = args[2];
+ auto cpu_array = args[2].cast<Array<String>>();
for (auto cpu : cpu_array) {
ICHECK(IsNumber(cpu)) << "The CPU core information '" << cpu << "' is not a number.";
cpus.push_back(std::stoi(cpu));
diff --git a/src/runtime/vulkan/vulkan_wrapped_func.cc b/src/runtime/vulkan/vulkan_wrapped_func.cc
index 41c3945d53..15f41afe91 100644
--- a/src/runtime/vulkan/vulkan_wrapped_func.cc
+++ b/src/runtime/vulkan/vulkan_wrapped_func.cc
@@ -54,7 +54,7 @@ void VulkanWrappedFunc::operator()(TVMArgs args, TVMRetValue* rv,
std::vector<VkDescriptorBufferInfo> descriptor_buffers;
descriptor_buffers.resize(num_buffer_args_);
for (size_t i = 0; i < num_buffer_args_; ++i) {
- void* buf = args[static_cast<int>(i)];
+ void* buf = args[static_cast<int>(i)].cast<void*>();
VkDescriptorBufferInfo binfo;
binfo.buffer = static_cast<VulkanBuffer*>(buf)->buffer;
binfo.offset = 0;
diff --git a/src/script/ir_builder/ir/ir.cc b/src/script/ir_builder/ir/ir.cc
index 66979f1eb6..2fb8da1e9f 100644
--- a/src/script/ir_builder/ir/ir.cc
+++ b/src/script/ir_builder/ir/ir.cc
@@ -103,7 +103,7 @@ Optional<ObjectRef> ModuleGetAttr(const String& key) {
if (IRBuilder::IsInScope()) {
IRModuleFrame frame = FindModuleFrame();
if (frame->attrs.find(key) != frame->attrs.end()) {
- return frame->attrs[key];
+ return frame->attrs[key].cast<ObjectRef>();
}
}
return NullOpt;
diff --git a/src/script/ir_builder/tir/ir.cc b/src/script/ir_builder/tir/ir.cc
index b44d0c29a2..1d99a1b852 100644
--- a/src/script/ir_builder/tir/ir.cc
+++ b/src/script/ir_builder/tir/ir.cc
@@ -489,7 +489,7 @@ AttrFrame Attr(ffi::Any node, String attr_key, PrimExpr value) {
node = node.as<PrimExpr>().value();
}
ObjectPtr<AttrFrameNode> n = make_object<AttrFrameNode>();
- n->node = node;
+ n->node = node.cast<ObjectRef>();
n->attr_key = attr_key;
n->value = value;
return AttrFrame(n);
diff --git a/src/script/printer/legacy_repr.cc b/src/script/printer/legacy_repr.cc
index 0316c3b27d..27301482cc 100644
--- a/src/script/printer/legacy_repr.cc
+++ b/src/script/printer/legacy_repr.cc
@@ -82,7 +82,7 @@ TVM_STATIC_IR_FUNCTOR(ReprLegacyPrinter, vtable)
if (i != 0) {
(*p) << ", ";
}
- p->Print(op->at(i));
+ p->Print(op->at(i).cast<ObjectRef>());
}
(*p) << ']';
});
@@ -98,10 +98,10 @@ TVM_STATIC_IR_FUNCTOR(ReprLegacyPrinter, vtable)
if (it->first.as<ffi::StringObj>()) {
(*p) << '\"' << Downcast<ffi::String>(it->first) << "\": ";
} else {
- p->Print(it->first);
+ p->Print(it->first.cast<ObjectRef>());
(*p) << ": ";
}
- p->Print(it->second);
+ p->Print(it->second.cast<ObjectRef>());
}
(*p) << '}';
});
diff --git a/src/script/printer/tir/expr.cc b/src/script/printer/tir/expr.cc
index 63c79af55c..b81c8ef5af 100644
--- a/src/script/printer/tir/expr.cc
+++ b/src/script/printer/tir/expr.cc
@@ -74,6 +74,7 @@ Doc PrintVar(const tir::Var& var, const ObjectPath& var_p, const IRDocsifier& d)
return doc.value();
}
LOG(FATAL) << "IndexError: Variable is not defined in the environment: " << var->name_hint;
+ TVM_FFI_UNREACHABLE();
}
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) //
@@ -264,7 +265,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
for (int i = 0; i < n_args; ++i) {
if ((i == 0) && (f_llvm_lookup_intrinsic_name)) {
- String name = (*f_llvm_lookup_intrinsic_name)(id);
+ String name = (*f_llvm_lookup_intrinsic_name)(id).cast<String>();
args.push_back(LiteralDoc::Str(name.c_str(), call_p->Attr("args")->ArrayIndex(i)));
} else {
args.push_back(d->AsDoc<ExprDoc>(call->args[i], call_p->Attr("args")->ArrayIndex(i)));
diff --git a/src/support/array.h b/src/support/array.h
index 25acc1ddc5..057e1668b3 100644
--- a/src/support/array.h
+++ b/src/support/array.h
@@ -167,7 +167,7 @@ struct AsVectorImpl<TSrcObjectRef, int> {
inline std::vector<int> operator()(const Array<TSrcObjectRef>& array) const {
TVMRetValue ret_value;
ret_value = array;
- Array<IntImm> as_int_vec = ret_value;
+ Array<IntImm> as_int_vec = ret_value.cast<Array<IntImm>>();
std::vector<int> results;
for (const auto& value : as_int_vec) {
@@ -182,7 +182,7 @@ struct AsVectorImpl<TSrcObjectRef, int64_t> {
inline std::vector<int64_t> operator()(const Array<TSrcObjectRef>& array) const {
TVMRetValue ret_value;
ret_value = array;
- Array<IntImm> as_int_vec = ret_value;
+ Array<IntImm> as_int_vec = ret_value.cast<Array<IntImm>>();
std::vector<int64_t> results;
for (const auto& value : as_int_vec) {
@@ -197,7 +197,7 @@ struct AsVectorImpl<TSrcObjectRef, double> {
inline std::vector<double> operator()(const Array<TSrcObjectRef>& array) const {
TVMRetValue ret_value;
ret_value = array;
- Array<FloatImm> as_int_vec = ret_value;
+ Array<FloatImm> as_int_vec = ret_value.cast<Array<FloatImm>>();
std::vector<double> results;
for (const auto& value : as_int_vec) {
@@ -230,7 +230,7 @@ struct AsArrayImpl<int, TDstObjectRef> {
for (auto x : vec) {
TVMRetValue ret_value;
ret_value = x;
- result.push_back(ret_value);
+ result.push_back(ret_value.cast<TDstObjectRef>());
}
return result;
}
@@ -244,7 +244,7 @@ struct AsArrayImpl<int64_t, TDstObjectRef> {
for (auto x : vec) {
TVMRetValue ret_value;
ret_value = x;
- result.push_back(ret_value);
+ result.push_back(ret_value.cast<TDstObjectRef>());
}
return result;
}
@@ -258,7 +258,7 @@ struct AsArrayImpl<double, TDstObjectRef> {
for (auto x : vec) {
TVMRetValue ret_value;
ret_value = x;
- result.push_back(ret_value);
+ result.push_back(ret_value.cast<TDstObjectRef>());
}
return result;
}
diff --git a/src/support/ffi_testing.cc b/src/support/ffi_testing.cc
index 5c59c08792..87447d1b5f 100644
--- a/src/support/ffi_testing.cc
+++ b/src/support/ffi_testing.cc
@@ -55,13 +55,13 @@ TVM_REGISTER_NODE_TYPE(TestAttrs);
TVM_REGISTER_GLOBAL("testing.test_wrap_callback")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- PackedFunc pf = args[0];
+ PackedFunc pf = args[0].cast<PackedFunc>();
*ret = runtime::TypedPackedFunc<void()>([pf]() { pf(); });
});
TVM_REGISTER_GLOBAL("testing.test_wrap_callback_suppress_err")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- PackedFunc pf = args[0];
+ PackedFunc pf = args[0].cast<PackedFunc>();
auto result = runtime::TypedPackedFunc<void()>([pf]() {
try {
pf();
@@ -73,21 +73,21 @@ TVM_REGISTER_GLOBAL("testing.test_wrap_callback_suppress_err")
TVM_REGISTER_GLOBAL("testing.test_raise_error_callback")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- std::string msg = args[0];
+ auto msg = args[0].cast<std::string>();
*ret = runtime::TypedPackedFunc<void()>([msg]() { LOG(FATAL) << msg; });
});
TVM_REGISTER_GLOBAL("testing.test_check_eq_callback")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- std::string msg = args[0];
+ auto msg = args[0].cast<std::string>();
*ret = runtime::TypedPackedFunc<void(int x, int y)>(
[msg](int x, int y) { CHECK_EQ(x, y) << msg; });
});
TVM_REGISTER_GLOBAL("testing.device_test").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- Device dev = args[0];
- int dtype = args[1];
- int did = args[2];
+ auto dev = args[0].cast<Device>();
+ int dtype = args[1].cast<int>();
+ int did = args[2].cast<int>();
CHECK_EQ(static_cast<int>(dev.device_type), dtype);
CHECK_EQ(static_cast<int>(dev.device_id), did);
*ret = dev;
@@ -115,7 +115,7 @@ TVM_REGISTER_GLOBAL("testing.ErrorTest").set_body_typed(ErrorTest);
// internal function used for debug and testing purposes
TVM_REGISTER_GLOBAL("testing.object_use_count").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- runtime::ObjectRef obj = args[0];
+ auto obj = args[0].cast<ObjectRef>();
// substract the current one because we always copy
// and get another value.
*ret = (obj.use_count() - 1);
@@ -270,7 +270,7 @@ class TestingEventLogger {
TVM_REGISTER_GLOBAL("testing.record_event").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
if (args.size() != 0 && args[0].as<String>()) {
- TestingEventLogger::ThreadLocal()->Record(args[0]);
+ TestingEventLogger::ThreadLocal()->Record(args[0].cast<String>());
} else {
TestingEventLogger::ThreadLocal()->Record("X");
}
diff --git a/src/target/codegen.cc b/src/target/codegen.cc
index 87c294050f..d7db5cc8e9 100644
--- a/src/target/codegen.cc
+++ b/src/target/codegen.cc
@@ -53,7 +53,7 @@ runtime::Module Build(IRModule mod, Target target) {
std::string build_f_name = "target.build." + target->kind->name;
const auto bf = tvm::ffi::Function::GetGlobal(build_f_name);
ICHECK(bf.has_value()) << build_f_name << " is not enabled";
- return (*bf)(mod, target);
+ return (*bf)(mod, target).cast<runtime::Module>();
}
/*! \brief Helper class to serialize module */
@@ -357,7 +357,8 @@ runtime::Module PackImportsToLLVM(const runtime::Module& mod, bool system_lib,
// the codegen function.
const auto codegen_f = tvm::ffi::Function::GetGlobal(codegen_f_name);
ICHECK(codegen_f.has_value()) << "codegen.codegen_blob is not presented.";
- return (*codegen_f)(ffi::Bytes(blob), system_lib, llvm_target_string, c_symbol_prefix);
+ return (*codegen_f)(ffi::Bytes(blob), system_lib, llvm_target_string, c_symbol_prefix)
+ .cast<runtime::Module>();
}
TVM_REGISTER_GLOBAL("target.Build").set_body_typed(Build);
diff --git a/src/target/datatype/registry.cc b/src/target/datatype/registry.cc
index e308b3cf35..11b40a6726 100644
--- a/src/target/datatype/registry.cc
+++ b/src/target/datatype/registry.cc
@@ -28,22 +28,23 @@ using runtime::TVMRetValue;
TVM_REGISTER_GLOBAL("dtype.register_custom_type")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- datatype::Registry::Global()->Register(args[0], static_cast<uint8_t>(args[1].operator int()));
+ datatype::Registry::Global()->Register(args[0].cast<std::string>(),
+ static_cast<uint8_t>(args[1].cast<int>()));
});
TVM_REGISTER_GLOBAL("dtype.get_custom_type_code")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- *ret = datatype::Registry::Global()->GetTypeCode(args[0]);
+ *ret = datatype::Registry::Global()->GetTypeCode(args[0].cast<std::string>());
});
TVM_REGISTER_GLOBAL("dtype.get_custom_type_name")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- *ret = Registry::Global()->GetTypeName(args[0].operator int());
+ *ret = Registry::Global()->GetTypeName(args[0].cast<int>());
});
TVM_REGISTER_GLOBAL("runtime._datatype_get_type_registered")
.set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- *ret = Registry::Global()->GetTypeRegistered(args[0].operator int());
+ *ret = Registry::Global()->GetTypeRegistered(args[0].cast<int>());
});
Registry* Registry::Global() {
@@ -128,7 +129,7 @@ uint64_t ConvertConstScalar(uint8_t type_code, double value) {
ss << "tvm.datatype.convertconstscalar.float.";
ss << datatype::Registry::Global()->GetTypeName(type_code);
auto make_const_scalar_func = tvm::ffi::Function::GetGlobal(ss.str());
- return (*make_const_scalar_func)(value).operator uint64_t();
+ return (*make_const_scalar_func)(value).cast<uint64_t>();
}
} // namespace datatype
diff --git a/src/target/llvm/codegen_amdgpu.cc b/src/target/llvm/codegen_amdgpu.cc
index 70ecdb5494..4cd833d2da 100644
--- a/src/target/llvm/codegen_amdgpu.cc
+++ b/src/target/llvm/codegen_amdgpu.cc
@@ -70,10 +70,10 @@ static inline int DetectROCMmaxThreadsPerBlock() {
if (api != nullptr) {
TVMRetValue val;
api->GetAttr(tvm_dev, tvm::runtime::kExist, &val);
- if (val.operator int() == 1) {
+ if (val.cast<int>() == 1) {
tvm::runtime::DeviceAPI::Get(tvm_dev)->GetAttr(tvm_dev, tvm::runtime::kMaxThreadsPerBlock,
&val);
- return val.operator int();
+ return val.cast<int>();
}
}
LOG(WARNING) << "Cannot get maximum number of threads for AMD codegen";
@@ -279,7 +279,7 @@ runtime::Module BuildAMDGPU(IRModule mod, Target target) {
llvm::TargetMachine* tm = llvm_target->GetOrCreateTargetMachine();
auto fbitcode = tvm::ffi::Function::GetGlobalRequired("tvm_callback_rocm_bitcode_path");
- Array<runtime::String> bitcode_files = fbitcode();
+ auto bitcode_files = fbitcode().cast<Array<runtime::String>>();
for (auto& bitcode_path : bitcode_files) {
std::unique_ptr<llvm::Module> mlib = llvm_instance.LoadIR(bitcode_path);
@@ -351,7 +351,7 @@ runtime::Module BuildAMDGPU(IRModule mod, Target target) {
arr.data = &obj[0];
arr.size = obj.length();
- std::string hsaco = (*flink)(&arr);
+ std::string hsaco = (*flink)(&arr).cast<std::string>();
std::string ll(data_ll.begin(), data_ll.end());
return ROCMModuleCreate(hsaco, "hsaco", ExtractFuncInfo(mod), ll, assembly);
}
diff --git a/src/target/llvm/codegen_hexagon.cc b/src/target/llvm/codegen_hexagon.cc
index d188f7e5a8..9b1025aff9 100644
--- a/src/target/llvm/codegen_hexagon.cc
+++ b/src/target/llvm/codegen_hexagon.cc
@@ -583,7 +583,7 @@ runtime::Module BuildHexagon(IRModule mod, Target target) {
<< "unexpected -mcpu value in target:" << mcpu;
extra_args.Set("hex_arch", llvm::StringRef(mcpu).drop_front(strlen("hexagon")).str());
}
- int rc = (*f)(so_name, o_names, extra_args);
+ int rc = (*f)(so_name, o_names, extra_args).cast<int>();
ICHECK(rc == 0) << "Failed to link " << so_name;
return HexagonModuleCreate(so_name, "so", ExtractFuncInfo(mod), asm_str, obj_str, ir_str, bc_str);
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 68a1a4ee68..f77e32efd5 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -123,9 +123,9 @@ std::unique_ptr<CodeGenLLVM> CodeGenLLVM::Create(LLVMTarget* llvm_target) {
std::string factory_template = "tvm.codegen.llvm.target_";
void* handle = nullptr;
if (auto f = tvm::ffi::Function::GetGlobal(factory_template + target)) {
- handle = (*f)();
+ handle = (*f)().cast<void*>();
} else if (auto f = tvm::ffi::Function::GetGlobal(factory_template + "cpu")) {
- handle = (*f)();
+ handle = (*f)().cast<void*>();
} else {
LOG(FATAL) << "no factory function for codegen for target " << target;
}
diff --git a/src/target/llvm/codegen_nvptx.cc b/src/target/llvm/codegen_nvptx.cc
index bb903522e1..2e2f34ba8a 100644
--- a/src/target/llvm/codegen_nvptx.cc
+++ b/src/target/llvm/codegen_nvptx.cc
@@ -331,7 +331,7 @@ runtime::Module BuildNVPTX(IRModule mod, Target target) {
llvm::TargetMachine* tm = llvm_target->GetOrCreateTargetMachine();
const auto flibdevice_path = tvm::ffi::Function::GetGlobal("tvm_callback_libdevice_path");
if (flibdevice_path.has_value()) {
- std::string path = (*flibdevice_path)(compute_ver);
+ std::string path = (*flibdevice_path)(compute_ver).cast<std::string>();
if (path.length() != 0) {
std::unique_ptr<llvm::Module> mlib = llvm_instance.LoadIR(path);
mlib->setTargetTriple(llvm_target->GetTargetTriple());
diff --git a/src/target/llvm/llvm_instance.cc b/src/target/llvm/llvm_instance.cc
index 8d54a1df3b..a83a146826 100644
--- a/src/target/llvm/llvm_instance.cc
+++ b/src/target/llvm/llvm_instance.cc
@@ -347,7 +347,7 @@ LLVMTargetInfo::LLVMTargetInfo(LLVMInstance& instance, const TargetJSON& target)
auto maybe_level = target.Get("opt-level");
#if TVM_LLVM_VERSION <= 170
if (maybe_level.has_value()) {
- int level = maybe_level.value().operator int();
+ int level = maybe_level.value().cast<int>();
if (level <= 0) {
opt_level_ = llvm::CodeGenOpt::None;
} else if (level == 1) {
@@ -385,7 +385,7 @@ LLVMTargetInfo::LLVMTargetInfo(LLVMInstance& instance, const TargetJSON& target)
auto GetBoolFlag = [&target](llvm::StringRef name) -> bool {
if (auto flag = target.Get(name.str())) {
- return flag.value().operator bool();
+ return flag.value().cast<bool>();
} else {
return false;
}
diff --git a/src/target/opt/build_cuda_on.cc b/src/target/opt/build_cuda_on.cc
index 449125ff15..58f80858d0 100644
--- a/src/target/opt/build_cuda_on.cc
+++ b/src/target/opt/build_cuda_on.cc
@@ -150,14 +150,14 @@ runtime::Module BuildCUDA(IRModule mod, Target target) {
std::string code = cg.Finish();
if (auto f = ffi::Function::GetGlobal("tvm_callback_cuda_postproc")) {
- code = (*f)(code, target).operator std::string();
+ code = (*f)(code, target).cast<std::string>();
}
std::string fmt = "ptx";
std::string ptx;
auto f_enter = ffi::Function::GetGlobal("target.TargetEnterScope");
(*f_enter)(target);
if (auto f = ffi::Function::GetGlobal("tvm_callback_cuda_compile")) {
- ptx = (*f)(code, target).operator std::string();
+ ptx = (*f)(code, target).cast<std::string>();
// Dirty matching to check PTX vs cubin.
// TODO(tqchen) more reliable checks
if (ptx[0] != '/') fmt = "cubin";
diff --git a/src/target/parsers/cpu.cc b/src/target/parsers/cpu.cc
index cccc53fd05..c28d80133f 100644
--- a/src/target/parsers/cpu.cc
+++ b/src/target/parsers/cpu.cc
@@ -33,7 +33,7 @@ Optional<String> DetectSystemTriple() {
auto pf = tvm::ffi::Function::GetGlobal("target.llvm_get_system_triple");
ICHECK(pf.has_value()) << "The target llvm_get_system_triple was not found, "
"please compile with USE_LLVM = ON";
- return (*pf)();
+ return (*pf)().cast<String>();
#endif
return {};
}
diff --git a/src/target/source/codegen_metal.cc b/src/target/source/codegen_metal.cc
index ecb5962d95..d115916def 100644
--- a/src/target/source/codegen_metal.cc
+++ b/src/target/source/codegen_metal.cc
@@ -458,7 +458,7 @@ runtime::Module BuildMetal(IRModule mod, Target target) {
std::string fsource = cg.Finish();
source_maker << fsource << "\n";
if (fmetal_compile) {
- fsource = (*fmetal_compile)(fsource, target).operator std::string();
+ fsource = (*fmetal_compile)(fsource, target).cast<std::string>();
}
smap[func_name] = fsource;
}
diff --git a/src/target/source/codegen_opencl.cc b/src/target/source/codegen_opencl.cc
index f62a8847f7..9814696b37 100644
--- a/src/target/source/codegen_opencl.cc
+++ b/src/target/source/codegen_opencl.cc
@@ -664,7 +664,7 @@ runtime::Module BuildOpenCL(IRModule mod, Target target) {
cg.AddFunction(gvar, prim_func);
std::string fsource = cg.Finish();
if (fpostproc) {
- fsource = (*fpostproc)(fsource, target).operator std::string();
+ fsource = (*fpostproc)(fsource, target).cast<std::string>();
}
code << fsource;
}
diff --git a/src/target/spirv/spirv_utils.cc b/src/target/spirv/spirv_utils.cc
index 180f8f31a6..6afd087e5d 100644
--- a/src/target/spirv/spirv_utils.cc
+++ b/src/target/spirv/spirv_utils.cc
@@ -160,7 +160,7 @@ std::pair<std::unordered_map<std::string, runtime::SPIRVShader>, std::string> Lo
TVMFFIByteArray arr;
arr.data = reinterpret_cast<const char*>(dmlc::BeginPtr(shader.data));
arr.size = shader.data.size() * sizeof(uint32_t);
- std::string transformed = (*postproc)(&arr, target);
+ std::string transformed = (*postproc)(&arr, target).cast<std::string>();
ICHECK_EQ(transformed.length() % 4U, 0U);
shader.data.resize(transformed.size() / 4U);
std::copy(transformed.begin(), transformed.end(),
diff --git a/src/target/target.cc b/src/target/target.cc
index 4aa4db3e1e..143436541e 100644
--- a/src/target/target.cc
+++ b/src/target/target.cc
@@ -399,7 +399,7 @@ Any TargetInternal::ParseType(const std::string& str, const TargetKindNode::Valu
std::vector<ObjectRef> result;
for (const std::string& substr : SplitString(interp_str, ',')) {
try {
- ObjectRef parsed = TargetInternal::ParseType(substr, *info.key);
+ ObjectRef parsed = TargetInternal::ParseType(substr, *info.key).cast<ObjectRef>();
result.push_back(parsed);
} catch (const Error& e) {
std::string index = "[" + std::to_string(result.size()) + "]";
@@ -446,7 +446,7 @@ Any TargetInternal::ParseType(const Any& obj, const TargetKindNode::ValueTypeInf
std::vector<ObjectRef> result;
for (const Any& e : *array) {
try {
- result.push_back(TargetInternal::ParseType(e, *info.key));
+ result.push_back(TargetInternal::ParseType(e, *info.key).cast<ObjectRef>());
} catch (const Error& e) {
std::string index = '[' + std::to_string(result.size()) + ']';
throw Error(e->kind, index + e->message, e->traceback);
@@ -771,8 +771,8 @@ void TargetInternal::ConstructorDispatcher(TVMArgs args, TVMRetValue* rv) {
return;
} else if (args.size() == 2) {
if (args[0].as<Target>().has_value() && args[1].as<Target>().has_value()) {
- Target target = args[0];
- Target host = args[1];
+ Target target = args[0].cast<Target>();
+ Target host = args[1].cast<Target>();
*rv = Target(target, host);
} else {
LOG(FATAL) << "ValueError: Invalid type of arguments. Expect 2 Target arguments.";
@@ -798,7 +798,7 @@ ObjectPtr<Object> TargetInternal::FromConfigString(const String& config_str) {
ICHECK(loader.has_value())
<< "AttributeError: \"target._load_config_dict\" is not registered. Please check "
"if the python module is properly loaded";
- Optional<Map<String, ffi::Any>> config = (*loader)(config_str);
+ auto config = (*loader)(config_str).cast<Optional<Map<String, ffi::Any>>>();
if (!config.defined()) {
TVM_FFI_THROW(ValueError) << "Cannot load config dict with python JSON loader";
}
@@ -923,7 +923,7 @@ ObjectPtr<Object> TargetInternal::FromConfig(Map<String, ffi::Any> config) {
}
// parse host
if (config.count(kHost)) {
- target->host = PackedFunc(ConstructorDispatcher)(config[kHost]).operator Target();
+ target->host = PackedFunc(ConstructorDispatcher)(config[kHost]).cast<Target>();
config.erase(kHost);
} else {
target->host = NullOpt;
@@ -964,7 +964,8 @@ ObjectPtr<Object> TargetInternal::FromConfig(Map<String, ffi::Any> config) {
}
// do extra pre-processing
if (target->kind->preprocessor != nullptr) {
- target->attrs = target->kind->preprocessor(Map<String, ffi::Any>(attrs));
+ target->attrs =
+ target->kind->preprocessor(Map<String, ffi::Any>(attrs)).cast<Map<String, ffi::Any>>();
} else {
target->attrs = attrs;
}
@@ -988,7 +989,7 @@ std::unordered_map<String, ffi::Any> TargetInternal::QueryDevice(int device_id,
TVMRetValue ret;
api->GetAttr(device, runtime::kExist, &ret);
- bool device_exists = ret;
+ bool device_exists = ret.cast<bool>();
if (!device_exists) {
ICHECK(device_exists) << "Requested reading the parameters for " << target->kind->name
<< " from device_id " << device_id << ", but device_id " << device_id
diff --git a/src/target/target_info.cc b/src/target/target_info.cc
index b047d3a96b..a63e45a81a 100644
--- a/src/target/target_info.cc
+++ b/src/target/target_info.cc
@@ -45,7 +45,7 @@ MemoryInfo GetMemoryInfo(const std::string& scope) {
LOG(WARNING) << "MemoryInfo for scope = " << scope << " is undefined";
return MemoryInfo();
} else {
- return (*f)();
+ return (*f)().cast<MemoryInfo>();
}
}
diff --git a/src/target/target_kind.cc b/src/target/target_kind.cc
index e51a4c21ff..bea7791aad 100644
--- a/src/target/target_kind.cc
+++ b/src/target/target_kind.cc
@@ -131,7 +131,7 @@ static bool DetectDeviceFlag(Device device, runtime::DeviceAttrKind flag, TVMRet
}
// Check if the device exists
api->GetAttr(device, runtime::kExist, val);
- int exists = *val;
+ int exists = val->cast<int>();
if (!exists) {
return false;
}
@@ -173,7 +173,7 @@ TargetJSON UpdateCUDAAttrs(TargetJSON target) {
LOG(WARNING) << "Unable to detect CUDA version, default to \"-arch=sm_50\" instead";
archInt = 50;
} else {
- archInt = std::stod(version.operator std::string()) * 10 + 0.1;
+ archInt = std::stod(version.cast<std::string>()) * 10 + 0.1;
}
target.Set("arch", String("sm_") + std::to_string(archInt));
}
@@ -201,7 +201,7 @@ TargetJSON UpdateNVPTXAttrs(TargetJSON target) {
LOG(WARNING) << "Unable to detect CUDA version, default to \"-mcpu=sm_50\" instead";
arch = 50;
} else {
- arch = std::stod(version.operator std::string()) * 10 + 0.1;
+ arch = std::stod(version.cast<std::string>()) * 10 + 0.1;
}
target.Set("mcpu", String("sm_") + std::to_string(arch));
}
@@ -224,7 +224,7 @@ TargetJSON UpdateROCmAttrs(TargetJSON target) {
} else {
TVMRetValue val;
if (const auto f_get_rocm_arch = tvm::ffi::Function::GetGlobal("tvm_callback_rocm_get_arch")) {
- arch = (*f_get_rocm_arch)().operator std::string();
+ arch = (*f_get_rocm_arch)().cast<std::string>();
}
target.Set("mcpu", String(arch));
}
@@ -238,7 +238,7 @@ TargetJSON UpdateROCmAttrs(TargetJSON target) {
LOG(WARNING) << "Unable to detect ROCm version, assuming >= 3.5";
version = 305;
} else {
- version = val.operator int();
+ version = val.cast<int>();
}
if (version < 305) {
Array<String> mattr;
diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc
index 186cd28811..7f0f90302e 100644
--- a/src/te/operation/create_primfunc.cc
+++ b/src/te/operation/create_primfunc.cc
@@ -744,7 +744,7 @@ PrimFunc GenerateAndCompletePrimFunc(const Array<te::Tensor>& arg_list,
{{"global_symbol", String("main")}, {"tir.noalias", true}});
const auto fcomplete = tvm::ffi::Function::GetGlobal("script.Complete");
ICHECK(fcomplete.has_value());
- func = (*fcomplete)(std::move(func), info->root_alloc);
+ func = (*fcomplete)(std::move(func), info->root_alloc).cast<PrimFunc>();
return func;
}
@@ -785,11 +785,11 @@ PrimFunc CreatePrimFunc(const Array<te::Tensor>& arg_list,
}
TVM_REGISTER_GLOBAL("te.CreatePrimFunc").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- Array<ObjectRef> arg_list = args[0];
+ Array<ObjectRef> arg_list = args[0].cast<Array<ObjectRef>>();
std::optional<DataType> index_dtype_override{std::nullopt};
// Add conversion to make std::optional compatible with FFI.
if (args[1] != nullptr) {
- index_dtype_override = args[1].operator DataType();
+ index_dtype_override = args[1].cast<DataType>();
}
*ret = CreatePrimFunc(arg_list, index_dtype_override);
});
@@ -818,7 +818,7 @@ PrimFunc GenerateAndCompletePrimFunc(const Array<ObjectRef>& arg_tir_var_list,
{{"global_symbol", String("main")}, {"tir.noalias", true}});
const auto fcomplete = tvm::ffi::Function::GetGlobal("script.Complete");
ICHECK(fcomplete.has_value());
- func = (*fcomplete)(std::move(func), info->root_alloc);
+ func = (*fcomplete)(std::move(func), info->root_alloc).cast<PrimFunc>();
return func;
}
diff --git a/src/tir/analysis/var_use_def_analysis.cc b/src/tir/analysis/var_use_def_analysis.cc
index 1ca5b670f6..d2d0313298 100644
--- a/src/tir/analysis/var_use_def_analysis.cc
+++ b/src/tir/analysis/var_use_def_analysis.cc
@@ -202,9 +202,9 @@ Array<Var> UndefinedVars(const PrimExpr& expr, const Array<Var>& args) {
TVM_REGISTER_GLOBAL("tir.analysis.UndefinedVars")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
if (auto opt_stmt = args[0].as<Stmt>()) {
- *rv = UndefinedVars(opt_stmt.value(), args[1]);
+ *rv = UndefinedVars(opt_stmt.value(), args[1].cast<Array<Var>>());
} else if (auto opt_expr = args[0].as<PrimExpr>()) {
- *rv = UndefinedVars(opt_expr.value(), args[1]);
+ *rv = UndefinedVars(opt_expr.value(), args[1].cast<Array<Var>>());
} else {
LOG(FATAL) << "either UndefinedVars(stmt, args) or UndefinedVars(expr, args) is expected";
}
diff --git a/src/tir/ir/buffer.cc b/src/tir/ir/buffer.cc
index 4c16c9805a..e1626c08a8 100644
--- a/src/tir/ir/buffer.cc
+++ b/src/tir/ir/buffer.cc
@@ -642,10 +642,20 @@ TVM_REGISTER_NODE_TYPE(BufferNode);
TVM_REGISTER_GLOBAL("tir.Buffer").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
ICHECK_EQ(args.size(), 11);
- auto buffer_type = args[8].operator String();
+ auto buffer_type = args[8].cast<String>();
BufferType type = (buffer_type == "auto_broadcast") ? kAutoBroadcast : kDefault;
- *ret = Buffer(args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7], type,
- args[9], args[10]);
+ auto data = args[0].cast<Var>();
+ auto dtype = args[1].cast<DataType>();
+ auto shape = args[2].cast<Array<PrimExpr>>();
+ auto strides = args[3].cast<Array<PrimExpr>>();
+ auto elem_offset = args[4].cast<PrimExpr>();
+ auto name = args[5].cast<String>();
+ auto data_alignment = args[6].cast<int>();
+ auto offset_factor = args[7].cast<int>();
+ auto axis_separators = args[9].cast<Array<IntImm>>();
+ auto span = args[10].cast<Span>();
+ *ret = Buffer(data, dtype, shape, strides, elem_offset, name, data_alignment, offset_factor, type,
+ axis_separators, span);
});
TVM_REGISTER_GLOBAL("tir.BufferAccessPtr").set_body_method(&Buffer::access_ptr);
diff --git a/src/tir/ir/data_type_rewriter.cc b/src/tir/ir/data_type_rewriter.cc
index 3f1576ab92..8a496d06be 100644
--- a/src/tir/ir/data_type_rewriter.cc
+++ b/src/tir/ir/data_type_rewriter.cc
@@ -371,7 +371,7 @@ Map<String, ffi::Any> IndexDataTypeRewriter::VisitBlockAnnotations(
const Map<String, ffi::Any>& annotations) {
auto new_annotations = annotations;
- std::function<ObjectRef(const Any&)> f_mutate_obj =
+ std::function<ObjectRef(const ObjectRef&)> f_mutate_obj =
[this, &f_mutate_obj](const ObjectRef& obj) -> ObjectRef {
if (!obj.defined()) {
return obj;
@@ -387,9 +387,9 @@ Map<String, ffi::Any> IndexDataTypeRewriter::VisitBlockAnnotations(
return obj;
};
for (const auto& [key, value] : annotations) {
- if (auto* object_ptr = value.as<Object>()) {
- auto new_value = f_mutate_obj(GetRef<ObjectRef>(object_ptr));
- if (!new_value.same_as(value)) {
+ if (auto opt_object_ref = value.as<ObjectRef>()) {
+ auto new_value = f_mutate_obj(*opt_object_ref);
+ if (!new_value.same_as(*opt_object_ref)) {
new_annotations.Set(key, new_value);
}
}
diff --git a/src/tir/ir/expr.cc b/src/tir/ir/expr.cc
index 7270c014cc..f88e749306 100644
--- a/src/tir/ir/expr.cc
+++ b/src/tir/ir/expr.cc
@@ -130,9 +130,9 @@ Var Var::copy_with_dtype(DataType dtype) const {
TVM_REGISTER_GLOBAL("tir.Var").set_body_typed([](String name_hint, runtime::TVMArgValue type,
Span span) {
if (type.as<Type>()) {
- return Var(name_hint, type.operator Type(), span);
+ return Var(name_hint, type.cast<Type>(), span);
} else {
- return Var(name_hint, type.operator DataType(), span);
+ return Var(name_hint, type.cast<DataType>(), span);
}
});
diff --git a/src/tir/ir/stmt_functor.cc b/src/tir/ir/stmt_functor.cc
index 3572b082b2..cf09ea306e 100644
--- a/src/tir/ir/stmt_functor.cc
+++ b/src/tir/ir/stmt_functor.cc
@@ -616,12 +616,12 @@ class IRTransformer final : public StmtExprMutator {
return fmutate(node);
}
if (f_preorder_ != nullptr) {
- T pre = f_preorder_(node);
+ T pre = f_preorder_(node).template cast<T>();
if (pre.defined()) return pre;
}
T new_node = fmutate(node);
if (f_postorder_ != nullptr) {
- T post = f_postorder_(new_node);
+ T post = f_postorder_(new_node).template cast<T>();
if (post.defined()) return post;
}
return new_node;
@@ -899,7 +899,7 @@ TVM_REGISTER_GLOBAL("tir.PostOrderVisit").set_body_typed([](ObjectRef node, Pack
});
TVM_REGISTER_GLOBAL("tir.PreOrderVisit").set_body_typed([](ObjectRef node, PackedFunc f) {
- tir::PreOrderVisit(node, [f](const ObjectRef& n) { return f(n); });
+ tir::PreOrderVisit(node, [f](const ObjectRef& n) { return f(n).cast<bool>(); });
});
TVM_REGISTER_GLOBAL("tir.Substitute")
diff --git a/src/tir/ir/utils.cc b/src/tir/ir/utils.cc
index de9d69e846..65495bd8dd 100644
--- a/src/tir/ir/utils.cc
+++ b/src/tir/ir/utils.cc
@@ -30,7 +30,7 @@ namespace tir {
ffi::Any NormalizeAttributeObject(ffi::Any obj) {
if (obj.type_index() == ffi::TypeIndex::kTVMFFIBool) {
- return Bool(obj.operator bool());
+ return Bool(obj.cast<bool>());
} else if (auto opt_int = obj.as<int>()) {
return Integer(opt_int.value());
} else if (auto opt_float = obj.as<double>()) {
@@ -42,7 +42,7 @@ ffi::Any NormalizeAttributeObject(ffi::Any obj) {
bool is_same = true;
for (const auto& [key, obj] : opt_map.value()) {
- ObjectRef new_obj = NormalizeAttributeObject(obj);
+ ObjectRef new_obj = NormalizeAttributeObject(obj.cast<ObjectRef>()).cast<ObjectRef>();
is_same = is_same && obj.same_as(new_obj);
new_map.Set(key, new_obj);
}
diff --git a/src/tir/op/op.cc b/src/tir/op/op.cc
index ecfca1730c..4c064185fd 100644
--- a/src/tir/op/op.cc
+++ b/src/tir/op/op.cc
@@ -293,7 +293,7 @@ PrimExpr min_value(const DataType& dtype, Span span) {
ICHECK(f) << "No minimum function registered for custom dtype " << (unsigned int)dtype.code();
// TODO(@hypercubestart) Document this change (and others associated with the overflowing
// floatimm min bug)
- return (*f)(dtype.bits());
+ return (*f)(dtype.bits()).cast<PrimExpr>();
} else if (dtype.is_int()) {
if (dtype.bits() == 64) {
return IntImm(dtype, std::numeric_limits<int64_t>::lowest(), span);
@@ -1073,9 +1073,9 @@ TVM_TIR_REGISTER_OP("TVMBackendFreeWorkspace")
// expose basic functions to node namespace
TVM_REGISTER_GLOBAL("node._const").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
if (auto opt = args[0].as<int64_t>()) {
- *ret = tir::make_const(args[1], opt.value(), args[2]);
+ *ret = tir::make_const(args[1].cast<DataType>(), *opt, args[2].cast<Span>());
} else if (auto opt = args[0].as<double>()) {
- *ret = tir::make_const(args[1], opt.value(), args[2]);
+ *ret = tir::make_const(args[1].cast<DataType>(), *opt, args[2].cast<Span>());
} else {
LOG(FATAL) << "First argument to tvm.tir.const must be int, float, or bool, "
<< "but instead received argument with type code " << args[0].GetTypeKey();
@@ -1120,17 +1120,17 @@ TVM_REGISTER_GLOBAL("tir.reinterpret").set_body_typed(tvm::reinterpret);
return (Func(a, b, span)); \
})
-#define REGISTER_MAKE_BIT_OP(Node, Func) \
- TVM_REGISTER_GLOBAL("tir." #Node).set_body_packed([](TVMArgs args, TVMRetValue* ret) { \
- bool lhs_is_int = args[0].type_index() == ffi::TypeIndex::kTVMFFIInt; \
- bool rhs_is_int = args[1].type_index() == ffi::TypeIndex::kTVMFFIInt; \
- if (lhs_is_int) { \
- *ret = (Func(args[0].operator int(), args[1].operator PrimExpr(), args[2])); \
- } else if (rhs_is_int) { \
- *ret = (Func(args[0].operator PrimExpr(), args[1].operator int(), args[2])); \
- } else { \
- *ret = (Func(args[0].operator PrimExpr(), args[1].operator PrimExpr(), args[2])); \
- } \
+#define REGISTER_MAKE_BIT_OP(Node, Func) \
+ TVM_REGISTER_GLOBAL("tir." #Node).set_body_packed([](TVMArgs args, TVMRetValue* ret) { \
+ bool lhs_is_int = args[0].type_index() == ffi::TypeIndex::kTVMFFIInt; \
+ bool rhs_is_int = args[1].type_index() == ffi::TypeIndex::kTVMFFIInt; \
+ if (lhs_is_int) { \
+ *ret = (Func(args[0].cast<int>(), args[1].cast<PrimExpr>(), args[2].cast<Span>())); \
+ } else if (rhs_is_int) { \
+ *ret = (Func(args[0].cast<PrimExpr>(), args[1].cast<int>(), args[2].cast<Span>())); \
+ } else { \
+ *ret = (Func(args[0].cast<PrimExpr>(), args[1].cast<PrimExpr>(), args[2].cast<Span>())); \
+ } \
})
REGISTER_MAKE_BINARY_OP(_OpAdd, add);
diff --git a/src/tir/schedule/instruction_traits.h b/src/tir/schedule/instruction_traits.h
index 924b2ddb00..d292bd3dc7 100644
--- a/src/tir/schedule/instruction_traits.h
+++ b/src/tir/schedule/instruction_traits.h
@@ -354,8 +354,7 @@ String UnpackedInstTraits<TTraits>::AsPython(const Array<Any>& inputs, const Arr
});
ffi::Any rv;
pf.CallPacked(ffi::PackedArgs(packed_args, kNumArgs), &rv);
- String result = rv;
- return result;
+ return rv.cast<String>();
}
template <class TTraits>
@@ -409,7 +408,7 @@ TVM_ALWAYS_INLINE Array<Any> UnpackedInstTraits<TTraits>::_ConvertOutputs(const
} else if (is_single_obj) {
return {rv};
} else if (is_array) {
- return rv;
+ return rv.cast<runtime::Array<Any>>();
}
}
diff --git a/src/tir/schedule/primitive/blockize_tensorize.cc b/src/tir/schedule/primitive/blockize_tensorize.cc
index c057a3d4fe..60ac7bf422 100644
--- a/src/tir/schedule/primitive/blockize_tensorize.cc
+++ b/src/tir/schedule/primitive/blockize_tensorize.cc
@@ -833,7 +833,7 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int
block->body = impl_block->body;
block->match_buffers = std::move(match_buffer_regions);
for (const auto& [key, val] : impl_block->annotations) {
- if (block->annotations.count(key) && block->annotations[key] != val) {
+ if (block->annotations.count(key) && !ffi::AnyEqual()(block->annotations[key], val)) {
LOG(WARNING) << "Conflict of annotation \"" << key << "\". Tensor intrinsic and schedule "
<< "has different values : " << block->annotations[key] << " vs " << val << " "
<< "The value from tensor intrinsic is skipped.";
diff --git a/src/tir/schedule/trace.cc b/src/tir/schedule/trace.cc
index 13bb93a798..cde7a6fc05 100644
--- a/src/tir/schedule/trace.cc
+++ b/src/tir/schedule/trace.cc
@@ -112,7 +112,7 @@ Array<Any> TranslateInputRVs(
} else if (input.as<BlockRVNode>() || // RV: block
input.as<LoopRVNode>() || // RV: loop
input.as<VarNode>()) { // RV: var
- auto it = rv_names.find(input);
+ auto it = rv_names.find(input.cast<ObjectRef>());
if (it != rv_names.end()) {
// Case 1. BlockRV, LoopRV, VarRV
results.push_back(it->second);
@@ -233,8 +233,9 @@ Array<String> TranslateAddOutputRVs(
results.reserve(outputs.size());
for (const Any& output : outputs) {
int i = rv_names->size();
- ICHECK(!rv_names->count(output))
- << "ValueError: The random variable has been produced once: " << rv_names->at(output);
+ ICHECK(!rv_names->count(output.cast<ObjectRef>()))
+ << "ValueError: The random variable has been produced once: "
+ << rv_names->at(output.cast<ObjectRef>());
String result{ObjectPtr<StringObj>{nullptr}};
if (output == nullptr) {
result = "_";
@@ -250,7 +251,7 @@ Array<String> TranslateAddOutputRVs(
throw;
}
results.push_back(result);
- rv_names->emplace(output, std::move(result));
+ rv_names->emplace(output.cast<ObjectRef>(), std::move(result));
}
return results;
}
@@ -260,7 +261,7 @@ void TranslateAddOutputRVs(const Array<String>& old_outputs, const Array<Any>& n
ICHECK_EQ(old_outputs.size(), new_outputs.size());
int n = old_outputs.size();
for (int i = 0; i < n; ++i) {
- named_rvs->emplace(Downcast<String>(old_outputs[i]), new_outputs[i]);
+ named_rvs->emplace(Downcast<String>(old_outputs[i]), new_outputs[i].cast<ObjectRef>());
}
}
@@ -333,7 +334,7 @@ ObjectRef TraceNode::AsJSON(bool remove_postproc) const {
: ObjectRef(inst->attrs),
/* 3: outputs */ TranslateAddOutputRVs(inst->outputs, &rv_names),
});
- if (Optional<ObjectRef> decision = this->GetDecision(inst)) {
+ if (auto decision = this->GetDecision(inst).cast<Optional<ObjectRef>>()) {
json_decisions.push_back(Array<ObjectRef>{
/* 0: index */ Integer(i),
/* 1: decision */ decision.value(),
@@ -414,7 +415,7 @@ void Trace::ApplyJSONToSchedule(ObjectRef json, Schedule sch) {
// Parse `json_insts`
std::unordered_map<std::string, ObjectRef> named_rvs{{"None", ObjectRef{nullptr}}};
int i = 0;
- for (const ObjectRef& inst_entry : json_insts) {
+ for (const Any& inst_entry : json_insts) {
InstructionKind kind{nullptr};
Array<Any> inputs{nullptr};
Array<Any> attrs{nullptr};
@@ -425,9 +426,9 @@ void Trace::ApplyJSONToSchedule(ObjectRef json, Schedule sch) {
ICHECK(arr && arr->size() == 4);
const auto* arr0 = arr->at(0).as<StringObj>();
kind = InstructionKind::Get(arr0->data);
- inputs = arr->at(1);
- attrs = arr->at(2);
- outputs = arr->at(3);
+ inputs = arr->at(1).cast<Array<Any>>();
+ attrs = arr->at(2).cast<Array<Any>>();
+ outputs = arr->at(3).cast<Array<String>>();
} catch (const tvm::Error& e) {
LOG(FATAL) << "ValueError: Each entry of a json instruction should be a tuple [inst_name, "
"inputs, attrs, outputs], but gets: "
@@ -498,8 +499,8 @@ Trace TraceNode::Simplified(bool remove_postproc) const {
} else if (obj.as<BlockRVNode>() || obj.as<LoopRVNode>() || obj.as<VarNode>()) {
used_rvs.insert(obj.as<Object>());
continue;
- } else if (obj.as<PrimExprNode>()) {
- PostOrderVisit(obj, [&used_rvs](const ObjectRef& obj) -> void {
+ } else if (auto prim_expr = obj.as<PrimExpr>()) {
+ PostOrderVisit(*prim_expr, [&used_rvs](const ObjectRef& obj) -> void {
if (obj.as<VarNode>()) {
used_rvs.insert(obj.get());
}
diff --git a/src/tir/schedule/traced_schedule.cc b/src/tir/schedule/traced_schedule.cc
index 780227a4e2..52a012b2a3 100644
--- a/src/tir/schedule/traced_schedule.cc
+++ b/src/tir/schedule/traced_schedule.cc
@@ -274,7 +274,7 @@ void TracedScheduleNode::Reorder(const Array<LoopRV>& ordered_loop_rvs) {
static const InstructionKind& kind = InstructionKind::Get("Reorder");
trace_->Append(/*inst=*/Instruction(/*kind=*/kind,
- /*inputs=*/Any(ordered_loop_rvs),
+ /*inputs=*/ordered_loop_rvs,
/*attrs=*/{},
/*outputs=*/{}));
}
diff --git a/src/tir/transforms/lower_custom_datatypes.cc b/src/tir/transforms/lower_custom_datatypes.cc
index 3e2dc130e7..19f529103b 100644
--- a/src/tir/transforms/lower_custom_datatypes.cc
+++ b/src/tir/transforms/lower_custom_datatypes.cc
@@ -56,7 +56,7 @@ class CustomDatatypesLowerer : public StmtExprMutator {
ICHECK(lower) << "Cast lowering function for target " << target_ << " destination type "
<< static_cast<unsigned>(type_code) << " source type "
<< static_cast<unsigned>(src_type_code) << " not found";
- return (*lower)(expr);
+ return (*lower)(expr).cast<PrimExpr>();
}
return expr;
}
@@ -68,7 +68,7 @@ class CustomDatatypesLowerer : public StmtExprMutator {
auto lower = datatype::GetFloatImmLowerFunc(target_, type_code);
ICHECK(lower) << "FloatImm lowering function for target " << target_ << " type "
<< static_cast<unsigned>(type_code) << " not found";
- return (*lower)(e);
+ return (*lower)(e).cast<PrimExpr>();
}
return e;
}
@@ -191,7 +191,7 @@ class CustomDatatypesLowerer : public StmtExprMutator {
ICHECK(lower) << "Intrinsic lowering function for target " << target_ << ", intrinsic name "
<< op->name << ", type " << static_cast<unsigned>(call->dtype.code())
<< " not found";
- return (*lower)(expr);
+ return (*lower)(expr).cast<PrimExpr>();
}
return expr;
}
@@ -206,7 +206,7 @@ class CustomDatatypesLowerer : public StmtExprMutator {
auto lower = datatype::Get##OP##LowerFunc(target_, type_code); \
ICHECK(lower) << #OP " lowering function for target " << target_ << " type " \
<< static_cast<unsigned>(type_code) << " not found"; \
- return (*lower)(expr); \
+ return (*lower)(expr).cast<PrimExpr>(); \
} \
return expr; \
}
diff --git a/src/tir/transforms/unsupported_dtype_legalize.cc b/src/tir/transforms/unsupported_dtype_legalize.cc
index 29ac31617f..047e80ea0e 100644
--- a/src/tir/transforms/unsupported_dtype_legalize.cc
+++ b/src/tir/transforms/unsupported_dtype_legalize.cc
@@ -738,9 +738,9 @@ bool CheckDataTypeSupport(const Target& target, const std::string& support_func_
bool has_native_support = false;
if (target->kind->name == "cuda") {
if (auto get_cv = tvm::ffi::Function::GetGlobal("tvm.contrib.nvcc.get_compute_version")) {
- std::string compute_version = (*get_cv)(target);
+ std::string compute_version = (*get_cv)(target).cast<std::string>();
if (auto check_support = tvm::ffi::Function::GetGlobal(support_func_name)) {
- has_native_support = (*check_support)(compute_version);
+ has_native_support = (*check_support)(compute_version).cast<bool>();
}
}
}
diff --git a/src/topi/broadcast.cc b/src/topi/broadcast.cc
index cb34617f65..660cc241e5 100644
--- a/src/topi/broadcast.cc
+++ b/src/topi/broadcast.cc
@@ -32,19 +32,19 @@ namespace topi {
using namespace tvm;
using namespace tvm::runtime;
-#define TOPI_REGISTER_BCAST_OP(OpName, Op) \
- TVM_REGISTER_GLOBAL(OpName).set_body_packed([](TVMArgs args, TVMRetValue* rv) { \
- bool lhs_is_tensor = args[0].as<tvm::te::Tensor>().has_value(); \
- bool rhs_is_tensor = args[1].as<tvm::te::Tensor>().has_value(); \
- if (lhs_is_tensor && rhs_is_tensor) { \
- *rv = Op(args[0].operator tvm::te::Tensor(), args[1].operator tvm::te::Tensor()); \
- } else if (!lhs_is_tensor && rhs_is_tensor) { \
- *rv = Op(args[0].operator tvm::PrimExpr(), args[1].operator tvm::te::Tensor()); \
- } else if (lhs_is_tensor && !rhs_is_tensor) { \
- *rv = Op(args[0].operator tvm::te::Tensor(), args[1].operator tvm::PrimExpr()); \
- } else if (!lhs_is_tensor && !rhs_is_tensor) { \
- *rv = Op(args[0].operator tvm::PrimExpr(), args[1].operator tvm::PrimExpr()); \
- } \
+#define TOPI_REGISTER_BCAST_OP(OpName, Op) \
+ TVM_REGISTER_GLOBAL(OpName).set_body_packed([](TVMArgs args, TVMRetValue* rv) { \
+ bool lhs_is_tensor = args[0].as<tvm::te::Tensor>().has_value(); \
+ bool rhs_is_tensor = args[1].as<tvm::te::Tensor>().has_value(); \
+ if (lhs_is_tensor && rhs_is_tensor) { \
+ *rv = Op(args[0].cast<tvm::te::Tensor>(), args[1].cast<tvm::te::Tensor>()); \
+ } else if (!lhs_is_tensor && rhs_is_tensor) { \
+ *rv = Op(args[0].cast<tvm::PrimExpr>(), args[1].cast<tvm::te::Tensor>()); \
+ } else if (lhs_is_tensor && !rhs_is_tensor) { \
+ *rv = Op(args[0].cast<tvm::te::Tensor>(), args[1].cast<tvm::PrimExpr>()); \
+ } else if (!lhs_is_tensor && !rhs_is_tensor) { \
+ *rv = Op(args[0].cast<tvm::PrimExpr>(), args[1].cast<tvm::PrimExpr>()); \
+ } \
});
TOPI_REGISTER_BCAST_OP("topi.add", topi::add);
@@ -74,7 +74,7 @@ TOPI_REGISTER_BCAST_OP("topi.greater_equal", topi::greater_equal);
TOPI_REGISTER_BCAST_OP("topi.less_equal", topi::less_equal);
TVM_REGISTER_GLOBAL("topi.broadcast_to").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = broadcast_to(args[0], args[1]);
+ *rv = broadcast_to(args[0].cast<te::Tensor>(), args[1].cast<Array<PrimExpr>>());
});
} // namespace topi
diff --git a/src/topi/einsum.cc b/src/topi/einsum.cc
index 47e5aeb4d2..8f9b9d953a 100644
--- a/src/topi/einsum.cc
+++ b/src/topi/einsum.cc
@@ -356,7 +356,7 @@ Array<PrimExpr> InferEinsumShape(const std::string& subscripts,
}
TVM_REGISTER_GLOBAL("topi.einsum").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = einsum(args[0], args[1]);
+ *rv = einsum(args[0].cast<std::string>(), args[1].cast<Array<Tensor>>());
});
} // namespace topi
diff --git a/src/topi/elemwise.cc b/src/topi/elemwise.cc
index a7d554ca09..ce7f7f4eb1 100644
--- a/src/topi/elemwise.cc
+++ b/src/topi/elemwise.cc
@@ -32,139 +32,139 @@ using namespace tvm;
using namespace tvm::runtime;
TVM_REGISTER_GLOBAL("topi.acos").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = acos(args[0]);
+ *rv = acos(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.acosh").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = acosh(args[0]);
+ *rv = acosh(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.asin").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = asin(args[0]);
+ *rv = asin(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.asinh").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = asinh(args[0]);
+ *rv = asinh(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.atanh").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = atanh(args[0]);
+ *rv = atanh(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.exp").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = exp(args[0]);
+ *rv = exp(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.fast_exp").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = fast_exp(args[0]);
+ *rv = fast_exp(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.erf").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = erf(args[0]);
+ *rv = erf(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.fast_erf").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = fast_erf(args[0]);
+ *rv = fast_erf(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.tan").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = tan(args[0]);
+ *rv = tan(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.cos").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = cos(args[0]);
+ *rv = cos(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.cosh").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = cosh(args[0]);
+ *rv = cosh(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.sin").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = sin(args[0]);
+ *rv = sin(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.sinh").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = sinh(args[0]);
+ *rv = sinh(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.tanh").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = tanh(args[0]);
+ *rv = tanh(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.fast_tanh").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = fast_tanh(args[0]);
+ *rv = fast_tanh(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.atan").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = atan(args[0]);
+ *rv = atan(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.sigmoid").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = sigmoid(args[0]);
+ *rv = sigmoid(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.sqrt").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = sqrt(args[0]);
+ *rv = sqrt(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.rsqrt").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = rsqrt(args[0]);
+ *rv = rsqrt(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.log").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = log(args[0]);
+ *rv = log(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.log2").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = log2(args[0]);
+ *rv = log2(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.log10").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = log10(args[0]);
+ *rv = log10(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.identity").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = identity(args[0]);
+ *rv = identity(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.negative").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = negative(args[0]);
+ *rv = negative(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.clip").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = clip(args[0], args[1], args[2]);
+ *rv = clip(args[0].cast<te::Tensor>(), args[1].cast<PrimExpr>(), args[2].cast<PrimExpr>());
});
TVM_REGISTER_GLOBAL("topi.cast").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = cast(args[0], args[1]);
+ *rv = cast(args[0].cast<te::Tensor>(), args[1].cast<DataType>());
});
TVM_REGISTER_GLOBAL("topi.reinterpret").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = reinterpret(args[0], args[1]);
+ *rv = reinterpret(args[0].cast<te::Tensor>(), args[1].cast<DataType>());
});
TVM_REGISTER_GLOBAL("topi.elemwise_sum").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = elemwise_sum(args[0]);
+ *rv = elemwise_sum(args[0].cast<Array<te::Tensor>>());
});
TVM_REGISTER_GLOBAL("topi.sign").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = sign(args[0]);
+ *rv = sign(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.full").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = full(args[0], args[1], args[2]);
+ *rv = full(args[0].cast<Array<PrimExpr>>(), args[1].cast<DataType>(), args[2].cast<PrimExpr>());
});
TVM_REGISTER_GLOBAL("topi.full_like").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = full_like(args[0], args[1]);
+ *rv = full_like(args[0].cast<te::Tensor>(), args[1].cast<PrimExpr>());
});
TVM_REGISTER_GLOBAL("topi.logical_not").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = logical_not(args[0]);
+ *rv = logical_not(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.bitwise_not").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = bitwise_not(args[0]);
+ *rv = bitwise_not(args[0].cast<te::Tensor>());
});
} // namespace topi
diff --git a/src/topi/nn.cc b/src/topi/nn.cc
index 331e3737ec..58ef57f508 100644
--- a/src/topi/nn.cc
+++ b/src/topi/nn.cc
@@ -46,145 +46,173 @@ using namespace tvm::runtime;
/* Ops from nn.h */
TVM_REGISTER_GLOBAL("topi.nn.relu").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = relu<float>(args[0]);
+ *rv = relu<float>(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.nn.leaky_relu").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = leaky_relu(args[0], args[1]);
+ *rv = leaky_relu(args[0].cast<te::Tensor>(), args[1].cast<double>());
});
TVM_REGISTER_GLOBAL("topi.nn.prelu").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = prelu(args[0], args[1], args[2]);
+ *rv = prelu(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(), args[2].cast<int>());
});
TVM_REGISTER_GLOBAL("topi.nn.pad").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = pad(args[0], args[1], args[2], args[3]);
+ *rv = pad(args[0].cast<te::Tensor>(), args[1].cast<Array<PrimExpr>>(),
+ args[2].cast<Array<PrimExpr>>(), args[3].cast<PrimExpr>());
});
TVM_REGISTER_GLOBAL("topi.nn.space_to_batch_nd").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = space_to_batch_nd(args[0], args[1], args[2], args[3], args[4]);
+ *rv = space_to_batch_nd(args[0].cast<te::Tensor>(), args[1].cast<Array<Integer>>(),
+ args[2].cast<Array<PrimExpr>>(), args[3].cast<Array<PrimExpr>>(),
+ args[4].cast<PrimExpr>());
});
TVM_REGISTER_GLOBAL("topi.nn.batch_to_space_nd").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = batch_to_space_nd(args[0], args[1], args[2], args[3]);
+ *rv = batch_to_space_nd(args[0].cast<te::Tensor>(), args[1].cast<Array<Integer>>(),
+ args[2].cast<Array<PrimExpr>>(), args[3].cast<Array<PrimExpr>>(),
+ args[4].cast<std::string>());
});
TVM_REGISTER_GLOBAL("topi.nn.nll_loss").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nll_loss(args[0], args[1], args[2], args[3], args[4]);
+ *rv = nll_loss(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(), args[2].cast<te::Tensor>(),
+ args[3].cast<std::string>(), args[4].cast<int>());
});
/* Ops from nn/dense.h */
TVM_REGISTER_GLOBAL("topi.nn.dense").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::dense(args[0], args[1], args[2], args[3]);
+ *rv = nn::dense(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(),
+ args[2].cast<te::Tensor>(), args[3].cast<DataType>());
});
/* Ops from nn/bias_add.h */
TVM_REGISTER_GLOBAL("topi.nn.bias_add").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::bias_add(args[0], args[1], args[2]);
+ *rv = nn::bias_add(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(), args[2].cast<int>());
});
/* Ops from nn/dilate.h */
TVM_REGISTER_GLOBAL("topi.nn.dilate").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::dilate(args[0], args[1], args[2]);
+ *rv = nn::dilate(args[0].cast<te::Tensor>(), args[1].cast<Array<Integer>>(),
+ args[2].cast<double>());
});
/* Ops from nn/flatten.h */
TVM_REGISTER_GLOBAL("topi.nn.flatten").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::flatten(args[0]);
+ *rv = nn::flatten(args[0].cast<te::Tensor>());
});
/* Ops from nn/mapping.h */
TVM_REGISTER_GLOBAL("topi.nn.scale_shift_nchw").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::scale_shift_nchw(args[0], args[1], args[2]);
+ *rv = nn::scale_shift_nchw(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(),
+ args[2].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.nn.scale_shift_nhwc").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::scale_shift_nhwc(args[0], args[1], args[2]);
+ *rv = nn::scale_shift_nhwc(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(),
+ args[2].cast<te::Tensor>());
});
/* Ops from nn/pooling.h */
TVM_REGISTER_GLOBAL("topi.nn.pool_grad").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::pool_grad(args[0], args[1], args[2], args[3], args[4],
- static_cast<nn::PoolType>(static_cast<int>(args[5])), args[6], args[7],
- args[8]);
+ *rv =
+ nn::pool_grad(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(),
+ args[2].cast<Array<PrimExpr>>(), args[3].cast<Array<PrimExpr>>(),
+ args[4].cast<Array<PrimExpr>>(), static_cast<nn::PoolType>(args[5].cast<int>()),
+ args[6].cast<bool>(), args[7].cast<std::string>(), args[8].cast<bool>());
});
TVM_REGISTER_GLOBAL("topi.nn.global_pool").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::global_pool(args[0], static_cast<nn::PoolType>(static_cast<int>(args[1])), args[2]);
+ *rv = nn::global_pool(args[0].cast<te::Tensor>(), static_cast<nn::PoolType>(args[1].cast<int>()),
+ args[2].cast<std::string>());
});
TVM_REGISTER_GLOBAL("topi.nn.adaptive_pool1d").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::adaptive_pool1d(args[0], args[1], static_cast<nn::PoolType>(static_cast<int>(args[2])),
- args[3]);
+ *rv = nn::adaptive_pool1d(args[0].cast<te::Tensor>(), args[1].cast<Array<PrimExpr>>(),
+ static_cast<nn::PoolType>(args[2].cast<int>()),
+ args[3].cast<std::string>());
});
TVM_REGISTER_GLOBAL("topi.nn.adaptive_pool").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::adaptive_pool(args[0], args[1], static_cast<nn::PoolType>(static_cast<int>(args[2])),
- args[3]);
+ *rv = nn::adaptive_pool(args[0].cast<te::Tensor>(), args[1].cast<Array<PrimExpr>>(),
+ static_cast<nn::PoolType>(args[2].cast<int>()),
+ args[3].cast<std::string>());
});
TVM_REGISTER_GLOBAL("topi.nn.adaptive_pool3d").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::adaptive_pool3d(args[0], args[1], static_cast<nn::PoolType>(static_cast<int>(args[2])),
- args[3]);
+ *rv = nn::adaptive_pool3d(args[0].cast<te::Tensor>(), args[1].cast<Array<PrimExpr>>(),
+ static_cast<nn::PoolType>(args[2].cast<int>()),
+ args[3].cast<std::string>());
});
TVM_REGISTER_GLOBAL("topi.nn.pool1d").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::pool1d(args[0], args[1], args[2], args[3], args[4],
- static_cast<nn::PoolType>(static_cast<int>(args[5])), args[6], args[7], args[8]);
+ *rv = nn::pool1d(args[0].cast<te::Tensor>(), args[1].cast<Array<PrimExpr>>(),
+ args[2].cast<Array<PrimExpr>>(), args[3].cast<Array<PrimExpr>>(),
+ args[4].cast<Array<PrimExpr>>(), static_cast<nn::PoolType>(args[5].cast<int>()),
+ args[6].cast<bool>(), args[7].cast<std::string>(), args[8].cast<bool>());
});
TVM_REGISTER_GLOBAL("topi.nn.pool2d").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::pool2d(args[0], args[1], args[2], args[3], args[4],
- static_cast<nn::PoolType>(static_cast<int>(args[5])), args[6], args[7], args[8]);
+ *rv = nn::pool2d(args[0].cast<te::Tensor>(), args[1].cast<Array<PrimExpr>>(),
+ args[2].cast<Array<PrimExpr>>(), args[3].cast<Array<PrimExpr>>(),
+ args[4].cast<Array<PrimExpr>>(), static_cast<nn::PoolType>(args[5].cast<int>()),
+ args[6].cast<bool>(), args[7].cast<std::string>(), args[8].cast<bool>());
});
TVM_REGISTER_GLOBAL("topi.nn.pool3d").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::pool3d(args[0], args[1], args[2], args[3], args[4],
- static_cast<nn::PoolType>(static_cast<int>(args[5])), args[6], args[7], args[8]);
+ *rv = nn::pool3d(args[0].cast<te::Tensor>(), args[1].cast<Array<PrimExpr>>(),
+ args[2].cast<Array<PrimExpr>>(), args[3].cast<Array<PrimExpr>>(),
+ args[4].cast<Array<PrimExpr>>(), static_cast<nn::PoolType>(args[5].cast<int>()),
+ args[6].cast<bool>(), args[7].cast<std::string>(), args[8].cast<bool>());
});
/* Ops from nn/softmax.h */
TVM_REGISTER_GLOBAL("topi.nn.softmax").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::softmax(args[0], args[1]);
+ *rv = nn::softmax(args[0].cast<te::Tensor>(), args[1].cast<int>());
});
TVM_REGISTER_GLOBAL("topi.nn.log_softmax").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::log_softmax(args[0]);
+ *rv = nn::log_softmax(args[0].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.nn.lrn").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::lrn(args[0], args[1], args[2], static_cast<double>(args[3]),
- static_cast<double>(args[4]), static_cast<double>(args[5]));
+ *rv = nn::lrn(args[0].cast<te::Tensor>(), args[1].cast<int>(), args[2].cast<int>(),
+ args[3].cast<double>(), args[4].cast<double>(), args[5].cast<double>());
});
/* Ops from nn/bnn.h */
TVM_REGISTER_GLOBAL("topi.nn.binarize_pack").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::binarize_pack(args[0], args[1]);
+ *rv = nn::binarize_pack(args[0].cast<te::Tensor>(), args[1].cast<int>());
});
TVM_REGISTER_GLOBAL("topi.nn.binary_dense").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::binary_dense(args[0], args[1]);
+ *rv = nn::binary_dense(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>());
});
/* Ops from nn/layer_norm.h */
TVM_REGISTER_GLOBAL("topi.nn.layer_norm").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::layer_norm(args[0], args[1], args[2], args[3], static_cast<double>(args[4]));
+ *rv = nn::layer_norm(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(),
+ args[2].cast<te::Tensor>(), args[3].cast<Array<Integer>>(),
+ args[4].cast<double>());
});
/* Ops from nn/group_norm.h */
TVM_REGISTER_GLOBAL("topi.nn.group_norm").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::group_norm(args[0], args[1], args[2], static_cast<int>(args[3]),
- static_cast<int>(args[4]), args[5], static_cast<double>(args[6]));
+ *rv = nn::group_norm(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(),
+ args[2].cast<te::Tensor>(), args[3].cast<int>(), args[4].cast<int>(),
+ args[5].cast<Array<Integer>>(), args[6].cast<double>());
});
/* Ops from nn/instance_norm.h */
TVM_REGISTER_GLOBAL("topi.nn.instance_norm").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::instance_norm(args[0], args[1], args[2], args[3], static_cast<double>(args[4]));
+ *rv = nn::instance_norm(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(),
+ args[2].cast<te::Tensor>(), args[3].cast<Array<Integer>>(),
+ args[4].cast<double>());
});
/* Ops from nn/rms_norm.h */
TVM_REGISTER_GLOBAL("topi.nn.rms_norm").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = nn::rms_norm(args[0], args[1], args[2], static_cast<double>(args[3]));
+ *rv = nn::rms_norm(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(),
+ args[2].cast<Array<Integer>>(), args[3].cast<double>());
});
} // namespace topi
diff --git a/src/topi/reduction.cc b/src/topi/reduction.cc
index 2400534b4f..b9daa1deea 100644
--- a/src/topi/reduction.cc
+++ b/src/topi/reduction.cc
@@ -33,39 +33,41 @@ using namespace tvm;
using namespace tvm::runtime;
TVM_REGISTER_GLOBAL("topi.sum").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = topi::sum(args[0], ArrayOrInt(args[1]), args[2]);
+ *rv = topi::sum(args[0].cast<te::Tensor>(), ArrayOrInt(args[1]), args[2].cast<bool>());
});
TVM_REGISTER_GLOBAL("topi.min").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = topi::min(args[0], ArrayOrInt(args[1]), args[2]);
+ *rv = topi::min(args[0].cast<te::Tensor>(), ArrayOrInt(args[1]), args[2].cast<bool>());
});
TVM_REGISTER_GLOBAL("topi.max").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = topi::max(args[0], ArrayOrInt(args[1]), args[2]);
+ *rv = topi::max(args[0].cast<te::Tensor>(), ArrayOrInt(args[1]), args[2].cast<bool>());
});
TVM_REGISTER_GLOBAL("topi.argmin").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = topi::argmin(args[0], ArrayOrInt(args[1]), args[2], false, args[3]);
+ *rv = topi::argmin(args[0].cast<te::Tensor>(), ArrayOrInt(args[1]), args[2].cast<bool>(), false,
+ args[3].cast<bool>());
});
TVM_REGISTER_GLOBAL("topi.argmax").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = topi::argmax(args[0], ArrayOrInt(args[1]), args[2], false, args[3]);
+ *rv = topi::argmax(args[0].cast<te::Tensor>(), ArrayOrInt(args[1]), args[2].cast<bool>(), false,
+ args[3].cast<bool>());
});
TVM_REGISTER_GLOBAL("topi.prod").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = topi::prod(args[0], ArrayOrInt(args[1]), args[2]);
+ *rv = topi::prod(args[0].cast<te::Tensor>(), ArrayOrInt(args[1]), args[2].cast<bool>());
});
TVM_REGISTER_GLOBAL("topi.all").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = topi::all(args[0], ArrayOrInt(args[1]), args[2]);
+ *rv = topi::all(args[0].cast<te::Tensor>(), ArrayOrInt(args[1]), args[2].cast<bool>());
});
TVM_REGISTER_GLOBAL("topi.any").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = topi::any(args[0], ArrayOrInt(args[1]), args[2]);
+ *rv = topi::any(args[0].cast<te::Tensor>(), ArrayOrInt(args[1]), args[2].cast<bool>());
});
TVM_REGISTER_GLOBAL("topi.collapse_sum").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = topi::collapse_sum(args[0], args[1]);
+ *rv = topi::collapse_sum(args[0].cast<te::Tensor>(), args[1].cast<Array<Integer>>());
});
} // namespace topi
diff --git a/src/topi/transform.cc b/src/topi/transform.cc
index f56b2978c8..fdd2af0573 100644
--- a/src/topi/transform.cc
+++ b/src/topi/transform.cc
@@ -38,129 +38,136 @@ using namespace tvm;
using namespace tvm::runtime;
TVM_REGISTER_GLOBAL("topi.expand_dims").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = expand_dims(args[0], args[1], args[2]);
+ *rv = expand_dims(args[0].cast<te::Tensor>(), args[1].cast<int>(), args[2].cast<int>());
});
TVM_REGISTER_GLOBAL("topi.transpose").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = transpose(args[0], args[1]);
+ *rv = transpose(args[0].cast<te::Tensor>(), args[1].cast<Array<Integer>>());
});
TVM_REGISTER_GLOBAL("topi.flip").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
// pass empty seq_lengths tensor to reverse_sequence
- *rv = reverse_sequence(args[0], Tensor(), args[1]);
+ *rv = reverse_sequence(args[0].cast<te::Tensor>(), Tensor(), args[1].cast<int>());
});
TVM_REGISTER_GLOBAL("topi.reverse_sequence").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = reverse_sequence(args[0], args[1], args[2], args[3]);
+ *rv =
+ reverse_sequence(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(), args[2].cast<int>());
});
TVM_REGISTER_GLOBAL("topi.reshape").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = reshape(args[0], args[1]);
+ *rv = reshape(args[0].cast<te::Tensor>(), args[1].cast<Array<PrimExpr>>());
});
TVM_REGISTER_GLOBAL("topi.sliding_window").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = sliding_window(args[0], args[1], args[2], args[3]);
+ *rv = sliding_window(args[0].cast<te::Tensor>(), args[1].cast<int>(),
+ args[2].cast<Array<Integer>>(), args[3].cast<Array<Integer>>());
});
TVM_REGISTER_GLOBAL("topi.squeeze").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = squeeze(args[0], ArrayOrInt(args[1]));
+ *rv = squeeze(args[0].cast<te::Tensor>(), ArrayOrInt(args[1]));
});
TVM_REGISTER_GLOBAL("topi.concatenate").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = concatenate(args[0], args[1]);
+ *rv = concatenate(args[0].cast<Array<te::Tensor>>(), args[1].cast<int>());
});
TVM_REGISTER_GLOBAL("topi.stack").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = stack(args[0], args[1]);
+ *rv = stack(args[0].cast<Array<te::Tensor>>(), args[1].cast<int>());
});
TVM_REGISTER_GLOBAL("topi.shape").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = shape(args[0], args[1]);
+ *rv = shape(args[0].cast<te::Tensor>(), args[1].cast<DataType>());
});
TVM_REGISTER_GLOBAL("topi.ndarray_size").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = ndarray_size(args[0], args[1]);
+ *rv = ndarray_size(args[0].cast<te::Tensor>(), args[1].cast<DataType>());
});
TVM_REGISTER_GLOBAL("topi.split").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
if (args[1].as<int>()) {
- *rv = split_n_sections(args[0], args[1], args[2]);
+ *rv = split_n_sections(args[0].cast<te::Tensor>(), args[1].cast<int>(), args[2].cast<int>());
} else {
- *rv = split_indices_array(args[0], args[1], args[2]);
+ *rv = split_indices_array(args[0].cast<te::Tensor>(), args[1].cast<Array<Integer>>(),
+ args[2].cast<int>());
}
});
TVM_REGISTER_GLOBAL("topi.layout_transform").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = layout_transform(args[0], args[1], args[2], args[3]);
+ *rv = layout_transform(args[0].cast<te::Tensor>(), args[1].cast<std::string>(),
+ args[2].cast<std::string>(), args[3].cast<std::string>());
});
TVM_REGISTER_GLOBAL("topi.take").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
if (args.size() == 4) {
- std::string mode = args[3];
- int batch_dims = args[2];
- *rv = take(args[0], args[1], batch_dims, mode);
+ auto mode = args[3].cast<std::string>();
+ int batch_dims = args[2].cast<int>();
+ *rv = take(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(), batch_dims, mode);
} else {
ICHECK_EQ(args.size(), 5) << "topi.take expects 4 or 5 arguments";
- int batch_dims = args[2];
- int axis = args[3];
- std::string mode = args[4];
- *rv = take(args[0], args[1], batch_dims, axis, mode);
+ int batch_dims = args[2].cast<int>();
+ int axis = args[3].cast<int>();
+ auto mode = args[4].cast<std::string>();
+ *rv = take(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(), batch_dims, axis, mode);
}
});
TVM_REGISTER_GLOBAL("topi.sequence_mask").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- double pad_val = args[2];
- int axis = args[3];
- *rv = sequence_mask(args[0], args[1], pad_val, axis);
+ double pad_val = args[2].cast<double>();
+ int axis = args[3].cast<int>();
+ *rv = sequence_mask(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(), pad_val, axis);
});
TVM_REGISTER_GLOBAL("topi.where").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = where(args[0], args[1], args[2]);
+ *rv = where(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(), args[2].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.arange").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = arange(args[0], args[1], args[2], args[3]);
+ *rv = arange(args[0].cast<PrimExpr>(), args[1].cast<PrimExpr>(), args[2].cast<PrimExpr>(),
+ args[3].cast<DataType>());
});
TVM_REGISTER_GLOBAL("topi.meshgrid").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = meshgrid(args[0], args[1]);
+ *rv = meshgrid(args[0].cast<Array<te::Tensor>>(), args[1].cast<std::string>());
});
TVM_REGISTER_GLOBAL("topi.repeat").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = repeat(args[0], args[1], args[2]);
+ *rv = repeat(args[0].cast<te::Tensor>(), args[1].cast<int>(), args[2].cast<int>());
});
TVM_REGISTER_GLOBAL("topi.tile").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = tile(args[0], args[1]);
+ *rv = tile(args[0].cast<te::Tensor>(), args[1].cast<Array<Integer>>());
});
TVM_REGISTER_GLOBAL("topi.gather").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = gather(args[0], args[1], args[2]);
+ *rv = gather(args[0].cast<te::Tensor>(), args[1].cast<int>(), args[2].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.gather_nd").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- int batch_dims = args[2];
- *rv = gather_nd(args[0], args[1], batch_dims);
+ int batch_dims = args[2].cast<int>();
+ *rv = gather_nd(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(), batch_dims);
});
TVM_REGISTER_GLOBAL("topi.unravel_index").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = unravel_index(args[0], args[1]);
+ *rv = unravel_index(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>());
});
TVM_REGISTER_GLOBAL("topi.sparse_to_dense").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = sparse_to_dense(args[0], args[1], args[2], args[3]);
+ *rv = sparse_to_dense(args[0].cast<te::Tensor>(), args[1].cast<Array<PrimExpr>>(),
+ args[2].cast<te::Tensor>(), args[3].cast<PrimExpr>());
});
TVM_REGISTER_GLOBAL("topi.matmul").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
switch (args.size()) {
case 2:
- *rv = matmul(args[0], args[1]);
+ *rv = matmul(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>());
break;
case 3:
- *rv = matmul(args[0], args[1], args[2]);
+ *rv = matmul(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(), args[2].cast<bool>());
break;
case 4:
- *rv = matmul(args[0], args[1], args[2], args[3]);
+ *rv = matmul(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(), args[2].cast<bool>(),
+ args[3].cast<bool>());
break;
default:
ICHECK(0) << "topi.matmul expects 2, 3 or 4 arguments";
@@ -169,28 +176,29 @@ TVM_REGISTER_GLOBAL("topi.matmul").set_body_packed([](TVMArgs args, TVMRetValue*
TVM_REGISTER_GLOBAL("topi.tensordot").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
if (args.size() == 2) {
- *rv = tensordot(args[0], args[1]);
+ *rv = tensordot(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>());
} else if (args.size() == 3) {
- *rv = tensordot(args[0], args[1], args[2]);
+ *rv = tensordot(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(), args[2].cast<int>());
} else {
- Array<PrimExpr> axes = args[3];
- *rv = tensordot(args[0], args[1], args[2], axes);
+ Array<PrimExpr> axes = args[3].cast<Array<PrimExpr>>();
+ *rv = tensordot(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(),
+ args[2].cast<Array<PrimExpr>>(), axes);
}
});
TVM_REGISTER_GLOBAL("topi.strided_slice").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- Tensor x = args[0];
- Array<PrimExpr> begin = args[1];
- Array<PrimExpr> end = args[2];
- Array<PrimExpr> strides = args[3];
- Array<Integer> axes = args[4];
- bool assume_inbound = args[6];
+ Tensor x = args[0].cast<te::Tensor>();
+ Array<PrimExpr> begin = args[1].cast<Array<PrimExpr>>();
+ Array<PrimExpr> end = args[2].cast<Array<PrimExpr>>();
+ Array<PrimExpr> strides = args[3].cast<Array<PrimExpr>>();
+ Array<Integer> axes = args[4].cast<Array<Integer>>();
+ bool assume_inbound = args[6].cast<bool>();
if (IsConstIntArray(begin) && IsConstIntArray(end) && IsConstIntArray(strides) &&
IsConstIntArray(x->shape)) {
- Array<Integer> begin_static = args[1];
- Array<Integer> end_static = args[2];
- Array<Integer> strides_static = args[3];
- std::string slice_mode = args[5];
+ Array<Integer> begin_static = args[1].cast<Array<Integer>>();
+ Array<Integer> end_static = args[2].cast<Array<Integer>>();
+ Array<Integer> strides_static = args[3].cast<Array<Integer>>();
+ auto slice_mode = args[5].cast<std::string>();
if (axes.size()) {
*rv = strided_slice_with_axes(x, begin_static, end_static, strides_static, axes, slice_mode);
} else {
@@ -207,38 +215,37 @@ TVM_REGISTER_GLOBAL("topi.strided_slice").set_body_packed([](TVMArgs args, TVMRe
TVM_REGISTER_GLOBAL("topi.dynamic_strided_slice")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- te::Tensor begin = args[1];
- te::Tensor end = args[2];
- te::Tensor strides = args[3];
- *rv = dynamic_strided_slice(args[0], begin, end, strides);
+ te::Tensor begin = args[1].cast<te::Tensor>();
+ te::Tensor end = args[2].cast<te::Tensor>();
+ te::Tensor strides = args[3].cast<te::Tensor>();
+ *rv = dynamic_strided_slice(args[0].cast<te::Tensor>(), begin, end, strides);
});
TVM_REGISTER_GLOBAL("topi.relax_dynamic_strided_slice")
- .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- te::Tensor begin = args[1];
- te::Tensor end = args[2];
- te::Tensor strides = args[3];
- Array<PrimExpr> output_shape = args[4];
- *rv = relax::dynamic_strided_slice(args[0], begin, end, strides, output_shape);
+ .set_body_typed([](te::Tensor x, te::Tensor begin, te::Tensor end, te::Tensor strides,
+ Array<PrimExpr> output_shape) {
+ return relax::dynamic_strided_slice(x, begin, end, strides, output_shape);
});
TVM_REGISTER_GLOBAL("topi.one_hot").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- int depth = args[3];
- int axis = args[4];
- DataType dtype = args[5];
- *rv = one_hot(args[0], args[1], args[2], depth, axis, dtype);
+ int depth = args[3].cast<int>();
+ int axis = args[4].cast<int>();
+ DataType dtype = args[5].cast<DataType>();
+ *rv = one_hot(args[0].cast<te::Tensor>(), args[1].cast<PrimExpr>(), args[2].cast<PrimExpr>(),
+ depth, axis, dtype);
});
TVM_REGISTER_GLOBAL("topi.matrix_set_diag").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- int k1 = args[2];
- int k2 = args[3];
- bool super_diag_right_align = args[4];
- bool sub_diag_right_align = args[5];
- *rv = matrix_set_diag(args[0], args[1], k1, k2, super_diag_right_align, sub_diag_right_align);
+ int k1 = args[2].cast<int>();
+ int k2 = args[3].cast<int>();
+ bool super_diag_right_align = args[4].cast<bool>();
+ bool sub_diag_right_align = args[5].cast<bool>();
+ *rv = matrix_set_diag(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(), k1, k2,
+ super_diag_right_align, sub_diag_right_align);
});
-TVM_REGISTER_GLOBAL("topi.adv_index").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = adv_index(args[0], args[1]);
+TVM_REGISTER_GLOBAL("topi.adv_index").set_body_typed([](te::Tensor x, Array<te::Tensor> indices) {
+ return adv_index(x, indices);
});
} // namespace topi
diff --git a/src/topi/utils.cc b/src/topi/utils.cc
index c3ad01b1a5..a4084b570c 100644
--- a/src/topi/utils.cc
+++ b/src/topi/utils.cc
@@ -29,17 +29,21 @@
namespace tvm {
namespace topi {
TVM_REGISTER_GLOBAL("topi.utils.is_empty_shape").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = topi::detail::is_empty_shape(args[0]);
+ *rv = topi::detail::is_empty_shape(args[0].cast<Array<PrimExpr>>());
});
TVM_REGISTER_GLOBAL("topi.utils.bilinear_sample_nchw")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = detail::bilinear_sample_nchw(args[0], args[1], args[2], args[3]);
+ *rv =
+ detail::bilinear_sample_nchw(args[0].cast<te::Tensor>(), args[1].cast<Array<PrimExpr>>(),
+ args[2].cast<PrimExpr>(), args[3].cast<PrimExpr>());
});
TVM_REGISTER_GLOBAL("topi.utils.bilinear_sample_nhwc")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = detail::bilinear_sample_nhwc(args[0], args[1], args[2], args[3]);
+ *rv =
+ detail::bilinear_sample_nhwc(args[0].cast<te::Tensor>(), args[1].cast<Array<PrimExpr>>(),
+ args[2].cast<PrimExpr>(), args[3].cast<PrimExpr>());
});
} // namespace topi
diff --git a/src/topi/vision.cc b/src/topi/vision.cc
index 68eab43fc2..62644e584f 100644
--- a/src/topi/vision.cc
+++ b/src/topi/vision.cc
@@ -32,7 +32,7 @@ using namespace tvm;
using namespace tvm::runtime;
TVM_REGISTER_GLOBAL("topi.vision.reorg").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- *rv = vision::reorg(args[0], args[1]);
+ *rv = vision::reorg(args[0].cast<te::Tensor>(), args[1].cast<int>());
});
} // namespace topi
diff --git a/tests/python/contrib/test_hexagon/README_RPC.md b/tests/python/contrib/test_hexagon/README_RPC.md
index 5beffd800b..652237d4f5 100644
--- a/tests/python/contrib/test_hexagon/README_RPC.md
+++ b/tests/python/contrib/test_hexagon/README_RPC.md
@@ -81,9 +81,9 @@ Which eventually jumps to the following line in C++, which creates a RPC client
```cpp
TVM_REGISTER_GLOBAL("rpc.Connect").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- std::string url = args[0];
- int port = args[1];
- std::string key = args[2];
+ auto url = args[0].cast<std::string>();
+ int port = args[1].cast<int>();
+ auto key = args[2].cast<std::string>();
*rv = RPCClientConnect(url, port, key,
TVMArgs(args.values + 3, args.type_codes + 3, args.size() - 3));
});
@@ -96,8 +96,8 @@ TVM_REGISTER_GLOBAL("rpc.Connect").set_body_packed([](TVMArgs args, TVMRetValue*
```cpp
TVM_REGISTER_GLOBAL("tvm.contrib.hexagon.create_hexagon_session")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
- std::string session_name = args[0];
- int remote_stack_size_bytes = args[1];
+ auto session_name = args[0].cast<std::string>();
+ int remote_stack_size_bytes = args[1].cast<int>();
HexagonTransportChannel* hexagon_channel =
new HexagonTransportChannel(hexagon_rpc_URI CDSP_DOMAIN, remote_stack_size_bytes);
std::unique_ptr<RPCChannel> channel(hexagon_channel);
@@ -229,8 +229,8 @@ The handler is passed to the following function
```cpp
void RPCCopyAmongRemote(RPCSession* handler, TVMArgs args, TVMRetValue* rv) {
- DLTensor* from = args[0];
- DLTensor* to = args[1];
+ auto from = args[0].cast<DLTensor*>();
+ auto to = args[1].cast<DLTensor*>();
...
handler->GetDeviceAPI(dev)->CopyDataFromTo(from, to, stream);
}
diff --git a/web/emcc/tvmjs_support.cc b/web/emcc/tvmjs_support.cc
index 620095391c..2c4ba7da08 100644
--- a/web/emcc/tvmjs_support.cc
+++ b/web/emcc/tvmjs_support.cc
@@ -134,7 +134,7 @@ class AsyncLocalSession : public LocalSession {
auto it = async_func_set_.find(func);
if (it != async_func_set_.end()) {
PackedFunc packed_callback([callback, this](TVMArgs args, TVMRetValue*) {
- int code = args[0];
+ int code = args[0].cast<int>();
TVMRetValue rv;
rv = args[1];
if (code == static_cast<int>(RPCCode::kReturn)) {
@@ -224,7 +224,7 @@ class AsyncLocalSession : public LocalSession {
}
CHECK(async_wait_.has_value());
PackedFunc packed_callback([on_complete](TVMArgs args, TVMRetValue*) {
- int code = args[0];
+ int code = args[0].cast<int>();
on_complete(static_cast<RPCCode>(code), args.Slice(1));
});
(*async_wait_)(packed_callback);
diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc
index 1939a48be6..0873101bbb 100644
--- a/web/emcc/wasm_runtime.cc
+++ b/web/emcc/wasm_runtime.cc
@@ -111,15 +111,15 @@ TVM_REGISTER_GLOBAL("testing.call").set_body_packed([](TVMArgs args, TVMRetValue
});
TVM_REGISTER_GLOBAL("testing.ret_string").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- *ret = args[0].operator String();
+ *ret = args[0].cast<String>();
});
TVM_REGISTER_GLOBAL("testing.log_info_str").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- LOG(INFO) << args[0].operator String();
+ LOG(INFO) << args[0].cast<String>();
});
TVM_REGISTER_GLOBAL("testing.log_fatal_str").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- LOG(FATAL) << args[0].operator String();
+ LOG(FATAL) << args[0].cast<String>();
});
TVM_REGISTER_GLOBAL("testing.add_one").set_body_typed([](int x) { return x + 1; });
@@ -131,7 +131,7 @@ TVM_REGISTER_GLOBAL("testing.wrap_callback").set_body_packed([](TVMArgs args, TV
// internal function used for debug and testing purposes
TVM_REGISTER_GLOBAL("testing.object_use_count").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
- runtime::ObjectRef obj = args[0];
+ auto obj = args[0].cast<ffi::ObjectRef>();
// subtract the current one because we always copy
// and get another value.
*ret = (obj.use_count() - 1);
diff --git a/web/emcc/webgpu_runtime.cc b/web/emcc/webgpu_runtime.cc
index 0b36f392c5..5b064cfc16 100644
--- a/web/emcc/webgpu_runtime.cc
+++ b/web/emcc/webgpu_runtime.cc
@@ -176,14 +176,14 @@ class WebGPUModuleNode final : public runtime::ModuleNode {
});
} else if (name == "webgpu.get_shader") {
return PackedFunc([this](TVMArgs args, TVMRetValue* rv) {
- std::string name = args[0];
+ auto name = args[0].cast<std::string>();
auto it = smap_.find(name);
ICHECK(it != smap_.end()) << "Cannot find code " << name;
*rv = it->second;
});
} else if (name == "webgpu.update_prebuild") {
return PackedFunc([this](TVMArgs args, TVMRetValue* rv) {
- std::string name = args[0];
+ auto name = args[0].cast<std::string>();
PackedFunc func = args[1];
prebuild_[name] = func;
});