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:29 UTC
(tvm) branch refactor-s2 updated (d5e4fa026a -> 7126c65695)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a change to branch refactor-s2
in repository https://gitbox.apache.org/repos/asf/tvm.git
from d5e4fa026a [MACOS] enable backtrace for macos
new 184a813e77 [WIN] Initial stab to make win ffi compile
new 25772b0c97 [FFI] Refactor to make Any require explicit cast to value
new 7126c65695 Update to enable explicit Any cast
The 3 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
apps/cpp_rpc/rpc_server.cc | 4 +-
apps/ios_rpc/tvmrpc/TVMRuntime.mm | 4 +-
docs/arch/runtime.rst | 4 +-
ffi/CMakeLists.txt | 3 +
ffi/cmake/Utils/CxxWarning.cmake | 2 +-
ffi/cmake/Utils/Library.cmake | 11 ++
ffi/include/tvm/ffi/any.h | 7 +-
ffi/include/tvm/ffi/base_details.h | 12 +-
ffi/include/tvm/ffi/cast.h | 6 +-
ffi/include/tvm/ffi/container/container_details.h | 2 +-
ffi/include/tvm/ffi/container/map.h | 11 +-
ffi/include/tvm/ffi/container/shape.h | 27 ++--
ffi/include/tvm/ffi/container/variant.h | 4 +-
ffi/include/tvm/ffi/error.h | 12 +-
ffi/include/tvm/ffi/function.h | 13 +-
ffi/include/tvm/ffi/function_details.h | 20 ++-
ffi/include/tvm/ffi/reflection/reflection.h | 2 +-
ffi/include/tvm/ffi/type_traits.h | 8 +-
ffi/src/ffi/container.cc | 53 ++++----
ffi/src/ffi/dtype.cc | 10 +-
ffi/src/ffi/ndarray.cc | 12 +-
ffi/src/ffi/testing.cc | 14 +-
ffi/src/ffi/traceback_win.cc | 128 +++++++++++++++++
ffi/tests/cpp/CMakeLists.txt | 1 +
ffi/tests/cpp/test_any.cc | 52 +++----
ffi/tests/cpp/test_array.cc | 28 ++--
ffi/tests/cpp/test_dtype.cc | 4 +-
ffi/tests/cpp/test_function.cc | 59 ++++----
ffi/tests/cpp/test_map.cc | 34 ++---
ffi/tests/cpp/test_optional.cc | 20 +--
ffi/tests/cpp/test_reflection.cc | 4 +-
ffi/tests/cpp/test_rvalue_ref.cc | 10 +-
ffi/tests/cpp/test_shape.cc | 4 +-
ffi/tests/cpp/test_tuple.cc | 12 +-
ffi/tests/cpp/test_variant.cc | 10 +-
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 +-
216 files changed, 1570 insertions(+), 1322 deletions(-)
create mode 100644 ffi/src/ffi/traceback_win.cc
(tvm) 01/03: [WIN] Initial stab to make win ffi compile
Posted by tq...@apache.org.
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 184a813e7783741d14063ea582fbd6dd2bc73e7d
Author: tqchen <tq...@gmail.com>
AuthorDate: Wed Apr 30 11:23:35 2025 -0400
[WIN] Initial stab to make win ffi compile
Still need followup refactors
---
ffi/CMakeLists.txt | 3 +
ffi/cmake/Utils/CxxWarning.cmake | 2 +-
ffi/cmake/Utils/Library.cmake | 11 ++
ffi/include/tvm/ffi/any.h | 1 +
ffi/include/tvm/ffi/base_details.h | 12 +-
ffi/include/tvm/ffi/container/container_details.h | 2 +-
ffi/include/tvm/ffi/container/map.h | 11 +-
ffi/include/tvm/ffi/error.h | 12 +-
ffi/include/tvm/ffi/function.h | 6 +-
ffi/include/tvm/ffi/function_details.h | 13 ++-
ffi/include/tvm/ffi/type_traits.h | 8 +-
ffi/src/ffi/container.cc | 53 ++++-----
ffi/src/ffi/dtype.cc | 6 +-
ffi/src/ffi/testing.cc | 12 +-
ffi/src/ffi/traceback_win.cc | 128 ++++++++++++++++++++++
ffi/tests/cpp/CMakeLists.txt | 1 +
16 files changed, 216 insertions(+), 65 deletions(-)
diff --git a/ffi/CMakeLists.txt b/ffi/CMakeLists.txt
index 5ba8fd0f2f..11e380e46f 100644
--- a/ffi/CMakeLists.txt
+++ b/ffi/CMakeLists.txt
@@ -59,6 +59,7 @@ target_link_libraries(tvm_ffi_header INTERFACE dlpack_header)
########## Target: `tvm_ffi` ##########
add_library(tvm_ffi_objs OBJECT
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/traceback.cc"
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/traceback_win.cc"
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/object.cc"
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/function.cc"
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/ndarray.cc"
@@ -110,6 +111,8 @@ target_link_libraries(tvm_ffi_static PUBLIC tvm_ffi_header)
install(TARGETS tvm_ffi_static DESTINATION lib${LIB_SUFFIX})
install(TARGETS tvm_ffi_shared DESTINATION lib${LIB_SUFFIX})
+add_msvc_compact_defs(tvm_ffi_objs)
+
########## Adding tests ##########
if (${PROJECT_NAME} STREQUAL ${CMAKE_PROJECT_NAME})
diff --git a/ffi/cmake/Utils/CxxWarning.cmake b/ffi/cmake/Utils/CxxWarning.cmake
index 649ce2911a..c272bfdf7b 100644
--- a/ffi/cmake/Utils/CxxWarning.cmake
+++ b/ffi/cmake/Utils/CxxWarning.cmake
@@ -23,7 +23,7 @@ function(add_cxx_warning target_name)
endif()
# MSVC
if(MSVC)
- target_compile_options(${target_name} PRIVATE "/W4" "/WX")
+ # target_compile_options(${target_name} PRIVATE "/W4" "/WX")
return()
endif()
message(FATAL_ERROR "Unsupported compiler: ${CMAKE_CXX_COMPILER_ID}")
diff --git a/ffi/cmake/Utils/Library.cmake b/ffi/cmake/Utils/Library.cmake
index 9f91c53ce6..6afea393d6 100644
--- a/ffi/cmake/Utils/Library.cmake
+++ b/ffi/cmake/Utils/Library.cmake
@@ -28,6 +28,17 @@ function(add_dsymutil target_name)
endif()
endfunction()
+function(add_msvc_compact_defs target_name)
+ # running if we are under msvc
+ if(MSVC)
+ target_compile_definitions(${target_name} PUBLIC -DWIN32_LEAN_AND_MEAN)
+ target_compile_definitions(${target_name} PUBLIC -D_CRT_SECURE_NO_WARNINGS)
+ target_compile_definitions(${target_name} PUBLIC -D_SCL_SECURE_NO_WARNINGS)
+ target_compile_definitions(${target_name} PUBLIC -D_ENABLE_EXTENDED_ALIGNED_STORAGE)
+ target_compile_definitions(${target_name} PUBLIC -DNOMINMAX)
+ endif()
+endfunction()
+
function(add_target_from_obj target_name obj_target_name)
add_library(${target_name}_static STATIC $<TARGET_OBJECTS:${obj_target_name}>)
set_target_properties(
diff --git a/ffi/include/tvm/ffi/any.h b/ffi/include/tvm/ffi/any.h
index 049c07066c..3f787e7571 100644
--- a/ffi/include/tvm/ffi/any.h
+++ b/ffi/include/tvm/ffi/any.h
@@ -304,6 +304,7 @@ class Any {
}
return *std::move(opt);
}
+
/*
* \brief Check if the two Any are same type and value in shallow comparison.
* \param other The other Any
diff --git a/ffi/include/tvm/ffi/base_details.h b/ffi/include/tvm/ffi/base_details.h
index 9debc5b8ee..cc54be77bb 100644
--- a/ffi/include/tvm/ffi/base_details.h
+++ b/ffi/include/tvm/ffi/base_details.h
@@ -31,6 +31,10 @@
#include <cstddef>
#include <utility>
+#if defined(_MSC_VER)
+#include <windows.h>
+#endif
+
#if defined(_MSC_VER)
#define TVM_FFI_INLINE __forceinline
#else
@@ -95,8 +99,8 @@
#define TVM_FFI_LOG_EXCEPTION_CALL_END(Name) \
} \
catch (const std::exception& err) { \
- std::cout << "Exception caught during " << #Name << ":\n" << err.what() << std::endl; \
- throw err; \
+ std::cerr << "Exception caught during " << #Name << ":\n" << err.what() << std::endl; \
+ exit(-1); \
}
/*!
@@ -197,7 +201,7 @@ TVM_FFI_INLINE uint64_t StableHashBytes(const char* data, size_t size) {
const char* end = it + size;
uint64_t result = 0;
for (; it + 8 <= end; it += 8) {
- if (TVM_FFI_IO_NO_ENDIAN_SWAP) {
+ if constexpr (TVM_FFI_IO_NO_ENDIAN_SWAP) {
u.a[0] = it[0];
u.a[1] = it[1];
u.a[2] = it[2];
@@ -240,7 +244,7 @@ TVM_FFI_INLINE uint64_t StableHashBytes(const char* data, size_t size) {
it += 1;
a += 1;
}
- if (!TVM_FFI_IO_NO_ENDIAN_SWAP) {
+ if constexpr (!TVM_FFI_IO_NO_ENDIAN_SWAP) {
std::swap(u.a[0], u.a[7]);
std::swap(u.a[1], u.a[6]);
std::swap(u.a[2], u.a[5]);
diff --git a/ffi/include/tvm/ffi/container/container_details.h b/ffi/include/tvm/ffi/container/container_details.h
index 9f54b32ecb..39ebcb4a58 100644
--- a/ffi/include/tvm/ffi/container/container_details.h
+++ b/ffi/include/tvm/ffi/container/container_details.h
@@ -107,7 +107,7 @@ class InplaceArrayBase {
* \brief Destroy the Inplace Array Base object
*/
~InplaceArrayBase() {
- if (!(std::is_standard_layout<ElemType>::value && std::is_trivial<ElemType>::value)) {
+ if constexpr (!(std::is_standard_layout<ElemType>::value && std::is_trivial<ElemType>::value)) {
size_t size = Self()->GetSize();
for (size_t i = 0; i < size; ++i) {
ElemType* fp = reinterpret_cast<ElemType*>(AddressOf(i));
diff --git a/ffi/include/tvm/ffi/container/map.h b/ffi/include/tvm/ffi/container/map.h
index 26eeee7925..7d805d61e9 100644
--- a/ffi/include/tvm/ffi/container/map.h
+++ b/ffi/include/tvm/ffi/container/map.h
@@ -292,7 +292,10 @@ class SmallMapObj : public MapObj, public details::InplaceArrayBase<SmallMapObj,
}
KVType* begin = static_cast<KVType*>(AddressOf(0));
// call destructor to destroy the item in `begin + index`
- (begin + index)->KVType::~KVType();
+ // Explicit call Any::~Any() to destroy the Any object
+ // Favor this over ~KVType as MSVC may not support ~KVType (need the original name)
+ (begin + index)->first.Any::~Any();
+ (begin + index)->second.Any::~Any();
// IMPORTANT: We do direct raw memmove to bring later items to the current position
// to preserve the order of insertion.
// This works because direct memory copy preserves the Any's move semantics.
@@ -1008,8 +1011,10 @@ class DenseMapObj : public MapObj {
void SetEmpty() const { Meta() = uint8_t(kEmptySlot); }
/*! \brief Destruct the item in the entry */
void DestructData() const {
- // explict call destructor to destroy the item
- (&Data())->KVType::~KVType();
+ // explicit call destructor to destroy the item
+ // Favor this over ~KVType as MSVC may not support ~KVType (need the original name)
+ (&Data())->first.Any::~Any();
+ (&Data())->second.Any::~Any();
}
/*! \brief Set the entry to be protected */
void SetProtected() const { Meta() = uint8_t(kProtectedSlot); }
diff --git a/ffi/include/tvm/ffi/error.h b/ffi/include/tvm/ffi/error.h
index 7473853c6c..8c97ebc8e9 100644
--- a/ffi/include/tvm/ffi/error.h
+++ b/ffi/include/tvm/ffi/error.h
@@ -84,8 +84,8 @@ class ErrorObj : public Object, public TVMFFIErrorInfo {
* \brief Update the traceback of the error object.
* \param traceback The traceback to update.
*/
- void UpdateTraceback(const char* traceback) {
- this->traceback_data_ = traceback;
+ void UpdateTraceback(const char* traceback_str) {
+ this->traceback_data_ = traceback_str;
this->traceback = this->traceback_data_.c_str();
this->what_data_ = (std::string("Traceback (most recent call last):\n") + this->traceback +
this->kind + ": " + this->message + '\n');
@@ -135,11 +135,6 @@ class ErrorBuilder {
explicit ErrorBuilder(std::string kind, std::string traceback, bool log_before_throw)
: kind_(kind), traceback_(traceback), log_before_throw_(log_before_throw) {}
-// MSVC disable warning in error builder as it is exepected
-#ifdef _MSC_VER
-#pragma disagnostic push
-#pragma warning(disable : 4722)
-#endif
// avoid inline to reduce binary size, error throw path do not need to be fast
[[noreturn]] ~ErrorBuilder() noexcept(false) {
::tvm::ffi::Error error(std::move(kind_), stream_.str(), std::move(traceback_));
@@ -148,9 +143,6 @@ class ErrorBuilder {
}
throw error;
}
-#ifdef _MSC_VER
-#pragma disagnostic pop
-#endif
std::ostringstream& stream() { return stream_; }
diff --git a/ffi/include/tvm/ffi/function.h b/ffi/include/tvm/ffi/function.h
index d5e31372c6..22cc553a61 100644
--- a/ffi/include/tvm/ffi/function.h
+++ b/ffi/include/tvm/ffi/function.h
@@ -346,9 +346,9 @@ class Function : public ObjectRef {
static Function ImportFromExternDLL(Function other) {
const FunctionObj* other_func = static_cast<const FunctionObj*>(other.get());
// the other function comes from the same dll, no action needed
- if (other_func->safe_call == FunctionObj::SafeCall ||
- other_func->safe_call == details::ImportedFunctionObjImpl::SafeCall ||
- other_func->safe_call == details::ExternCFunctionObjImpl::SafeCall) {
+ if (other_func->safe_call == &(FunctionObj::SafeCall) ||
+ other_func->safe_call == &(details::ImportedFunctionObjImpl::SafeCall) ||
+ other_func->safe_call == &(details::ExternCFunctionObjImpl::SafeCall)) {
return other;
}
// the other function coems from a different library
diff --git a/ffi/include/tvm/ffi/function_details.h b/ffi/include/tvm/ffi/function_details.h
index f877a45272..0244398efe 100644
--- a/ffi/include/tvm/ffi/function_details.h
+++ b/ffi/include/tvm/ffi/function_details.h
@@ -39,7 +39,10 @@ namespace details {
/*!
* \brief Empty struct, used to reduce unused variable in global static initialization.
*/
-struct EmptyStruct {};
+struct EmptyStruct {
+ /*! \brief at least one byte to ensure address */
+ uint8_t pad;
+};
template <typename ArgType>
struct Arg2Str {
@@ -203,8 +206,12 @@ TVM_FFI_INLINE void unpack_call(const std::string* optional_name, const F& f, co
int32_t num_args, Any* rv) {
using FuncInfo = FunctionInfo<F>;
FGetFuncSignature f_sig = FuncInfo::Sig;
- static_assert(FuncInfo::unpacked_supported,
- "The function signature cannot support unpacked call");
+
+ // somehow MSVC does not support the static constexpr member in this case, function is fine
+#ifndef _MSC_VER
+ static_assert(FuncInfo::unpacked_supported, "The function signature do not support unpacked");
+#endif
+
if (nargs != num_args) {
TVM_FFI_THROW(TypeError) << "Mismatched number of arguments when calling: `"
<< (optional_name == nullptr ? "" : *optional_name)
diff --git a/ffi/include/tvm/ffi/type_traits.h b/ffi/include/tvm/ffi/type_traits.h
index f8a24ac177..d350aea82a 100644
--- a/ffi/include/tvm/ffi/type_traits.h
+++ b/ffi/include/tvm/ffi/type_traits.h
@@ -140,9 +140,7 @@ struct TypeTraits<std::nullptr_t> : public TypeTraitsBase {
return nullptr;
}
- static TVM_FFI_INLINE std::nullptr_t MoveFromAnyStorageAfterCheck(TVMFFIAny* src) {
- return nullptr;
- }
+ static TVM_FFI_INLINE std::nullptr_t MoveFromAnyStorageAfterCheck(TVMFFIAny*) { return nullptr; }
static TVM_FFI_INLINE std::optional<std::nullptr_t> TryConvertFromAnyView(const TVMFFIAny* src) {
if (src->type_index == TypeIndex::kTVMFFINone) {
@@ -256,7 +254,7 @@ struct TypeTraits<Int, std::enable_if_t<std::is_integral_v<Int>>> : public TypeT
return src->type_index == TypeIndex::kTVMFFIInt;
}
- static TVM_FFI_INLINE int CopyFromAnyStorageAfterCheck(const TVMFFIAny* src) {
+ static TVM_FFI_INLINE Int CopyFromAnyStorageAfterCheck(const TVMFFIAny* src) {
return static_cast<Int>(src->v_int64);
}
@@ -406,7 +404,7 @@ struct TypeTraits<DLTensor*> : public TypeTraitsBase {
result->v_ptr = src;
}
- static TVM_FFI_INLINE void MoveToAny(DLTensor* src, TVMFFIAny* result) {
+ static TVM_FFI_INLINE void MoveToAny(DLTensor*, TVMFFIAny*) {
TVM_FFI_THROW(RuntimeError)
<< "DLTensor* cannot be held in Any as it does not retain ownership, use NDArray instead";
}
diff --git a/ffi/src/ffi/container.cc b/ffi/src/ffi/container.cc
index 885e8395ed..563505b74c 100644
--- a/ffi/src/ffi/container.cc
+++ b/ffi/src/ffi/container.cc
@@ -39,7 +39,6 @@ TVM_FFI_REGISTER_GLOBAL("ffi.ArrayGetItem")
TVM_FFI_REGISTER_GLOBAL("ffi.ArraySize").set_body_typed([](const ffi::ArrayObj* n) -> int64_t {
return static_cast<int64_t>(n->size());
});
-
// Map
TVM_FFI_REGISTER_GLOBAL("ffi.Map").set_body_packed([](ffi::PackedArgs args, Any* ret) {
TVM_FFI_ICHECK_EQ(args.size() % 2, 0);
@@ -60,33 +59,35 @@ TVM_FFI_REGISTER_GLOBAL("ffi.MapGetItem")
TVM_FFI_REGISTER_GLOBAL("ffi.MapCount")
.set_body_typed([](const ffi::MapObj* n, const Any& k) -> int64_t { return n->count(k); });
+// Favor struct outside function scope as MSVC may have bug for in fn scope struct.
+class MapForwardIterFunctor {
+ public:
+ MapForwardIterFunctor(ffi::MapObj::iterator iter, ffi::MapObj::iterator end)
+ : iter_(iter), end_(end) {}
+ // 0 get current key
+ // 1 get current value
+ // 2 move to next: return true if success, false if end
+ Any operator()(int command) const {
+ if (command == 0) {
+ return (*iter_).first;
+ } else if (command == 1) {
+ return (*iter_).second;
+ } else {
+ ++iter_;
+ if (iter_ == end_) {
+ return false;
+ }
+ return true;
+ }
+ }
+
+ private:
+ mutable ffi::MapObj::iterator iter_;
+ ffi::MapObj::iterator end_;
+};
+
TVM_FFI_REGISTER_GLOBAL("ffi.MapForwardIterFunctor")
.set_body_typed([](const ffi::MapObj* n) -> ffi::Function {
- class MapForwardIterFunctor {
- public:
- MapForwardIterFunctor(ffi::MapObj::iterator iter, ffi::MapObj::iterator end)
- : iter_(iter), end_(end) {}
- // 0 get current key
- // 1 get current value
- // 2 move to next: return true if success, false if end
- Any operator()(int command) const {
- if (command == 0) {
- return (*iter_).first;
- } else if (command == 1) {
- return (*iter_).second;
- } else {
- ++iter_;
- if (iter_ == end_) {
- return false;
- }
- return true;
- }
- }
-
- private:
- mutable ffi::MapObj::iterator iter_;
- ffi::MapObj::iterator end_;
- };
return ffi::Function::FromUnpacked(MapForwardIterFunctor(n->begin(), n->end()));
});
diff --git a/ffi/src/ffi/dtype.cc b/ffi/src/ffi/dtype.cc
index 25691b9350..a065c70389 100644
--- a/ffi/src/ffi/dtype.cc
+++ b/ffi/src/ffi/dtype.cc
@@ -207,8 +207,8 @@ inline DLDataType StringViewToDLDataType_(std::string_view str) {
const char* scan;
auto parse_float = [&](const std::string_view& str, int offset, int code, int bits) {
- dtype.code = code;
- dtype.bits = bits;
+ dtype.code = static_cast<uint8_t>(code);
+ dtype.bits = static_cast<uint8_t>(bits);
scan = str.data() + offset;
char* endpt = nullptr;
if (*scan == 'x') {
@@ -288,7 +288,7 @@ inline DLDataType StringViewToDLDataType_(std::string_view str) {
dtype.bits = 16;
scan = str.data() + 6;
} else if (str.compare(0, 6, "custom") == 0) {
- dtype.code = details::ParseCustomDataTypeCode(str, &scan);
+ dtype.code = static_cast<uint8_t>(details::ParseCustomDataTypeCode(str, &scan));
} else {
scan = str.data();
TVM_FFI_THROW(ValueError) << "unknown dtype `" << str << '`';
diff --git a/ffi/src/ffi/testing.cc b/ffi/src/ffi/testing.cc
index 73f247f242..b4dc5b718a 100644
--- a/ffi/src/ffi/testing.cc
+++ b/ffi/src/ffi/testing.cc
@@ -26,6 +26,12 @@
namespace tvm {
namespace ffi {
+void TestRaiseError(String kind, String msg) {
+ throw ffi::Error(kind, msg, TVM_FFI_TRACEBACK_HERE);
+}
+
+TVM_FFI_REGISTER_GLOBAL("testing.test_raise_error").set_body_typed(TestRaiseError);
+
TVM_FFI_REGISTER_GLOBAL("testing.nop").set_body_packed([](PackedArgs args, Any* ret) {
*ret = args[0];
});
@@ -41,12 +47,6 @@ TVM_FFI_REGISTER_GLOBAL("testing.apply").set_body_packed([](PackedArgs args, Any
TestApply(f, args.Slice(1), ret);
});
-void TestRaiseError(String kind, String msg) {
- throw ffi::Error(kind, msg, TVM_FFI_TRACEBACK_HERE);
-}
-
-TVM_FFI_REGISTER_GLOBAL("testing.test_raise_error").set_body_typed(TestRaiseError);
-
TVM_FFI_REGISTER_GLOBAL("testing.run_check_signal").set_body_typed([](int nsec) {
for (int i = 0; i < nsec; ++i) {
if (TVMFFIEnvCheckSignals() != 0) {
diff --git a/ffi/src/ffi/traceback_win.cc b/ffi/src/ffi/traceback_win.cc
new file mode 100644
index 0000000000..786b70447a
--- /dev/null
+++ b/ffi/src/ffi/traceback_win.cc
@@ -0,0 +1,128 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file traceback_win.cc
+ * \brief Traceback implementation on windows platform
+ * \note We use the term "traceback" to be consistent with python naming convention.
+ */
+#ifdef _MSC_VER
+
+#include <tvm/ffi/c_api.h>
+#include <tvm/ffi/error.h>
+
+#include "./traceback.h"
+
+namespace tvm {
+namespace ffi {
+namespace {
+
+#include <dbghelp.h>
+#include <windows.h>
+
+#include <iostream>
+#include <vector>
+
+#pragma comment(lib, "DbgHelp.lib")
+
+std::string Traceback() {
+ TracebackStorage traceback;
+ HANDLE process = GetCurrentProcess();
+ HANDLE thread = GetCurrentThread();
+
+ SymInitialize(process, NULL, TRUE);
+ CONTEXT context = {};
+ RtlCaptureContext(&context);
+
+ STACKFRAME64 stack = {};
+ DWORD machine_type;
+
+#if defined(_M_X64)
+ machine_type = IMAGE_FILE_MACHINE_AMD64;
+ stack.AddrPC.Offset = context.Rip;
+ stack.AddrFrame.Offset = context.Rbp;
+ stack.AddrStack.Offset = context.Rsp;
+#elif defined(_M_IX86)
+ machine_type = IMAGE_FILE_MACHINE_I386;
+ stack.AddrPC.Offset = context.Eip;
+ stack.AddrFrame.Offset = context.Ebp;
+ stack.AddrStack.Offset = context.Esp;
+#else
+#error "Platform not supported!"
+#endif
+
+ stack.AddrPC.Mode = AddrModeFlat;
+ stack.AddrFrame.Mode = AddrModeFlat;
+ stack.AddrStack.Mode = AddrModeFlat;
+
+ while (!traceback.ExceedTracebackLimit()) {
+ if (!StackWalk64(machine_type, process, thread, &stack, &context, nullptr,
+ SymFunctionTableAccess64, SymGetModuleBase64, nullptr)) {
+ break;
+ }
+
+ if (stack.AddrPC.Offset == 0) {
+ break;
+ }
+ const char* filename = nullptr;
+ const char* symbol = "<unknown>";
+ int lineno = 0;
+ // Get file and line number
+ IMAGEHLP_LINE64 line_info = {};
+ line_info.SizeOfStruct = sizeof(IMAGEHLP_LINE64);
+ DWORD displacement32 = 0;
+
+ if (SymGetLineFromAddr64(process, stack.AddrPC.Offset, &displacement32, &line_info)) {
+ filename = line_info.FileName;
+ lineno = line_info.LineNumber;
+ }
+ // allocate symbol info that aligns to the SYMBOL_INFO
+ // we use u64 here to be safe
+ size_t total_symbol_bytes = sizeof(SYMBOL_INFO) + MAX_SYM_NAME * sizeof(TCHAR);
+ size_t total_u64_words = (total_symbol_bytes + 7) / 8;
+ static_assert(8 % alignof(SYMBOL_INFO) == 0);
+ std::vector<uint64_t> symbol_buffer(total_u64_words, 0);
+ PSYMBOL_INFO symbol_info = reinterpret_cast<PSYMBOL_INFO>(symbol_buffer.data());
+ symbol_info->SizeOfStruct = sizeof(SYMBOL_INFO);
+ symbol_info->MaxNameLen = MAX_SYM_NAME;
+ DWORD64 displacement = 0;
+ if (SymFromAddr(process, stack.AddrPC.Offset, &displacement, symbol_info)) {
+ symbol = symbol_info->Name;
+ }
+
+ if (ShouldStopTraceback(filename, symbol)) {
+ break;
+ }
+ if (ShouldExcludeFrame(filename, symbol)) {
+ continue;
+ }
+ traceback.Append(filename, symbol, lineno);
+ }
+ SymCleanup(process);
+ return traceback.GetTraceback();
+}
+} // namespace
+} // namespace ffi
+} // namespace tvm
+
+const char* TVMFFITraceback(const char*, int, const char*) {
+ static thread_local std::string traceback_str;
+ traceback_str = ::tvm::ffi::Traceback();
+ return traceback_str.c_str();
+}
+#endif // _MSC_VER
diff --git a/ffi/tests/cpp/CMakeLists.txt b/ffi/tests/cpp/CMakeLists.txt
index ce0b318d41..3ae2df73c8 100644
--- a/ffi/tests/cpp/CMakeLists.txt
+++ b/ffi/tests/cpp/CMakeLists.txt
@@ -17,6 +17,7 @@ set_target_properties(
add_cxx_warning(tvm_ffi_tests)
add_sanitizer_address(tvm_ffi_tests)
add_dsymutil(tvm_ffi_tests)
+add_msvc_compact_defs(tvm_ffi_tests)
target_link_libraries(tvm_ffi_tests PRIVATE tvm_ffi_shared)
add_googletest(tvm_ffi_tests)
(tvm) 03/03: Update to enable explicit Any cast
Posted by tq...@apache.org.
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;
});
(tvm) 02/03: [FFI] Refactor to make Any require explicit cast to value
Posted by tq...@apache.org.
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 25772b0c97f3f7e4701807aae4123543aaa954a5
Author: tqchen <ti...@gmail.com>
AuthorDate: Wed Apr 30 21:41:08 2025 -0400
[FFI] Refactor to make Any require explicit cast to value
---
ffi/include/tvm/ffi/any.h | 6 +--
ffi/include/tvm/ffi/cast.h | 6 +--
ffi/include/tvm/ffi/container/shape.h | 27 ++++++-------
ffi/include/tvm/ffi/container/variant.h | 4 +-
ffi/include/tvm/ffi/function.h | 7 ++--
ffi/include/tvm/ffi/function_details.h | 7 +++-
ffi/include/tvm/ffi/reflection/reflection.h | 2 +-
ffi/src/ffi/dtype.cc | 4 +-
ffi/src/ffi/ndarray.cc | 12 +++++-
ffi/src/ffi/testing.cc | 2 +-
ffi/tests/cpp/test_any.cc | 52 ++++++++++++-------------
ffi/tests/cpp/test_array.cc | 28 +++++++-------
ffi/tests/cpp/test_dtype.cc | 4 +-
ffi/tests/cpp/test_function.cc | 59 +++++++++++++++--------------
ffi/tests/cpp/test_map.cc | 34 ++++++++---------
ffi/tests/cpp/test_optional.cc | 20 +++++-----
ffi/tests/cpp/test_reflection.cc | 4 +-
ffi/tests/cpp/test_rvalue_ref.cc | 10 ++---
ffi/tests/cpp/test_shape.cc | 4 +-
ffi/tests/cpp/test_tuple.cc | 12 +++---
ffi/tests/cpp/test_variant.cc | 10 ++---
21 files changed, 166 insertions(+), 148 deletions(-)
diff --git a/ffi/include/tvm/ffi/any.h b/ffi/include/tvm/ffi/any.h
index 3f787e7571..48b74fc924 100644
--- a/ffi/include/tvm/ffi/any.h
+++ b/ffi/include/tvm/ffi/any.h
@@ -105,7 +105,7 @@ class AnyView {
}
template <typename T, typename = std::enable_if_t<TypeTraits<T>::convert_enabled>>
- TVM_FFI_INLINE operator T() const {
+ TVM_FFI_INLINE T cast() const {
std::optional<T> opt = TypeTraits<T>::TryConvertFromAnyView(&data_);
if (!opt.has_value()) {
TVM_FFI_THROW(TypeError) << "Cannot convert from type `"
@@ -280,7 +280,7 @@ class Any {
}
template <typename T, typename = std::enable_if_t<TypeTraits<T>::convert_enabled>>
- TVM_FFI_INLINE operator T() const& {
+ TVM_FFI_INLINE T cast() const& {
std::optional<T> opt = TypeTraits<T>::TryConvertFromAnyView(&data_);
if (!opt.has_value()) {
TVM_FFI_THROW(TypeError) << "Cannot convert from type `"
@@ -291,7 +291,7 @@ class Any {
}
template <typename T, typename = std::enable_if_t<TypeTraits<T>::storage_enabled>>
- TVM_FFI_INLINE operator T() && {
+ TVM_FFI_INLINE T cast() && {
if (TypeTraits<T>::CheckAnyStorage(&data_)) {
return TypeTraits<T>::MoveFromAnyStorageAfterCheck(&data_);
}
diff --git a/ffi/include/tvm/ffi/cast.h b/ffi/include/tvm/ffi/cast.h
index c6e89c044d..99069fb13b 100644
--- a/ffi/include/tvm/ffi/cast.h
+++ b/ffi/include/tvm/ffi/cast.h
@@ -115,7 +115,7 @@ inline T Downcast(const Any& ref) {
if constexpr (std::is_same_v<T, Any>) {
return ref;
} else {
- return ref.operator T();
+ return ref.cast<T>();
}
}
@@ -131,7 +131,7 @@ inline T Downcast(Any&& ref) {
if constexpr (std::is_same_v<T, Any>) {
return std::move(ref);
} else {
- return std::move(ref).operator T();
+ return std::move(ref).cast<T>();
}
}
@@ -148,7 +148,7 @@ inline OptionalType Downcast(const std::optional<Any>& ref) {
if constexpr (std::is_same_v<OptionalType, Any>) {
return ref.value();
} else {
- return ref.value().operator OptionalType();
+ return ref.value().cast<OptionalType>();
}
} else {
return OptionalType(std::nullopt);
diff --git a/ffi/include/tvm/ffi/container/shape.h b/ffi/include/tvm/ffi/container/shape.h
index 14a46b76df..6681d1b2df 100644
--- a/ffi/include/tvm/ffi/container/shape.h
+++ b/ffi/include/tvm/ffi/container/shape.h
@@ -68,10 +68,16 @@ class ShapeObjStdImpl : public ShapeObj {
std::vector<int64_t> data_;
};
-TVM_FFI_INLINE ObjectPtr<ShapeObj> MakeEmptyShape() {
- ObjectPtr<ShapeObj> p = make_object<ShapeObj>();
- p->data = nullptr;
- p->size = 0;
+TVM_FFI_INLINE ObjectPtr<ShapeObj> MakeEmptyShape(size_t length, int64_t** mutable_data) {
+ ObjectPtr<ShapeObj> p = make_inplace_array_object<ShapeObj, int64_t>(length);
+ static_assert(alignof(ShapeObj) % alignof(int64_t) == 0);
+ static_assert(sizeof(ShapeObj) % alignof(int64_t) == 0);
+ int64_t* data = reinterpret_cast<int64_t*>(reinterpret_cast<char*>(p.get()) + sizeof(ShapeObj));
+ if (mutable_data) {
+ *mutable_data = data;
+ }
+ p->data = data;
+ p->size = length;
return p;
}
@@ -79,14 +85,9 @@ TVM_FFI_INLINE ObjectPtr<ShapeObj> MakeEmptyShape() {
template <typename IterType>
TVM_FFI_INLINE ObjectPtr<ShapeObj> MakeInplaceShape(IterType begin, IterType end) {
size_t length = std::distance(begin, end);
- ObjectPtr<ShapeObj> p = make_inplace_array_object<ShapeObj, int64_t>(length);
- static_assert(alignof(ShapeObj) % alignof(int64_t) == 0);
- static_assert(sizeof(ShapeObj) % alignof(int64_t) == 0);
- int64_t* dest_data =
- reinterpret_cast<int64_t*>(reinterpret_cast<char*>(p.get()) + sizeof(ShapeObj));
- p->data = dest_data;
- p->size = length;
- std::copy(begin, end, dest_data);
+ int64_t* mutable_data;
+ ObjectPtr<ShapeObj> p = MakeEmptyShape(length, &mutable_data);
+ std::copy(begin, end, mutable_data);
return p;
}
@@ -101,7 +102,7 @@ class Shape : public ObjectRef {
using index_type = ShapeObj::index_type;
/*! \brief Default constructor */
- Shape() : ObjectRef(details::MakeEmptyShape()) {}
+ Shape() : ObjectRef(details::MakeEmptyShape(0, nullptr)) {}
/*!
* \brief Constructor from iterator
diff --git a/ffi/include/tvm/ffi/container/variant.h b/ffi/include/tvm/ffi/container/variant.h
index 4ce044bcfa..6c34a8a15c 100644
--- a/ffi/include/tvm/ffi/container/variant.h
+++ b/ffi/include/tvm/ffi/container/variant.h
@@ -95,12 +95,12 @@ class Variant {
template <typename T, typename = enable_if_variant_contains_t<T>>
TVM_FFI_INLINE T Get() const& {
- return data_.operator T();
+ return data_.cast<T>();
}
template <typename T, typename = enable_if_variant_contains_t<T>>
TVM_FFI_INLINE T Get() && {
- return std::move(data_).operator T();
+ return std::move(data_).cast<T>();
}
TVM_FFI_INLINE std::string GetTypeKey() const { return data_.GetTypeKey(); }
diff --git a/ffi/include/tvm/ffi/function.h b/ffi/include/tvm/ffi/function.h
index 22cc553a61..f6baca5999 100644
--- a/ffi/include/tvm/ffi/function.h
+++ b/ffi/include/tvm/ffi/function.h
@@ -429,11 +429,12 @@ class Function : public ObjectRef {
* \note This function do not depend on Array so core do not have container dep.
*/
static std::vector<String> ListGlobalNames() {
- Function fname_functor = GetGlobalRequired("ffi.FunctionListGlobalNamesFunctor")();
+ Function fname_functor =
+ GetGlobalRequired("ffi.FunctionListGlobalNamesFunctor")().cast<Function>();
std::vector<String> names;
- int len = fname_functor(-1);
+ int len = fname_functor(-1).cast<int>();
for (int i = 0; i < len; ++i) {
- names.push_back(fname_functor(i));
+ names.push_back(fname_functor(i).cast<String>());
}
return names;
}
diff --git a/ffi/include/tvm/ffi/function_details.h b/ffi/include/tvm/ffi/function_details.h
index 0244398efe..19425f4fda 100644
--- a/ffi/include/tvm/ffi/function_details.h
+++ b/ffi/include/tvm/ffi/function_details.h
@@ -236,7 +236,12 @@ template <typename R>
struct typed_packed_call_dispatcher {
template <typename F, typename... Args>
TVM_FFI_INLINE static R run(const F& f, Args&&... args) {
- return f(std::forward<Args>(args)...);
+ Any res = f(std::forward<Args>(args)...);
+ if constexpr (std::is_same_v<R, Any>) {
+ return res;
+ } else {
+ return std::move(res).cast<R>();
+ }
}
};
diff --git a/ffi/include/tvm/ffi/reflection/reflection.h b/ffi/include/tvm/ffi/reflection/reflection.h
index 2d03b56b92..b1088c1eda 100644
--- a/ffi/include/tvm/ffi/reflection/reflection.h
+++ b/ffi/include/tvm/ffi/reflection/reflection.h
@@ -93,7 +93,7 @@ class ReflectionDef {
template <typename T>
static int FieldSetter(void* field, const TVMFFIAny* value) {
TVM_FFI_SAFE_CALL_BEGIN();
- *reinterpret_cast<T*>(field) = AnyView::CopyFromTVMFFIAny(*value);
+ *reinterpret_cast<T*>(field) = AnyView::CopyFromTVMFFIAny(*value).cast<T>();
TVM_FFI_SAFE_CALL_END();
}
diff --git a/ffi/src/ffi/dtype.cc b/ffi/src/ffi/dtype.cc
index a065c70389..cad38b3240 100644
--- a/ffi/src/ffi/dtype.cc
+++ b/ffi/src/ffi/dtype.cc
@@ -29,7 +29,7 @@ namespace details {
*/
inline String DLDataTypeCodeGetCustomTypeName(DLDataTypeCode type_code) {
static Function fget_custom_type_name = Function::GetGlobalRequired("dtype.get_custom_type_name");
- return fget_custom_type_name(static_cast<int>(type_code)).operator String();
+ return fget_custom_type_name(static_cast<int>(type_code)).cast<String>();
}
/*!
@@ -65,7 +65,7 @@ inline int ParseCustomDataTypeCode(const std::string_view& str, const char** sca
auto type_name = str.substr(7, custom_name_len);
TVM_FFI_ICHECK(str.data() == tmp);
static Function fget_custom_type_code = Function::GetGlobalRequired("dtype.get_custom_type_code");
- return fget_custom_type_code(std::string(type_name));
+ return fget_custom_type_code(std::string(type_name)).cast<int>();
}
/*
diff --git a/ffi/src/ffi/ndarray.cc b/ffi/src/ffi/ndarray.cc
index 0b20189532..d4c1470566 100644
--- a/ffi/src/ffi/ndarray.cc
+++ b/ffi/src/ffi/ndarray.cc
@@ -26,9 +26,19 @@
namespace tvm {
namespace ffi {
+
// Shape
TVM_FFI_REGISTER_GLOBAL("ffi.Shape").set_body_packed([](ffi::PackedArgs args, Any* ret) {
- *ret = Shape(args.data(), args.data() + args.size());
+ int64_t* mutable_data;
+ ObjectPtr<ShapeObj> shape = details::MakeEmptyShape(args.size(), &mutable_data);
+ for (int i = 0; i < args.size(); ++i) {
+ if (auto opt_int = args[i].as<int64_t>()) {
+ mutable_data[i] = *opt_int;
+ } else {
+ TVM_FFI_THROW(ValueError) << "Expect shape to take list of int arguments";
+ }
+ }
+ *ret = Shape(shape);
});
} // namespace ffi
} // namespace tvm
diff --git a/ffi/src/ffi/testing.cc b/ffi/src/ffi/testing.cc
index b4dc5b718a..bdca9c970c 100644
--- a/ffi/src/ffi/testing.cc
+++ b/ffi/src/ffi/testing.cc
@@ -43,7 +43,7 @@ TVM_FFI_REGISTER_GLOBAL("testing.echo").set_body_packed([](PackedArgs args, Any*
void TestApply(Function f, PackedArgs args, Any* ret) { f.CallPacked(args, ret); }
TVM_FFI_REGISTER_GLOBAL("testing.apply").set_body_packed([](PackedArgs args, Any* ret) {
- Function f = args[0];
+ auto f = args[0].cast<Function>();
TestApply(f, args.Slice(1), ret);
});
diff --git a/ffi/tests/cpp/test_any.cc b/ffi/tests/cpp/test_any.cc
index d289e9a305..496ae9469d 100644
--- a/ffi/tests/cpp/test_any.cc
+++ b/ffi/tests/cpp/test_any.cc
@@ -37,7 +37,7 @@ TEST(Any, Int) {
EXPECT_THROW(
{
try {
- [[maybe_unused]] int64_t v0 = view0;
+ [[maybe_unused]] auto v0 = view0.cast<int>();
} catch (const Error& error) {
EXPECT_STREQ(error->kind, "TypeError");
std::string what = error.what();
@@ -51,7 +51,7 @@ TEST(Any, Int) {
EXPECT_EQ(view1.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFIInt);
EXPECT_EQ(view1.CopyToTVMFFIAny().v_int64, 1);
- int32_t int_v1 = view1;
+ auto int_v1 = view1.cast<int>();
EXPECT_EQ(int_v1, 1);
int64_t v1 = 2;
@@ -68,7 +68,7 @@ TEST(Any, bool) {
EXPECT_THROW(
{
try {
- [[maybe_unused]] bool v0 = view0;
+ [[maybe_unused]] auto v0 = view0.cast<bool>();
} catch (const Error& error) {
EXPECT_STREQ(error->kind, "TypeError");
std::string what = error.what();
@@ -82,7 +82,7 @@ TEST(Any, bool) {
EXPECT_EQ(view1.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFIBool);
EXPECT_EQ(view1.CopyToTVMFFIAny().v_int64, 1);
- int32_t int_v1 = view1;
+ auto int_v1 = view1.cast<int>();
EXPECT_EQ(int_v1, 1);
bool v1 = false;
@@ -120,7 +120,7 @@ TEST(Any, Float) {
EXPECT_THROW(
{
try {
- [[maybe_unused]] double v0 = view0;
+ [[maybe_unused]] auto v0 = view0.cast<double>();
} catch (const Error& error) {
EXPECT_STREQ(error->kind, "TypeError");
std::string what = error.what();
@@ -131,7 +131,7 @@ TEST(Any, Float) {
::tvm::ffi::Error);
AnyView view1_int = 1;
- float float_v1 = view1_int;
+ auto float_v1 = view1_int.cast<float>();
EXPECT_EQ(float_v1, 1);
AnyView view2 = 2.2;
@@ -154,7 +154,7 @@ TEST(Any, Device) {
EXPECT_THROW(
{
try {
- [[maybe_unused]] DLDevice v0 = view0;
+ [[maybe_unused]] auto v0 = view0.cast<DLDevice>();
} catch (const Error& error) {
EXPECT_STREQ(error->kind, "TypeError");
std::string what = error.what();
@@ -167,7 +167,7 @@ TEST(Any, Device) {
DLDevice device{kDLCUDA, 1};
AnyView view1_device = device;
- DLDevice dtype_v1 = view1_device;
+ auto dtype_v1 = view1_device.cast<DLDevice>();
EXPECT_EQ(dtype_v1.device_type, kDLCUDA);
EXPECT_EQ(dtype_v1.device_id, 1);
@@ -187,7 +187,7 @@ TEST(Any, DLTensor) {
EXPECT_THROW(
{
try {
- [[maybe_unused]] DLTensor* v0 = view0;
+ [[maybe_unused]] auto v0 = view0.cast<DLTensor*>();
} catch (const Error& error) {
EXPECT_STREQ(error->kind, "TypeError");
std::string what = error.what();
@@ -200,7 +200,7 @@ TEST(Any, DLTensor) {
DLTensor dltensor;
AnyView view1_dl = &dltensor;
- DLTensor* dl_v1 = view1_dl;
+ auto dl_v1 = view1_dl.cast<DLTensor*>();
EXPECT_EQ(dl_v1, &dltensor);
}
@@ -225,7 +225,7 @@ TEST(Any, Object) {
EXPECT_EQ(v1.use_count(), 2);
// convert to weak raw object ptr
- const TIntObj* v1_ptr = view2;
+ const TIntObj* v1_ptr = view2.cast<const TIntObj*>();
EXPECT_EQ(v1.use_count(), 2);
EXPECT_EQ(v1_ptr->value, 11);
Any any2 = v1_ptr;
@@ -239,7 +239,7 @@ TEST(Any, Object) {
// convert to ObjectRef
{
- TNumber v1_obj_ref = view2;
+ auto v1_obj_ref = view2.cast<TNumber>();
EXPECT_EQ(v1.use_count(), 3);
any2 = v1_obj_ref;
EXPECT_EQ(v1.use_count(), 4);
@@ -251,7 +251,7 @@ TEST(Any, Object) {
EXPECT_THROW(
{
try {
- [[maybe_unused]] TFloat v0 = view1;
+ [[maybe_unused]] auto v0 = view1.cast<TFloat>();
} catch (const Error& error) {
EXPECT_STREQ(error->kind, "TypeError");
std::string what = error.what();
@@ -263,13 +263,13 @@ TEST(Any, Object) {
},
::tvm::ffi::Error);
// Try to convert to number
- TNumber number0 = any1;
+ auto number0 = any1.cast<TNumber>();
EXPECT_EQ(v1.use_count(), 3);
EXPECT_TRUE(number0.as<TIntObj>());
EXPECT_EQ(number0.as<TIntObj>()->value, 11);
EXPECT_TRUE(!any1.as<int>().has_value());
- TInt int1 = view2;
+ auto int1 = view2.cast<TInt>();
EXPECT_EQ(v1.use_count(), 4);
any1.reset();
EXPECT_EQ(v1.use_count(), 3);
@@ -278,63 +278,63 @@ TEST(Any, Object) {
TEST(Any, ObjectRefWithFallbackTraits) {
// Test case for TPrimExpr fallback from Any
Any any1 = TPrimExpr("float32", 3.14);
- TPrimExpr v0 = any1;
+ auto v0 = any1.cast<TPrimExpr>();
EXPECT_EQ(v0->value, 3.14);
EXPECT_EQ(v0->dtype, "float32");
any1 = true;
- TPrimExpr v1 = any1;
+ auto v1 = any1.cast<TPrimExpr>();
EXPECT_EQ(v1->value, 1);
EXPECT_EQ(v1->dtype, "bool");
any1 = int64_t(42);
- TPrimExpr v2 = any1;
+ auto v2 = any1.cast<TPrimExpr>();
EXPECT_EQ(v2->value, 42);
EXPECT_EQ(v2->dtype, "int64");
any1 = 2.718;
- TPrimExpr v3 = any1;
+ auto v3 = any1.cast<TPrimExpr>();
EXPECT_EQ(v3->value, 2.718);
EXPECT_EQ(v3->dtype, "float32");
// Test case for TPrimExpr fallback from AnyView
TPrimExpr texpr1("float32", 3.14);
AnyView view1 = texpr1;
- TPrimExpr v4 = view1;
+ auto v4 = view1.cast<TPrimExpr>();
EXPECT_EQ(v4->value, 3.14);
EXPECT_EQ(v4->dtype, "float32");
view1 = true;
- TPrimExpr v5 = view1;
+ auto v5 = view1.cast<TPrimExpr>();
EXPECT_EQ(v5->value, 1);
EXPECT_EQ(v5->dtype, "bool");
view1 = int64_t(42);
- TPrimExpr v6 = view1;
+ auto v6 = view1.cast<TPrimExpr>();
EXPECT_EQ(v6->value, 42);
EXPECT_EQ(v6->dtype, "int64");
view1 = 2.718;
- TPrimExpr v7 = view1;
+ auto v7 = view1.cast<TPrimExpr>();
EXPECT_EQ(v7->value, 2.718);
EXPECT_EQ(v7->dtype, "float32");
// Test case for TPrimExpr fallback from Any with String
any1 = std::string("test_string");
- TPrimExpr v8 = any1;
+ auto v8 = any1.cast<TPrimExpr>();
EXPECT_EQ(v8->dtype, "test_string");
EXPECT_EQ(v8->value, 0);
// Test case for TPrimExpr fallback from AnyView with String
view1 = "test_string";
- TPrimExpr v9 = view1;
+ auto v9 = view1.cast<TPrimExpr>();
EXPECT_EQ(v9->dtype, "test_string");
EXPECT_EQ(v9->value, 0);
}
TEST(Any, ObjectMove) {
Any any1 = TPrimExpr("float32", 3.14);
- TPrimExpr v0 = std::move(any1);
+ auto v0 = std::move(any1).cast<TPrimExpr>();
EXPECT_EQ(v0->value, 3.14);
EXPECT_EQ(v0.use_count(), 1);
}
diff --git a/ffi/tests/cpp/test_array.cc b/ffi/tests/cpp/test_array.cc
index d782267129..15ed0bff06 100644
--- a/ffi/tests/cpp/test_array.cc
+++ b/ffi/tests/cpp/test_array.cc
@@ -185,8 +185,8 @@ TEST(Array, InsertEraseRange) {
TEST(Array, FuncArrayAnyArg) {
Function fadd_one =
- Function::FromUnpacked([](Array<Any> a) -> Any { return a[0].operator int() + 1; });
- EXPECT_EQ(fadd_one(Array<Any>{1}).operator int(), 2);
+ Function::FromUnpacked([](Array<Any> a) -> Any { return a[0].cast<int>() + 1; });
+ EXPECT_EQ(fadd_one(Array<Any>{1}).cast<int>(), 2);
}
TEST(Array, MapUniquePropogation) {
@@ -200,10 +200,10 @@ TEST(Array, MapUniquePropogation) {
TEST(Array, AnyImplicitConversion) {
Array<Any> arr0_mixed = {11.1, 1};
- EXPECT_EQ(arr0_mixed[1].operator int(), 1);
+ EXPECT_EQ(arr0_mixed[1].cast<int>(), 1);
AnyView view0 = arr0_mixed;
- Array<double> arr0_float = view0;
+ auto arr0_float = view0.cast<Array<double>>();
// they are not the same because arr_mixed
// stores arr_mixed[1] as int but we need to convert to float
EXPECT_TRUE(!arr0_float.same_as(arr0_mixed));
@@ -211,23 +211,23 @@ TEST(Array, AnyImplicitConversion) {
Any any1 = arr0_float;
// if storage check passes, the same array get returned
- Array<double> arr1_float = any1;
+ auto arr1_float = any1.cast<Array<double>>();
EXPECT_TRUE(arr1_float.same_as(arr0_float));
// total count equals 3 include any1
EXPECT_EQ(arr1_float.use_count(), 3);
// convert to Array<Any> do not need any conversion
- Array<Any> arr1_mixed = any1;
+ auto arr1_mixed = any1.cast<Array<Any>>();
EXPECT_TRUE(arr1_mixed.same_as(arr1_float));
EXPECT_EQ(arr1_float.use_count(), 4);
}
TEST(Array, AnyConvertCheck) {
Array<Any> arr = {11.1, 1};
- EXPECT_EQ(arr[1].operator int(), 1);
+ EXPECT_EQ(arr[1].cast<int>(), 1);
AnyView view0 = arr;
- Array<double> arr1 = view0;
+ auto arr1 = view0.cast<Array<double>>();
EXPECT_EQ(arr1[0], 11.1);
EXPECT_EQ(arr1[1], 1.0);
@@ -236,7 +236,7 @@ TEST(Array, AnyConvertCheck) {
EXPECT_THROW(
{
try {
- [[maybe_unused]] Array<int> arr2 = any1;
+ [[maybe_unused]] auto arr2 = any1.cast<Array<int>>();
} catch (const Error& error) {
EXPECT_STREQ(error->kind, "TypeError");
std::string what = error.what();
@@ -249,13 +249,13 @@ TEST(Array, AnyConvertCheck) {
Array<Array<TNumber>> arr_nested = {{}, {TInt(1), TFloat(2)}};
any1 = arr_nested;
- Array<Array<TNumber>> arr1_nested = any1;
+ auto arr1_nested = any1.cast<Array<Array<TNumber>>>();
EXPECT_EQ(arr1_nested.use_count(), 3);
EXPECT_THROW(
{
try {
- [[maybe_unused]] Array<Array<int>> arr2 = any1;
+ [[maybe_unused]] auto arr2 = any1.cast<Array<Array<int>>>();
} catch (const Error& error) {
EXPECT_STREQ(error->kind, "TypeError");
std::string what = error.what();
@@ -270,9 +270,9 @@ TEST(Array, AnyConvertCheck) {
TEST(Array, Upcast) {
Array<int> a0 = {1, 2, 3};
Array<Any> a1 = a0;
- EXPECT_EQ(a1[0].operator int(), 1);
- EXPECT_EQ(a1[1].operator int(), 2);
- EXPECT_EQ(a1[2].operator int(), 3);
+ EXPECT_EQ(a1[0].cast<int>(), 1);
+ EXPECT_EQ(a1[1].cast<int>(), 2);
+ EXPECT_EQ(a1[2].cast<int>(), 3);
Array<Array<int>> a2 = {a0};
Array<Array<Any>> a3 = a2;
diff --git a/ffi/tests/cpp/test_dtype.cc b/ffi/tests/cpp/test_dtype.cc
index 6373f5a4ae..5dce997a9a 100644
--- a/ffi/tests/cpp/test_dtype.cc
+++ b/ffi/tests/cpp/test_dtype.cc
@@ -84,7 +84,7 @@ TEST(DataType, AnyConversion) {
EXPECT_THROW(
{
try {
- [[maybe_unused]] DLDataType v0 = view0;
+ [[maybe_unused]] auto v0 = view0.cast<DLDataType>();
} catch (const Error& error) {
EXPECT_STREQ(error->kind, "TypeError");
std::string what = error.what();
@@ -97,7 +97,7 @@ TEST(DataType, AnyConversion) {
DLDataType dtype{kDLFloat, 32, 1};
AnyView view1_dtype = dtype;
- DLDataType dtype_v1 = view1_dtype;
+ auto dtype_v1 = view1_dtype.cast<DLDataType>();
EXPECT_EQ(dtype_v1.code, kDLFloat);
EXPECT_EQ(dtype_v1.bits, 32);
EXPECT_EQ(dtype_v1.lanes, 1);
diff --git a/ffi/tests/cpp/test_function.cc b/ffi/tests/cpp/test_function.cc
index 9c793e2493..d8895da934 100644
--- a/ffi/tests/cpp/test_function.cc
+++ b/ffi/tests/cpp/test_function.cc
@@ -32,50 +32,50 @@ using namespace tvm::ffi::testing;
TEST(Func, FromPacked) {
Function fadd1 = Function::FromPacked([](const AnyView* args, int32_t num_args, Any* rv) {
EXPECT_EQ(num_args, 1);
- int32_t a = args[0];
+ int32_t a = args[0].cast<int32_t>();
*rv = a + 1;
});
- int b = fadd1(1);
+ int b = fadd1(1).cast<int>();
EXPECT_EQ(b, 2);
Function fadd2 = Function::FromPacked([](const AnyView* args, int32_t num_args, Any* rv) {
EXPECT_EQ(num_args, 1);
- TInt a = args[0];
+ auto a = args[0].cast<TInt>();
EXPECT_EQ(a.use_count(), 2);
*rv = a->value + 1;
});
- EXPECT_EQ(fadd2(TInt(12)).operator int(), 13);
+ EXPECT_EQ(fadd2(TInt(12)).cast<int>(), 13);
}
TEST(Func, PackedArgs) {
Function fadd1 = Function::FromPacked([](PackedArgs args, Any* rv) {
EXPECT_EQ(args.size(), 1);
- int32_t a = args[0];
+ int32_t a = args[0].cast<int>();
*rv = a + 1;
});
- int b = fadd1(1);
+ int b = fadd1(1).cast<int>();
EXPECT_EQ(b, 2);
Function fadd2 = Function::FromPacked([](PackedArgs args, Any* rv) {
EXPECT_EQ(args.size(), 1);
- TInt a = args[0];
+ TInt a = args[0].cast<TInt>();
EXPECT_EQ(a.use_count(), 2);
*rv = a->value + 1;
});
- EXPECT_EQ(fadd2(TInt(12)).operator int(), 13);
+ EXPECT_EQ(fadd2(TInt(12)).cast<int>(), 13);
TInt v(12);
AnyView data[3];
PackedArgs::Fill(data, 3, 1, v);
- EXPECT_EQ(data[0].operator int(), 3);
- EXPECT_EQ(data[1].operator int(), 1);
- EXPECT_EQ(data[2].operator TInt()->value, 12);
+ EXPECT_EQ(data[0].cast<int>(), 3);
+ EXPECT_EQ(data[1].cast<int>(), 1);
+ EXPECT_EQ(data[2].cast<TInt>()->value, 12);
}
TEST(Func, FromUnpacked) {
// try decution
Function fadd1 = Function::FromUnpacked([](const int32_t& a) -> int { return a + 1; });
- int b = fadd1(1);
+ int b = fadd1(1).cast<int>();
EXPECT_EQ(b, 2);
// convert that triggers error
@@ -120,8 +120,8 @@ TEST(Func, FromUnpacked) {
},
"fpass_and_return");
TInt a(11);
- Function fret = fpass_and_return(std::move(a), 11, 11);
- EXPECT_EQ(fret(12).operator int(), 23);
+ auto fret = fpass_and_return(std::move(a), 11, 11).cast<Function>();
+ EXPECT_EQ(fret(12).cast<int>(), 23);
EXPECT_THROW(
{
@@ -140,28 +140,29 @@ TEST(Func, FromUnpacked) {
Function fconcact =
Function::FromUnpacked([](const String& a, const String& b) -> String { return a + b; });
- EXPECT_EQ(fconcact("abc", "def").operator String(), "abcdef");
+ EXPECT_EQ(fconcact("abc", "def").cast<String>(), "abcdef");
}
TEST(Func, PassReturnAny) {
- Function fadd_one = Function::FromUnpacked([](Any a) -> Any { return a.operator int() + 1; });
- EXPECT_EQ(fadd_one(1).operator int(), 2);
+ Function fadd_one = Function::FromUnpacked([](Any a) -> Any { return a.cast<int>() + 1; });
+ EXPECT_EQ(fadd_one(1).cast<int>(), 2);
}
TEST(Func, Global) {
Function::SetGlobal("testing.add1",
Function::FromUnpacked([](const int32_t& a) -> int { return a + 1; }));
auto fadd1 = Function::GetGlobalRequired("testing.add1");
- int b = fadd1(1);
+ int b = fadd1(1).cast<int>();
EXPECT_EQ(b, 2);
auto fnot_exist = Function::GetGlobal("testing.not_existing_func");
EXPECT_TRUE(!fnot_exist);
- Function fname_functor = Function::GetGlobal("ffi.FunctionListGlobalNamesFunctor").value()();
+ auto fname_functor =
+ Function::GetGlobal("ffi.FunctionListGlobalNamesFunctor").value()().cast<Function>();
Array<String> names;
- int len = fname_functor(-1);
+ int len = fname_functor(-1).cast<int>();
for (int i = 0; i < len; ++i) {
- names.push_back(fname_functor(i));
+ names.push_back(fname_functor(i).cast<String>());
}
EXPECT_TRUE(std::find(names.begin(), names.end(), "testing.add1") != names.end());
}
@@ -172,7 +173,7 @@ TEST(Func, TypedFunction) {
TypedFunction<int(int)> fadd2([](int a) -> int { return a + 2; });
EXPECT_EQ(fadd2(1), 3);
- EXPECT_EQ(fadd2.packed()(1).operator int(), 3);
+ EXPECT_EQ(fadd2.packed()(1).cast<int>(), 3);
TypedFunction<void(int)> fcheck_int;
EXPECT_TRUE(fcheck_int == nullptr);
@@ -184,14 +185,14 @@ TEST(Func, TypedFunctionAsAny) {
TypedFunction<int(int)> fadd1 = [](int a) -> int { return a + 1; };
Any fany(std::move(fadd1));
EXPECT_TRUE(fadd1 == nullptr);
- TypedFunction<int(int)> fadd1_dup = fany;
+ auto fadd1_dup = fany.cast<TypedFunction<int(int)>>();
EXPECT_EQ(fadd1_dup(1), 2);
}
TEST(Func, TypedFunctionAsAnyView) {
TypedFunction<int(int)> fadd2 = [](int a) -> int { return a + 2; };
AnyView fview(fadd2);
- TypedFunction<int(int)> fadd2_dup = fview;
+ auto fadd2_dup = fview.cast<TypedFunction<int(int)>>();
EXPECT_EQ(fadd2_dup(1), 3);
}
@@ -200,22 +201,22 @@ TEST(Func, ObjectRefWithFallbackTraits) {
// through TPrimExpr
Function freturn_primexpr = Function::FromUnpacked([](TPrimExpr a) -> TPrimExpr { return a; });
- TPrimExpr result_int = freturn_primexpr(1);
+ auto result_int = freturn_primexpr(1).cast<TPrimExpr>();
EXPECT_EQ(result_int->dtype, "int64");
EXPECT_EQ(result_int->value, 1);
// Test case for float
- TPrimExpr result_float = freturn_primexpr(2.5);
+ auto result_float = freturn_primexpr(2.5).cast<TPrimExpr>();
EXPECT_EQ(result_float->dtype, "float32");
EXPECT_EQ(result_float->value, 2.5);
// Test case for bool
- TPrimExpr result_bool = freturn_primexpr(true);
+ auto result_bool = freturn_primexpr(true).cast<TPrimExpr>();
EXPECT_EQ(result_bool->dtype, "bool");
EXPECT_EQ(result_bool->value, 1);
// Test case for string
- TPrimExpr result_string = freturn_primexpr("test_string");
+ auto result_string = freturn_primexpr("test_string").cast<TPrimExpr>();
EXPECT_EQ(result_string->dtype, "test_string");
EXPECT_EQ(result_string->value, 0);
@@ -240,6 +241,6 @@ TVM_FFI_REGISTER_GLOBAL("testing.Int_GetValue").set_body_method(&TIntObj::GetVal
TEST(Func, Register) {
Function fget_value = Function::GetGlobalRequired("testing.Int_GetValue");
TInt a(12);
- EXPECT_EQ(fget_value(a).operator int(), 12);
+ EXPECT_EQ(fget_value(a).cast<int>(), 12);
}
} // namespace
diff --git a/ffi/tests/cpp/test_map.cc b/ffi/tests/cpp/test_map.cc
index 58d81143cc..c579f47442 100644
--- a/ffi/tests/cpp/test_map.cc
+++ b/ffi/tests/cpp/test_map.cc
@@ -54,7 +54,7 @@ TEST(Map, PODKey) {
auto it = map0.find(1.1);
EXPECT_TRUE(it != map0.end());
- EXPECT_EQ((*it).second.operator int(), 3);
+ EXPECT_EQ((*it).second.cast<int>(), 3);
}
TEST(Map, Object) {
@@ -160,39 +160,39 @@ TEST(Map, AnyImplicitConversion) {
// check will trigger copy
AnyView view0 = map0;
- Map<int, double> map1 = view0;
+ auto map1 = view0.cast<Map<int, double>>();
EXPECT_TRUE(!map1.same_as(map0));
EXPECT_EQ(map1[1], 2);
EXPECT_EQ(map1[2], 3.1);
EXPECT_EQ(map1.use_count(), 1);
- Map<int, Any> map2 = view0;
+ auto map2 = view0.cast<Map<int, Any>>();
EXPECT_TRUE(map2.same_as(map0));
EXPECT_EQ(map2.use_count(), 2);
- Map<Any, double> map3 = view0;
+ auto map3 = view0.cast<Map<Any, double>>();
EXPECT_TRUE(!map3.same_as(map0));
EXPECT_EQ(map3.use_count(), 1);
Map<Any, Any> map4{{"yes", 1.1}, {"no", 2.2}};
Any any1 = map4;
- Map<String, double> map5 = any1;
+ auto map5 = any1.cast<Map<String, double>>();
EXPECT_TRUE(map5.same_as(map4));
EXPECT_EQ(map5.use_count(), 3);
- Map<String, Any> map6 = any1;
+ auto map6 = any1.cast<Map<String, Any>>();
EXPECT_TRUE(map6.same_as(map4));
EXPECT_EQ(map6.use_count(), 4);
- EXPECT_EQ(map6["yes"].operator double(), 1.1);
- EXPECT_EQ(map6["no"].operator double(), 2.2);
+ EXPECT_EQ(map6["yes"].cast<double>(), 1.1);
+ EXPECT_EQ(map6["no"].cast<double>(), 2.2);
- Map<Any, Any> map7 = any1;
+ auto map7 = any1.cast<Map<Any, Any>>();
EXPECT_TRUE(map7.same_as(map4));
EXPECT_EQ(map7.use_count(), 5);
- Map<Any, TPrimExpr> map8 = any1;
+ auto map8 = any1.cast<Map<Any, TPrimExpr>>();
EXPECT_TRUE(!map8.same_as(map4));
EXPECT_EQ(map8.use_count(), 1);
EXPECT_EQ(map8["yes"]->value, 1.1);
@@ -201,10 +201,10 @@ TEST(Map, AnyImplicitConversion) {
TEST(Map, AnyConvertCheck) {
Map<Any, Any> map = {{11, 1.1}};
- EXPECT_EQ(map[11].operator double(), 1.1);
+ EXPECT_EQ(map[11].cast<double>(), 1.1);
AnyView view0 = map;
- Map<int64_t, double> arr1 = view0;
+ auto arr1 = view0.cast<Map<int64_t, double>>();
EXPECT_EQ(arr1[11], 1.1);
Any any1 = map;
@@ -213,7 +213,7 @@ TEST(Map, AnyConvertCheck) {
EXPECT_THROW(
{
try {
- [[maybe_unused]] WrongMap arr2 = any1;
+ [[maybe_unused]] auto arr2 = any1.cast<WrongMap>();
} catch (const Error& error) {
EXPECT_STREQ(error->kind, "TypeError");
std::string what = error.what();
@@ -230,7 +230,7 @@ TEST(Map, AnyConvertCheck) {
EXPECT_THROW(
{
try {
- [[maybe_unused]] WrongMap2 arr2 = any1;
+ [[maybe_unused]] auto arr2 = any1.cast<WrongMap2>();
} catch (const Error& error) {
EXPECT_STREQ(error->kind, "TypeError");
std::string what = error.what();
@@ -249,14 +249,14 @@ TEST(Map, PackedFuncGetItem) {
Map<String, int64_t> map{{"x", 1}, {"y", 2}};
Any k("x");
Any v = f(map, k);
- EXPECT_EQ(v.operator int(), 1);
+ EXPECT_EQ(v.cast<int>(), 1);
}
TEST(Map, Upcast) {
Map<int, int> m0 = {{1, 2}, {3, 4}};
Map<Any, Any> m1 = m0;
- EXPECT_EQ(m1[1].operator int(), 2);
- EXPECT_EQ(m1[3].operator int(), 4);
+ EXPECT_EQ(m1[1].cast<int>(), 2);
+ EXPECT_EQ(m1[3].cast<int>(), 4);
static_assert(details::type_contains_v<Map<Any, Any>, Map<String, int>>);
Map<String, Array<int>> m2 = {{"x", {1}}, {"y", {2}}};
diff --git a/ffi/tests/cpp/test_optional.cc b/ffi/tests/cpp/test_optional.cc
index 65e4042f8d..a45a391904 100644
--- a/ffi/tests/cpp/test_optional.cc
+++ b/ffi/tests/cpp/test_optional.cc
@@ -42,11 +42,11 @@ TEST(Optional, TInt) {
Any z_any = std::move(y);
EXPECT_TRUE(z_any != nullptr);
- EXPECT_EQ((z_any.operator TInt())->value, 11);
+ EXPECT_EQ((z_any.cast<TInt>())->value, 11);
EXPECT_TRUE(!y.has_value());
// move from any to optional
- Optional<TInt> y2 = std::move(z_any);
+ auto y2 = std::move(z_any).cast<Optional<TInt>>();
EXPECT_EQ(y2.use_count(), 1);
EXPECT_TRUE(y2.has_value());
EXPECT_EQ(y2.value_or(TInt(12))->value, 11);
@@ -73,14 +73,14 @@ TEST(Optional, AnyConvert_int) {
EXPECT_TRUE(opt_v0.has_value());
AnyView view0 = opt_v0;
- EXPECT_EQ(view0.operator int(), 1);
+ EXPECT_EQ(view0.cast<int>(), 1);
Any any1;
- Optional<int> opt_v1 = std::move(any1);
+ auto opt_v1 = std::move(any1).cast<Optional<int>>();
EXPECT_TRUE(!opt_v1.has_value());
Optional<int> opt_v2 = 11;
Any any2 = std::move(opt_v2);
- EXPECT_EQ(any2.operator int(), 11);
+ EXPECT_EQ(any2.cast<int>(), 11);
}
TEST(Optional, AnyConvert_Array) {
@@ -88,21 +88,21 @@ TEST(Optional, AnyConvert_Array) {
Array<Array<TNumber>> arr_nested = {{}, {TInt(1), TFloat(2)}};
view0 = arr_nested;
- Optional<Array<Array<TNumber>>> opt_arr = view0;
+ auto opt_arr = view0.cast<Optional<Array<Array<TNumber>>>>();
EXPECT_EQ(arr_nested.use_count(), 2);
- Optional<Array<Array<TNumber>>> arr1 = view0;
+ auto arr1 = view0.cast<Optional<Array<Array<TNumber>>>>();
EXPECT_EQ(arr_nested.use_count(), 3);
EXPECT_EQ(arr1.value()[1][1].as<TFloatObj>()->value, 2);
Any any1;
- Optional<Array<Array<TNumber>>> arr2 = any1;
+ auto arr2 = any1.cast<Optional<Array<Array<TNumber>>>>();
EXPECT_TRUE(!arr2.has_value());
EXPECT_THROW(
{
try {
- [[maybe_unused]] Optional<Array<Array<int>>> arr2 = view0;
+ [[maybe_unused]] auto arr2 = view0.cast<Optional<Array<Array<int>>>>();
} catch (const Error& error) {
EXPECT_STREQ(error->kind, "TypeError");
std::string what = error.what();
@@ -167,7 +167,7 @@ TEST(Optional, OptionalInArray) {
}
Any any = arr;
- Array<Optional<Array<TInt>>> opt_arr = any;
+ auto opt_arr = any.cast<Array<Optional<Array<TInt>>>>();
EXPECT_EQ(opt_arr[0].value()[0]->value, 0);
}
} // namespace
diff --git a/ffi/tests/cpp/test_reflection.cc b/ffi/tests/cpp/test_reflection.cc
index 56ef77525e..fec167d257 100644
--- a/ffi/tests/cpp/test_reflection.cc
+++ b/ffi/tests/cpp/test_reflection.cc
@@ -42,11 +42,11 @@ TEST(Reflection, GetFieldByteOffset) {
TEST(Reflection, FieldGetter) {
ObjectRef a = TInt(10);
details::ReflectionFieldGetter getter(details::GetReflectionFieldInfo("test.Int", "value"));
- EXPECT_EQ(getter(a).operator int(), 10);
+ EXPECT_EQ(getter(a).cast<int>(), 10);
ObjectRef b = TFloat(10.0);
details::ReflectionFieldGetter getter_float(
details::GetReflectionFieldInfo("test.Float", "value"));
- EXPECT_EQ(getter_float(b).operator double(), 10.0);
+ EXPECT_EQ(getter_float(b).cast<double>(), 10.0);
}
} // namespace
diff --git a/ffi/tests/cpp/test_rvalue_ref.cc b/ffi/tests/cpp/test_rvalue_ref.cc
index a16ef1ac7e..65d86d764a 100644
--- a/ffi/tests/cpp/test_rvalue_ref.cc
+++ b/ffi/tests/cpp/test_rvalue_ref.cc
@@ -37,12 +37,12 @@ TEST(RValueRef, Basic) {
arr.push_back(val);
return arr;
});
- Array<int> a = append(RValueRef(Array<int>({1, 2})), 3, true);
+ auto a = append(RValueRef(Array<int>({1, 2})), 3, true).cast<Array<int>>();
EXPECT_EQ(a.size(), 3);
- a = append(RValueRef(std::move(a)), 4, true);
+ a = append(RValueRef(std::move(a)), 4, true).cast<Array<int>>();
EXPECT_EQ(a.size(), 4);
// pass in lvalue instead, the append still will succeed but array will not be unique
- a = append(a, 5, false);
+ a = append(a, 5, false).cast<Array<int>>();
EXPECT_EQ(a.size(), 5);
}
@@ -90,8 +90,8 @@ TEST(RValueRef, ParamChecking) {
TPrimExpr expr = *std::move(a);
return expr->dtype;
});
- EXPECT_EQ(func3(RValueRef(String("int32"))).operator String(), "int32");
+ EXPECT_EQ(func3(RValueRef(String("int32"))).cast<String>(), "int32");
// triggered a lvalue based conversion
- EXPECT_EQ(func3(String("int32")).operator String(), "int32");
+ EXPECT_EQ(func3(String("int32")).cast<String>(), "int32");
}
} // namespace
diff --git a/ffi/tests/cpp/test_shape.cc b/ffi/tests/cpp/test_shape.cc
index c6d8d5dbd8..0ccba7820a 100644
--- a/ffi/tests/cpp/test_shape.cc
+++ b/ffi/tests/cpp/test_shape.cc
@@ -55,7 +55,7 @@ TEST(Shape, AnyConvert) {
Shape shape0 = Shape({1, 2, 3});
Any any0 = shape0;
- Shape shape1 = any0;
+ auto shape1 = any0.cast<Shape>();
EXPECT_EQ(shape1.size(), 3);
EXPECT_EQ(shape1[0], 1);
EXPECT_EQ(shape1[1], 2);
@@ -63,7 +63,7 @@ TEST(Shape, AnyConvert) {
Array<Any> arr({1, 2});
AnyView any_view0 = arr;
- Shape shape2 = any_view0;
+ auto shape2 = any_view0.cast<Shape>();
EXPECT_EQ(shape2.size(), 2);
EXPECT_EQ(shape2[0], 1);
EXPECT_EQ(shape2[1], 2);
diff --git a/ffi/tests/cpp/test_tuple.cc b/ffi/tests/cpp/test_tuple.cc
index 5512db80e6..9d24e60056 100644
--- a/ffi/tests/cpp/test_tuple.cc
+++ b/ffi/tests/cpp/test_tuple.cc
@@ -73,12 +73,12 @@ TEST(Tuple, AnyConvert) {
EXPECT_EQ(arr0[1].as<TInt>().value()->value, 2);
// directly reuse the underlying storage.
- Tuple<int, TInt> tuple1 = view0;
+ auto tuple1 = view0.cast<Tuple<int, TInt>>();
EXPECT_TRUE(tuple0.same_as(tuple1));
Any any0 = view0;
// trigger a copy due to implict conversion
- Tuple<TPrimExpr, TInt> tuple2 = any0;
+ auto tuple2 = any0.cast<Tuple<TPrimExpr, TInt>>();
EXPECT_TRUE(!tuple0.same_as(tuple2));
EXPECT_EQ(tuple2.Get<0>()->value, 1);
EXPECT_EQ(tuple2.Get<1>()->value, 2);
@@ -89,10 +89,10 @@ TEST(Tuple, FromUnpacked) {
Function fadd1 = Function::FromUnpacked([](const Tuple<int, TPrimExpr>& a) -> int {
return a.Get<0>() + static_cast<int>(a.Get<1>()->value);
});
- int b = fadd1(Tuple<int, float>(1, 2));
+ int b = fadd1(Tuple<int, float>(1, 2)).cast<int>();
EXPECT_EQ(b, 3);
- int c = fadd1(Array<Any>({1, 2}));
+ int c = fadd1(Array<Any>({1, 2})).cast<int>();
EXPECT_EQ(c, 3);
// convert that triggers error
@@ -130,8 +130,8 @@ TEST(Tuple, FromUnpacked) {
TEST(Tuple, Upcast) {
Tuple<int, float> t0(1, 2.0f);
Tuple<Any, Any> t1 = t0;
- EXPECT_EQ(t1.Get<0>().operator int(), 1);
- EXPECT_EQ(t1.Get<1>().operator float(), 2.0f);
+ EXPECT_EQ(t1.Get<0>().cast<int>(), 1);
+ EXPECT_EQ(t1.Get<1>().cast<float>(), 2.0f);
static_assert(details::type_contains_v<Tuple<Any, Any>, Tuple<int, float>>);
static_assert(details::type_contains_v<Tuple<Any, float>, Tuple<int, float>>);
static_assert(details::type_contains_v<Tuple<TNumber, float>, Tuple<TInt, float>>);
diff --git a/ffi/tests/cpp/test_variant.cc b/ffi/tests/cpp/test_variant.cc
index 107ca55c89..77e4b8eee1 100644
--- a/ffi/tests/cpp/test_variant.cc
+++ b/ffi/tests/cpp/test_variant.cc
@@ -47,14 +47,14 @@ TEST(Variant, AnyConvert) {
// implicit convert to variant
Any any0 = 1;
- Variant<TPrimExpr, Array<TPrimExpr>> v1 = any0;
+ auto v1 = any0.cast<Variant<TPrimExpr, Array<TPrimExpr>>>();
EXPECT_EQ(v1.Get<TPrimExpr>()->value, 1);
// move from any to variant
Variant<TInt, int> v2 = TInt(1);
Any any1 = std::move(v2);
- Variant<TInt, int> v3 = std::move(any1);
- TInt v4 = std::move(v3).Get<TInt>();
+ auto v3 = std::move(any1).cast<Variant<TInt, int>>();
+ auto v4 = std::move(v3).Get<TInt>();
EXPECT_EQ(v4->value, 1);
EXPECT_EQ(v4.use_count(), 1);
}
@@ -81,7 +81,7 @@ TEST(Variant, FromUnpacked) {
return a.Get<TInt>()->value + 1;
}
});
- int b = fadd1(1);
+ int b = fadd1(1).cast<int>();
EXPECT_EQ(b, 2);
// convert that triggers error
@@ -107,7 +107,7 @@ TEST(Variant, FromUnpacked) {
return a[0].Get<TInt>()->value + 1;
}
});
- int c = fadd2(Array<Any>({1, 2}));
+ int c = fadd2(Array<Any>({1, 2})).cast<int>();
EXPECT_EQ(c, 2);
// convert that triggers error