chore: import upstream snapshot with attribution

python/tvm/tirx/__init__.py

@@ -0,0 +1,114 @@
+# isort: skip_file
+# 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.
+# pylint: disable=unused-import, redefined-builtin
+"""Namespace for Tensor-level IR"""
+
+import tvm.script
+
+tvm.script.register_dialect("tirx", "tvm.tirx.script")
+
+
+from tvm.ir import Expr
+from tvm.runtime import const
+
+from .buffer import Buffer, decl_buffer, DataProducer
+from .expr import convert
+from .expr import Var, Reduce, FloatImm, IntImm, StringImm, Cast
+from .expr import Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod
+from .expr import Min, Max, EQ, NE, LT, LE, GT, GE, And, Or, Not
+from .expr import Select, BufferLoad, ProducerLoad, Ramp, Broadcast, Shuffle
+from .expr import CallEffectKind, Let, IterVar, CommReducer
+
+from .stmt import Stmt, Bind, AssertStmt, ForKind, For, While
+
+# Legacy alias: LetStmt was folded into Bind (which now accepts an optional body)
+LetStmt = Bind
+
+from .stmt import BufferStore, AllocBuffer, AttrStmt, DeclBuffer
+
+from .stmt import SeqStmt
+from .stmt import IfThenElse, Evaluate, stmt_seq, stmt_list
+from .stmt import BufferRegion, MatchBufferRegion, SBlock, SBlockRealize
+from .stmt import TilePrimitiveCall, ScopeIdDefStmt
+
+from .function import PrimFunc, TensorIntrin, IndexMap
+
+from .op import call_packed_lowered, call_cpacked_lowered, call_tir
+from .op import call_packed, call_cpacked, call_intrin, call_pure_extern, call_extern
+from .op import call_llvm_intrin, call_llvm_pure_intrin, ret, all, any, min_value, max_value, trace
+from .op import tvm_stack_alloca, tvm_stack_make_shape, tvm_stack_make_array
+from .op import tvm_tuple, handle_add_byte_offset, tvm_struct_get, tvm_struct_set
+from .op import address_of, lookup_param, assume, undef
+from .op import continue_loop, break_loop
+from .op import tvm_thread_allreduce, type_annotation, tvm_access_ptr, ptr_byte_offset
+from .op import tvm_throw_last_error
+from .op import (
+    tvm_load_matrix_sync,
+    tvm_store_matrix_sync,
+    tvm_mma_sync,
+    tvm_bmma_sync,
+    tvm_fill_fragment,
+)
+from .op import vectorlow, vectorhigh, vectorcombine
+from .op import infinity, reinterpret
+from .op import exp, exp2, exp10, log, log2, log10, log1p, ldexp, clz
+from .op import sin, sinh, asin, asinh
+from .op import cos, cosh, acos, acosh
+from .op import tan, tanh, atan, atan2, atanh
+from .op import bitwise_and, bitwise_not, bitwise_or, bitwise_xor
+from .op import erf, sigmoid, sqrt, rsqrt, floor, ceil, hypot
+from .op import trunc, abs, round, nextafter, nearbyint, power, pow, popcount, fmod, if_then_else
+from .op import likely, isnan, isnullptr, isfinite, isinf, copysign
+from .op import div, indexdiv, indexmod, truncdiv, truncmod, floordiv, floormod, ceildiv, logaddexp
+from .op import comm_reducer, min, max, sum
+from .op import q_multiply_shift, q_multiply_shift_per_axis, shift_left, shift_right
+from .op import TVMBackendAllocWorkspace, TVMBackendFreeWorkspace
+from .op import start_profile_intrinsic, end_profile_intrinsic
+from .op import vscale, get_active_lane_mask, get_vscale_expr
+from .op import dp4a
+from .op import ignore_loop_partition
+
+# TIRX-specific imports (must come before subpackage imports to avoid circular imports)
+from .exec_scope import ExecScope, ScopeIdDef
+from .layout import TileLayout, Layout, SwizzleLayout, ComposeLayout
+from .predicate import Predicate
+from .expr_functor import ExprFunctor
+
+from . import transform
+from . import analysis
+from . import backend
+from . import stmt_functor
+
+from .functor import PyStmtExprVisitor, PyStmtExprMutator
+
+# Compiler-only submodules. Skip under `TVM_USE_RUNTIME_LIB=1` since they
+# perform compiler-side FFI at module load (schema engine looks up
+# `ir.RegisterOp`; codegen registry hooks the build pipeline).
+from tvm.base import _RUNTIME_ONLY as _RUNTIME_ONLY_TIRX  # pylint: disable=wrong-import-position
+
+if not _RUNTIME_ONLY_TIRX:
+    from .build import build
+    from .compilation_pipeline import (
+        get_tir_pipeline,
+        get_default_tir_pipeline,
+        register_tir_pipeline,
+    )
+
+import tvm.script
+
+tvm.script.register_dialect("tirx", "tvm.tirx.script")