This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 1d735e19c0 [REFACTOR] Clean up tensor helpers and parent library
visibility (#20770)
1d735e19c0 is described below
commit 1d735e19c0dbf56c3c6ced7914d3e90d2c922a46
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Oct 16 22:31:54 2026 -0400
[REFACTOR] Clean up tensor helpers and parent library visibility (#20770)
Use canonical ir.Var construction for tensor variables and remove unused
Python replacement and FFI helpers while retaining operator traversal
coverage.
Move compiler and runtime-extra visibility policies into parent-owned
headers with explicit consumer includes, keeping standalone TIRx
independent. Generated C host modules use their own FFI export alias.
---
include/tvm/base.h | 51 ++++++++++++
include/tvm/relax/analysis.h | 1 +
include/tvm/relax/backend.h | 1 +
include/tvm/relax/backend/adreno/transform.h | 1 +
include/tvm/relax/binding_rewrite.h | 1 +
include/tvm/relax/block_builder.h | 1 +
include/tvm/relax/dataflow_matcher.h | 1 +
include/tvm/relax/dataflow_pattern.h | 1 +
include/tvm/relax/distributed/global_info.h | 1 +
include/tvm/relax/distributed/transform.h | 1 +
include/tvm/relax/distributed/type.h | 1 +
include/tvm/relax/exec_builder.h | 1 +
include/tvm/relax/expr.h | 1 +
include/tvm/relax/global_info.h | 1 +
include/tvm/relax/script/ir_builder/ir.h | 1 +
include/tvm/relax/tir_pattern.h | 1 +
include/tvm/relax/transform.h | 1 +
include/tvm/relax/type.h | 1 +
include/tvm/relax/type_functor.h | 1 +
include/tvm/relax/utils.h | 1 +
include/tvm/runtime/disco/builtin.h | 1 +
include/tvm/runtime/disco/cuda_ipc_memory.h | 1 +
include/tvm/runtime/disco/disco_worker.h | 1 +
include/tvm/runtime/disco/session.h | 1 +
include/tvm/runtime/extra_base.h | 49 ++++++++++++
include/tvm/runtime/timing.h | 1 +
include/tvm/runtime/vm/executable.h | 1 +
include/tvm/runtime/vm/tensor_cache_support.h | 1 +
include/tvm/s_tir/analysis.h | 1 +
include/tvm/s_tir/backend/adreno/transform.h | 1 +
include/tvm/s_tir/data_layout.h | 1 +
include/tvm/s_tir/iter_var.h | 1 +
include/tvm/s_tir/meta_schedule/arg_info.h | 1 +
include/tvm/s_tir/meta_schedule/builder.h | 1 +
include/tvm/s_tir/meta_schedule/cost_model.h | 1 +
include/tvm/s_tir/meta_schedule/database.h | 1 +
.../tvm/s_tir/meta_schedule/feature_extractor.h | 1 +
include/tvm/s_tir/meta_schedule/measure_callback.h | 1 +
.../tvm/s_tir/meta_schedule/measure_candidate.h | 1 +
include/tvm/s_tir/meta_schedule/mutator.h | 1 +
include/tvm/s_tir/meta_schedule/postproc.h | 1 +
include/tvm/s_tir/meta_schedule/runner.h | 1 +
include/tvm/s_tir/meta_schedule/schedule_rule.h | 1 +
include/tvm/s_tir/meta_schedule/search_strategy.h | 1 +
include/tvm/s_tir/meta_schedule/space_generator.h | 1 +
include/tvm/s_tir/meta_schedule/task_scheduler.h | 1 +
include/tvm/s_tir/meta_schedule/tune_context.h | 1 +
include/tvm/s_tir/sblock_dependence_info.h | 1 +
include/tvm/s_tir/sblock_scope.h | 1 +
include/tvm/s_tir/schedule/schedule.h | 1 +
include/tvm/s_tir/schedule/state.h | 1 +
include/tvm/s_tir/stmt.h | 1 +
include/tvm/s_tir/stmt_functor.h | 1 +
include/tvm/s_tir/stmt_functor_with_analyzer.h | 1 +
include/tvm/s_tir/tensor_intrin.h | 1 +
include/tvm/s_tir/transform.h | 1 +
include/tvm/te/operation.h | 1 +
include/tvm/te/reduction.h | 1 +
include/tvm/te/tensor.h | 1 +
src/relax/ir/py_expr_functor.cc | 1 +
src/relax/ir/transform.cc | 1 +
src/relax/training/utils.h | 1 +
src/relax/transform/dataflow_inplace.cc | 1 +
src/relax/transform/utils.h | 1 +
src/runtime/extra/contrib/dnnl/dnnl_kernel.h | 1 +
src/runtime/extra/disco/builtin.cc | 1 +
src/runtime/extra/disco/disco_worker.cc | 1 +
src/runtime/rpc/rpc_socket_impl.cc | 1 +
src/runtime/vm/tensor_cache_support.cc | 1 +
src/s_tir/schedule/primitive.h | 1 +
src/s_tir/schedule/primitive/sampling.cc | 2 +
src/s_tir/schedule/state.cc | 1 +
src/s_tir/support/parallel_for.h | 1 +
tirx/include/tvm/runtime/base.h | 56 +------------
tirx/python/tvm_tirx/tirx/tensor.py | 2 +-
tirx/python/tvm_tirx/tirx/transform/common.py | 91 +---------------------
tirx/src/codegen/c/codegen_c_host.cc | 2 +-
tirx/src/tirx/ir/tensor.cc | 7 +-
tirx/tests/python/tirx/test_op.py | 20 +----
79 files changed, 182 insertions(+), 168 deletions(-)
diff --git a/include/tvm/base.h b/include/tvm/base.h
new file mode 100644
index 0000000000..feecde856b
--- /dev/null
+++ b/include/tvm/base.h
@@ -0,0 +1,51 @@
+/*
+ * 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 tvm/base.h
+ * \brief Parent compiler symbol visibility.
+ */
+#ifndef TVM_BASE_H_
+#define TVM_BASE_H_
+
+// TVM_EXPORTS selects compiler exports on MSVC; consumers import them.
+// TVM_DLL_EXPORT always exports. Embedders may override TVM_DLL.
+#ifdef __EMSCRIPTEN__
+#include <emscripten/emscripten.h>
+#endif
+
+// --- TVM_DLL family (libtvm_compiler) ---
+#if !defined(TVM_DLL) && defined(__EMSCRIPTEN__)
+#define TVM_DLL EMSCRIPTEN_KEEPALIVE
+#define TVM_DLL_EXPORT EMSCRIPTEN_KEEPALIVE
+#endif
+#if !defined(TVM_DLL) && defined(_MSC_VER)
+#ifdef TVM_EXPORTS
+#define TVM_DLL __declspec(dllexport)
+#else
+#define TVM_DLL __declspec(dllimport)
+#endif
+#define TVM_DLL_EXPORT __declspec(dllexport)
+#endif
+#ifndef TVM_DLL
+#define TVM_DLL __attribute__((visibility("default")))
+#define TVM_DLL_EXPORT __attribute__((visibility("default")))
+#endif
+
+#endif // TVM_BASE_H_
diff --git a/include/tvm/relax/analysis.h b/include/tvm/relax/analysis.h
index 8aa335d9ef..5daf98b081 100644
--- a/include/tvm/relax/analysis.h
+++ b/include/tvm/relax/analysis.h
@@ -24,6 +24,7 @@
#ifndef TVM_RELAX_ANALYSIS_H_
#define TVM_RELAX_ANALYSIS_H_
+#include <tvm/base.h>
#include <tvm/ir/module.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/op_attr_types.h>
diff --git a/include/tvm/relax/backend.h b/include/tvm/relax/backend.h
index e7d13c47b2..de9aa70387 100644
--- a/include/tvm/relax/backend.h
+++ b/include/tvm/relax/backend.h
@@ -24,6 +24,7 @@
#ifndef TVM_RELAX_BACKEND_H_
#define TVM_RELAX_BACKEND_H_
+#include <tvm/base.h>
#include <tvm/relax/transform.h>
namespace tvm {
diff --git a/include/tvm/relax/backend/adreno/transform.h
b/include/tvm/relax/backend/adreno/transform.h
index 3a57c061cc..985d7e0236 100644
--- a/include/tvm/relax/backend/adreno/transform.h
+++ b/include/tvm/relax/backend/adreno/transform.h
@@ -24,6 +24,7 @@
#ifndef TVM_RELAX_BACKEND_ADRENO_TRANSFORM_H_
#define TVM_RELAX_BACKEND_ADRENO_TRANSFORM_H_
+#include <tvm/base.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/transform.h>
namespace tvm {
diff --git a/include/tvm/relax/binding_rewrite.h
b/include/tvm/relax/binding_rewrite.h
index b5884d19aa..d116478570 100644
--- a/include/tvm/relax/binding_rewrite.h
+++ b/include/tvm/relax/binding_rewrite.h
@@ -24,6 +24,7 @@
#ifndef TVM_RELAX_BINDING_REWRITE_H_
+#include <tvm/base.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/unique_name_supply.h>
#include <tvm/relax/analysis.h>
diff --git a/include/tvm/relax/block_builder.h
b/include/tvm/relax/block_builder.h
index 8f430d83ec..9649b8c43c 100644
--- a/include/tvm/relax/block_builder.h
+++ b/include/tvm/relax/block_builder.h
@@ -24,6 +24,7 @@
#ifndef TVM_RELAX_BLOCK_BUILDER_H_
#define TVM_RELAX_BLOCK_BUILDER_H_
+#include <tvm/base.h>
#include <tvm/ir/unique_name_supply.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/utils.h>
diff --git a/include/tvm/relax/dataflow_matcher.h
b/include/tvm/relax/dataflow_matcher.h
index 8a834d1fcd..37d43aa905 100644
--- a/include/tvm/relax/dataflow_matcher.h
+++ b/include/tvm/relax/dataflow_matcher.h
@@ -24,6 +24,7 @@
#ifndef TVM_RELAX_DATAFLOW_MATCHER_H_
#define TVM_RELAX_DATAFLOW_MATCHER_H_
+#include <tvm/base.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/optional.h>
#include <tvm/relax/dataflow_pattern.h>
diff --git a/include/tvm/relax/dataflow_pattern.h
b/include/tvm/relax/dataflow_pattern.h
index 8a5cd8f7a0..1c7c688403 100644
--- a/include/tvm/relax/dataflow_pattern.h
+++ b/include/tvm/relax/dataflow_pattern.h
@@ -24,6 +24,7 @@
#ifndef TVM_RELAX_DATAFLOW_PATTERN_H_
#define TVM_RELAX_DATAFLOW_PATTERN_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/registry.h>
diff --git a/include/tvm/relax/distributed/global_info.h
b/include/tvm/relax/distributed/global_info.h
index 51f677a498..36d7d3b303 100644
--- a/include/tvm/relax/distributed/global_info.h
+++ b/include/tvm/relax/distributed/global_info.h
@@ -25,6 +25,7 @@
#ifndef TVM_RELAX_DISTRIBUTED_GLOBAL_INFO_H_
#define TVM_RELAX_DISTRIBUTED_GLOBAL_INFO_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/shape.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/module.h>
diff --git a/include/tvm/relax/distributed/transform.h
b/include/tvm/relax/distributed/transform.h
index 12a87614cf..ceec9300f1 100644
--- a/include/tvm/relax/distributed/transform.h
+++ b/include/tvm/relax/distributed/transform.h
@@ -24,6 +24,7 @@
#ifndef TVM_RELAX_DISTRIBUTED_TRANSFORM_H_
#define TVM_RELAX_DISTRIBUTED_TRANSFORM_H_
+#include <tvm/base.h>
#include <tvm/ir/transform.h>
#include <tvm/relax/dataflow_pattern.h>
#include <tvm/relax/expr.h>
diff --git a/include/tvm/relax/distributed/type.h
b/include/tvm/relax/distributed/type.h
index 121fc4d1ec..a6a1316dd5 100644
--- a/include/tvm/relax/distributed/type.h
+++ b/include/tvm/relax/distributed/type.h
@@ -25,6 +25,7 @@
#ifndef TVM_RELAX_DISTRIBUTED_TYPE_H_
#define TVM_RELAX_DISTRIBUTED_TYPE_H_
+#include <tvm/base.h>
#include <tvm/relax/distributed/global_info.h>
#include <tvm/relax/type.h>
diff --git a/include/tvm/relax/exec_builder.h b/include/tvm/relax/exec_builder.h
index e71e0607b9..bb1d2ca848 100644
--- a/include/tvm/relax/exec_builder.h
+++ b/include/tvm/relax/exec_builder.h
@@ -23,6 +23,7 @@
#ifndef TVM_RELAX_EXEC_BUILDER_H_
#define TVM_RELAX_EXEC_BUILDER_H_
+#include <tvm/base.h>
#include <tvm/ffi/extra/dataclass.h>
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/extra/structural_hash.h>
diff --git a/include/tvm/relax/expr.h b/include/tvm/relax/expr.h
index fa16d746a9..60de2c70d0 100644
--- a/include/tvm/relax/expr.h
+++ b/include/tvm/relax/expr.h
@@ -19,6 +19,7 @@
#ifndef TVM_RELAX_EXPR_H_
#define TVM_RELAX_EXPR_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/extra/structural_equal.h>
diff --git a/include/tvm/relax/global_info.h b/include/tvm/relax/global_info.h
index 820717c926..6ffc691822 100644
--- a/include/tvm/relax/global_info.h
+++ b/include/tvm/relax/global_info.h
@@ -24,6 +24,7 @@
#ifndef TVM_RELAX_GLOBAL_INFO_H_
#define TVM_RELAX_GLOBAL_INFO_H_
+#include <tvm/base.h>
#include <tvm/ir/global_info.h>
#include <tvm/runtime/memory_scope.h>
#include <tvm/target/target.h>
diff --git a/include/tvm/relax/script/ir_builder/ir.h
b/include/tvm/relax/script/ir_builder/ir.h
index 37c3e63cf3..25555bb1a6 100644
--- a/include/tvm/relax/script/ir_builder/ir.h
+++ b/include/tvm/relax/script/ir_builder/ir.h
@@ -19,6 +19,7 @@
#ifndef TVM_RELAX_SCRIPT_IR_BUILDER_IR_H_
#define TVM_RELAX_SCRIPT_IR_BUILDER_IR_H_
+#include <tvm/base.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/script/ir_builder/frame.h>
#include <tvm/relax/type.h>
diff --git a/include/tvm/relax/tir_pattern.h b/include/tvm/relax/tir_pattern.h
index 9be6fbbfc8..efc7a3ffd1 100644
--- a/include/tvm/relax/tir_pattern.h
+++ b/include/tvm/relax/tir_pattern.h
@@ -25,6 +25,7 @@
#ifndef TVM_RELAX_TIR_PATTERN_H_
#define TVM_RELAX_TIR_PATTERN_H_
+#include <tvm/base.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tirx/function.h>
diff --git a/include/tvm/relax/transform.h b/include/tvm/relax/transform.h
index f048e56c2a..03d88ee513 100644
--- a/include/tvm/relax/transform.h
+++ b/include/tvm/relax/transform.h
@@ -24,6 +24,7 @@
#ifndef TVM_RELAX_TRANSFORM_H_
#define TVM_RELAX_TRANSFORM_H_
+#include <tvm/base.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/transform.h>
#include <tvm/relax/dataflow_pattern.h>
diff --git a/include/tvm/relax/type.h b/include/tvm/relax/type.h
index 70f0bad3ae..c2c4095a11 100644
--- a/include/tvm/relax/type.h
+++ b/include/tvm/relax/type.h
@@ -24,6 +24,7 @@
#ifndef TVM_RELAX_TYPE_H_
#define TVM_RELAX_TYPE_H_
+#include <tvm/base.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/attrs.h>
diff --git a/include/tvm/relax/type_functor.h b/include/tvm/relax/type_functor.h
index 4ee6359c1f..9374c8c849 100644
--- a/include/tvm/relax/type_functor.h
+++ b/include/tvm/relax/type_functor.h
@@ -24,6 +24,7 @@
#ifndef TVM_RELAX_TYPE_FUNCTOR_H_
#define TVM_RELAX_TYPE_FUNCTOR_H_
+#include <tvm/base.h>
#include <tvm/ir/object_functor.h>
#include <tvm/relax/distributed/type.h>
#include <tvm/relax/type.h>
diff --git a/include/tvm/relax/utils.h b/include/tvm/relax/utils.h
index 2a905c2624..0107b78b46 100644
--- a/include/tvm/relax/utils.h
+++ b/include/tvm/relax/utils.h
@@ -24,6 +24,7 @@
#ifndef TVM_RELAX_UTILS_H_
#define TVM_RELAX_UTILS_H_
+#include <tvm/base.h>
#include <tvm/ffi/error.h>
#include <tvm/ir/module.h>
#include <tvm/relax/expr.h>
diff --git a/include/tvm/runtime/disco/builtin.h
b/include/tvm/runtime/disco/builtin.h
index ff1a7f4e3b..03c9e8c35a 100644
--- a/include/tvm/runtime/disco/builtin.h
+++ b/include/tvm/runtime/disco/builtin.h
@@ -21,6 +21,7 @@
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/extra/module.h>
+#include <tvm/runtime/extra_base.h>
#include <tvm/runtime/tensor.h>
#include <string>
diff --git a/include/tvm/runtime/disco/cuda_ipc_memory.h
b/include/tvm/runtime/disco/cuda_ipc_memory.h
index cc2b023421..84676696c4 100644
--- a/include/tvm/runtime/disco/cuda_ipc_memory.h
+++ b/include/tvm/runtime/disco/cuda_ipc_memory.h
@@ -21,6 +21,7 @@
#define TVM_RUNTIME_DISCO_CUDA_IPC_MEMORY_H_
#include <tvm/runtime/base.h>
+#include <tvm/runtime/extra_base.h>
#include <tvm/runtime/tensor.h>
#include <vector>
diff --git a/include/tvm/runtime/disco/disco_worker.h
b/include/tvm/runtime/disco/disco_worker.h
index 4d5c8a3acc..d0bf688a2e 100644
--- a/include/tvm/runtime/disco/disco_worker.h
+++ b/include/tvm/runtime/disco/disco_worker.h
@@ -27,6 +27,7 @@
#include <tvm/ffi/function.h>
#include <tvm/runtime/disco/session.h>
+#include <tvm/runtime/extra_base.h>
#include <vector>
diff --git a/include/tvm/runtime/disco/session.h
b/include/tvm/runtime/disco/session.h
index 7b1bd192b3..e1de0d73a1 100644
--- a/include/tvm/runtime/disco/session.h
+++ b/include/tvm/runtime/disco/session.h
@@ -75,6 +75,7 @@
#include <tvm/ffi/container/shape.h>
#include <tvm/ffi/function.h>
+#include <tvm/runtime/extra_base.h>
#include <tvm/runtime/tensor.h>
#include <mutex>
diff --git a/include/tvm/runtime/extra_base.h b/include/tvm/runtime/extra_base.h
new file mode 100644
index 0000000000..6a193b6a4c
--- /dev/null
+++ b/include/tvm/runtime/extra_base.h
@@ -0,0 +1,49 @@
+/*
+ * 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 tvm/runtime/extra_base.h
+ * \brief Parent runtime-extra symbol visibility.
+ */
+#ifndef TVM_RUNTIME_EXTRA_BASE_H_
+#define TVM_RUNTIME_EXTRA_BASE_H_
+
+// TVM_RUNTIME_EXTRA_EXPORTS selects runtime-extra exports on MSVC.
+// Other libraries and downstream consumers import them. Embedders may
+// override TVM_RUNTIME_EXTRA_DLL independently of compiler/core runtime.
+#ifdef __EMSCRIPTEN__
+#include <emscripten/emscripten.h>
+#endif
+
+// --- TVM_RUNTIME_EXTRA_DLL (libtvm_runtime_extra) ---
+#ifndef TVM_RUNTIME_EXTRA_DLL
+#if defined(__EMSCRIPTEN__)
+#define TVM_RUNTIME_EXTRA_DLL EMSCRIPTEN_KEEPALIVE
+#elif defined(_MSC_VER)
+#ifdef TVM_RUNTIME_EXTRA_EXPORTS
+#define TVM_RUNTIME_EXTRA_DLL __declspec(dllexport)
+#else
+#define TVM_RUNTIME_EXTRA_DLL __declspec(dllimport)
+#endif
+#else
+#define TVM_RUNTIME_EXTRA_DLL __attribute__((visibility("default")))
+#endif
+#endif
+
+#endif // TVM_RUNTIME_EXTRA_BASE_H_
diff --git a/include/tvm/runtime/timing.h b/include/tvm/runtime/timing.h
index a19b0baeae..7ac8bbd574 100644
--- a/include/tvm/runtime/timing.h
+++ b/include/tvm/runtime/timing.h
@@ -26,6 +26,7 @@
#include <tvm/ffi/function.h>
#include <tvm/runtime/device_api.h>
+#include <tvm/runtime/extra_base.h>
#include <tvm/runtime/timer.h>
namespace tvm {
diff --git a/include/tvm/runtime/vm/executable.h
b/include/tvm/runtime/vm/executable.h
index 6b1eecd1f2..bebfd6154f 100644
--- a/include/tvm/runtime/vm/executable.h
+++ b/include/tvm/runtime/vm/executable.h
@@ -25,6 +25,7 @@
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
+#include <tvm/runtime/extra_base.h>
#include <tvm/runtime/memory_scope.h>
#include <tvm/support/io.h>
#include <tvm/support/serializer.h>
diff --git a/include/tvm/runtime/vm/tensor_cache_support.h
b/include/tvm/runtime/vm/tensor_cache_support.h
index 0293810deb..728ed0084f 100644
--- a/include/tvm/runtime/vm/tensor_cache_support.h
+++ b/include/tvm/runtime/vm/tensor_cache_support.h
@@ -21,6 +21,7 @@
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
+#include <tvm/runtime/extra_base.h>
#include <tvm/runtime/tensor.h>
#include <string>
diff --git a/include/tvm/s_tir/analysis.h b/include/tvm/s_tir/analysis.h
index bf7b20bbd5..8372485317 100644
--- a/include/tvm/s_tir/analysis.h
+++ b/include/tvm/s_tir/analysis.h
@@ -24,6 +24,7 @@
#ifndef TVM_S_TIR_ANALYSIS_H_
#define TVM_S_TIR_ANALYSIS_H_
+#include <tvm/base.h>
#include <tvm/ir/module.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/transform.h>
diff --git a/include/tvm/s_tir/backend/adreno/transform.h
b/include/tvm/s_tir/backend/adreno/transform.h
index 66174815ad..ac9cffe11b 100644
--- a/include/tvm/s_tir/backend/adreno/transform.h
+++ b/include/tvm/s_tir/backend/adreno/transform.h
@@ -24,6 +24,7 @@
#ifndef TVM_S_TIR_BACKEND_ADRENO_TRANSFORM_H_
#define TVM_S_TIR_BACKEND_ADRENO_TRANSFORM_H_
+#include <tvm/base.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/transform.h>
#include <tvm/s_tir/transform.h>
diff --git a/include/tvm/s_tir/data_layout.h b/include/tvm/s_tir/data_layout.h
index ad36a76a07..f71e21b3e9 100644
--- a/include/tvm/s_tir/data_layout.h
+++ b/include/tvm/s_tir/data_layout.h
@@ -25,6 +25,7 @@
#ifndef TVM_S_TIR_DATA_LAYOUT_H_
#define TVM_S_TIR_DATA_LAYOUT_H_
+#include <tvm/base.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
diff --git a/include/tvm/s_tir/iter_var.h b/include/tvm/s_tir/iter_var.h
index 75277ad9ae..96ddb9ff61 100644
--- a/include/tvm/s_tir/iter_var.h
+++ b/include/tvm/s_tir/iter_var.h
@@ -23,6 +23,7 @@
#ifndef TVM_S_TIR_ITER_VAR_H_
#define TVM_S_TIR_ITER_VAR_H_
+#include <tvm/base.h>
#include <tvm/ffi/optional.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/expr.h>
diff --git a/include/tvm/s_tir/meta_schedule/arg_info.h
b/include/tvm/s_tir/meta_schedule/arg_info.h
index b53a90931b..a8da1bfc3f 100644
--- a/include/tvm/s_tir/meta_schedule/arg_info.h
+++ b/include/tvm/s_tir/meta_schedule/arg_info.h
@@ -19,6 +19,7 @@
#ifndef TVM_S_TIR_META_SCHEDULE_ARG_INFO_H_
#define TVM_S_TIR_META_SCHEDULE_ARG_INFO_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/shape.h>
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/reflection/registry.h>
diff --git a/include/tvm/s_tir/meta_schedule/builder.h
b/include/tvm/s_tir/meta_schedule/builder.h
index fc12630ac1..b639aaa645 100644
--- a/include/tvm/s_tir/meta_schedule/builder.h
+++ b/include/tvm/s_tir/meta_schedule/builder.h
@@ -19,6 +19,7 @@
#ifndef TVM_S_TIR_META_SCHEDULE_BUILDER_H_
#define TVM_S_TIR_META_SCHEDULE_BUILDER_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/function.h>
diff --git a/include/tvm/s_tir/meta_schedule/cost_model.h
b/include/tvm/s_tir/meta_schedule/cost_model.h
index a6b5e901dc..e9840e4e47 100644
--- a/include/tvm/s_tir/meta_schedule/cost_model.h
+++ b/include/tvm/s_tir/meta_schedule/cost_model.h
@@ -20,6 +20,7 @@
#ifndef TVM_S_TIR_META_SCHEDULE_COST_MODEL_H_
#define TVM_S_TIR_META_SCHEDULE_COST_MODEL_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
diff --git a/include/tvm/s_tir/meta_schedule/database.h
b/include/tvm/s_tir/meta_schedule/database.h
index c8fb38debe..2748989e41 100644
--- a/include/tvm/s_tir/meta_schedule/database.h
+++ b/include/tvm/s_tir/meta_schedule/database.h
@@ -19,6 +19,7 @@
#ifndef TVM_S_TIR_META_SCHEDULE_DATABASE_H_
#define TVM_S_TIR_META_SCHEDULE_DATABASE_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
diff --git a/include/tvm/s_tir/meta_schedule/feature_extractor.h
b/include/tvm/s_tir/meta_schedule/feature_extractor.h
index d55c205195..072c8d9ac9 100644
--- a/include/tvm/s_tir/meta_schedule/feature_extractor.h
+++ b/include/tvm/s_tir/meta_schedule/feature_extractor.h
@@ -20,6 +20,7 @@
#ifndef TVM_S_TIR_META_SCHEDULE_FEATURE_EXTRACTOR_H_
#define TVM_S_TIR_META_SCHEDULE_FEATURE_EXTRACTOR_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
diff --git a/include/tvm/s_tir/meta_schedule/measure_callback.h
b/include/tvm/s_tir/meta_schedule/measure_callback.h
index 9a690e0a1b..7b16d386de 100644
--- a/include/tvm/s_tir/meta_schedule/measure_callback.h
+++ b/include/tvm/s_tir/meta_schedule/measure_callback.h
@@ -20,6 +20,7 @@
#ifndef TVM_S_TIR_META_SCHEDULE_MEASURE_CALLBACK_H_
#define TVM_S_TIR_META_SCHEDULE_MEASURE_CALLBACK_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
diff --git a/include/tvm/s_tir/meta_schedule/measure_candidate.h
b/include/tvm/s_tir/meta_schedule/measure_candidate.h
index c0af677157..0ee3b81934 100644
--- a/include/tvm/s_tir/meta_schedule/measure_candidate.h
+++ b/include/tvm/s_tir/meta_schedule/measure_candidate.h
@@ -20,6 +20,7 @@
#ifndef TVM_S_TIR_META_SCHEDULE_MEASURE_CANDIDATE_H_
#define TVM_S_TIR_META_SCHEDULE_MEASURE_CANDIDATE_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
diff --git a/include/tvm/s_tir/meta_schedule/mutator.h
b/include/tvm/s_tir/meta_schedule/mutator.h
index 85205213d8..312b0dbc72 100644
--- a/include/tvm/s_tir/meta_schedule/mutator.h
+++ b/include/tvm/s_tir/meta_schedule/mutator.h
@@ -20,6 +20,7 @@
#ifndef TVM_S_TIR_META_SCHEDULE_MUTATOR_H_
#define TVM_S_TIR_META_SCHEDULE_MUTATOR_H_
+#include <tvm/base.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/registry.h>
diff --git a/include/tvm/s_tir/meta_schedule/postproc.h
b/include/tvm/s_tir/meta_schedule/postproc.h
index 3f235f9a36..e4b5a8f6b8 100644
--- a/include/tvm/s_tir/meta_schedule/postproc.h
+++ b/include/tvm/s_tir/meta_schedule/postproc.h
@@ -20,6 +20,7 @@
#ifndef TVM_S_TIR_META_SCHEDULE_POSTPROC_H_
#define TVM_S_TIR_META_SCHEDULE_POSTPROC_H_
+#include <tvm/base.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
diff --git a/include/tvm/s_tir/meta_schedule/runner.h
b/include/tvm/s_tir/meta_schedule/runner.h
index d32425d1fb..4fc6bdbac7 100644
--- a/include/tvm/s_tir/meta_schedule/runner.h
+++ b/include/tvm/s_tir/meta_schedule/runner.h
@@ -19,6 +19,7 @@
#ifndef TVM_S_TIR_META_SCHEDULE_RUNNER_H_
#define TVM_S_TIR_META_SCHEDULE_RUNNER_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/optional.h>
diff --git a/include/tvm/s_tir/meta_schedule/schedule_rule.h
b/include/tvm/s_tir/meta_schedule/schedule_rule.h
index cf0f47dbf9..6a57565115 100644
--- a/include/tvm/s_tir/meta_schedule/schedule_rule.h
+++ b/include/tvm/s_tir/meta_schedule/schedule_rule.h
@@ -20,6 +20,7 @@
#ifndef TVM_S_TIR_META_SCHEDULE_SCHEDULE_RULE_H_
#define TVM_S_TIR_META_SCHEDULE_SCHEDULE_RULE_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/function.h>
diff --git a/include/tvm/s_tir/meta_schedule/search_strategy.h
b/include/tvm/s_tir/meta_schedule/search_strategy.h
index 4fb9ea2ec7..4e51d6bfa5 100644
--- a/include/tvm/s_tir/meta_schedule/search_strategy.h
+++ b/include/tvm/s_tir/meta_schedule/search_strategy.h
@@ -19,6 +19,7 @@
#ifndef TVM_S_TIR_META_SCHEDULE_SEARCH_STRATEGY_H_
#define TVM_S_TIR_META_SCHEDULE_SEARCH_STRATEGY_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/optional.h>
diff --git a/include/tvm/s_tir/meta_schedule/space_generator.h
b/include/tvm/s_tir/meta_schedule/space_generator.h
index 5f3cb93842..9eb8d3c4c5 100644
--- a/include/tvm/s_tir/meta_schedule/space_generator.h
+++ b/include/tvm/s_tir/meta_schedule/space_generator.h
@@ -19,6 +19,7 @@
#ifndef TVM_S_TIR_META_SCHEDULE_SPACE_GENERATOR_H_
#define TVM_S_TIR_META_SCHEDULE_SPACE_GENERATOR_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
diff --git a/include/tvm/s_tir/meta_schedule/task_scheduler.h
b/include/tvm/s_tir/meta_schedule/task_scheduler.h
index aa62284c9f..63d85de700 100644
--- a/include/tvm/s_tir/meta_schedule/task_scheduler.h
+++ b/include/tvm/s_tir/meta_schedule/task_scheduler.h
@@ -19,6 +19,7 @@
#ifndef TVM_S_TIR_META_SCHEDULE_TASK_SCHEDULER_H_
#define TVM_S_TIR_META_SCHEDULE_TASK_SCHEDULER_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/optional.h>
diff --git a/include/tvm/s_tir/meta_schedule/tune_context.h
b/include/tvm/s_tir/meta_schedule/tune_context.h
index c9468749fb..5c4b8600ad 100644
--- a/include/tvm/s_tir/meta_schedule/tune_context.h
+++ b/include/tvm/s_tir/meta_schedule/tune_context.h
@@ -19,6 +19,7 @@
#ifndef TVM_S_TIR_META_SCHEDULE_TUNE_CONTEXT_H_
#define TVM_S_TIR_META_SCHEDULE_TUNE_CONTEXT_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/function.h>
diff --git a/include/tvm/s_tir/sblock_dependence_info.h
b/include/tvm/s_tir/sblock_dependence_info.h
index 4d1c8309eb..815d85c6b9 100644
--- a/include/tvm/s_tir/sblock_dependence_info.h
+++ b/include/tvm/s_tir/sblock_dependence_info.h
@@ -31,6 +31,7 @@
#ifndef TVM_S_TIR_SBLOCK_DEPENDENCE_INFO_H_
#define TVM_S_TIR_SBLOCK_DEPENDENCE_INFO_H_
+#include <tvm/base.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/sblock_scope.h>
diff --git a/include/tvm/s_tir/sblock_scope.h b/include/tvm/s_tir/sblock_scope.h
index 6a7c480df3..67281c3f1f 100644
--- a/include/tvm/s_tir/sblock_scope.h
+++ b/include/tvm/s_tir/sblock_scope.h
@@ -25,6 +25,7 @@
#ifndef TVM_S_TIR_SBLOCK_SCOPE_H_
#define TVM_S_TIR_SBLOCK_SCOPE_H_
+#include <tvm/base.h>
#include <tvm/ir/module.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
diff --git a/include/tvm/s_tir/schedule/schedule.h
b/include/tvm/s_tir/schedule/schedule.h
index c090eab07a..46fcea90f8 100644
--- a/include/tvm/s_tir/schedule/schedule.h
+++ b/include/tvm/s_tir/schedule/schedule.h
@@ -19,6 +19,7 @@
#ifndef TVM_S_TIR_SCHEDULE_SCHEDULE_H_
#define TVM_S_TIR_SCHEDULE_SCHEDULE_H_
+#include <tvm/base.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/random_engine.h>
#include <tvm/s_tir/schedule/state.h>
diff --git a/include/tvm/s_tir/schedule/state.h
b/include/tvm/s_tir/schedule/state.h
index 60f5b09827..27d3d0d2a3 100644
--- a/include/tvm/s_tir/schedule/state.h
+++ b/include/tvm/s_tir/schedule/state.h
@@ -23,6 +23,7 @@
#ifndef TVM_S_TIR_SCHEDULE_STATE_H_
#define TVM_S_TIR_SCHEDULE_STATE_H_
+#include <tvm/base.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/module.h>
#include <tvm/ir/prim/expr.h>
diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h
index 30e868164d..f534ffc6a2 100644
--- a/include/tvm/s_tir/stmt.h
+++ b/include/tvm/s_tir/stmt.h
@@ -27,6 +27,7 @@
#ifndef TVM_S_TIR_STMT_H_
#define TVM_S_TIR_STMT_H_
+#include <tvm/base.h>
#include <tvm/ffi/container/tuple.h>
#include <tvm/s_tir/iter_var.h>
#include <tvm/tirx/stmt.h>
diff --git a/include/tvm/s_tir/stmt_functor.h b/include/tvm/s_tir/stmt_functor.h
index 7e8b9f2758..cf089d6b53 100644
--- a/include/tvm/s_tir/stmt_functor.h
+++ b/include/tvm/s_tir/stmt_functor.h
@@ -23,6 +23,7 @@
#ifndef TVM_S_TIR_STMT_FUNCTOR_H_
#define TVM_S_TIR_STMT_FUNCTOR_H_
+#include <tvm/base.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/stmt_functor.h>
diff --git a/include/tvm/s_tir/stmt_functor_with_analyzer.h
b/include/tvm/s_tir/stmt_functor_with_analyzer.h
index aee173e357..d1452c561c 100644
--- a/include/tvm/s_tir/stmt_functor_with_analyzer.h
+++ b/include/tvm/s_tir/stmt_functor_with_analyzer.h
@@ -25,6 +25,7 @@
#ifndef TVM_S_TIR_STMT_FUNCTOR_WITH_ANALYZER_H_
#define TVM_S_TIR_STMT_FUNCTOR_WITH_ANALYZER_H_
+#include <tvm/base.h>
#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/stmt_functor_with_analyzer.h>
diff --git a/include/tvm/s_tir/tensor_intrin.h
b/include/tvm/s_tir/tensor_intrin.h
index 3943555a93..c583ec87fb 100644
--- a/include/tvm/s_tir/tensor_intrin.h
+++ b/include/tvm/s_tir/tensor_intrin.h
@@ -24,6 +24,7 @@
#ifndef TVM_S_TIR_TENSOR_INTRIN_H_
#define TVM_S_TIR_TENSOR_INTRIN_H_
+#include <tvm/base.h>
#include <tvm/tirx/function.h>
namespace tvm {
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index 2ec062e39a..dcb177878a 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -24,6 +24,7 @@
#ifndef TVM_S_TIR_TRANSFORM_H_
#define TVM_S_TIR_TRANSFORM_H_
+#include <tvm/base.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/transform.h>
#include <tvm/runtime/memory_scope.h>
diff --git a/include/tvm/te/operation.h b/include/tvm/te/operation.h
index ce5004c588..f6cb430af8 100644
--- a/include/tvm/te/operation.h
+++ b/include/tvm/te/operation.h
@@ -24,6 +24,7 @@
#ifndef TVM_TE_OPERATION_H_
#define TVM_TE_OPERATION_H_
+#include <tvm/base.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/prim/expr.h>
diff --git a/include/tvm/te/reduction.h b/include/tvm/te/reduction.h
index 29f4516970..353255ab93 100644
--- a/include/tvm/te/reduction.h
+++ b/include/tvm/te/reduction.h
@@ -23,6 +23,7 @@
#ifndef TVM_TE_REDUCTION_H_
#define TVM_TE_REDUCTION_H_
+#include <tvm/base.h>
#include <tvm/s_tir/iter_var.h>
namespace tvm::prim {
diff --git a/include/tvm/te/tensor.h b/include/tvm/te/tensor.h
index ec477f337d..cb9c113614 100644
--- a/include/tvm/te/tensor.h
+++ b/include/tvm/te/tensor.h
@@ -24,6 +24,7 @@
#ifndef TVM_TE_TENSOR_H_
#define TVM_TE_TENSOR_H_
+#include <tvm/base.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/prim/op.h>
diff --git a/src/relax/ir/py_expr_functor.cc b/src/relax/ir/py_expr_functor.cc
index 3f011d2008..c3910fbeb8 100644
--- a/src/relax/ir/py_expr_functor.cc
+++ b/src/relax/ir/py_expr_functor.cc
@@ -21,6 +21,7 @@
* \file src/relax/py_expr_functor.cc
* \brief The backbone of PyExprVisitor/PyExprMutator.
*/
+#include <tvm/base.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/expr_functor.h>
diff --git a/src/relax/ir/transform.cc b/src/relax/ir/transform.cc
index 91bf21683b..a20564f5c9 100644
--- a/src/relax/ir/transform.cc
+++ b/src/relax/ir/transform.cc
@@ -21,6 +21,7 @@
* \file relax/ir/transform.cc
* \brief Relax specific transformation passes.
*/
+#include <tvm/base.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/dataclass.h>
#include <tvm/ffi/function.h>
diff --git a/src/relax/training/utils.h b/src/relax/training/utils.h
index c22588804d..32c073f8fd 100644
--- a/src/relax/training/utils.h
+++ b/src/relax/training/utils.h
@@ -24,6 +24,7 @@
#ifndef TVM_RELAX_TRAINING_UTILS_H_
#define TVM_RELAX_TRAINING_UTILS_H_
+#include <tvm/base.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/transform.h>
diff --git a/src/relax/transform/dataflow_inplace.cc
b/src/relax/transform/dataflow_inplace.cc
index 9b320bb37f..f735a0a49c 100644
--- a/src/relax/transform/dataflow_inplace.cc
+++ b/src/relax/transform/dataflow_inplace.cc
@@ -23,6 +23,7 @@
* into in-place versions.
*/
+#include <tvm/base.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/reflection/registry.h>
diff --git a/src/relax/transform/utils.h b/src/relax/transform/utils.h
index 921d10f0d3..67d8a18f2c 100644
--- a/src/relax/transform/utils.h
+++ b/src/relax/transform/utils.h
@@ -25,6 +25,7 @@
#define TVM_RELAX_TRANSFORM_UTILS_H_
#include <builtin_fp16.h>
+#include <tvm/base.h>
#include <tvm/ffi/cast.h>
#include <tvm/ir/expr_functor.h>
#include <tvm/ir/function.h>
diff --git a/src/runtime/extra/contrib/dnnl/dnnl_kernel.h
b/src/runtime/extra/contrib/dnnl/dnnl_kernel.h
index edac82b51f..ef76397d1c 100644
--- a/src/runtime/extra/contrib/dnnl/dnnl_kernel.h
+++ b/src/runtime/extra/contrib/dnnl/dnnl_kernel.h
@@ -28,6 +28,7 @@
#include <tvm/ffi/error.h>
#include <tvm/ffi/function.h>
#include <tvm/runtime/base.h>
+#include <tvm/runtime/extra_base.h>
#include <vector>
diff --git a/src/runtime/extra/disco/builtin.cc
b/src/runtime/extra/disco/builtin.cc
index 46b968e77d..c703823ea0 100644
--- a/src/runtime/extra/disco/builtin.cc
+++ b/src/runtime/extra/disco/builtin.cc
@@ -23,6 +23,7 @@
#include <tvm/runtime/disco/builtin.h>
#include <tvm/runtime/disco/disco_worker.h>
#include <tvm/runtime/disco/session.h>
+#include <tvm/runtime/extra_base.h>
#include <tvm/runtime/tensor.h>
#include <sstream>
diff --git a/src/runtime/extra/disco/disco_worker.cc
b/src/runtime/extra/disco/disco_worker.cc
index 5e3a8b6668..6d7eb4740d 100644
--- a/src/runtime/extra/disco/disco_worker.cc
+++ b/src/runtime/extra/disco/disco_worker.cc
@@ -20,6 +20,7 @@
#include <tvm/runtime/disco/builtin.h>
#include <tvm/runtime/disco/disco_worker.h>
#include <tvm/runtime/disco/session.h>
+#include <tvm/runtime/extra_base.h>
#include "../../../support/process_id.h"
#include "./protocol.h"
diff --git a/src/runtime/rpc/rpc_socket_impl.cc
b/src/runtime/rpc/rpc_socket_impl.cc
index cff92a17eb..d0dd4a3a81 100644
--- a/src/runtime/rpc/rpc_socket_impl.cc
+++ b/src/runtime/rpc/rpc_socket_impl.cc
@@ -23,6 +23,7 @@
*/
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/runtime/extra_base.h>
#include <tvm/support/io.h>
#include <tvm/support/serializer.h>
diff --git a/src/runtime/vm/tensor_cache_support.cc
b/src/runtime/vm/tensor_cache_support.cc
index 0c7559bcd6..a45e67917e 100644
--- a/src/runtime/vm/tensor_cache_support.cc
+++ b/src/runtime/vm/tensor_cache_support.cc
@@ -38,6 +38,7 @@
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/runtime/extra_base.h>
#include <tvm/runtime/tensor.h>
#include <tvm/runtime/vm/tensor_cache_support.h>
#include <tvm/support/file_utils.h>
diff --git a/src/s_tir/schedule/primitive.h b/src/s_tir/schedule/primitive.h
index e4092685b6..fd6b484531 100644
--- a/src/s_tir/schedule/primitive.h
+++ b/src/s_tir/schedule/primitive.h
@@ -19,6 +19,7 @@
#ifndef TVM_S_TIR_SCHEDULE_PRIMITIVE_H_
#define TVM_S_TIR_SCHEDULE_PRIMITIVE_H_
+#include <tvm/base.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/random_engine.h>
#include <tvm/s_tir/schedule/state.h>
diff --git a/src/s_tir/schedule/primitive/sampling.cc
b/src/s_tir/schedule/primitive/sampling.cc
index 5b74c40e9b..6dc7ff1586 100644
--- a/src/s_tir/schedule/primitive/sampling.cc
+++ b/src/s_tir/schedule/primitive/sampling.cc
@@ -17,6 +17,8 @@
* under the License.
*/
+#include <tvm/base.h>
+
#include <random>
#include "../utils.h"
diff --git a/src/s_tir/schedule/state.cc b/src/s_tir/schedule/state.cc
index ee313c696f..88d7762ed0 100644
--- a/src/s_tir/schedule/state.cc
+++ b/src/s_tir/schedule/state.cc
@@ -16,6 +16,7 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/base.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/reflection/registry.h>
diff --git a/src/s_tir/support/parallel_for.h b/src/s_tir/support/parallel_for.h
index 22b2073e56..9bcbea7d7c 100644
--- a/src/s_tir/support/parallel_for.h
+++ b/src/s_tir/support/parallel_for.h
@@ -24,6 +24,7 @@
#ifndef TVM_S_TIR_SUPPORT_PARALLEL_FOR_H_
#define TVM_S_TIR_SUPPORT_PARALLEL_FOR_H_
+#include <tvm/base.h>
#include <tvm/ffi/error.h>
#include <tvm/runtime/base.h>
diff --git a/tirx/include/tvm/runtime/base.h b/tirx/include/tvm/runtime/base.h
index 0848fff40c..c50ca9d84f 100644
--- a/tirx/include/tvm/runtime/base.h
+++ b/tirx/include/tvm/runtime/base.h
@@ -35,48 +35,13 @@
#define TVM_TIRX_VERSION "0.1.0a0"
#endif
-// The compiler, core runtime and runtime-extra libraries expose independent
-// DLL policies so downstream embedders can override their visibility
separately.
-//
-// TVM_DLL / TVM_DLL_EXPORT: symbols in libtvm_compiler.
-// - TVM_DLL is dllexport when TVM_EXPORTS is defined (compiler build),
-// dllimport otherwise (downstream consumers, runtime TUs).
-// - TVM_DLL_EXPORT is always dllexport.
-//
-// TVM_RUNTIME_DLL / TVM_RUNTIME_DLL_EXPORT: symbols in libtvm_runtime.
-// - TVM_RUNTIME_DLL is dllexport when TVM_RUNTIME_EXPORTS is defined
-// (runtime build), dllimport otherwise.
-// - TVM_RUNTIME_DLL_EXPORT is always dllexport.
-//
-// TVM_RUNTIME_EXTRA_DLL: symbols in libtvm_runtime_extra.
-// - dllexport when TVM_RUNTIME_EXTRA_EXPORTS is defined (extra build),
-// dllimport otherwise, including core-runtime and compiler builds.
-//
-// On non-MSVC platforms the import/export decision is made by the dynamic
-// loader, so the macros expand to visibility("default"). Under
-// Emscripten they expand to EMSCRIPTEN_KEEPALIVE.
+// Core runtime visibility is independent of TIRx compiler visibility.
+// TVM_RUNTIME_EXPORTS selects core-runtime exports on MSVC; downstream
+// consumers import them. Embedders may override each library policy.
#ifdef __EMSCRIPTEN__
#include <emscripten/emscripten.h>
#endif
-// --- TVM_DLL family (libtvm_compiler) ---
-#if !defined(TVM_DLL) && defined(__EMSCRIPTEN__)
-#define TVM_DLL EMSCRIPTEN_KEEPALIVE
-#define TVM_DLL_EXPORT EMSCRIPTEN_KEEPALIVE
-#endif
-#if !defined(TVM_DLL) && defined(_MSC_VER)
-#ifdef TVM_EXPORTS
-#define TVM_DLL __declspec(dllexport)
-#else
-#define TVM_DLL __declspec(dllimport)
-#endif
-#define TVM_DLL_EXPORT __declspec(dllexport)
-#endif
-#ifndef TVM_DLL
-#define TVM_DLL __attribute__((visibility("default")))
-#define TVM_DLL_EXPORT __attribute__((visibility("default")))
-#endif
-
// Core IR, TIRx and shared compiler infrastructure live in
libtvm_tirx_compiler.
#ifndef TVM_TIRX_DLL
#if defined(__EMSCRIPTEN__)
@@ -110,19 +75,4 @@
#define TVM_RUNTIME_DLL_EXPORT __attribute__((visibility("default")))
#endif
-// --- TVM_RUNTIME_EXTRA_DLL (libtvm_runtime_extra) ---
-#ifndef TVM_RUNTIME_EXTRA_DLL
-#if defined(__EMSCRIPTEN__)
-#define TVM_RUNTIME_EXTRA_DLL EMSCRIPTEN_KEEPALIVE
-#elif defined(_MSC_VER)
-#ifdef TVM_RUNTIME_EXTRA_EXPORTS
-#define TVM_RUNTIME_EXTRA_DLL __declspec(dllexport)
-#else
-#define TVM_RUNTIME_EXTRA_DLL __declspec(dllimport)
-#endif
-#else
-#define TVM_RUNTIME_EXTRA_DLL __attribute__((visibility("default")))
-#endif
-#endif
-
#endif // TVM_RUNTIME_BASE_H_
diff --git a/tirx/python/tvm_tirx/tirx/tensor.py
b/tirx/python/tvm_tirx/tirx/tensor.py
index f497854fdd..aab69a485e 100644
--- a/tirx/python/tvm_tirx/tirx/tensor.py
+++ b/tirx/python/tvm_tirx/tirx/tensor.py
@@ -354,7 +354,7 @@ def decl_tensor(
layout,
loc,
)
- return _ffi_api.TensorVar(name, tensor_type, loc) # type: ignore
+ return Var(name, tensor_type, loc=loc)
def tensor_data_ptr(tensor, *, ty=None, loc: Location = UNKNOWN_LOC):
diff --git a/tirx/python/tvm_tirx/tirx/transform/common.py
b/tirx/python/tvm_tirx/tirx/transform/common.py
index 433b3e7769..2a9e682e90 100644
--- a/tirx/python/tvm_tirx/tirx/transform/common.py
+++ b/tirx/python/tvm_tirx/tirx/transform/common.py
@@ -18,96 +18,7 @@
import tvm_ffi
-from tvm_tirx.ir import Call, Evaluate, Op, Stmt, Var
-from tvm_tirx.tirx import Expr, decl_tensor, is_tensor_var
-from tvm_tirx.tirx.layout import Iter, TileLayout
-
-
-class TensorReplacer:
- """
- Replace tensor variables with other tensor variables.
- Tensor variables are ordinary Vars, so the same mapping also rewrites
- ``tensor_data_ptr`` projections.
- """
-
- def __init__(
- self, tensor_map: dict[Var, Var] | None = None, var_map: dict[Var,
Var] | None = None
- ):
- super().__init__()
- self.tensor_map = tensor_map if tensor_map is not None else {}
- self.var_map = var_map if var_map is not None else {}
- for old_tensor, new_tensor in self.tensor_map.items():
- self.var_map[old_tensor] = new_tensor
-
- def __call__(self, node):
- def replace_var(op: Var):
- if is_tensor_var(op):
- return self._mutate_tensor(op)
- return self.var_map.get(op, op)
-
- return tvm_ffi.structural_map(
- node,
- (Var, replace_var),
- order="post",
- )
-
- def _replace_expr(self, expr: Expr):
- return tvm_ffi.structural_map(
- expr,
- (Var, lambda var: self.var_map.get(var, var)),
- order="post",
- )
-
- def _mutate_tensor(self, tensor: Var):
- if tensor in self.tensor_map:
- return self.tensor_map[tensor]
-
- new_shape = [self._replace_expr(expr) for expr in tensor.ty.shape]
- new_strides = [self._replace_expr(expr) for expr in tensor.ty.strides]
- new_elem_offset = (
- self._replace_expr(tensor.ty.elem_offset) if tensor.ty.elem_offset
is not None else None
- )
- if isinstance(tensor.ty.layout, TileLayout):
- new_shard = [
- Iter(self._replace_expr(it.extent),
self._replace_expr(it.stride), it.axis)
- for it in tensor.ty.layout.shard
- ]
- new_replicate = [
- Iter(self._replace_expr(it.extent),
self._replace_expr(it.stride), it.axis)
- for it in tensor.ty.layout.replica
- ]
- new_layout = TileLayout.from_iters(
- new_shard,
- new_replicate,
- offset=tensor.ty.layout.offset,
- )
- else:
- new_layout = tensor.ty.layout
-
- unchanged = (
- all(old is new for old, new in zip(tensor.ty.shape, new_shape))
- and all(old is new for old, new in zip(tensor.ty.strides,
new_strides))
- and tensor.ty.elem_offset is new_elem_offset
- and tensor.ty.layout is new_layout
- )
- if unchanged:
- return tensor
-
- new_tensor = decl_tensor(
- new_shape,
- tensor.ty.dtype,
- tensor.name,
- None,
- new_strides,
- new_elem_offset,
- tensor.scope(),
- tensor.ty.data_alignment,
- tensor.ty.offset_factor,
- layout=new_layout,
- )
- self.tensor_map[tensor] = new_tensor
- self.var_map[tensor] = new_tensor
- return new_tensor
+from tvm_tirx.ir import Call, Evaluate, Op, Stmt
def seek_kernel_replace_point(stmt: Stmt, body: Stmt) -> Stmt:
diff --git a/tirx/src/codegen/c/codegen_c_host.cc
b/tirx/src/codegen/c/codegen_c_host.cc
index 708079aa46..0d0097bcda 100644
--- a/tirx/src/codegen/c/codegen_c_host.cc
+++ b/tirx/src/codegen/c/codegen_c_host.cc
@@ -49,11 +49,11 @@ void CodeGenCHost::Init(bool output_ssa, bool emit_asserts,
bool emit_fwd_func_d
emit_fwd_func_decl_ = emit_fwd_func_decl;
declared_globals_.clear();
decl_stream << "// tvm target: " << target_str << "\n";
- decl_stream << "#define TVM_EXPORTS\n";
decl_stream << "#include \"tvm/ffi/extra/c_env_api.h\"\n";
decl_stream << "#include \"tvm/runtime/base.h\"\n";
decl_stream << "#include \"tvm/runtime/c_backend_api.h\"\n";
decl_stream << "#include \"tvm/ffi/c_api.h\"\n";
+ decl_stream << "#define TVM_DLL TVM_FFI_DLL_EXPORT\n";
decl_stream << "#include <math.h>\n";
decl_stream << "#include <stdbool.h>\n";
CodeGenCHost::InitGlobalContext();
diff --git a/tirx/src/tirx/ir/tensor.cc b/tirx/src/tirx/ir/tensor.cc
index 9cda90b20d..8864a345bb 100644
--- a/tirx/src/tirx/ir/tensor.cc
+++ b/tirx/src/tirx/ir/tensor.cc
@@ -444,13 +444,8 @@ bool TensorVar::IsScalar(bool alloc_or_decl) const {
return ty()->IsScalar(alloc
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
- .def("tirx.TensorVar",
- [](ffi::String name, TensorType type, Location loc) {
- return TensorVar(std::move(name), std::move(type),
std::move(loc));
- })
.def_method("tirx.TensorStorageScope", &TensorVar::scope)
- .def_method("tirx.TensorIsScalar", &TensorVar::IsScalar)
- .def_method("tirx.TensorData", &TensorVar::data);
+ .def_method("tirx.TensorIsScalar", &TensorVar::IsScalar);
}
} // namespace tirx
diff --git a/tirx/tests/python/tirx/test_op.py
b/tirx/tests/python/tirx/test_op.py
index 4e39fbf10b..c0be7f08fe 100644
--- a/tirx/tests/python/tirx/test_op.py
+++ b/tirx/tests/python/tirx/test_op.py
@@ -25,7 +25,6 @@ from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx import Var, decl_tensor
from tvm_tirx.tirx.analysis import undefined_vars
-from tvm_tirx.tirx.transform.common import TensorReplacer
def test_instruction_is_a_void_call():
@@ -91,7 +90,9 @@ def test_runtime_operands_are_seen_and_replaced():
for value in (src, dst, mbar, row, mask, policy):
assert any(v.same_as(value) for v in free)
new_row = Var("new_row", "int32")
- replacement = TensorReplacer(var_map={row: new_row})(Evaluate(call))
+ replacement = tvm_ffi.structural_map(
+ Evaluate(call), (Var, lambda var: new_row if var.same_as(row) else
var), order="post"
+ )
new_free = undefined_vars(replacement)
assert not any(v.same_as(row) for v in new_free)
assert any(v.same_as(new_row) for v in new_free)
@@ -132,18 +133,3 @@ def test_direct_call_validation_cannot_bypass_signature():
]:
with pytest.raises((TypeError, ValueError, RuntimeError)):
call.validate()
-
-
-def test_tensor_replacer_no_shared_default():
- r1, r2 = TensorReplacer(), TensorReplacer()
- a, b = decl_tensor((64,), "float32"), decl_tensor((64,), "float32")
- r1.tensor_map[a] = b
- assert not r2.tensor_map
-
-
-def test_tensor_replacer_replaces_strides_and_elem_offset():
- n, m = Var("n", "int32"), Var("m", "int32")
- a = decl_tensor((64,), "float32", strides=[n], elem_offset=n)
- new = TensorReplacer(var_map={n: m})(tvm_tirx.ir.TensorStore(a, [0], 1.0))
- assert new.dest.strides[0].same_as(m)
- assert new.dest.elem_offset.same_as(m)
