wehub-resource-sync
115 lines
4.6 KiB
Python
115 lines
4.6 KiB
Python
# 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")
|