This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx-standalone-root in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 533e0b6660dc467b28dc91cb7e6d807fb7f4b7b2 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. --- 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 +----- 6 files changed, 10 insertions(+), 168 deletions(-) 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)
