# 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")