chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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")
|
||||
@@ -0,0 +1,21 @@
|
||||
# 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.
|
||||
"""FFI APIs for tvm.tirx"""
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
tvm_ffi.init_ffi_api("tirx", __name__)
|
||||
@@ -0,0 +1,20 @@
|
||||
# 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.
|
||||
"""Namespace of all TIR analysis utils."""
|
||||
# pylint: disable=wildcard-import, invalid-name
|
||||
|
||||
from .analysis import *
|
||||
@@ -0,0 +1,21 @@
|
||||
# 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.
|
||||
"""FFI APIs for tvm.tirx.analysis"""
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
tvm_ffi.init_ffi_api("tirx.analysis", __name__)
|
||||
@@ -0,0 +1,160 @@
|
||||
# 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.
|
||||
"""Wrapping existing analysis utils."""
|
||||
|
||||
# pylint: disable=invalid-name
|
||||
|
||||
from tvm.ir import IRModule
|
||||
from tvm.tirx.expr import Var
|
||||
from tvm.tirx.stmt import Expr
|
||||
|
||||
from .. import Stmt
|
||||
from ..function import PrimFunc
|
||||
from . import _ffi_api
|
||||
|
||||
|
||||
def expr_deep_equal(lhs: Expr, rhs: Expr) -> bool:
|
||||
"""Deeply compare two nested expressions.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
lhs : Expr
|
||||
The left operand.
|
||||
|
||||
rhs : Expr
|
||||
The right operand.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : bool
|
||||
The comparison result
|
||||
|
||||
Note
|
||||
----
|
||||
|
||||
This function does not remap variable bindings, it will not
|
||||
return true for (let x = 1 in x + 1) vs (let y = 1 in y + 1), unless x.same_as(y).
|
||||
Use py:func:`tvm_ffi.structural_equal` to handle structural variable remapping.
|
||||
|
||||
Due to the restriction of not remapping variables, this function can run
|
||||
faster than StructuralEqual and can be used as a utility function during arithmetic
|
||||
simplifications.
|
||||
|
||||
Always consider py:func:`tvm_ffi.structural_equal` first, which handles
|
||||
the structural remapping.
|
||||
|
||||
See Also
|
||||
--------
|
||||
tvm_ffi.structural_equal
|
||||
"""
|
||||
return _ffi_api.expr_deep_equal(lhs, rhs) # type: ignore
|
||||
|
||||
|
||||
def verify_ssa(func: PrimFunc) -> bool:
|
||||
"""Verify if the func is in SSA form.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func: tvm.tirx.PrimFunc
|
||||
The module to be verified.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : bool
|
||||
The result of verification.
|
||||
"""
|
||||
return _ffi_api.verify_ssa(func) # type: ignore
|
||||
|
||||
|
||||
def verify_memory(func: PrimFunc) -> bool:
|
||||
"""Verify if func contains illegal host side direct memory access.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func: tvm.tirx.PrimFunc
|
||||
The module to be verified.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : bool
|
||||
The result of verification.
|
||||
"""
|
||||
return _ffi_api.verify_memory(func) # type: ignore
|
||||
|
||||
|
||||
def undefined_vars(node: Stmt | Expr, defs: list[Var] | None = None) -> list[Var]:
|
||||
"""Find undefined vars in a TIR statement or expression.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node: Union[Stmt, Expr]
|
||||
The TIR statement or expression to be checked.
|
||||
|
||||
defs: Optional[List[Var]]
|
||||
The vars that is defined
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : List[Var]
|
||||
The undefined vars.
|
||||
"""
|
||||
defs = defs or []
|
||||
return _ffi_api.UndefinedVars(node, defs) # type: ignore # pylint: disable=no-member
|
||||
|
||||
|
||||
def verify_well_formed(obj: PrimFunc | IRModule, assert_mode: bool = True) -> bool:
|
||||
"""Verify if the given TIR is well-formed. The verification includes:
|
||||
- Check if expressions not contain vars that is defined outside the block.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
obj: Union[tvm.tirx.PrimFunc, tvm.ir.IRModule]
|
||||
The function or module to be verified.
|
||||
|
||||
assert_mode: bool
|
||||
The indicator if it raises an error when the function is not well-formed.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result: bool
|
||||
Whether it is a well-formed TIR function.
|
||||
"""
|
||||
return _ffi_api.VerifyWellFormed(obj, assert_mode) # type: ignore # pylint: disable=no-member
|
||||
|
||||
|
||||
def verify_tirx_well_formed(
|
||||
obj: PrimFunc | IRModule, assert_mode: bool = True, device_func: bool = False
|
||||
) -> bool:
|
||||
"""Verify if the given TIRX is well-formed.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
obj: Union[tvm.tirx.PrimFunc, tvm.ir.IRModule]
|
||||
The function or module to be verified.
|
||||
|
||||
assert_mode: bool
|
||||
The indicator if it raises an error when the function is not well-formed.
|
||||
|
||||
device_func: bool
|
||||
The indicator if it is a device function.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result: bool
|
||||
Whether it is a well-formed TIRX function.
|
||||
"""
|
||||
return _ffi_api.VerifyTIRxWellFormed(obj, assert_mode, device_func) # type: ignore # pylint: disable=no-member
|
||||
@@ -0,0 +1,17 @@
|
||||
# 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.
|
||||
"""TIRx backend compatibility package."""
|
||||
@@ -0,0 +1,803 @@
|
||||
# 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.
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from enum import Enum
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import triton.profiler as proton
|
||||
import tvm_ffi
|
||||
|
||||
import tvm
|
||||
from tvm.script import tirx as T
|
||||
from tvm.support import nvcc
|
||||
|
||||
|
||||
def is_running_under_pytest():
|
||||
"""Check if the code is being executed within a pytest session."""
|
||||
return "PYTEST_CURRENT_TEST" in os.environ
|
||||
|
||||
|
||||
def setup():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--dump-ptx", type=str, help="Dump PTX code to specified file")
|
||||
parser.add_argument("--dump-source", action="store_true", help="Dump source code")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.dump_ptx:
|
||||
|
||||
@tvm_ffi.register_global_func("tvm_callback_cuda_compile", override=True)
|
||||
def tvm_callback_cuda_compile(code, target):
|
||||
ptx = nvcc.compile_cuda(code, target_format="ptx")
|
||||
with open(args.dump_ptx, "w", encoding="utf-8") as f:
|
||||
f.write(ptx.decode())
|
||||
return ptx
|
||||
|
||||
return args
|
||||
|
||||
|
||||
_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m")
|
||||
|
||||
|
||||
# proton-viewer -m avg_time/us prints average kernel time in microseconds (see
|
||||
# triton/profiler/viewer.py avg_time_factor_dict). Store microseconds as-is.
|
||||
PROTON_AVG_TIME_METRIC = "avg_time/us"
|
||||
|
||||
|
||||
def _parse_proton_tree(text, *, kernel: str = ""):
|
||||
"""Parse proton-viewer tree output into {impl: time_us}.
|
||||
|
||||
Accepts ALL depth-1 nodes (no KNOWN_IMPLS filter). For each depth-1 impl,
|
||||
takes the slowest depth-2 child kernel time.
|
||||
|
||||
Tree numbers are microseconds when ProtonContext uses avg_time/us.
|
||||
|
||||
Returns (impl_times, baseline_errors) where:
|
||||
impl_times: {str: float} — impl name to avg time in microseconds
|
||||
baseline_errors: {str: str} — impl name to error message
|
||||
"""
|
||||
_ = kernel # kept for callers; unit does not depend on workload
|
||||
impl = None
|
||||
results = {}
|
||||
baseline_errors = {}
|
||||
for raw in text.splitlines():
|
||||
line = _ANSI_RE.sub("", raw).rstrip()
|
||||
if not line:
|
||||
continue
|
||||
if line.startswith("BASELINE_ERROR:"):
|
||||
parts = line.split(":", 2)
|
||||
if len(parts) >= 3:
|
||||
baseline_errors[parts[1].strip()] = parts[2].strip()
|
||||
continue
|
||||
# Depth-1 impl header: starts with tree drawing chars
|
||||
if line and line[0] in "\u251c\u2514": # ├ └
|
||||
parts = line.split("\u2500", 1)[-1].split() # split on ─
|
||||
if len(parts) >= 2:
|
||||
impl = parts[1]
|
||||
else:
|
||||
impl = None
|
||||
continue
|
||||
# Depth-2 kernel: contains tree drawing chars at deeper indent
|
||||
if impl and ("\u251c\u2500" in line or "\u2514\u2500" in line): # ├─ └─
|
||||
parts = line.split("\u2500", 1)[-1].split()
|
||||
if len(parts) >= 2:
|
||||
name = parts[1]
|
||||
if (
|
||||
"vectorized_elementwise_kernel" in name
|
||||
or "elementwise_kernel_with_index" in name
|
||||
):
|
||||
continue
|
||||
try:
|
||||
t = float(parts[0])
|
||||
results[impl] = max(results.get(impl, 0), t)
|
||||
except ValueError:
|
||||
pass
|
||||
return results, baseline_errors
|
||||
|
||||
|
||||
class ProtonContext:
|
||||
"""Context manager for Proton profiling sessions.
|
||||
|
||||
Always captures proton-viewer output and parses impl times so that
|
||||
get_impl_times() / get_baseline_errors() work after exiting the context.
|
||||
|
||||
The proton tree is printed to **stdout** by default (visible on screen
|
||||
when running kernels interactively). When the environment variable
|
||||
``TIRX_BENCH_JSON=1`` is set (done automatically by ``--json`` mode),
|
||||
the tree goes to **stderr** instead so it does not corrupt the JSON on
|
||||
stdout.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name="kernel",
|
||||
hook="triton",
|
||||
debug=False,
|
||||
nsight=False,
|
||||
metric=PROTON_AVG_TIME_METRIC,
|
||||
kernel="",
|
||||
):
|
||||
self.name = name
|
||||
self.hook = hook
|
||||
self.debug = debug
|
||||
self.nsight = nsight
|
||||
self.metric = metric
|
||||
self.kernel = kernel
|
||||
self._impl_times = {}
|
||||
self._baseline_errors = {}
|
||||
|
||||
def __enter__(self):
|
||||
if not is_running_under_pytest() and not self.debug and not self.nsight:
|
||||
proton.start(self.name, hook=self.hook)
|
||||
proton.deactivate()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
if not is_running_under_pytest() and not self.debug and not self.nsight:
|
||||
proton.finalize()
|
||||
|
||||
hatchet = f"{self.name}.hatchet"
|
||||
result = subprocess.run(
|
||||
["proton-viewer", "-m", self.metric, hatchet],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
self._impl_times, self._baseline_errors = _parse_proton_tree(
|
||||
result.stdout, kernel=self.kernel
|
||||
)
|
||||
out = sys.stderr if os.environ.get("TIRX_BENCH_JSON") else sys.stdout
|
||||
print(f"# proton {PROTON_AVG_TIME_METRIC} (microseconds)\n", file=out, end="")
|
||||
print(result.stdout, file=out, end="")
|
||||
else:
|
||||
print(
|
||||
f"proton-viewer failed (rc={result.returncode}): {result.stderr}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
if os.path.exists(hatchet):
|
||||
os.remove(hatchet)
|
||||
|
||||
def get_impl_times(self):
|
||||
"""Return {impl_name: avg_time_us} parsed from proton-viewer output."""
|
||||
return dict(self._impl_times)
|
||||
|
||||
def get_baseline_errors(self):
|
||||
"""Return {impl_name: error_message} from BASELINE_ERROR lines."""
|
||||
return dict(self._baseline_errors)
|
||||
|
||||
|
||||
def _get_l2_cache_bytes():
|
||||
"""Query L2 cache size from the current CUDA device, fallback to 128MB."""
|
||||
try:
|
||||
props = torch.cuda.get_device_properties(torch.cuda.current_device())
|
||||
if hasattr(props, "l2_cache_size") and props.l2_cache_size > 0:
|
||||
return props.l2_cache_size
|
||||
except Exception:
|
||||
pass
|
||||
return 128 * 1024 * 1024 # 128MB default (B200)
|
||||
|
||||
|
||||
def _tensor_bytes(args, _seen=None):
|
||||
"""Sum the byte size of all torch/tvm tensors in a nested value."""
|
||||
if _seen is None:
|
||||
_seen = set()
|
||||
total = 0
|
||||
if isinstance(args, list | tuple):
|
||||
for a in args:
|
||||
total += _tensor_bytes(a, _seen)
|
||||
elif isinstance(args, Mapping):
|
||||
for a in args.values():
|
||||
total += _tensor_bytes(a, _seen)
|
||||
elif isinstance(args, torch.Tensor):
|
||||
key = ("torch", args.device.type, args.device.index, int(args.data_ptr()))
|
||||
if key not in _seen:
|
||||
_seen.add(key)
|
||||
total += args.nelement() * args.element_size()
|
||||
elif hasattr(args, "numpy"): # tvm.runtime.NDArray
|
||||
try:
|
||||
key = ("tvm", int(args.handle.value))
|
||||
except Exception:
|
||||
key = ("tvm", id(args))
|
||||
if key not in _seen:
|
||||
_seen.add(key)
|
||||
try:
|
||||
total += int(np.prod(args.shape)) * np.dtype(str(args.dtype)).itemsize
|
||||
except Exception:
|
||||
total += args.numpy().nbytes
|
||||
return total
|
||||
|
||||
|
||||
def tensor_bytes(*values):
|
||||
"""Return unique torch/tvm tensor bytes for kernel-owned byte accounting.
|
||||
|
||||
The benchmark driver does not use this implicitly. Kernel benchmark
|
||||
factories may call it when their invocation footprint is exactly the set of
|
||||
tensors in ``values``.
|
||||
"""
|
||||
if len(values) == 1:
|
||||
return _tensor_bytes(values[0])
|
||||
return _tensor_bytes(values)
|
||||
|
||||
|
||||
def _compute_group_count(input_bytes, l2_bytes=None):
|
||||
"""Return TK-style input-group count from one invocation's byte footprint."""
|
||||
if input_bytes <= 0:
|
||||
return 1
|
||||
if l2_bytes is None:
|
||||
l2_bytes = _get_l2_cache_bytes()
|
||||
threshold = l2_bytes * 3
|
||||
if input_bytes >= threshold:
|
||||
return 1
|
||||
return int(threshold // input_bytes) + 1
|
||||
|
||||
|
||||
def _make_bench_input(input_factory):
|
||||
value = input_factory()
|
||||
if not isinstance(value, tuple) or len(value) != 2:
|
||||
raise TypeError("input_factory must return (case, input_bytes)")
|
||||
|
||||
case, input_bytes = value
|
||||
try:
|
||||
input_bytes = int(input_bytes)
|
||||
except (TypeError, ValueError) as err:
|
||||
raise TypeError("input_factory input_bytes must be an integer") from err
|
||||
if input_bytes < 0:
|
||||
raise ValueError("input_factory input_bytes must be non-negative")
|
||||
return case, input_bytes
|
||||
|
||||
|
||||
def prepare_input_groups(input_factory, l2_bytes=None):
|
||||
"""Materialize TK-style input groups from a single-group factory.
|
||||
|
||||
``input_factory`` must return ``(case, input_bytes)``. ``case`` is passed
|
||||
back to every benchmark function unchanged. ``input_bytes`` defines one
|
||||
invocation's L2-eviction footprint and is intentionally owned by the kernel
|
||||
benchmark harness instead of inferred here.
|
||||
"""
|
||||
if not callable(input_factory):
|
||||
raise TypeError("input_factory must be callable")
|
||||
if l2_bytes is None:
|
||||
l2_bytes = _get_l2_cache_bytes()
|
||||
|
||||
sample, input_bytes = _make_bench_input(input_factory)
|
||||
num_groups = _compute_group_count(input_bytes, l2_bytes)
|
||||
groups = [sample]
|
||||
for _ in range(num_groups - 1):
|
||||
case, _ = _make_bench_input(input_factory)
|
||||
groups.append(case)
|
||||
|
||||
return groups, {
|
||||
"num_groups": num_groups,
|
||||
"input_bytes": input_bytes,
|
||||
"l2_bytes": l2_bytes,
|
||||
"l2_eviction_factor": 3,
|
||||
"flush_l2": False,
|
||||
}
|
||||
|
||||
|
||||
def _bench_event_groups(funcs, groups, warmup, repeat, cooldown_s):
|
||||
num_groups = len(groups)
|
||||
results = {}
|
||||
|
||||
for idx, (name, func) in enumerate(funcs.items()):
|
||||
if idx > 0:
|
||||
time.sleep(cooldown_s)
|
||||
|
||||
for i in range(warmup):
|
||||
func(groups[i % num_groups])
|
||||
|
||||
start_event = torch.cuda.Event(enable_timing=True)
|
||||
end_event = torch.cuda.Event(enable_timing=True)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
start_event.record()
|
||||
for i in range(repeat):
|
||||
func(groups[i % num_groups])
|
||||
end_event.record()
|
||||
|
||||
torch.cuda.synchronize()
|
||||
results[name] = start_event.elapsed_time(end_event) / repeat * 1000.0
|
||||
|
||||
time.sleep(cooldown_s)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def _bench_proton_groups(
|
||||
funcs, groups, warmup, repeat, cooldown_s, proton_name, debug, nsight, *, kernel=""
|
||||
):
|
||||
num_groups = len(groups)
|
||||
with ProtonContext(proton_name, debug=debug, nsight=nsight, kernel=kernel) as ctx:
|
||||
for idx, (name, func) in enumerate(funcs.items()):
|
||||
if idx > 0:
|
||||
time.sleep(cooldown_s)
|
||||
|
||||
for i in range(warmup):
|
||||
func(groups[i % num_groups])
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if not is_running_under_pytest() and not debug and not nsight:
|
||||
proton.activate()
|
||||
with proton.scope(name, metrics={}):
|
||||
for i in range(repeat):
|
||||
func(groups[i % num_groups])
|
||||
proton.deactivate()
|
||||
else:
|
||||
for i in range(repeat):
|
||||
func(groups[i % num_groups])
|
||||
torch.cuda.synchronize()
|
||||
|
||||
time.sleep(cooldown_s)
|
||||
|
||||
return ctx.get_impl_times(), ctx.get_baseline_errors()
|
||||
|
||||
|
||||
def _flush_l2_legacy(flush_l2_size):
|
||||
if flush_l2_size > 0:
|
||||
torch.empty(flush_l2_size, dtype=torch.int, device="cuda").zero_()
|
||||
|
||||
|
||||
def _bench_legacy_callable(func, warmup, repeat, proton_name, debug, nsight, flush_l2_size):
|
||||
start_event = torch.cuda.Event(enable_timing=True)
|
||||
end_event = torch.cuda.Event(enable_timing=True)
|
||||
|
||||
def timed_loop():
|
||||
start_event.record()
|
||||
for _ in range(repeat):
|
||||
_flush_l2_legacy(flush_l2_size)
|
||||
func()
|
||||
end_event.record()
|
||||
|
||||
for _ in range(warmup):
|
||||
_flush_l2_legacy(flush_l2_size)
|
||||
func()
|
||||
torch.cuda.synchronize()
|
||||
if not is_running_under_pytest() and not debug and not nsight:
|
||||
proton.activate()
|
||||
with proton.scope(proton_name, metrics={}):
|
||||
timed_loop()
|
||||
proton.deactivate()
|
||||
else:
|
||||
timed_loop()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
return start_event.elapsed_time(end_event) / repeat * 1000.0
|
||||
|
||||
|
||||
# Labels identifying our own kernel (vs external reference impls). Must match
|
||||
# OUR_IMPLS in bench_suite's ratio_diff.py. Used by the TIRX_BENCH_IMPLS filter.
|
||||
OURS_IMPLS = frozenset({"tir", "tirx"})
|
||||
|
||||
|
||||
def bench_impls_mode():
|
||||
"""Current impl-selection mode: 'all' (default), 'ours', or 'baseline'.
|
||||
|
||||
Set via the ``TIRX_BENCH_IMPLS`` env var (by bench_suite ``run.py --impls``).
|
||||
A kernel's ``run_bench`` can use this to skip *building/warming* reference
|
||||
impls (e.g. flashinfer autotune, deepgemm/cublas ext setup) that ``bench``'s
|
||||
filter alone cannot avoid, since that setup runs before ``bench`` is called.
|
||||
"""
|
||||
mode = os.environ.get("TIRX_BENCH_IMPLS", "all").lower()
|
||||
if mode not in {"all", "ours", "baseline"}:
|
||||
raise ValueError(f"TIRX_BENCH_IMPLS must be 'all', 'ours', or 'baseline', got {mode!r}")
|
||||
return mode
|
||||
|
||||
|
||||
def bench_only_ours():
|
||||
"""True when only our own kernel should be benched (reference setup skippable)."""
|
||||
return bench_impls_mode() == "ours"
|
||||
|
||||
|
||||
def filter_impls(funcs):
|
||||
"""Filter a ``{label: callable}`` impl map per the current ``bench_impls_mode``.
|
||||
|
||||
Call this right after building ``funcs`` so any subsequent
|
||||
``if "<ref>" in funcs:`` reference-setup blocks are skipped in 'ours' mode.
|
||||
"""
|
||||
mode = bench_impls_mode()
|
||||
if mode == "ours":
|
||||
return {n: f for n, f in funcs.items() if n in OURS_IMPLS}
|
||||
if mode == "baseline":
|
||||
return {n: f for n, f in funcs.items() if n not in OURS_IMPLS}
|
||||
return funcs
|
||||
|
||||
|
||||
def bench(
|
||||
funcs,
|
||||
input_factory=None,
|
||||
warmup=500,
|
||||
repeat=100,
|
||||
cooldown_s=1.0,
|
||||
timer="proton",
|
||||
proton_name="kernel",
|
||||
l2_bytes=None,
|
||||
debug=False,
|
||||
nsight=False,
|
||||
flush_l2_size=int(8e8 // 4),
|
||||
references=None,
|
||||
rounds=1,
|
||||
round_cooldown_s=1.0,
|
||||
validate_case=None,
|
||||
):
|
||||
"""Benchmark implementations with a factory-owned input footprint.
|
||||
|
||||
This is the single TIRx benchmark API. It follows the ThunderKittens-style
|
||||
multi-input protocol for L2 eviction and supports either Proton/CUPTI or
|
||||
CUDA-event timing. The benchmark driver never infers which tensors belong
|
||||
to a workload; ``input_factory`` owns that definition by returning
|
||||
``(case, input_bytes)``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
funcs : dict[str, callable]
|
||||
Map of implementation name to callable. Each callable receives one
|
||||
``case`` returned by ``input_factory``. This should hold only *our*
|
||||
kernel(s); external baselines go in ``references``.
|
||||
references : dict[str, callable], optional
|
||||
Map of reference-impl name to a no-arg *builder* that does the heavy
|
||||
import/setup and returns the run callable. Builders run lazily and only
|
||||
when that impl will actually be benched (skipped entirely under
|
||||
``--impls ours``); a builder that raises is recorded as a
|
||||
``BASELINE_ERROR`` instead of failing the workload.
|
||||
input_factory : callable
|
||||
Factory returning ``(case, input_bytes)`` for one benchmark group.
|
||||
warmup : int
|
||||
Number of untimed warmup iterations per implementation.
|
||||
repeat : int
|
||||
Number of timed iterations per round.
|
||||
cooldown_s : float
|
||||
Seconds to sleep between impls for thermal cooldown.
|
||||
rounds : int
|
||||
Independent benchmark rounds (compile + inputs once; each round runs
|
||||
warmup + repeat for every selected impl).
|
||||
round_cooldown_s : float
|
||||
Seconds to sleep between rounds (ignored when ``rounds == 1``).
|
||||
validate_case : callable, optional
|
||||
Called once on the first prepared ``case`` (after ``prepare_input_groups``,
|
||||
before warmup/repeat rounds). Under bench_suite, ``run_kernel_bench`` holds
|
||||
the per-GPU lock for the whole ``run_bench()`` call.
|
||||
timer : {"event", "proton"}
|
||||
Timing backend.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
``{"impls": {name: us}, "round_samples": {name: [us, ...]}, ...}``.
|
||||
Times are stored in microseconds (same unit as pinned bench_suite baselines).
|
||||
"""
|
||||
if repeat <= 0:
|
||||
raise ValueError("repeat must be positive")
|
||||
if warmup < 0:
|
||||
raise ValueError("warmup must be non-negative")
|
||||
if rounds < 1:
|
||||
raise ValueError("rounds must be >= 1")
|
||||
if round_cooldown_s < 0:
|
||||
raise ValueError("round_cooldown_s must be non-negative")
|
||||
if timer not in {"event", "proton"}:
|
||||
raise ValueError(f"unsupported timer {timer!r}; expected event or proton")
|
||||
|
||||
if callable(funcs) and input_factory is None:
|
||||
return _bench_legacy_callable(
|
||||
funcs,
|
||||
warmup=warmup,
|
||||
repeat=repeat,
|
||||
proton_name=proton_name,
|
||||
debug=debug,
|
||||
nsight=nsight,
|
||||
flush_l2_size=flush_l2_size,
|
||||
)
|
||||
|
||||
if input_factory is None:
|
||||
raise TypeError("input_factory is required when funcs is a mapping")
|
||||
if not isinstance(funcs, Mapping) or not funcs:
|
||||
raise TypeError("funcs must be a non-empty mapping of name to callable")
|
||||
for name, func in funcs.items():
|
||||
if not isinstance(name, str):
|
||||
raise TypeError("func names must be strings")
|
||||
if not callable(func):
|
||||
raise TypeError(f"funcs[{name!r}] must be callable")
|
||||
|
||||
# Select impls for this run. ``funcs`` holds our own kernel(s); external
|
||||
# baselines are passed as ``references`` (name -> no-arg builder). A builder
|
||||
# is invoked here ONLY when its impl will actually be benched, so --impls
|
||||
# ours skips all reference setup, --impls baseline skips ours, and a builder
|
||||
# that fails is recorded as a BASELINE_ERROR rather than failing the
|
||||
# workload. Legacy kernels that put references directly in ``funcs`` are
|
||||
# still handled by the label filter (filter_impls) below.
|
||||
funcs = filter_impls(funcs)
|
||||
build_errors: dict[str, str] = {}
|
||||
if bench_impls_mode() != "ours":
|
||||
for ref_name, builder in (references or {}).items():
|
||||
if not isinstance(ref_name, str) or not callable(builder):
|
||||
raise TypeError("references must map a name to a no-arg builder callable")
|
||||
try:
|
||||
ref_fn = builder()
|
||||
except Exception as e:
|
||||
build_errors[ref_name] = str(e)
|
||||
print(f"BASELINE_ERROR: {ref_name}: {e}", file=sys.stderr)
|
||||
continue
|
||||
if ref_fn is None:
|
||||
continue
|
||||
if not callable(ref_fn):
|
||||
raise TypeError(f"references[{ref_name!r}] builder must return a callable")
|
||||
funcs = {**funcs, ref_name: ref_fn}
|
||||
|
||||
if not funcs:
|
||||
# Nothing to bench in this mode (e.g. 'ours' on a kernel with no tir
|
||||
# impl, or 'baseline' with no references). Return an empty-but-valid
|
||||
# result so the workload is a no-op.
|
||||
return {
|
||||
"impls": {},
|
||||
"round_samples": {},
|
||||
"errors": build_errors,
|
||||
"timer": timer,
|
||||
"benchmark_protocol": {
|
||||
"warmup": warmup,
|
||||
"repeat": repeat,
|
||||
"cooldown_s": cooldown_s,
|
||||
"rounds": rounds,
|
||||
"round_cooldown_s": round_cooldown_s,
|
||||
"order": [],
|
||||
},
|
||||
}
|
||||
|
||||
inputs, protocol = prepare_input_groups(input_factory, l2_bytes=l2_bytes)
|
||||
num_groups = len(inputs)
|
||||
if num_groups == 0:
|
||||
return {
|
||||
"impls": {},
|
||||
"round_samples": {},
|
||||
"errors": build_errors,
|
||||
"timer": timer,
|
||||
"benchmark_protocol": {
|
||||
**protocol,
|
||||
"warmup": warmup,
|
||||
"repeat": repeat,
|
||||
"cooldown_s": cooldown_s,
|
||||
"rounds": rounds,
|
||||
"round_cooldown_s": round_cooldown_s,
|
||||
"order": list(funcs.keys()),
|
||||
},
|
||||
}
|
||||
|
||||
if validate_case is not None:
|
||||
validate_case(inputs[0])
|
||||
|
||||
errors = dict(build_errors)
|
||||
round_samples: dict[str, list[float]] = {}
|
||||
for round_idx in range(rounds):
|
||||
if round_idx > 0:
|
||||
time.sleep(round_cooldown_s)
|
||||
if timer == "event":
|
||||
impls = _bench_event_groups(funcs, inputs, warmup, repeat, cooldown_s)
|
||||
proton_errors = {}
|
||||
else:
|
||||
impls, proton_errors = _bench_proton_groups(
|
||||
funcs,
|
||||
inputs,
|
||||
warmup,
|
||||
repeat,
|
||||
cooldown_s,
|
||||
proton_name,
|
||||
debug,
|
||||
nsight,
|
||||
kernel=proton_name,
|
||||
)
|
||||
errors.update(proton_errors)
|
||||
for impl, sec in impls.items():
|
||||
round_samples.setdefault(impl, []).append(sec)
|
||||
|
||||
if not round_samples:
|
||||
aggregated = {}
|
||||
else:
|
||||
import statistics
|
||||
|
||||
aggregated = {impl: statistics.mean(samples) for impl, samples in round_samples.items()}
|
||||
|
||||
return {
|
||||
"impls": aggregated,
|
||||
"round_samples": round_samples,
|
||||
"errors": errors,
|
||||
"timer": timer,
|
||||
"benchmark_protocol": {
|
||||
**protocol,
|
||||
"warmup": warmup,
|
||||
"repeat": repeat,
|
||||
"cooldown_s": cooldown_s,
|
||||
"rounds": rounds,
|
||||
"round_cooldown_s": round_cooldown_s,
|
||||
"order": list(funcs.keys()),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# utils for tg4perfetto profiler, adapted from https://github.com/flashinfer-ai/flashinfer
|
||||
|
||||
|
||||
class EventType(Enum):
|
||||
kBegin = 0
|
||||
kEnd = 1
|
||||
kInstant = 2
|
||||
kFinalize = 3
|
||||
|
||||
|
||||
def decode_tag(tag, num_groups):
|
||||
block_group_tag = tag >> 12
|
||||
event_idx = (tag >> 2) & 0x3FF
|
||||
event_type = tag & 0x3
|
||||
return (block_group_tag // num_groups, block_group_tag % num_groups, event_idx, event_type)
|
||||
|
||||
|
||||
def export_to_perfetto_trace(
|
||||
profiler_buffer: np.ndarray, file_name: str, event_type_names: list[str]
|
||||
) -> None:
|
||||
if is_running_under_pytest():
|
||||
return
|
||||
|
||||
import torch
|
||||
|
||||
# pip install git+https://github.com/ihavnoid/tg4perfetto.git
|
||||
from tg4perfetto import TraceGenerator
|
||||
|
||||
profiler_buffer_host = torch.tensor(profiler_buffer)
|
||||
num_blocks, num_groups = profiler_buffer_host[:1].view(dtype=torch.int32)
|
||||
num_blocks = int(num_blocks)
|
||||
num_groups = int(num_groups)
|
||||
tgen = TraceGenerator(file_name)
|
||||
|
||||
tid_map = {}
|
||||
track_map = {}
|
||||
finish_idx = set()
|
||||
for block_idx in range(num_blocks):
|
||||
pid = tgen.create_group(f"block_{block_idx}")
|
||||
for group_idx in range(num_groups):
|
||||
tid = pid.create_group(f"group_{group_idx}")
|
||||
tid_map[(block_idx, group_idx)] = tid
|
||||
|
||||
for i in range(1, len(profiler_buffer_host)):
|
||||
if profiler_buffer_host[i] == 0:
|
||||
continue
|
||||
tag, timestamp = profiler_buffer_host[i : i + 1].view(dtype=torch.uint32)
|
||||
tag = int(tag)
|
||||
timestamp = int(timestamp)
|
||||
block_idx, group_idx, event_idx, event_type = decode_tag(tag, num_groups)
|
||||
|
||||
if event_type == EventType.kFinalize.value:
|
||||
finish_idx.add((block_idx, group_idx))
|
||||
if len(finish_idx) == num_blocks * num_groups:
|
||||
break
|
||||
else:
|
||||
if (block_idx, group_idx) in finish_idx:
|
||||
continue
|
||||
|
||||
event = event_type_names[event_idx]
|
||||
tid = tid_map[(block_idx, group_idx)]
|
||||
|
||||
if (block_idx, group_idx, event_idx) in track_map:
|
||||
track = track_map[(block_idx, group_idx, event_idx)]
|
||||
else:
|
||||
track = tid.create_track()
|
||||
track_map[(block_idx, group_idx, event_idx)] = track
|
||||
|
||||
if event_type == EventType.kBegin.value:
|
||||
track.open(timestamp, event)
|
||||
elif event_type == EventType.kEnd.value:
|
||||
track.close(timestamp)
|
||||
elif event_type == EventType.kInstant.value:
|
||||
track.instant(timestamp, event)
|
||||
|
||||
tgen.flush()
|
||||
|
||||
|
||||
@T.meta_class
|
||||
class CudaProfiler:
|
||||
"""A lightweight wrapper around T.timer_* CUDA intrinsics.
|
||||
|
||||
Stores repeated arguments used by timer_init/start/end/finalize so users can
|
||||
call concise methods in kernels. Intended to mirror Pipeline/TileScheduler helpers.
|
||||
|
||||
When ``profiler_enabled`` is False (or a false-y Expr), calls to
|
||||
``init/start/end/finalize`` become no-ops. This allows constructing a
|
||||
profiler unconditionally and eliminating external ``if PROFILER_ON:`` guards.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
profiler_buffer: T.Buffer,
|
||||
write_stride: int,
|
||||
num_groups: int,
|
||||
default_leader: None | tvm.tirx.Expr | bool = None,
|
||||
profiler_enabled: bool | tvm.tirx.Expr = True,
|
||||
):
|
||||
self.buffer = profiler_buffer
|
||||
self.write_stride = write_stride
|
||||
self.num_groups = num_groups
|
||||
self.default_leader = default_leader
|
||||
# Accept either a Python bool or a Expr; normalize simple bools to T.bool
|
||||
# so we can use it uniformly inside macros for conditional emission.
|
||||
if isinstance(profiler_enabled, bool | np.bool_):
|
||||
self.profiler_enabled = T.bool(bool(profiler_enabled))
|
||||
else:
|
||||
# Assume Expr-like input; use as-is
|
||||
self.profiler_enabled = profiler_enabled # type: ignore[assignment]
|
||||
|
||||
self.profiler_tag = T.alloc_buffer([1], "uint64", scope="local", align=8)
|
||||
self.profiler_write_offset = T.alloc_buffer([1], "uint32", scope="local", align=8)
|
||||
|
||||
def _leader(self, leader: None | tvm.tirx.Expr | bool):
|
||||
if leader is not None:
|
||||
if isinstance(leader, bool | np.bool_):
|
||||
return T.bool(bool(leader))
|
||||
return leader
|
||||
if self.default_leader is not None:
|
||||
return self.default_leader
|
||||
return T.bool(True)
|
||||
|
||||
@T.inline
|
||||
def init(self, group_id: tvm.tirx.Expr):
|
||||
if self.profiler_enabled:
|
||||
T.cuda.timer_init(
|
||||
self.buffer.data,
|
||||
self.profiler_tag.data,
|
||||
self.profiler_write_offset.data,
|
||||
self.num_groups,
|
||||
group_id,
|
||||
)
|
||||
|
||||
@T.inline
|
||||
def start(self, event_type: Enum, leader: None | tvm.tirx.Expr | bool = None):
|
||||
if self.profiler_enabled:
|
||||
T.cuda.timer_start(
|
||||
event_type,
|
||||
self.buffer.data,
|
||||
self.profiler_tag.data,
|
||||
self.profiler_write_offset.data,
|
||||
self.write_stride,
|
||||
self._leader(leader),
|
||||
)
|
||||
|
||||
@T.inline
|
||||
def end(self, event_type: Enum, leader: None | tvm.tirx.Expr | bool = None):
|
||||
if self.profiler_enabled:
|
||||
T.cuda.timer_end(
|
||||
event_type,
|
||||
self.buffer.data,
|
||||
self.profiler_tag.data,
|
||||
self.profiler_write_offset.data,
|
||||
self.write_stride,
|
||||
self._leader(leader),
|
||||
)
|
||||
|
||||
@T.inline
|
||||
def finalize(self, leader: None | tvm.tirx.Expr | bool = None):
|
||||
if self.profiler_enabled:
|
||||
T.cuda.timer_finalize(
|
||||
self.buffer.data,
|
||||
self.profiler_tag.data,
|
||||
self.profiler_write_offset.data,
|
||||
self.write_stride,
|
||||
self._leader(leader),
|
||||
)
|
||||
@@ -0,0 +1,577 @@
|
||||
# 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.
|
||||
"""Abstraction for array data structures."""
|
||||
|
||||
import functools
|
||||
from numbers import Integral
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
import tvm
|
||||
from tvm.ir import PointerType, PrimType, Range
|
||||
from tvm.runtime import Object, Scriptable, convert
|
||||
|
||||
from . import _ffi_api
|
||||
|
||||
|
||||
@tvm_ffi.register_object("tirx.Buffer")
|
||||
class Buffer(Object, Scriptable):
|
||||
"""Symbolic data buffer in TVM.
|
||||
|
||||
Buffer provide a way to represent data layout
|
||||
specialization of data structure in TVM.
|
||||
|
||||
Do not construct directly, use :py:func:`~decl_buffer` instead.
|
||||
See the documentation of :py:func:`decl_buffer` for more details.
|
||||
|
||||
See Also
|
||||
--------
|
||||
decl_buffer : Declare a buffer
|
||||
"""
|
||||
|
||||
READ = 1
|
||||
WRITE = 2
|
||||
|
||||
def access_ptr(self, access_mask, ptr_type="handle", content_lanes=1, offset=0, extent=None):
|
||||
"""Get an access pointer to the head of buffer.
|
||||
|
||||
This is the recommended method to get buffer data
|
||||
ptress when interacting with external functions.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
access_mask : int
|
||||
The access pattern MASK. Indicate whether the
|
||||
access will read or write to the data content.
|
||||
|
||||
ptr_type : str or tvm.ir.Type, optional
|
||||
The data type of the result pointer. Do not specify
|
||||
unless we want to cast pointer to specific type.
|
||||
|
||||
content_lanes: int, optional
|
||||
The number of lanes for the data type. This value
|
||||
is greater than one for vector types.
|
||||
|
||||
offset: Expr, optional
|
||||
The offset of pointer. We can use it to offset by
|
||||
the number of elements from the address of ptr.
|
||||
|
||||
extent: Expr, optional
|
||||
The extent of pointer.
|
||||
|
||||
Examples
|
||||
--------
|
||||
.. code-block:: python
|
||||
|
||||
# Get access ptr for read
|
||||
buffer.access_ptr("r")
|
||||
# Get access ptr for read/write with bitmask
|
||||
buffer.access_ptr(Buffer.READ | Buffer.WRITE)
|
||||
# Get access ptr for read/write with str flag
|
||||
buffer.access_ptr("rw")
|
||||
# Get access ptr for read with offset
|
||||
buffer.access_ptr("r", offset = 100)
|
||||
# Get access ptr for read with extent
|
||||
buffer.access_ptr("r", extent = 100)
|
||||
"""
|
||||
if isinstance(access_mask, str):
|
||||
mask = 0
|
||||
for value in access_mask:
|
||||
if value == "r":
|
||||
mask = mask | Buffer.READ
|
||||
elif value == "w":
|
||||
mask = mask | Buffer.WRITE
|
||||
else:
|
||||
raise ValueError(f"Unknown access_mask {access_mask}")
|
||||
access_mask = mask
|
||||
if isinstance(ptr_type, str):
|
||||
ptr_type = (
|
||||
PointerType(PrimType("void"))
|
||||
if ptr_type == "handle"
|
||||
else PointerType(PrimType(ptr_type))
|
||||
)
|
||||
elif isinstance(ptr_type, PrimType):
|
||||
ptr_type = PointerType(ptr_type)
|
||||
offset = convert(offset)
|
||||
extent = convert(extent)
|
||||
return _ffi_api.BufferAccessPtr(
|
||||
self,
|
||||
access_mask,
|
||||
ptr_type,
|
||||
content_lanes,
|
||||
offset,
|
||||
extent, # type: ignore
|
||||
)
|
||||
|
||||
def vload(self, begin, dtype=None, predicate=None):
|
||||
"""Generate an Expr that loads dtype from begin index.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
begin : Array of Expr
|
||||
The beginning index in unit of Buffer.dtype
|
||||
|
||||
dtype : str
|
||||
The data type to be loaded,
|
||||
can be vector type which have lanes that is multiple of Buffer.dtype
|
||||
|
||||
predicate : Optional[Expr]
|
||||
A vector mask of boolean values indicating which lanes of a vector are to be
|
||||
loaded. The number lanes of the mask must be equal to the number of lanes being loaded.
|
||||
|
||||
Returns
|
||||
-------
|
||||
load : Expr
|
||||
The corresponding load expression.
|
||||
"""
|
||||
begin = (begin,) if isinstance(begin, int) or tvm.ir.is_prim_expr(begin) else begin
|
||||
dtype = dtype if dtype else self.dtype
|
||||
return _ffi_api.BufferVLoad(self, begin, dtype, predicate) # type: ignore
|
||||
|
||||
def vstore(self, begin, value, predicate=None):
|
||||
"""Generate a Stmt that store value into begin index.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
begin : Array of Expr
|
||||
The beginning index in unit of Buffer.dtype
|
||||
|
||||
value : Expr
|
||||
The value to be stored.
|
||||
|
||||
predicate : Optional[Expr]
|
||||
A vector mask of boolean values indicating which lanes of a vector are to be
|
||||
stored. The number lanes of the mask must be equal to the number of lanes in
|
||||
value.
|
||||
|
||||
Returns
|
||||
-------
|
||||
store : Stmt
|
||||
The corresponding store stmt.
|
||||
"""
|
||||
begin = (begin,) if isinstance(begin, int) or tvm.ir.is_prim_expr(begin) else begin
|
||||
return _ffi_api.BufferVStore(self, begin, value, predicate) # type: ignore
|
||||
|
||||
def scope(self):
|
||||
"""Return the storage scope associated with this buffer.
|
||||
Returns
|
||||
-------
|
||||
scope : str
|
||||
The storage scope associated with this buffer.
|
||||
"""
|
||||
return _ffi_api.BufferStorageScope(self) # type: ignore
|
||||
|
||||
def get_flattened_buffer(self):
|
||||
"""Generate a Buffer that is a flattened version of this buffer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
flattened : Buffer
|
||||
The corresponding flat buffer.
|
||||
"""
|
||||
return _ffi_api.BufferGetFlattenedBuffer(self) # type: ignore
|
||||
|
||||
def with_allocated_addr(self, allocated_addr):
|
||||
"""Return a new buffer with the allocated address."""
|
||||
return _ffi_api.BufferWithAllocatedAddr(self, allocated_addr) # type: ignore
|
||||
|
||||
def with_dtype(self, dtype):
|
||||
"""Return a new buffer with the dtype."""
|
||||
return _ffi_api.BufferWithDtype(self, dtype) # type: ignore
|
||||
|
||||
def with_data(self, data):
|
||||
"""Return a new buffer with the data."""
|
||||
return _ffi_api.BufferWithData(self, data) # type: ignore
|
||||
|
||||
def offset_of(self, indices):
|
||||
"""Determine the offset of the provided indices in the flattened buffer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
indices : Union[Expr, List[Expr]]
|
||||
|
||||
The indices of the element in the original buffer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
flattened_indices: List[Expr]
|
||||
|
||||
The offset indices of the element in the flattened buffer.
|
||||
"""
|
||||
return _ffi_api.BufferOffsetOf(self, indices) # type: ignore
|
||||
|
||||
@property
|
||||
def byte_offset(self):
|
||||
"""Get the byte offset of the buffer."""
|
||||
return self.elem_offset * tvm.DataType(self.dtype).bits // 8
|
||||
|
||||
def elem_offset_of(self, indices, inner=True):
|
||||
"""Get the element offset of the buffer at the given indices.
|
||||
Note that indices subject to buffer's layout mapping.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
indices : Union[Expr, List[Expr]]
|
||||
The indices of the element in the original buffer.
|
||||
|
||||
inner : bool, optional
|
||||
If False, the offset is relative to the original buffer.
|
||||
Default is True.
|
||||
|
||||
Returns
|
||||
-------
|
||||
offset: Expr
|
||||
The element offset of the buffer at the given indices.
|
||||
"""
|
||||
if inner:
|
||||
return _ffi_api.BufferOffsetOfp(self, indices)
|
||||
return self.elem_offset + _ffi_api.BufferOffsetOfp(self, indices)
|
||||
|
||||
def byte_offset_of(self, indices, inner=True):
|
||||
"""Get the byte offset of the buffer at the given indices.
|
||||
Note that indices subject to buffer's layout mapping.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
indices : Union[Expr, List[Expr]]
|
||||
The indices of the element in the original buffer.
|
||||
|
||||
inner : bool, optional
|
||||
If False, the offset is relative to the original buffer.
|
||||
Default is True.
|
||||
|
||||
Returns
|
||||
-------
|
||||
offset: Expr
|
||||
The byte offset of the buffer at the given indices.
|
||||
"""
|
||||
return self.elem_offset_of(indices, inner) * tvm.DataType(self.dtype).bits // 8
|
||||
|
||||
def is_scalar(self, alloc_or_decl=True):
|
||||
"""Check if the buffer is a scalar.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
alloc_or_decl : bool, optional
|
||||
Whether to consider alloc_scalar and decl_scalar as scalar. True for alloc_scalar,
|
||||
False for decl_scalar.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bool: True if the buffer is a scalar, False otherwise.
|
||||
"""
|
||||
return _ffi_api.BufferIsScalar(self, alloc_or_decl)
|
||||
|
||||
def ptr_to(self, indices):
|
||||
"""Get the pointer to the buffer at the given indices (logical indices).
|
||||
|
||||
Note that the bufferload inside requires LowerTIPp pass to apply the layout to get the physical indices.
|
||||
""" # noqa: E501
|
||||
assert len(indices) == len(self.shape), (
|
||||
f"The number of indices {indices} does not match the shape of the buffer {self.shape}"
|
||||
)
|
||||
return tvm.tirx.address_of(self[tuple(indices)])
|
||||
|
||||
def view(self, *args, **kwargs) -> "Buffer":
|
||||
"""Creates a new view of the buffer. (used by parser)
|
||||
|
||||
Supported signatures are ``view(*shape, layout=None)``, where shape can contain
|
||||
``-1`` to indicate that the dimension size is auto-inferred, and
|
||||
``view(dtype: Union[str, tvm.DataType])``.
|
||||
|
||||
Returns
|
||||
-------
|
||||
view : DeclBufferFrame
|
||||
The corresponding view buffer.
|
||||
"""
|
||||
|
||||
def _infer_shape(shape):
|
||||
shape = list(shape)
|
||||
if -1 in shape and shape.count(-1) == 1:
|
||||
size = functools.reduce(lambda x, y: x * y, self.shape)
|
||||
n_size = functools.reduce(lambda x, y: x * y, [s for s in shape if s != -1], 1)
|
||||
shape[shape.index(-1)] = size // n_size
|
||||
else:
|
||||
# Only validate the shape product when both old and new shapes
|
||||
# are fully concrete: a Expr `==` returns an `EQ` node, not
|
||||
# a Python bool, and `assert <Expr>` raises (no __bool__).
|
||||
if all(isinstance(s, int) for s in shape) and all(
|
||||
isinstance(s, int) for s in self.shape
|
||||
):
|
||||
assert functools.reduce(lambda x, y: x * y, shape) == functools.reduce(
|
||||
lambda x, y: x * y, self.shape
|
||||
), (
|
||||
"The shape of the buffer "
|
||||
+ str(self.shape)
|
||||
+ " and the new shape "
|
||||
+ str(shape)
|
||||
+ " are not compatible"
|
||||
)
|
||||
return shape
|
||||
|
||||
if len(args) == 1 and isinstance(args[0], str | tvm.DataType) and not kwargs:
|
||||
cast_dtype = tvm.DataType(args[0])
|
||||
cur_dtype = tvm.DataType(self.dtype)
|
||||
if cast_dtype.bits > cur_dtype.bits:
|
||||
# cast up
|
||||
assert cast_dtype.bits % cur_dtype.bits == 0
|
||||
ratio = cast_dtype.bits // cur_dtype.bits
|
||||
layout = self.layout.pack(ratio)
|
||||
shape = [s for s in self.shape[:-1]] + [self.shape[-1] // ratio]
|
||||
new_elem_offset = self.elem_offset // ratio
|
||||
else:
|
||||
# cast down
|
||||
assert cur_dtype.bits % cast_dtype.bits == 0
|
||||
ratio = cur_dtype.bits // cast_dtype.bits
|
||||
layout = self.layout.unpack(ratio)
|
||||
shape = [s for s in self.shape[:-1]] + [self.shape[-1] * ratio]
|
||||
new_elem_offset = self.elem_offset * ratio
|
||||
return tvm.tirx.script.builder.decl_buffer(
|
||||
shape,
|
||||
cast_dtype,
|
||||
self.data,
|
||||
self.strides,
|
||||
new_elem_offset,
|
||||
None,
|
||||
self.scope(),
|
||||
self.data_alignment,
|
||||
self.offset_factor,
|
||||
"",
|
||||
self.axis_separators,
|
||||
layout,
|
||||
)
|
||||
else:
|
||||
# --- Signature 1: view(*shape, **opts) ---
|
||||
# Check if all positional args are integers/PrimExprs with dtype int32 or int64 (the shape) # noqa: E501
|
||||
shape = args
|
||||
assert all(
|
||||
isinstance(arg, int)
|
||||
or (tvm.ir.is_prim_expr(arg) and arg.ty.dtype in ["int32", "int64"])
|
||||
for arg in shape
|
||||
), "shape must be a list of integers or PrimExprs with dtype int32 or int64"
|
||||
# Safely get optional keyword arguments
|
||||
layout = kwargs.get("layout", None)
|
||||
# Assert there are no other kwargs
|
||||
assert set(kwargs.keys()).issubset({"layout"}), (
|
||||
f"Unsupported kwargs for view: {set(kwargs.keys()) - {'layout'}}"
|
||||
)
|
||||
|
||||
if layout is None:
|
||||
shape = _infer_shape(shape)
|
||||
|
||||
return tvm.tirx.script.builder.decl_buffer(
|
||||
shape,
|
||||
self.dtype,
|
||||
self.data,
|
||||
self.strides,
|
||||
self.elem_offset,
|
||||
None,
|
||||
self.scope(),
|
||||
self.data_alignment,
|
||||
self.offset_factor,
|
||||
"",
|
||||
self.axis_separators,
|
||||
self.layout if layout is None else layout,
|
||||
)
|
||||
|
||||
def local(self, *shape, layout=None) -> "Buffer":
|
||||
"""Create a thread-local view of this buffer.
|
||||
|
||||
When called with no shape arguments, auto-infers a 1D shape from
|
||||
the layout's non-thread component (i.e. ``layout.storage().shard``).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
shape : tuple of Expr
|
||||
The shape of the local view for indexing. If omitted, a 1D
|
||||
shape is computed automatically.
|
||||
|
||||
layout : optional
|
||||
Override layout. If None, uses the storage layout
|
||||
(parent layout with thread axes removed).
|
||||
|
||||
Returns
|
||||
-------
|
||||
local : DeclBufferFrame
|
||||
The corresponding local buffer.
|
||||
"""
|
||||
if not shape:
|
||||
local_layout = self.layout.storage()
|
||||
total = functools.reduce(
|
||||
lambda x, y: x * y, [it.extent for it in local_layout.shard], 1
|
||||
)
|
||||
shape = (total,)
|
||||
return tvm.tirx.script.builder.decl_buffer(
|
||||
shape,
|
||||
self.dtype,
|
||||
self.data,
|
||||
self.strides,
|
||||
self.elem_offset,
|
||||
None,
|
||||
self.scope(),
|
||||
self.data_alignment,
|
||||
self.offset_factor,
|
||||
"",
|
||||
self.axis_separators,
|
||||
self.layout.storage() if layout is None else layout,
|
||||
)
|
||||
|
||||
def permute(self, *dims) -> "Buffer":
|
||||
"""Permute the dimensions of the buffer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dims : tuple of int
|
||||
The permutation of dimensions.
|
||||
|
||||
Returns
|
||||
-------
|
||||
permuted : DeclBufferFrame
|
||||
The buffer with permuted dimensions.
|
||||
"""
|
||||
new_shape = [self.shape[d] for d in dims]
|
||||
# Permute *logical* dims, not the layout's fine-grained shard iters: a
|
||||
# tcgen05/atom layout maps several shard iters to each logical axis, so
|
||||
# group by the current shape first and permute whole groups. ``group``
|
||||
# returns a regrouped layout (degenerate extent-1 iters folded away)
|
||||
# plus seps over *that* layout — permute the regrouped one, not
|
||||
# ``self.layout``. For a simple layout (one shard iter per axis) this
|
||||
# reduces to ``permute_dims(dims)``.
|
||||
grouped, seps = self.layout.group(list(self.shape))
|
||||
new_layout = grouped.permute_by_groups(seps, list(dims))
|
||||
return tvm.tirx.script.builder.decl_buffer(
|
||||
new_shape,
|
||||
self.dtype,
|
||||
self.data,
|
||||
self.strides,
|
||||
self.elem_offset,
|
||||
None,
|
||||
self.scope(),
|
||||
self.data_alignment,
|
||||
self.offset_factor,
|
||||
"",
|
||||
self.axis_separators,
|
||||
new_layout,
|
||||
)
|
||||
|
||||
def __getitem__(self, indices):
|
||||
from ..arith import Analyzer # pylint: disable=import-outside-toplevel
|
||||
from .expr import BufferLoad, Ramp # pylint: disable=import-outside-toplevel
|
||||
from .stmt import BufferRegion # pylint: disable=import-outside-toplevel
|
||||
|
||||
if not isinstance(indices, tuple | list):
|
||||
indices = [indices]
|
||||
has_slice = any(isinstance(i, slice) for i in indices)
|
||||
has_step = any(
|
||||
isinstance(i, slice) and (i.step is not None and i.step != 1) for i in indices
|
||||
)
|
||||
has_implicit_slice = len(indices) < len(self.shape)
|
||||
analyzer = Analyzer()
|
||||
if (has_slice and not has_step) or has_implicit_slice:
|
||||
region = []
|
||||
for i, index in enumerate(indices):
|
||||
if isinstance(index, slice):
|
||||
start = 0 if index.start is None else index.start
|
||||
stop = self.shape[i] if index.stop is None else index.stop
|
||||
region.append(Range.from_min_extent(start, analyzer.simplify(stop - start)))
|
||||
else:
|
||||
region.append(
|
||||
Range.from_min_extent(
|
||||
index,
|
||||
tvm.tirx.expr.IntImm(index.ty, 1) if tvm.ir.is_prim_expr(index) else 1,
|
||||
)
|
||||
)
|
||||
if has_implicit_slice:
|
||||
for i in range(len(indices), len(self.shape)):
|
||||
region.append(Range.from_min_extent(0, self.shape[i]))
|
||||
return BufferRegion(self, region)
|
||||
else:
|
||||
expr_indices = []
|
||||
for i, index in enumerate(indices):
|
||||
if isinstance(index, slice):
|
||||
start = 0 if index.start is None else index.start
|
||||
stop = self.shape[i] if index.stop is None else index.stop
|
||||
step = 1 if index.step is None else index.step
|
||||
# We should ensure the dtype of start is the same with that of step.
|
||||
if tvm.ir.is_prim_expr(start) and isinstance(step, int):
|
||||
step = tvm.tirx.expr.IntImm(start.ty, step)
|
||||
lanes = analyzer.simplify((stop - start + step - 1) // step)
|
||||
if lanes == 1:
|
||||
expr_indices.append(start)
|
||||
else:
|
||||
expr_indices.append(Ramp(start, step, int(lanes)))
|
||||
else:
|
||||
expr_indices.append(index)
|
||||
return BufferLoad(self, expr_indices)
|
||||
|
||||
|
||||
def decl_buffer(
|
||||
shape,
|
||||
dtype=None,
|
||||
name="buffer",
|
||||
data=None,
|
||||
strides=None,
|
||||
elem_offset=None,
|
||||
scope="",
|
||||
data_alignment=-1,
|
||||
offset_factor=0,
|
||||
buffer_type="",
|
||||
axis_separators=None,
|
||||
span=None,
|
||||
layout="default",
|
||||
):
|
||||
# pylint: disable=import-outside-toplevel
|
||||
from .expr import Var
|
||||
from .layout import S, TileLayout
|
||||
|
||||
shape = (shape,) if tvm.ir.is_prim_expr(shape) or isinstance(shape, Integral) else shape
|
||||
dtype = "float32" if dtype is None else dtype
|
||||
strides = () if strides is None else strides
|
||||
|
||||
if axis_separators is None:
|
||||
axis_separators = []
|
||||
|
||||
if layout == "default":
|
||||
layout = TileLayout(S[tuple(shape)]) if shape else None
|
||||
|
||||
if offset_factor != 0 and elem_offset is None:
|
||||
shape_ty = shape[0].ty if shape and tvm.ir.is_prim_expr(shape[0]) else "int32"
|
||||
elem_offset = Var(f"{name}_elem_offset", shape_ty)
|
||||
if data is None:
|
||||
# Bool is represented as uint1 in the IR, but stored as int8
|
||||
storage_type = dtype if isinstance(dtype, PrimType) else PrimType(dtype)
|
||||
storage_type = PrimType("int8") if storage_type.dtype == "bool" else storage_type
|
||||
data = Var(name, PointerType(storage_type, scope), span)
|
||||
return _ffi_api.Buffer( # type: ignore
|
||||
data,
|
||||
dtype,
|
||||
shape,
|
||||
strides,
|
||||
elem_offset,
|
||||
name,
|
||||
data_alignment,
|
||||
offset_factor,
|
||||
buffer_type,
|
||||
axis_separators,
|
||||
span,
|
||||
layout,
|
||||
)
|
||||
|
||||
|
||||
@tvm_ffi.register_object("tirx.DataProducer")
|
||||
class DataProducer(Object):
|
||||
pass
|
||||
@@ -0,0 +1,243 @@
|
||||
# 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=invalid-name
|
||||
"""The build utils in python."""
|
||||
|
||||
import tvm
|
||||
from tvm import ir
|
||||
from tvm.ir.module import IRModule
|
||||
from tvm.target import Target
|
||||
from tvm.tirx import PrimFunc
|
||||
|
||||
|
||||
def split_host_device_mods(mod: IRModule) -> tuple[IRModule, dict[Target, IRModule]]:
|
||||
"""Split an IRModule into host and device modules.
|
||||
|
||||
This function takes an IRModule containing functions with different target attributes
|
||||
and separates them into host (CPU) and device (GPU/accelerator) modules. Functions
|
||||
are categorized based on their target attribute in func_attr.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mod : tvm.IRModule
|
||||
The input module to split.
|
||||
The module should contain functions with target attributes in their func_attr.
|
||||
Functions with "cpu" in their target string are considered host functions,
|
||||
while others are considered device functions.
|
||||
|
||||
Returns
|
||||
-------
|
||||
host_mod : tvm.IRModule
|
||||
The module containing host functions (CPU-targeted functions)
|
||||
device_mod_dict : Dict[Target, tvm.IRModule]
|
||||
A dict mapping targets to device modules. Each device module contains
|
||||
functions targeting the same device (e.g., CUDA GPU, OpenCL, etc.)
|
||||
|
||||
Examples
|
||||
--------
|
||||
Given an IRModule with the following functions:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@I.ir_module
|
||||
class Module:
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def add(a: T.int32, b: T.int32) -> T.int32:
|
||||
T.func_attr({"target": T.target({"arch": "sm_90", "keys": ["cuda", "gpu"],
|
||||
"kind": "cuda", "max_num_threads": 1024}))
|
||||
return a + b
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def add_host(a: T.int32, b: T.int32) -> T.int32:
|
||||
T.func_attr({"target": T.target({"keys": ["cpu"], "kind": "c"}))
|
||||
return a + b
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def main_kernel(A: T.handle, B: T.handle, C: T.handle, length: T.int32):
|
||||
T.func_attr({"target": T.target({"arch": "sm_90", "keys": ["cuda", "gpu"],
|
||||
"kind": "cuda"}),
|
||||
"calling_conv": 2, # kDeviceKernelLaunch for device kernels
|
||||
"tirx.is_global_func": True})
|
||||
# ... kernel implementation
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(self_handle: T.handle, args: T.handle, num_args: T.int32, result: T.handle):
|
||||
T.func_attr({"target": T.target({"keys": ["cpu"], "kind": "c"}),
|
||||
"calling_conv": 1, # kCPackedFunc for entry functions
|
||||
"tirx.is_entry_func": True})
|
||||
# ... main function implementation
|
||||
|
||||
The function will return:
|
||||
- host_mod: Contains `add_host` and `main` functions (CPU targets)
|
||||
- device_mod_dict: Contains a CUDA module with `add` and `main_kernel` functions
|
||||
|
||||
Notes
|
||||
-----
|
||||
- Functions are categorized based on string matching of their target attribute
|
||||
- Functions with "cpu" in the target string are considered host functions
|
||||
- Device functions are grouped by their target to create separate modules
|
||||
- The function uses string-based target matching due to target hash limitations
|
||||
- All functions must have a `calling_conv` attribute in their func_attr:
|
||||
- Private helper functions (private=True): use `calling_conv: 0` (kDefault, by default)
|
||||
- Public entry functions: use `calling_conv: 1` (kCPackedFunc)
|
||||
- Device kernel functions: use `calling_conv: 2` (kDeviceKernelLaunch)
|
||||
"""
|
||||
|
||||
def is_host_func(f):
|
||||
target = f.attrs.get("target", tvm.target.Target("llvm"))
|
||||
return target.kind.name in ["llvm", "c"]
|
||||
|
||||
host_mod = tvm.tirx.transform.Filter(is_host_func)(mod)
|
||||
device_mod = tvm.tirx.transform.Filter(lambda f: not is_host_func(f))(mod)
|
||||
# TODO(syfeng): Here we use str as key since target hash is not correct
|
||||
target_str2target = {}
|
||||
device_func_dict = {}
|
||||
device_mod_dict: dict[Target, IRModule] = {}
|
||||
for gv, func in device_mod.functions.items():
|
||||
target = func.attrs.get("target", None)
|
||||
target_str = str(target) if target is not None else ""
|
||||
target_str2target[target_str] = target # This might be overridden by the last one
|
||||
device_func_dict.setdefault(target_str, dict()).update({gv: func})
|
||||
for target_str in target_str2target.keys():
|
||||
target = target_str2target[target_str]
|
||||
device_mod_dict[target] = tvm.IRModule(device_func_dict[target_str], attrs=device_mod.attrs)
|
||||
return host_mod, device_mod_dict
|
||||
|
||||
|
||||
def codegen_build(mod: IRModule, target: Target) -> tvm.runtime.Module:
|
||||
"""Build a runtime module from an IRModule and a Target."""
|
||||
if tvm.ir.transform.PassContext.current().config.get("tirx.disable_assert", False):
|
||||
mod = tvm.tirx.transform.SkipAssert()(mod)
|
||||
build_f_name = "target.build." + target.kind.name
|
||||
bf = tvm.get_global_func(build_f_name)
|
||||
if bf is None:
|
||||
raise ValueError(f"{build_f_name} is not enabled")
|
||||
return bf(mod, target)
|
||||
|
||||
|
||||
def tir_to_runtime(
|
||||
host_mod: IRModule, device_mod_dict: dict[Target, IRModule], target_host: Target
|
||||
):
|
||||
"""Convert a collection of TIR IRModules (keyed by Target) into a single runtime Module."""
|
||||
|
||||
# Get the first module to get the attributes
|
||||
# necessary for tests/python/codegen/test_target_codegen_blob.py::test_cuda_multi_lib
|
||||
mhost_all = ir.IRModule({}, attrs=host_mod.attrs)
|
||||
|
||||
mhost_all.update(host_mod)
|
||||
device_modules = []
|
||||
for target, device_mod in device_mod_dict.items():
|
||||
if len(device_mod.functions) != 0:
|
||||
device_modules.append(codegen_build(device_mod, target))
|
||||
|
||||
mhost = codegen_build(mhost_all, target_host)
|
||||
for dev_mod in device_modules:
|
||||
if dev_mod is not None:
|
||||
mhost.import_module(dev_mod)
|
||||
return mhost
|
||||
|
||||
|
||||
def build(
|
||||
mod: PrimFunc | IRModule,
|
||||
target: str | Target | None = None,
|
||||
pipeline: None | str | tvm.transform.Pass = "default",
|
||||
):
|
||||
"""Build a function with a signature, generating code for devices
|
||||
coupled with target information.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mod : Union[PrimFunc, IRModule]
|
||||
The input to be built.
|
||||
target : Optional[Union[str, Target]]
|
||||
The target for compilation.
|
||||
pipeline : Union[None, str, tvm.transform.Pass]
|
||||
The pipeline to use for compilation.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tvm.runtime.Module
|
||||
A module combining both host and device code.
|
||||
"""
|
||||
# Convert PrimFunc to IRModule
|
||||
if isinstance(mod, PrimFunc):
|
||||
mod = tvm.IRModule.from_expr(mod)
|
||||
else:
|
||||
assert isinstance(mod, tvm.IRModule)
|
||||
|
||||
# Step 0: Determine the target in environment
|
||||
# It's used to bind the PrimFunc without target attr to serve as a default target
|
||||
target_to_bind = Target.current() if target is None else target
|
||||
if target_to_bind is None:
|
||||
target_to_bind = "llvm"
|
||||
assert target_to_bind is not None
|
||||
target_to_bind = Target(target_to_bind)
|
||||
|
||||
# Step 1: Determine the target to search for tirx pipeline
|
||||
target = Target.current() if target is None else target
|
||||
if target is None:
|
||||
for func in mod.functions.values():
|
||||
f_target = func.attrs.get("target", None)
|
||||
if f_target is not None:
|
||||
target = f_target
|
||||
break
|
||||
if target is not None:
|
||||
target = Target(target)
|
||||
|
||||
# Step 2: Determine the host target
|
||||
target_host = "llvm" if tvm.runtime.enabled("llvm") else "c"
|
||||
if target is not None:
|
||||
if target.host is not None:
|
||||
target_host = target.host
|
||||
elif (
|
||||
tvm.device(target.kind.name, 0).dlpack_device_type() == tvm.cpu(0).dlpack_device_type()
|
||||
):
|
||||
target_host = target
|
||||
target_host = Target(target_host)
|
||||
target_to_bind = target_to_bind.with_host(target_host)
|
||||
|
||||
# Step 3: Bind the target to the input module
|
||||
mod = tvm.tirx.transform.BindTarget(target_to_bind)(mod)
|
||||
|
||||
# Step 4: Apply the tirx pipeline
|
||||
if pipeline is not None:
|
||||
# custom pipeline
|
||||
assert isinstance(pipeline, str)
|
||||
pipeline, finalize_host_passes, finalize_device_passes = tvm.tirx.get_tir_pipeline(pipeline)
|
||||
else:
|
||||
# default pipeline depends on the target
|
||||
pipeline, finalize_host_passes, finalize_device_passes = tvm.tirx.get_default_tir_pipeline(
|
||||
target
|
||||
)
|
||||
mod = pipeline(mod)
|
||||
|
||||
# Step 5: Get host and device modules
|
||||
host_mod, device_mod_dict = split_host_device_mods(mod)
|
||||
|
||||
# Step 6: Apply finalization passes
|
||||
host_mod = finalize_host_passes()(host_mod)
|
||||
device_mod_dict = {
|
||||
target: finalize_device_passes()(device_mod)
|
||||
for target, device_mod in device_mod_dict.items()
|
||||
}
|
||||
|
||||
# Convert TIR IRModules to runtime Module by calling target.build
|
||||
return tir_to_runtime(host_mod, device_mod_dict, target_host)
|
||||
|
||||
|
||||
tvm.register_global_func("tirx.build", build)
|
||||
@@ -0,0 +1,162 @@
|
||||
# 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=invalid-name
|
||||
"""The TIR backend compilation pipeline."""
|
||||
|
||||
import tvm
|
||||
from tvm import tirx
|
||||
|
||||
|
||||
def default_tir_pipeline():
|
||||
"""The default tirx pipeline used in tvm.tirx.build"""
|
||||
|
||||
@tvm.transform.module_pass(opt_level=0)
|
||||
def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
|
||||
"""The default lowering passes for TIR backend."""
|
||||
pass_ctx = tvm.transform.PassContext.current()
|
||||
config = pass_ctx.config
|
||||
passes = [
|
||||
tirx.transform.LowerInitBlock(),
|
||||
tvm.s_tir.transform.UnifyThreadBinding(),
|
||||
tirx.transform.StmtSimplify(),
|
||||
tirx.transform.FlattenBuffer(),
|
||||
tirx.transform.BF16ComputeLegalize(),
|
||||
tirx.transform.NarrowDataType(32),
|
||||
tirx.transform.VectorizeLoop(not bool(config.get("tir.disable_vectorize", False))),
|
||||
tirx.transform.UnrollLoop(),
|
||||
tirx.transform.StmtSimplify(),
|
||||
]
|
||||
if not bool(config.get("tir.disable_cse_tir", False)):
|
||||
passes.append(tirx.transform.CommonSubexprElim())
|
||||
passes.extend(
|
||||
[
|
||||
tirx.transform.FP8ComputeLegalize(),
|
||||
tirx.transform.VerifyMemory(),
|
||||
tirx.transform.AnnotateEntryFunc(),
|
||||
tirx.transform.SplitHostDevice(),
|
||||
tirx.transform.MakePackedAPI(),
|
||||
tirx.transform.FP8StorageLegalize(),
|
||||
tirx.transform.BF16StorageLegalize(),
|
||||
]
|
||||
)
|
||||
mod = tvm.ir.transform.Sequential(passes)(mod)
|
||||
return mod
|
||||
|
||||
return _pipeline, finalize_host_passes, finalize_device_passes
|
||||
|
||||
|
||||
def tirx_pipeline():
|
||||
"""The TIRX pipeline used in tvm.tirx.build"""
|
||||
|
||||
@tvm.transform.module_pass(opt_level=0)
|
||||
def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
|
||||
"""The default lowering passes for TIR backend."""
|
||||
pass_ctx = tvm.transform.PassContext.current()
|
||||
config = pass_ctx.config
|
||||
passes = [
|
||||
tirx.transform.LowerTIRx(),
|
||||
tvm.s_tir.transform.UnifyThreadBinding(),
|
||||
tirx.transform.StmtSimplify(),
|
||||
tirx.transform.LowerTIRxOpaque(),
|
||||
tirx.transform.FlattenBuffer(),
|
||||
tirx.transform.BF16ComputeLegalize(),
|
||||
tirx.transform.NarrowDataType(32),
|
||||
tirx.transform.VectorizeLoop(not bool(config.get("tir.disable_vectorize", False))),
|
||||
tirx.transform.UnrollLoop(),
|
||||
tirx.transform.StmtSimplify(),
|
||||
]
|
||||
if not bool(config.get("tir.disable_cse_tir", False)):
|
||||
passes.append(tirx.transform.CommonSubexprElim())
|
||||
passes.extend(
|
||||
[
|
||||
tirx.transform.FP8ComputeLegalize(),
|
||||
tirx.transform.VerifyMemory(),
|
||||
tirx.transform.AnnotateEntryFunc(),
|
||||
tirx.transform.SplitHostDevice(),
|
||||
tirx.transform.MakePackedAPI(),
|
||||
tirx.transform.FP8StorageLegalize(),
|
||||
tirx.transform.BF16StorageLegalize(),
|
||||
]
|
||||
)
|
||||
mod = tvm.ir.transform.Sequential(passes)(mod)
|
||||
return mod
|
||||
|
||||
return _pipeline, finalize_host_passes, finalize_device_passes
|
||||
|
||||
|
||||
def finalize_host_passes(): # pylint: disable=unused-argument
|
||||
"""The default finalization passes for TIR backend."""
|
||||
host_pass_list = [
|
||||
tirx.transform.LowerTVMBuiltin(),
|
||||
tirx.transform.LowerIntrin(),
|
||||
]
|
||||
return tvm.ir.transform.Sequential(host_pass_list)
|
||||
|
||||
|
||||
def finalize_device_passes(): # pylint: disable=unused-argument
|
||||
"""The default finalization passes for TIR backend."""
|
||||
device_pass_list = [
|
||||
tirx.transform.LowerWarpMemory(),
|
||||
tirx.transform.StmtSimplify(),
|
||||
tirx.transform.LowerIntrin(),
|
||||
]
|
||||
return tvm.ir.transform.Sequential(device_pass_list)
|
||||
|
||||
|
||||
def finalize_device_passes_tirx(): # pylint: disable=unused-argument
|
||||
"""The TIRx finalization passes for TIR backend."""
|
||||
device_pass_list = [tirx.transform.LowerIntrin()]
|
||||
return tvm.ir.transform.Sequential(device_pass_list)
|
||||
|
||||
|
||||
# global map of pre-built pipelines
|
||||
PIPELINE_MAP = {"default": default_tir_pipeline, "tirx": tirx_pipeline}
|
||||
|
||||
|
||||
def register_tir_pipeline(name: str, pipeline_factory) -> None:
|
||||
"""Register a named TIR pipeline factory."""
|
||||
|
||||
PIPELINE_MAP[name] = pipeline_factory
|
||||
|
||||
|
||||
def get_tir_pipeline(name: str | None = None, **kwargs) -> tvm.transform.Pass:
|
||||
"""Get pre-build pipeline by name
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : Optional[str]
|
||||
Name of the pipeline
|
||||
"""
|
||||
if name == "default":
|
||||
# for now, default to s_tir pipeline
|
||||
name = "s_tir"
|
||||
if name not in PIPELINE_MAP:
|
||||
raise ValueError(
|
||||
f"Unknown pre-built pipeline {name},candidates are {list(PIPELINE_MAP.keys())}"
|
||||
)
|
||||
return PIPELINE_MAP[name](**kwargs)
|
||||
|
||||
|
||||
def get_default_tir_pipeline(
|
||||
target: tvm.target.Target, # pylint: disable=unused-argument
|
||||
) -> tvm.transform.Pass:
|
||||
"""Get the default TIR pipeline for the given target."""
|
||||
if target.kind.name == "opencl" and "adreno" in target.keys:
|
||||
return get_tir_pipeline("adreno")
|
||||
else:
|
||||
return get_tir_pipeline("s_tir")
|
||||
@@ -0,0 +1,408 @@
|
||||
# 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.
|
||||
"""ExecContext: per-program-point active-thread state.
|
||||
|
||||
The active thread set is represented as a ``TileLayout``: active axes live in
|
||||
``layout.shard`` and per-axis lower bounds live in ``layout.offset``. Filters
|
||||
narrow that layout; scope switches derive the current ``inter``/``intra`` view.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from tvm.tirx.layout import Axis, Iter, TileLayout
|
||||
|
||||
WG_SIZE = 4
|
||||
|
||||
KERNEL = "kernel"
|
||||
CLUSTER = "cluster"
|
||||
CTA = "cta"
|
||||
WARPGROUP = "warpgroup"
|
||||
WARP = "warp"
|
||||
THREAD = "thread"
|
||||
|
||||
SCOPE_KINDS = (KERNEL, CLUSTER, CTA, WARPGROUP, WARP, THREAD)
|
||||
|
||||
LANE_FLAT = "flat"
|
||||
LANE_WG_OUTER = "wg_outer"
|
||||
LANE_W_INNER = "w_inner"
|
||||
LANE_CTA_THREAD = "cta_thread"
|
||||
LANE_WG_THREAD = "wg_thread"
|
||||
|
||||
|
||||
class ExecContextError(Exception):
|
||||
"""Raised on structural violations of the ExecContext model."""
|
||||
|
||||
|
||||
def _ceildiv(lhs: int, rhs: int) -> int:
|
||||
return -((-lhs) // rhs)
|
||||
|
||||
|
||||
def _gcd(lhs: int, rhs: int) -> int:
|
||||
while rhs:
|
||||
lhs, rhs = rhs, lhs % rhs
|
||||
return abs(lhs)
|
||||
|
||||
|
||||
def _extended_gcd(lhs: int, rhs: int) -> tuple[int, int, int]:
|
||||
if rhs == 0:
|
||||
return lhs, 1, 0
|
||||
gcd, x1, y1 = _extended_gcd(rhs, lhs % rhs)
|
||||
return gcd, y1, x1 - (lhs // rhs) * y1
|
||||
|
||||
|
||||
def _mod_inverse(value: int, modulus: int) -> int:
|
||||
if modulus == 1:
|
||||
return 0
|
||||
gcd, inv, _ = _extended_gcd(value % modulus, modulus)
|
||||
if gcd != 1:
|
||||
raise ExecContextError(f"{value} has no inverse modulo {modulus}")
|
||||
return inv % modulus
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AxisRange:
|
||||
"""An active slice offset + stride * [0, extent) on one TileLayout axis."""
|
||||
|
||||
extent: int
|
||||
offset: int = 0
|
||||
stride: int = 1
|
||||
|
||||
def intersect(self, lo: int, hi: int) -> AxisRange:
|
||||
i_lo = max(0, _ceildiv(lo - self.offset, self.stride))
|
||||
i_hi = min(self.extent, (hi - 1 - self.offset) // self.stride + 1)
|
||||
if i_hi <= i_lo:
|
||||
raise ExecContextError(
|
||||
f"filter produces empty range: current=[{self.offset},"
|
||||
f" {self.offset + self.extent}) ∩ [{lo}, {hi})"
|
||||
)
|
||||
return AxisRange(
|
||||
extent=i_hi - i_lo, offset=self.offset + self.stride * i_lo, stride=self.stride
|
||||
)
|
||||
|
||||
def modulo(self, modulus: int, residue: int) -> AxisRange:
|
||||
residue %= modulus
|
||||
rhs = (residue - self.offset) % modulus
|
||||
g = _gcd(self.stride, modulus)
|
||||
if rhs % g != 0:
|
||||
raise ExecContextError(
|
||||
f"modulo filter produces empty range: {self.offset} + {self.stride} * i"
|
||||
f" == {residue} mod {modulus}"
|
||||
)
|
||||
reduced_stride = self.stride // g
|
||||
reduced_rhs = rhs // g
|
||||
reduced_modulus = modulus // g
|
||||
period = reduced_modulus
|
||||
i0 = (reduced_rhs * _mod_inverse(reduced_stride, reduced_modulus)) % reduced_modulus
|
||||
if i0 >= self.extent:
|
||||
raise ExecContextError(
|
||||
f"modulo filter produces empty range: {self.offset} + {self.stride} * i"
|
||||
f" == {residue} mod {modulus}"
|
||||
)
|
||||
return AxisRange(
|
||||
extent=(self.extent - 1 - i0) // period + 1,
|
||||
offset=self.offset + self.stride * i0,
|
||||
stride=self.stride * period,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ActiveSet:
|
||||
"""Active thread set represented by a TileLayout."""
|
||||
|
||||
layout: TileLayout
|
||||
|
||||
@staticmethod
|
||||
def from_axes(axes: list[tuple[str, AxisRange]]) -> ActiveSet:
|
||||
shard = [Iter(axis_range.extent, axis_range.stride, name) for name, axis_range in axes]
|
||||
offset = {
|
||||
Axis.get(name): axis_range.offset for name, axis_range in axes if axis_range.offset != 0
|
||||
}
|
||||
return ActiveSet(TileLayout.from_iters(shard, [], offset))
|
||||
|
||||
@property
|
||||
def size(self) -> int:
|
||||
result = 1
|
||||
for it in self.layout.shard:
|
||||
result *= int(it.extent)
|
||||
return result
|
||||
|
||||
@property
|
||||
def axis_names(self) -> list[str]:
|
||||
return [str(it.axis.name) for it in self.layout.shard]
|
||||
|
||||
def axis(self, name: str) -> AxisRange:
|
||||
for it in self.layout.shard:
|
||||
if str(it.axis.name) != name:
|
||||
continue
|
||||
offset = 0
|
||||
for axis, value in self.layout.offset.items():
|
||||
if str(axis.name) == name:
|
||||
offset = int(value)
|
||||
break
|
||||
return AxisRange(int(it.extent), offset, int(it.stride))
|
||||
raise ValueError(f"unknown active-set axis: {name!r}")
|
||||
|
||||
def replace_axis(self, axis: str, axis_range: AxisRange) -> ActiveSet:
|
||||
axes: list[tuple[str, AxisRange]] = []
|
||||
found = False
|
||||
for name in self.axis_names:
|
||||
if name == axis:
|
||||
axes.append((name, axis_range))
|
||||
found = True
|
||||
else:
|
||||
axes.append((name, self.axis(name)))
|
||||
if not found:
|
||||
raise ValueError(f"unknown active-set axis: {axis!r}")
|
||||
return ActiveSet.from_axes(axes)
|
||||
|
||||
@property
|
||||
def laneid(self) -> AxisRange:
|
||||
return self.axis("laneid")
|
||||
|
||||
@property
|
||||
def warpid(self) -> AxisRange:
|
||||
return self.axis("warpid")
|
||||
|
||||
@property
|
||||
def cta_id(self) -> AxisRange:
|
||||
return self.axis("cta_id")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LaneBinding:
|
||||
"""Resolution of a user-declared ScopeIdDef Var to one active-set axis."""
|
||||
|
||||
axis: str
|
||||
kind: str
|
||||
declared_extent: int
|
||||
|
||||
|
||||
def initial_A(*, lane_ext: int = 32, warp_ext: int, cta_ext: int = 1) -> ActiveSet:
|
||||
"""Build A at PrimFunc device entry: all threads active, offsets all zero."""
|
||||
return ActiveSet.from_axes(
|
||||
[
|
||||
("laneid", AxisRange(lane_ext, 0)),
|
||||
("warpid", AxisRange(warp_ext, 0)),
|
||||
("cta_id", AxisRange(cta_ext, 0)),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def filter_narrow(A: ActiveSet, binding: LaneBinding, lo: int, hi: int) -> ActiveSet:
|
||||
"""Intersect A's binding axis with [lo, hi)."""
|
||||
if lo >= hi:
|
||||
raise ExecContextError(f"filter range [{lo}, {hi}) is empty or inverted")
|
||||
|
||||
if binding.kind == LANE_CTA_THREAD:
|
||||
new_warpid, new_laneid = _flat_product_range(A.warpid, A.laneid, lo, hi)
|
||||
return A.replace_axis("laneid", new_laneid).replace_axis("warpid", new_warpid)
|
||||
|
||||
if binding.kind == LANE_WG_THREAD:
|
||||
factored = _factor_warpid(A.warpid)
|
||||
if factored is None:
|
||||
raise ExecContextError(
|
||||
"filter on flat warpgroup-thread range requires factorable warpid axis"
|
||||
)
|
||||
wid_in_wg, wgid = factored
|
||||
new_wid_in_wg, new_laneid = _flat_product_range(wid_in_wg, A.laneid, lo, hi)
|
||||
if wgid.extent != 1:
|
||||
if new_wid_in_wg == wid_in_wg and new_laneid == A.laneid:
|
||||
return A
|
||||
raise ExecContextError(
|
||||
"flat warpgroup-thread range across multiple warpgroups is not representable"
|
||||
)
|
||||
new_warpid = AxisRange(
|
||||
extent=new_wid_in_wg.extent, offset=wgid.offset * WG_SIZE + new_wid_in_wg.offset
|
||||
)
|
||||
return A.replace_axis("laneid", new_laneid).replace_axis("warpid", new_warpid)
|
||||
|
||||
if binding.kind == LANE_FLAT:
|
||||
new_axis = A.axis(binding.axis).intersect(lo, hi)
|
||||
return A.replace_axis(binding.axis, new_axis)
|
||||
|
||||
if binding.axis != "warpid":
|
||||
raise ExecContextError(
|
||||
f"kind={binding.kind!r} only valid for axis='warpid'; got {binding.axis!r}"
|
||||
)
|
||||
|
||||
wp = A.warpid
|
||||
if wp.stride != 1:
|
||||
raise ExecContextError(
|
||||
f"kind={binding.kind!r} requires unit-stride warpid axis; got stride={wp.stride}"
|
||||
)
|
||||
if binding.kind == LANE_WG_OUTER:
|
||||
if wp.offset % WG_SIZE != 0 or wp.extent % WG_SIZE != 0:
|
||||
raise ExecContextError(
|
||||
f"filter on wg_outer requires warpid axis aligned to WG_SIZE={WG_SIZE};"
|
||||
f" got extent={wp.extent}, offset={wp.offset}"
|
||||
)
|
||||
cur_outer = AxisRange(extent=wp.extent // WG_SIZE, offset=wp.offset // WG_SIZE)
|
||||
new_outer = cur_outer.intersect(lo, hi)
|
||||
return A.replace_axis(
|
||||
"warpid",
|
||||
AxisRange(extent=new_outer.extent * WG_SIZE, offset=new_outer.offset * WG_SIZE),
|
||||
)
|
||||
|
||||
if binding.kind == LANE_W_INNER:
|
||||
cur_inner_off = wp.offset % WG_SIZE
|
||||
if wp.extent > WG_SIZE - cur_inner_off:
|
||||
raise ExecContextError(
|
||||
"filter on w_inner would break A's TileLayout box: warpid spans multiple"
|
||||
f" warpgroups (extent={wp.extent}, offset={wp.offset})"
|
||||
)
|
||||
cur_inner = AxisRange(extent=wp.extent, offset=cur_inner_off)
|
||||
new_inner = cur_inner.intersect(lo, hi)
|
||||
outer_base = (wp.offset // WG_SIZE) * WG_SIZE
|
||||
return A.replace_axis(
|
||||
"warpid", AxisRange(extent=new_inner.extent, offset=outer_base + new_inner.offset)
|
||||
)
|
||||
|
||||
raise ValueError(f"unknown axis kind: {binding.kind!r}")
|
||||
|
||||
|
||||
def filter_modulo(A: ActiveSet, axis: str, modulus: int, residue: int) -> ActiveSet:
|
||||
"""Intersect an active-set axis with ``axis % modulus == residue``."""
|
||||
if modulus <= 0:
|
||||
raise ExecContextError(f"modulus must be positive, got {modulus}")
|
||||
new_axis = A.axis(axis).modulo(modulus, residue)
|
||||
return A.replace_axis(axis, new_axis)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Split:
|
||||
"""A scope_switch split of A."""
|
||||
|
||||
inter: dict[str, AxisRange]
|
||||
intra: dict[str, AxisRange]
|
||||
|
||||
|
||||
def _factor_warpid(warp: AxisRange) -> tuple[AxisRange, AxisRange] | None:
|
||||
if warp.stride != 1:
|
||||
return None
|
||||
off = warp.offset
|
||||
ext = warp.extent
|
||||
wid_off = off % WG_SIZE
|
||||
wgid_off = off // WG_SIZE
|
||||
|
||||
if wid_off == 0 and ext % WG_SIZE == 0:
|
||||
return (
|
||||
AxisRange(extent=WG_SIZE, offset=0),
|
||||
AxisRange(extent=ext // WG_SIZE, offset=wgid_off),
|
||||
)
|
||||
if ext <= WG_SIZE - wid_off:
|
||||
return (AxisRange(extent=ext, offset=wid_off), AxisRange(extent=1, offset=wgid_off))
|
||||
return None
|
||||
|
||||
|
||||
def _flat_product_range(
|
||||
major: AxisRange, lane: AxisRange, lo: int, hi: int
|
||||
) -> tuple[AxisRange, AxisRange]:
|
||||
active_min = major.offset * 32 + lane.offset
|
||||
active_max = (
|
||||
(major.offset + major.stride * (major.extent - 1)) * 32
|
||||
+ lane.offset
|
||||
+ lane.stride * (lane.extent - 1)
|
||||
+ 1
|
||||
)
|
||||
if lo <= active_min and active_max <= hi:
|
||||
return major, lane
|
||||
|
||||
if major.stride != 1 or lane.stride != 1:
|
||||
raise ExecContextError("flat thread range narrowing requires unit-stride axes")
|
||||
|
||||
lane_hi = lane.offset + lane.extent
|
||||
major_hi = major.offset + major.extent
|
||||
hit_lo = max(major.offset, (lo - lane_hi) // 32 + 1)
|
||||
hit_hi = min(major_hi, _ceildiv(hi - lane.offset, 32))
|
||||
if hit_hi <= hit_lo:
|
||||
raise ExecContextError("flat thread range produces empty active set")
|
||||
|
||||
if hit_hi == hit_lo + 1:
|
||||
new_lane_lo = max(lane.offset, lo - hit_lo * 32)
|
||||
new_lane_hi = min(lane_hi, hi - hit_lo * 32)
|
||||
if new_lane_hi <= new_lane_lo:
|
||||
raise ExecContextError("flat thread range produces empty lane range")
|
||||
return AxisRange(1, hit_lo), AxisRange(new_lane_hi - new_lane_lo, new_lane_lo)
|
||||
|
||||
if lo <= hit_lo * 32 + lane.offset and (hit_hi - 1) * 32 + lane_hi <= hi:
|
||||
return AxisRange(hit_hi - hit_lo, hit_lo), lane
|
||||
|
||||
raise ExecContextError("flat thread range would require a non-rectangular lane/warp active set")
|
||||
|
||||
|
||||
def scope_switch(A: ActiveSet, scope_kind: str) -> Split:
|
||||
"""Split A into (inter, intra) for the target scope kind."""
|
||||
if scope_kind == THREAD:
|
||||
return Split(inter={"laneid": A.laneid, "warpid": A.warpid, "cta_id": A.cta_id}, intra={})
|
||||
if scope_kind == WARP:
|
||||
return Split(inter={"warpid": A.warpid, "cta_id": A.cta_id}, intra={"laneid": A.laneid})
|
||||
if scope_kind == CTA:
|
||||
return Split(inter={"cta_id": A.cta_id}, intra={"laneid": A.laneid, "warpid": A.warpid})
|
||||
if scope_kind == CLUSTER:
|
||||
return Split(inter={}, intra={"laneid": A.laneid, "warpid": A.warpid, "cta_id": A.cta_id})
|
||||
if scope_kind == WARPGROUP:
|
||||
factored = _factor_warpid(A.warpid)
|
||||
if factored is None:
|
||||
raise ExecContextError(
|
||||
"scope_switch(warpgroup) failed: warpid axis"
|
||||
f" (extent={A.warpid.extent}, offset={A.warpid.offset})"
|
||||
" crosses warpgroup boundary and is not aligned"
|
||||
)
|
||||
wid_in_wg, wgid = factored
|
||||
return Split(
|
||||
inter={"wgid": wgid, "cta_id": A.cta_id},
|
||||
intra={"laneid": A.laneid, "wid_in_wg": wid_in_wg},
|
||||
)
|
||||
if scope_kind == KERNEL:
|
||||
return Split(inter={"laneid": A.laneid, "warpid": A.warpid, "cta_id": A.cta_id}, intra={})
|
||||
raise ValueError(f"unknown scope kind: {scope_kind!r}")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExecContext:
|
||||
"""Per-program-point compiler state: active set + scope kind + split."""
|
||||
|
||||
A: ActiveSet
|
||||
scope_kind: str
|
||||
inter: dict[str, AxisRange]
|
||||
intra: dict[str, AxisRange]
|
||||
|
||||
@staticmethod
|
||||
def at_kernel_entry(*, lane_ext: int = 32, warp_ext: int, cta_ext: int = 1) -> ExecContext:
|
||||
A = initial_A(lane_ext=lane_ext, warp_ext=warp_ext, cta_ext=cta_ext)
|
||||
split = scope_switch(A, KERNEL)
|
||||
return ExecContext(A=A, scope_kind=KERNEL, inter=split.inter, intra=split.intra)
|
||||
|
||||
def with_filter(self, binding: LaneBinding, lo: int, hi: int) -> ExecContext:
|
||||
new_A = filter_narrow(self.A, binding, lo, hi)
|
||||
split = scope_switch(new_A, self.scope_kind)
|
||||
return ExecContext(
|
||||
A=new_A, scope_kind=self.scope_kind, inter=split.inter, intra=split.intra
|
||||
)
|
||||
|
||||
def with_cta_axis_modulo(self, axis: str, modulus: int, residue: int) -> ExecContext:
|
||||
new_A = filter_modulo(self.A, axis, modulus, residue)
|
||||
split = scope_switch(new_A, self.scope_kind)
|
||||
return ExecContext(
|
||||
A=new_A, scope_kind=self.scope_kind, inter=split.inter, intra=split.intra
|
||||
)
|
||||
|
||||
def with_scope_switch(self, scope_kind: str) -> ExecContext:
|
||||
split = scope_switch(self.A, scope_kind)
|
||||
return ExecContext(A=self.A, scope_kind=scope_kind, inter=split.inter, intra=split.intra)
|
||||
@@ -0,0 +1,101 @@
|
||||
# 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=no-member, super-init-not-called
|
||||
|
||||
"""Definition of execution scope."""
|
||||
|
||||
from tvm_ffi import register_object
|
||||
|
||||
from tvm.runtime import Object
|
||||
|
||||
from . import _ffi_api
|
||||
from .expr import Expr, Var
|
||||
|
||||
|
||||
@register_object("tirx.ScopeIdDef")
|
||||
class ScopeIdDef(Object):
|
||||
"""Definition of scope identifiers with their extents and parent-child relationships.
|
||||
|
||||
The constructor accepts ``parent`` and ``cur`` as scope-name strings; they
|
||||
are converted by the FFI into the closed ``ScopeBinding`` enum and stored
|
||||
on the ``scope`` field (an ``int`` value of that enum).
|
||||
|
||||
``extents=None`` defers the extent: the value is inferred from sibling
|
||||
ScopeIdDef relationships at LowerTIRx entry via the verifier's closure.
|
||||
Deferred form requires ``def_ids`` to contain exactly one Var.
|
||||
"""
|
||||
|
||||
def_ids: list[Var]
|
||||
extents: list[Expr] | None
|
||||
scope: int
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
def_ids: list[Var],
|
||||
extents: list[Expr] | None,
|
||||
parent: str,
|
||||
cur: str,
|
||||
preferred_extents: list[Expr] | None = None,
|
||||
):
|
||||
self.__init_handle_by_constructor__(
|
||||
_ffi_api.ScopeIdDef, def_ids, extents, parent, cur, preferred_extents
|
||||
)
|
||||
|
||||
|
||||
_SCOPE_KIND_TO_NAME = {
|
||||
2: "cluster",
|
||||
3: "cta",
|
||||
4: "warpgroup",
|
||||
5: "warp",
|
||||
6: "thread",
|
||||
}
|
||||
|
||||
|
||||
# Mirror of ``enum class ScopeBinding`` in tvm/tirx/exec_scope.h. Maps the
|
||||
# ``int`` value of ``ScopeIdDef.scope`` back to the ``(parent, cur)`` pair
|
||||
# that ``ScopeIdDef.__init__`` accepts — needed when Python code wants to
|
||||
# rebuild a ``ScopeIdDef`` from an existing one (e.g. a StmtMutator
|
||||
# walking and rewriting extents).
|
||||
_SCOPE_BINDING_TO_PARENT_CUR = {
|
||||
0: ("kernel", "cluster"),
|
||||
1: ("kernel", "cta"),
|
||||
2: ("cluster", "cta"),
|
||||
3: ("cta", "warpgroup"),
|
||||
4: ("cta", "warp"),
|
||||
5: ("warpgroup", "warp"),
|
||||
6: ("warp", "thread"),
|
||||
7: ("cta", "thread"),
|
||||
8: ("warpgroup", "thread"),
|
||||
9: ("cluster", "cta_pair"),
|
||||
}
|
||||
|
||||
|
||||
@register_object("tirx.ExecScope")
|
||||
class ExecScope(Object):
|
||||
"""An execution scope, identified by one of {cluster, cta, warpgroup, warp,
|
||||
thread}. The ctor FATALs on any other name."""
|
||||
|
||||
kind: int
|
||||
scope_id_def: list[ScopeIdDef]
|
||||
|
||||
def __init__(self, name: str):
|
||||
self.__init_handle_by_constructor__(_ffi_api.ExecScope, name)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""Human-readable name of this scope (derived from ``kind``)."""
|
||||
return _SCOPE_KIND_TO_NAME[self.kind]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,671 @@
|
||||
# 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.
|
||||
"""
|
||||
TIR expression functors in Python.
|
||||
|
||||
This module implements the visitor and mutator patterns for TIR expressions.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import TypeVar
|
||||
|
||||
import tvm
|
||||
from tvm.ir import Expr, Range
|
||||
from tvm.tirx import IterVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
def _visit_array(arr: list[T], callback: Callable[[T], None]) -> None:
|
||||
"""Visit elements in an array using a callback function.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
arr : List[T]
|
||||
The array to be visited
|
||||
callback : Callable[[T], None]
|
||||
The callback function
|
||||
"""
|
||||
for item in arr:
|
||||
callback(item)
|
||||
|
||||
|
||||
class ExprFunctor:
|
||||
"""An abstract visitor over Expr, with visiting function defined for each Expr type."""
|
||||
|
||||
def __init__(self):
|
||||
self._dispatch_map = {
|
||||
"tirx.Var": self.visit_var_,
|
||||
"tirx.BufferLoad": self.visit_buffer_load_,
|
||||
"tirx.ProducerLoad": self.visit_producer_load_,
|
||||
"tirx.Let": self.visit_let_,
|
||||
"tirx.Call": self.visit_call_,
|
||||
"tirx.Add": self.visit_add_,
|
||||
"tirx.Sub": self.visit_sub_,
|
||||
"tirx.Mul": self.visit_mul_,
|
||||
"tirx.Div": self.visit_div_,
|
||||
"tirx.Mod": self.visit_mod_,
|
||||
"tirx.FloorDiv": self.visit_floordiv_,
|
||||
"tirx.FloorMod": self.visit_floormod_,
|
||||
"tirx.Min": self.visit_min_,
|
||||
"tirx.Max": self.visit_max_,
|
||||
"tirx.EQ": self.visit_eq_,
|
||||
"tirx.NE": self.visit_ne_,
|
||||
"tirx.LT": self.visit_lt_,
|
||||
"tirx.LE": self.visit_le_,
|
||||
"tirx.GT": self.visit_gt_,
|
||||
"tirx.GE": self.visit_ge_,
|
||||
"tirx.And": self.visit_and_,
|
||||
"tirx.Or": self.visit_or_,
|
||||
"tirx.Reduce": self.visit_reduce_,
|
||||
"tirx.Cast": self.visit_cast_,
|
||||
"tirx.Not": self.visit_not_,
|
||||
"tirx.Select": self.visit_select_,
|
||||
"tirx.Ramp": self.visit_ramp_,
|
||||
"tirx.Broadcast": self.visit_broadcast_,
|
||||
"tirx.Shuffle": self.visit_shuffle_,
|
||||
"tirx.IntImm": self.visit_int_imm_,
|
||||
"tirx.FloatImm": self.visit_float_imm_,
|
||||
"tirx.StringImm": self.visit_string_imm_,
|
||||
}
|
||||
|
||||
def visit_expr(self, expr: Expr):
|
||||
"""Apply the visitor to an expression.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
expr : Expr
|
||||
The expression to be visited.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : Any
|
||||
The result of the visit.
|
||||
"""
|
||||
if expr is None:
|
||||
return None
|
||||
|
||||
key = expr.__class__.__name__
|
||||
if key.endswith("Node"):
|
||||
key = key[:-4]
|
||||
|
||||
key = "tirx." + key
|
||||
if key in self._dispatch_map:
|
||||
return self._dispatch_map[key](expr)
|
||||
|
||||
return self.visit_expr_default_(expr)
|
||||
|
||||
def visit_var_(self, op):
|
||||
"""Default visitor for Var node."""
|
||||
return None
|
||||
|
||||
def visit_buffer_load_(self, op):
|
||||
"""Default visitor for BufferLoad node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_producer_load_(self, op):
|
||||
"""Default visitor for ProducerLoad node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_let_(self, op):
|
||||
"""Default visitor for Let node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_call_(self, op):
|
||||
"""Default visitor for Call node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_add_(self, op):
|
||||
"""Default visitor for Add node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_sub_(self, op):
|
||||
"""Default visitor for Sub node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_mul_(self, op):
|
||||
"""Default visitor for Mul node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_div_(self, op):
|
||||
"""Default visitor for Div node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_mod_(self, op):
|
||||
"""Default visitor for Mod node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_floordiv_(self, op):
|
||||
"""Default visitor for FloorDiv node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_floormod_(self, op):
|
||||
"""Default visitor for FloorMod node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_min_(self, op):
|
||||
"""Default visitor for Min node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_max_(self, op):
|
||||
"""Default visitor for Max node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_eq_(self, op):
|
||||
"""Default visitor for EQ node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_ne_(self, op):
|
||||
"""Default visitor for NE node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_lt_(self, op):
|
||||
"""Default visitor for LT node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_le_(self, op):
|
||||
"""Default visitor for LE node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_gt_(self, op):
|
||||
"""Default visitor for GT node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_ge_(self, op):
|
||||
"""Default visitor for GE node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_and_(self, op):
|
||||
"""Default visitor for And node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_or_(self, op):
|
||||
"""Default visitor for Or node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_reduce_(self, op):
|
||||
"""Default visitor for Reduce node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_cast_(self, op):
|
||||
"""Default visitor for Cast node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_not_(self, op):
|
||||
"""Default visitor for Not node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_select_(self, op):
|
||||
"""Default visitor for Select node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_ramp_(self, op):
|
||||
"""Default visitor for Ramp node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_broadcast_(self, op):
|
||||
"""Default visitor for Broadcast node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_shuffle_(self, op):
|
||||
"""Default visitor for Shuffle node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_int_imm_(self, op):
|
||||
"""Default visitor for IntImm node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_float_imm_(self, op):
|
||||
"""Default visitor for FloatImm node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_string_imm_(self, op):
|
||||
"""Default visitor for StringImm node."""
|
||||
return self.visit_expr_default_(op)
|
||||
|
||||
def visit_expr_default_(self, op):
|
||||
"""Default visitor implementation."""
|
||||
raise NotImplementedError(f"Do not have a default for {op.__class__.__name__}")
|
||||
|
||||
def __call__(self, expr):
|
||||
"""Call visitor on expression.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
expr : Expr
|
||||
The expression.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : Any
|
||||
The result of visiting.
|
||||
"""
|
||||
return self.visit_expr(expr)
|
||||
|
||||
|
||||
class ExprVisitor(ExprFunctor):
|
||||
"""A visitor over Expr.
|
||||
|
||||
This is a visitor that recursively traverses an expression. Subclasses can
|
||||
override the visit methods to customize the behavior.
|
||||
"""
|
||||
|
||||
def visit_var_(self, op):
|
||||
"""Visitor implementation for Var."""
|
||||
pass
|
||||
|
||||
def visit_buffer_load_(self, op):
|
||||
"""Visitor implementation for BufferLoad."""
|
||||
|
||||
def _visit_indices(index):
|
||||
self.visit_expr(index)
|
||||
|
||||
_visit_array(op.indices, _visit_indices)
|
||||
|
||||
def visit_producer_load_(self, op):
|
||||
"""Visitor implementation for ProducerLoad."""
|
||||
|
||||
def _visit_indices(index):
|
||||
self.visit_expr(index)
|
||||
|
||||
_visit_array(op.indices, _visit_indices)
|
||||
|
||||
def visit_let_(self, op):
|
||||
"""Visitor implementation for Let."""
|
||||
self.visit_expr(op.value)
|
||||
self.visit_expr(op.body)
|
||||
|
||||
def visit_call_(self, op):
|
||||
"""Visitor implementation for Call."""
|
||||
|
||||
def _visit_arg(arg):
|
||||
self.visit_expr(arg)
|
||||
|
||||
_visit_array(op.args, _visit_arg)
|
||||
|
||||
def _visit_binary_op(self, op):
|
||||
"""Helper to visit binary operators."""
|
||||
self.visit_expr(op.a)
|
||||
self.visit_expr(op.b)
|
||||
|
||||
def visit_add_(self, op):
|
||||
"""Visitor implementation for Add."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_sub_(self, op):
|
||||
"""Visitor implementation for Sub."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_mul_(self, op):
|
||||
"""Visitor implementation for Mul."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_div_(self, op):
|
||||
"""Visitor implementation for Div."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_mod_(self, op):
|
||||
"""Visitor implementation for Mod."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_floordiv_(self, op):
|
||||
"""Visitor implementation for FloorDiv."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_floormod_(self, op):
|
||||
"""Visitor implementation for FloorMod."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_min_(self, op):
|
||||
"""Visitor implementation for Min."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_max_(self, op):
|
||||
"""Visitor implementation for Max."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_eq_(self, op):
|
||||
"""Visitor implementation for EQ."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_ne_(self, op):
|
||||
"""Visitor implementation for NE."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_lt_(self, op):
|
||||
"""Visitor implementation for LT."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_le_(self, op):
|
||||
"""Visitor implementation for LE."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_gt_(self, op):
|
||||
"""Visitor implementation for GT."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_ge_(self, op):
|
||||
"""Visitor implementation for GE."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_and_(self, op):
|
||||
"""Visitor implementation for And."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_or_(self, op):
|
||||
"""Visitor implementation for Or."""
|
||||
self._visit_binary_op(op)
|
||||
|
||||
def visit_int_imm_(self, op):
|
||||
"""Visitor implementation for IntImm."""
|
||||
pass
|
||||
|
||||
def visit_float_imm_(self, op):
|
||||
"""Visitor implementation for FloatImm."""
|
||||
pass
|
||||
|
||||
def visit_string_imm_(self, op):
|
||||
"""Visitor implementation for StringImm."""
|
||||
pass
|
||||
|
||||
def visit_reduce_(self, op):
|
||||
"""Visitor implementation for Reduce."""
|
||||
|
||||
def _visit_iter_var(iv):
|
||||
self.visit_expr(iv.dom.min)
|
||||
self.visit_expr(iv.dom.extent)
|
||||
|
||||
def _visit_source(source):
|
||||
self.visit_expr(source)
|
||||
|
||||
_visit_array(op.axis, _visit_iter_var)
|
||||
_visit_array(op.source, _visit_source)
|
||||
|
||||
if op.init:
|
||||
_visit_array(op.init, _visit_source)
|
||||
|
||||
self.visit_expr(op.condition)
|
||||
|
||||
def visit_cast_(self, op):
|
||||
"""Visitor implementation for Cast."""
|
||||
self.visit_expr(op.value)
|
||||
|
||||
def visit_not_(self, op):
|
||||
"""Visitor implementation for Not."""
|
||||
self.visit_expr(op.a)
|
||||
|
||||
def visit_select_(self, op):
|
||||
"""Visitor implementation for Select."""
|
||||
self.visit_expr(op.condition)
|
||||
self.visit_expr(op.true_value)
|
||||
self.visit_expr(op.false_value)
|
||||
|
||||
def visit_ramp_(self, op):
|
||||
"""Visitor implementation for Ramp."""
|
||||
self.visit_expr(op.base)
|
||||
self.visit_expr(op.stride)
|
||||
self.visit_expr(op.lanes)
|
||||
|
||||
def visit_shuffle_(self, op):
|
||||
"""Visitor implementation for Shuffle."""
|
||||
|
||||
def _visit_expr(expr):
|
||||
self.visit_expr(expr)
|
||||
|
||||
_visit_array(op.indices, _visit_expr)
|
||||
_visit_array(op.vectors, _visit_expr)
|
||||
|
||||
def visit_broadcast_(self, op):
|
||||
"""Visitor implementation for Broadcast."""
|
||||
self.visit_expr(op.value)
|
||||
self.visit_expr(op.lanes)
|
||||
|
||||
|
||||
class ExprMutator(ExprFunctor):
|
||||
"""A mutator over Expr.
|
||||
|
||||
This is a mutator that recursively transforms an expression. Subclasses can
|
||||
override the visit methods to customize the behavior.
|
||||
"""
|
||||
|
||||
def visit_var_(self, op):
|
||||
"""Mutator implementation for Var."""
|
||||
return op
|
||||
|
||||
def visit_buffer_load_(self, op):
|
||||
"""Mutator implementation for BufferLoad."""
|
||||
indices = [self.visit_expr(index) for index in op.indices]
|
||||
|
||||
if all(old_index is new_index for old_index, new_index in zip(op.indices, indices)):
|
||||
return op
|
||||
else:
|
||||
return tvm.tirx.BufferLoad(op.buffer, indices, op.predicate)
|
||||
|
||||
def visit_producer_load_(self, op):
|
||||
"""Mutator implementation for ProducerLoad."""
|
||||
indices = [self.visit_expr(index) for index in op.indices]
|
||||
|
||||
if all(old_index is new_index for old_index, new_index in zip(op.indices, indices)):
|
||||
return op
|
||||
else:
|
||||
return tvm.tirx.ProducerLoad(op.producer, indices)
|
||||
|
||||
def visit_let_(self, op):
|
||||
"""Mutator implementation for Let."""
|
||||
var = self.visit_var_(op.var)
|
||||
value = self.visit_expr(op.value)
|
||||
body = self.visit_expr(op.body)
|
||||
|
||||
if var is op.var and value is op.value and body is op.body:
|
||||
return op
|
||||
else:
|
||||
return tvm.tirx.Let(var, value, body)
|
||||
|
||||
def visit_call_(self, op):
|
||||
"""Mutator implementation for Call."""
|
||||
args = [self.visit_expr(arg) for arg in op.args]
|
||||
|
||||
if all(old_arg is new_arg for old_arg, new_arg in zip(op.args, args)):
|
||||
return op
|
||||
else:
|
||||
return tvm.ir.Call(op.op, args, attrs=op.attrs, span=op.span, ret_ty=op.ty)
|
||||
|
||||
def _mutate_binary_op(self, op_cls, op):
|
||||
"""Helper to mutate binary operators."""
|
||||
a = self.visit_expr(op.a)
|
||||
b = self.visit_expr(op.b)
|
||||
|
||||
if a is op.a and b is op.b:
|
||||
return op
|
||||
else:
|
||||
return op_cls(a, b)
|
||||
|
||||
def visit_add_(self, op):
|
||||
"""Mutator implementation for Add."""
|
||||
return self._mutate_binary_op(tvm.tirx.Add, op)
|
||||
|
||||
def visit_sub_(self, op):
|
||||
"""Mutator implementation for Sub."""
|
||||
return self._mutate_binary_op(tvm.tirx.Sub, op)
|
||||
|
||||
def visit_mul_(self, op):
|
||||
"""Mutator implementation for Mul."""
|
||||
return self._mutate_binary_op(tvm.tirx.Mul, op)
|
||||
|
||||
def visit_div_(self, op):
|
||||
"""Mutator implementation for Div."""
|
||||
return self._mutate_binary_op(tvm.tirx.Div, op)
|
||||
|
||||
def visit_mod_(self, op):
|
||||
"""Mutator implementation for Mod."""
|
||||
return self._mutate_binary_op(tvm.tirx.Mod, op)
|
||||
|
||||
def visit_floordiv_(self, op):
|
||||
"""Mutator implementation for FloorDiv."""
|
||||
return self._mutate_binary_op(tvm.tirx.FloorDiv, op)
|
||||
|
||||
def visit_floormod_(self, op):
|
||||
"""Mutator implementation for FloorMod."""
|
||||
return self._mutate_binary_op(tvm.tirx.FloorMod, op)
|
||||
|
||||
def visit_min_(self, op):
|
||||
"""Mutator implementation for Min."""
|
||||
return self._mutate_binary_op(tvm.tirx.Min, op)
|
||||
|
||||
def visit_max_(self, op):
|
||||
"""Mutator implementation for Max."""
|
||||
return self._mutate_binary_op(tvm.tirx.Max, op)
|
||||
|
||||
def visit_eq_(self, op):
|
||||
"""Mutator implementation for EQ."""
|
||||
return self._mutate_binary_op(tvm.tirx.EQ, op)
|
||||
|
||||
def visit_ne_(self, op):
|
||||
"""Mutator implementation for NE."""
|
||||
return self._mutate_binary_op(tvm.tirx.NE, op)
|
||||
|
||||
def visit_lt_(self, op):
|
||||
"""Mutator implementation for LT."""
|
||||
return self._mutate_binary_op(tvm.tirx.LT, op)
|
||||
|
||||
def visit_le_(self, op):
|
||||
"""Mutator implementation for LE."""
|
||||
return self._mutate_binary_op(tvm.tirx.LE, op)
|
||||
|
||||
def visit_gt_(self, op):
|
||||
"""Mutator implementation for GT."""
|
||||
return self._mutate_binary_op(tvm.tirx.GT, op)
|
||||
|
||||
def visit_ge_(self, op):
|
||||
"""Mutator implementation for GE."""
|
||||
return self._mutate_binary_op(tvm.tirx.GE, op)
|
||||
|
||||
def visit_and_(self, op):
|
||||
"""Mutator implementation for And."""
|
||||
return self._mutate_binary_op(tvm.tirx.And, op)
|
||||
|
||||
def visit_or_(self, op):
|
||||
"""Mutator implementation for Or."""
|
||||
return self._mutate_binary_op(tvm.tirx.Or, op)
|
||||
|
||||
def visit_int_imm_(self, op):
|
||||
"""Mutator implementation for IntImm."""
|
||||
return op
|
||||
|
||||
def visit_float_imm_(self, op):
|
||||
"""Mutator implementation for FloatImm."""
|
||||
return op
|
||||
|
||||
def visit_string_imm_(self, op):
|
||||
"""Mutator implementation for StringImm."""
|
||||
return op
|
||||
|
||||
def visit_reduce_(self, op):
|
||||
"""Mutator implementation for Reduce."""
|
||||
|
||||
def _mutate_iter_var(iv):
|
||||
old_dom = iv.dom
|
||||
new_min = self.visit_expr(old_dom.min)
|
||||
new_extent = self.visit_expr(old_dom.extent)
|
||||
|
||||
if new_min is old_dom.min and new_extent is old_dom.extent:
|
||||
return iv
|
||||
else:
|
||||
new_dom = Range.FromMinExtent(new_min, new_extent)
|
||||
return IterVar(new_dom, iv.var, iv.iter_type, iv.thread_tag)
|
||||
|
||||
axis = [_mutate_iter_var(iv) for iv in op.axis]
|
||||
source = [self.visit_expr(e) for e in op.source]
|
||||
init = [self.visit_expr(e) for e in op.init] if op.init else []
|
||||
condition = self.visit_expr(op.condition)
|
||||
|
||||
axis_unchanged = all(old_iv is new_iv for old_iv, new_iv in zip(op.axis, axis))
|
||||
source_unchanged = all(old_e is new_e for old_e, new_e in zip(op.source, source))
|
||||
init_unchanged = (
|
||||
True if not op.init else all(old_e is new_e for old_e, new_e in zip(op.init, init))
|
||||
)
|
||||
condition_unchanged = condition is op.condition
|
||||
|
||||
if axis_unchanged and source_unchanged and init_unchanged and condition_unchanged:
|
||||
return op
|
||||
else:
|
||||
return tvm.tirx.Reduce(op.combiner, source, axis, condition, op.value_index, init)
|
||||
|
||||
def visit_cast_(self, op):
|
||||
"""Mutator implementation for Cast."""
|
||||
value = self.visit_expr(op.value)
|
||||
|
||||
if value is op.value:
|
||||
return op
|
||||
else:
|
||||
return tvm.tirx.Cast(op.ty, value)
|
||||
|
||||
def visit_not_(self, op):
|
||||
"""Mutator implementation for Not."""
|
||||
a = self.visit_expr(op.a)
|
||||
|
||||
if a is op.a:
|
||||
return op
|
||||
else:
|
||||
return tvm.tirx.Not(a)
|
||||
|
||||
def visit_select_(self, op):
|
||||
"""Mutator implementation for Select."""
|
||||
condition = self.visit_expr(op.condition)
|
||||
true_value = self.visit_expr(op.true_value)
|
||||
false_value = self.visit_expr(op.false_value)
|
||||
|
||||
if (
|
||||
condition is op.condition
|
||||
and true_value is op.true_value
|
||||
and false_value is op.false_value
|
||||
):
|
||||
return op
|
||||
else:
|
||||
return tvm.tirx.Select(condition, true_value, false_value)
|
||||
|
||||
def visit_ramp_(self, op):
|
||||
"""Mutator implementation for Ramp."""
|
||||
base = self.visit_expr(op.base)
|
||||
stride = self.visit_expr(op.stride)
|
||||
lanes = self.visit_expr(op.lanes)
|
||||
|
||||
if base is op.base and stride is op.stride and lanes is op.lanes:
|
||||
return op
|
||||
else:
|
||||
return tvm.tirx.Ramp(base, stride, lanes)
|
||||
|
||||
def visit_broadcast_(self, op):
|
||||
"""Mutator implementation for Broadcast."""
|
||||
value = self.visit_expr(op.value)
|
||||
lanes = self.visit_expr(op.lanes)
|
||||
|
||||
if value is op.value and lanes is op.lanes:
|
||||
return op
|
||||
else:
|
||||
return tvm.tirx.Broadcast(value, lanes)
|
||||
|
||||
def visit_shuffle_(self, op):
|
||||
"""Mutator implementation for Shuffle."""
|
||||
vectors = [self.visit_expr(v) for v in op.vectors]
|
||||
|
||||
vectors_unchanged = all(old_v is new_v for old_v, new_v in zip(op.vectors, vectors))
|
||||
|
||||
if vectors_unchanged:
|
||||
return op
|
||||
else:
|
||||
return tvm.tirx.Shuffle(vectors, op.indices)
|
||||
@@ -0,0 +1,578 @@
|
||||
# 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=unrecognized-inline-option
|
||||
"""Function data types."""
|
||||
|
||||
import collections
|
||||
import inspect
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Optional
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
import tvm
|
||||
import tvm.runtime
|
||||
from tvm.ir import BaseFunc, Range
|
||||
from tvm.runtime import Object, Scriptable
|
||||
|
||||
from ..runtime._tensor import Tensor
|
||||
from . import _ffi_api
|
||||
from .buffer import Buffer
|
||||
from .expr import Expr, Var
|
||||
|
||||
|
||||
@tvm_ffi.register_object("tirx.PrimFunc")
|
||||
class PrimFunc(BaseFunc, Scriptable):
|
||||
"""A function declaration expression.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
params: List[Union[tvm.tirx.Var, tvm.tirx.Buffer]]
|
||||
List of input parameters to the function.
|
||||
|
||||
body: tvm.tirx.Stmt
|
||||
The body of the function.
|
||||
|
||||
ret_type: tvm.ir.Type
|
||||
The return type annotation of the function.
|
||||
|
||||
buffer_map : Map[tvm.tirx.Var, tvm.tirx.Buffer]
|
||||
The buffer binding map.
|
||||
|
||||
attrs: Optional[tvm.Attrs]
|
||||
Attributes of the function, can be None
|
||||
|
||||
span : Optional[Span]
|
||||
The location of this itervar in the source code.
|
||||
"""
|
||||
|
||||
def __init__(self, params, body, ret_type=None, buffer_map=None, attrs=None, span=None):
|
||||
# Legacy compatibility: expand body-carrying leaf stmt wrappers
|
||||
# (e.g. DeclBuffer/AllocBuffer forms) into SeqStmt form.
|
||||
from .stmt import _normalize_legacy_stmt
|
||||
|
||||
body = _normalize_legacy_stmt(body)
|
||||
if ret_type is None:
|
||||
ret_type = tvm.ir.Type.missing()
|
||||
param_list = []
|
||||
buffer_map = {} if buffer_map is None else buffer_map
|
||||
for x in params:
|
||||
x = tvm.runtime.convert(x) if not isinstance(x, Object) else x
|
||||
if isinstance(x, Buffer):
|
||||
var = Var(x.name, dtype="handle")
|
||||
param_list.append(var)
|
||||
buffer_map[var] = x
|
||||
elif isinstance(x, Var):
|
||||
param_list.append(x)
|
||||
else:
|
||||
raise TypeError("params can only contain Var or Buffer")
|
||||
|
||||
if attrs is None:
|
||||
attrs = tvm.ir.make_node("ir.DictAttrs")
|
||||
|
||||
self.__init_handle_by_constructor__(
|
||||
_ffi_api.PrimFunc,
|
||||
param_list,
|
||||
body,
|
||||
ret_type,
|
||||
buffer_map,
|
||||
attrs,
|
||||
span,
|
||||
) # type: ignore
|
||||
|
||||
def with_body(self, new_body, span=None):
|
||||
"""Create a new PrimFunc with the same set signatures but a new body.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
new_body : Stmt
|
||||
The new body.
|
||||
|
||||
span : Optional[Span]
|
||||
The location of this itervar in the source code.
|
||||
|
||||
Returns
|
||||
-------
|
||||
new_func : PrimFunc
|
||||
The created new function.
|
||||
"""
|
||||
return PrimFunc(
|
||||
self.params,
|
||||
new_body,
|
||||
self.ret_type,
|
||||
self.buffer_map,
|
||||
self.attrs,
|
||||
span,
|
||||
)
|
||||
|
||||
def specialize(self, param_map: Mapping[Var, Expr | Buffer]):
|
||||
"""Specialize parameters of PrimFunc
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
param_map : Mapping[Var, Union[Expr, Buffer]]
|
||||
The mapping from function params to the instance
|
||||
|
||||
Examples
|
||||
--------
|
||||
We can define a Meta TIR function with symbolic shape:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def mem_copy(a: T.handle, b: T.handle, m: T.int32, n: T.int32) -> None:
|
||||
A = T.match_buffer(a, (m, n), "float32")
|
||||
B = T.match_buffer(b, (m, n), "float32")
|
||||
|
||||
for i, j in T.grid(m, n):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
B[vi, vj] = A[vi, vj]
|
||||
|
||||
Then we can make it specialized with given shapes or buffers.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
a, _, m, n = mem_copy.params
|
||||
func = mem_copy.specialize({a: tirx.decl_buffer((16, 16))})
|
||||
# or
|
||||
func = mem_copy.specialize({n: 16, m: 16})
|
||||
|
||||
The specialized function:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def mem_copy_16_16(a: T.handle, b: T.handle) -> None:
|
||||
A = T.match_buffer(a, (16, 16), "float32")
|
||||
B = T.match_buffer(b, (16, 16), "float32")
|
||||
|
||||
for i, j in T.grid(16, 16):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
B[vi, vj] = A[vi, vj]
|
||||
|
||||
Returns
|
||||
-------
|
||||
func : PrimFunc
|
||||
The new function with parameter specialized
|
||||
"""
|
||||
return _ffi_api.Specialize(self, param_map) # type: ignore
|
||||
|
||||
|
||||
@tvm_ffi.register_object("tirx.TensorIntrin")
|
||||
class TensorIntrin(Object):
|
||||
"""A tensor intrinsic.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
desc : PrimFunc
|
||||
The function to describe the computation.
|
||||
|
||||
impl : PrimFunc
|
||||
The function of the implementation for the execution.
|
||||
"""
|
||||
|
||||
def __init__(self, desc, impl):
|
||||
self.__init_handle_by_constructor__(_ffi_api.TensorIntrin, desc, impl)
|
||||
|
||||
@staticmethod
|
||||
def register(name: str, desc: PrimFunc, impl: PrimFunc, override: bool = False):
|
||||
"""Register a tensor intrinsic with its name.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
The name of the TensorIntrin to register.
|
||||
desc : PrimFunc
|
||||
The function to describe the computation.
|
||||
impl : PrimFunc
|
||||
The function of the implementation for the execution.
|
||||
override: bool
|
||||
Whether override existing intrinsic.
|
||||
"""
|
||||
return _ffi_api.TensorIntrinRegister(name, TensorIntrin(desc, impl), override) # type: ignore
|
||||
|
||||
@staticmethod
|
||||
def get(name: str, allow_missing: bool = False) -> Optional["TensorIntrin"]:
|
||||
"""Look up a tensor intrinsic by its name.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
The name of the TensorIntrin to look up.
|
||||
|
||||
allow_missing : bool
|
||||
Whether to allow missing tensor intrin. If False, raise an error if the tensor intrin
|
||||
doesn't exist.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : Optional[TensorIntrin]
|
||||
The TensorIntrin with the specified name, or None if not found.
|
||||
"""
|
||||
return _ffi_api.TensorIntrinGet(name, allow_missing) # pylint: type: ignore
|
||||
|
||||
|
||||
@tvm_ffi.register_object("tirx.IndexMap")
|
||||
class IndexMap(Object):
|
||||
"""A mapping from multi-dimensional indices to another set of multi-dimensional indices
|
||||
|
||||
Parameters
|
||||
----------
|
||||
initial_indices : List[Var]
|
||||
Variables representing the indices prior to remapping.
|
||||
final_indices : List[Expr]
|
||||
Expressions defining the indices after remapping.
|
||||
inverse_index_map : Union[Callable, Optional[IndexMap]]
|
||||
The optional pre-defined inverse index map.
|
||||
When this is defined, IndexMap::Inverse will return the pre-defined inverse index map.
|
||||
Otherwise, the inverse index map will be computed on the fly.
|
||||
It is the user's responsibility to ensure the correctness of the pre-defined inverse
|
||||
index map.
|
||||
"""
|
||||
|
||||
initial_indices: list[Var]
|
||||
final_indices: list[Expr]
|
||||
|
||||
# Sentinel value used to indicate which groups of pre-flattening axes
|
||||
# should be used to post-flattening axes axes. See
|
||||
# Stage.transform_layout for more details.
|
||||
AXIS_SEPARATOR = "axis_separator"
|
||||
|
||||
def __init__(self, initial_indices, final_indices, inverse_index_map):
|
||||
if isinstance(inverse_index_map, Callable):
|
||||
inverse_index_map = IndexMap.from_func(inverse_index_map)
|
||||
self.__init_handle_by_constructor__(
|
||||
_ffi_api.IndexMap, initial_indices, final_indices, inverse_index_map
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def from_func(
|
||||
mapping_function: Callable,
|
||||
ndim: int | None = None,
|
||||
inverse_index_map: Callable | Optional["IndexMap"] = None,
|
||||
*,
|
||||
index_dtype: str = "int64",
|
||||
):
|
||||
"""Create an index map from a function
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mapping_function : Callable
|
||||
|
||||
The function to map from source indices to target indices.
|
||||
The function should accept `tirx.Var` parameters and return
|
||||
a either a `tirx.Expr`, or a list of `tirx.Expr`.
|
||||
Returning a `tirx.Expr` is equivalent to returning a
|
||||
list of length 1 containing that `tirx.Expr`.
|
||||
|
||||
ndim: Optional[int]
|
||||
|
||||
The dimensionality of the buffer to which this
|
||||
transformation should be applied. If mapping_function uses
|
||||
variadic argument `*args`, `ndim` must be specified. If
|
||||
mapping_function does not use variadic arguments, ndim is
|
||||
optional.
|
||||
|
||||
inverse_index_map : Union[Callable, Optional[IndexMap]]
|
||||
The optional pre-defined inverse index map.
|
||||
When this is defined, IndexMap::Inverse will return the pre-defined inverse index map.
|
||||
Otherwise, the inverse index map will be computed on the fly.
|
||||
It is the user's responsibility to ensure the correctness of the pre-defined inverse
|
||||
index map.
|
||||
|
||||
Returns
|
||||
-------
|
||||
index_map: IndexMap
|
||||
|
||||
Returns an IndexMap representing the `mapping_function`.
|
||||
|
||||
"""
|
||||
index_map, axis_separators = IndexMap.from_func_with_separators(
|
||||
mapping_function,
|
||||
ndim,
|
||||
inverse_index_map,
|
||||
index_dtype=index_dtype,
|
||||
)
|
||||
assert not axis_separators, (
|
||||
"The mapping_function provided to IndexMap.from_func "
|
||||
"may not return IndexMap.AXIS_SEPARATOR. "
|
||||
"If required, please use IndexMap.from_func_with_separators instead."
|
||||
)
|
||||
return index_map
|
||||
|
||||
@staticmethod
|
||||
def from_func_with_separators(
|
||||
mapping_function: Callable,
|
||||
ndim: int | None = None,
|
||||
inverse_index_map: Callable | Optional["IndexMap"] = None,
|
||||
*,
|
||||
index_dtype: str = "int64",
|
||||
):
|
||||
"""Create an index map from a function
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mapping_function : Callable
|
||||
|
||||
The function to map from source indices to target indices.
|
||||
The function should accept tirx.Var parameters and return
|
||||
either a `tirx.Expr` or a list. Each element of the
|
||||
returned list should be either a `tirx.Expr` or the
|
||||
object `IndexMap.AXIS_SEPARATOR`. Returning a
|
||||
`tirx.Expr` is equivalent to returning a list of length
|
||||
1 containing that `tirx.Expr`.
|
||||
|
||||
ndim: Optional[int]
|
||||
|
||||
The dimensionality of the buffer to which this
|
||||
transformation should be applied. If mapping_function uses
|
||||
variadic argument `*args`, ndim must be specified. If
|
||||
mapping_function does not use variadic arguments, ndim is
|
||||
optional.
|
||||
|
||||
inverse_index_map : Union[Callable, Optional[IndexMap]]
|
||||
The optional pre-defined inverse index map.
|
||||
When this is defined, IndexMap::Inverse will return the pre-defined inverse index map.
|
||||
Otherwise, the inverse index map will be computed on the fly.
|
||||
It is the user's responsibility to ensure the correctness of the pre-defined inverse
|
||||
index map.
|
||||
|
||||
index_dtype : str
|
||||
The default index dtype to use for input iters in the mapping function.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret: Tuple[IndexMap, List[int]]
|
||||
|
||||
Returns a tuple whose first element is an IndexMap
|
||||
representing the `mapping_function`, and whose second index
|
||||
is a list of indices at which `IndexMap.AXIS_SEPARATOR`
|
||||
occurred.
|
||||
|
||||
"""
|
||||
params = inspect.signature(mapping_function).parameters
|
||||
|
||||
args = []
|
||||
var_arg_name = None
|
||||
kwargs = collections.OrderedDict()
|
||||
|
||||
for name, param in params.items():
|
||||
if param.kind in [
|
||||
inspect.Parameter.POSITIONAL_ONLY,
|
||||
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||||
]:
|
||||
args.append(tvm.tirx.Var(name, index_dtype))
|
||||
|
||||
elif param.kind == inspect.Parameter.VAR_POSITIONAL:
|
||||
var_arg_name = name
|
||||
|
||||
elif param.kind == inspect.Parameter.KEYWORD_ONLY:
|
||||
kwargs[name] = tvm.tirx.Var(name, index_dtype)
|
||||
|
||||
else:
|
||||
raise ValueError("transform_layout mapping may not have *args")
|
||||
|
||||
# Now that all the named arguments have been collected,
|
||||
# everything that remains should go to the *args, if
|
||||
# specified.
|
||||
if var_arg_name is not None:
|
||||
assert ndim is not None, "ndim must be specified when *args is used"
|
||||
num_var_args = ndim - len(args) - len(kwargs)
|
||||
for i in range(num_var_args):
|
||||
args.append(tvm.tirx.Var(f"{var_arg_name}_{i}", index_dtype))
|
||||
|
||||
mapping = mapping_function(*args, **kwargs)
|
||||
|
||||
initial_indices = args + list(kwargs.values())
|
||||
|
||||
final_indices = []
|
||||
axis_separators = []
|
||||
|
||||
try:
|
||||
iter(mapping)
|
||||
is_iterable = True
|
||||
except TypeError:
|
||||
is_iterable = False
|
||||
|
||||
if is_iterable:
|
||||
for val in mapping:
|
||||
if tvm.ir.is_prim_expr(val):
|
||||
final_indices.append(val)
|
||||
elif val is IndexMap.AXIS_SEPARATOR:
|
||||
axis_separators.append(len(final_indices))
|
||||
else:
|
||||
raise TypeError(
|
||||
"Expected mapping function to return list of "
|
||||
"either tvm.ir.Expr or IndexMap.AXIS_SEPARATOR. "
|
||||
f"Instead received {val} of type {type(val)}."
|
||||
)
|
||||
else:
|
||||
final_indices.append(mapping)
|
||||
|
||||
return IndexMap(initial_indices, final_indices, inverse_index_map), axis_separators
|
||||
|
||||
def is_equivalent_to(self, other_map: "IndexMap", analyzer=None) -> bool:
|
||||
"""Return if the index maps are equivalent.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
other_map: IndexMap
|
||||
|
||||
The IndexMap to which the comparison should be made.
|
||||
|
||||
analyzer : Optional[tvm.arith.Analyzer]
|
||||
|
||||
The analyzer to use while comparing the mapped indices. When
|
||||
provided, its accumulated bindings and constraints are reused so
|
||||
that maps that are only equivalent under those bindings can be
|
||||
proven equal.
|
||||
|
||||
Returns
|
||||
-------
|
||||
is_equivalent: bool
|
||||
|
||||
True if the two mappings represent the same
|
||||
transformation, otherwise False
|
||||
"""
|
||||
if len(self.initial_indices) != len(other_map.initial_indices):
|
||||
return False
|
||||
if len(self.final_indices) != len(other_map.final_indices):
|
||||
return False
|
||||
|
||||
if analyzer is None:
|
||||
analyzer = tvm.arith.Analyzer()
|
||||
|
||||
mapped_other_final_indices = other_map.map_indices(self.initial_indices, analyzer=analyzer)
|
||||
for self_index, other_index in zip(self.final_indices, mapped_other_final_indices):
|
||||
if not analyzer.can_prove_equal(self_index, other_index):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def map_indices(self, indices: list[Expr], analyzer=None) -> list[Expr]:
|
||||
"""Apply the index map to a set of indices
|
||||
|
||||
Parameters
|
||||
----------
|
||||
indices : List[Expr]
|
||||
The indices to be mapped
|
||||
analyzer : Optional[tvm.arith.Analyzer]
|
||||
The analyzer to use while simplifying mapped indices.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : List[Expr]
|
||||
The mapped indices
|
||||
"""
|
||||
return _ffi_api.IndexMapMapIndices(self, indices, analyzer)
|
||||
|
||||
def map_shape(self, shape: list[Expr], analyzer=None) -> list[Expr]:
|
||||
"""Apply the index map to a buffer shape
|
||||
|
||||
Parameters
|
||||
----------
|
||||
shape : List[Expr]
|
||||
The buffer shape to be mapped
|
||||
analyzer : Optional[tvm.arith.Analyzer]
|
||||
The analyzer to use while simplifying mapped shape expressions.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : List[Expr]
|
||||
The mapped shape
|
||||
"""
|
||||
return _ffi_api.IndexMapMapShape(self, shape, analyzer)
|
||||
|
||||
def map_tensor(self, arr_src: Tensor) -> Tensor:
|
||||
"""Apply thie index map to transform the layout of the input Tensor
|
||||
|
||||
Parameters
|
||||
----------
|
||||
arr_src : runtime.Tensor
|
||||
The Tensor to be transformed
|
||||
|
||||
Returns
|
||||
-------
|
||||
arr_dst : runtime.Tensor
|
||||
The transformed Tensor
|
||||
"""
|
||||
return _ffi_api.IndexMapMapTensor(self, arr_src)
|
||||
|
||||
def inverse(self, shape: list[Range | Expr], analyzer=None) -> "IndexMap":
|
||||
"""Return the inverse of the map
|
||||
|
||||
Throws an error if the function is not bijective.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
shape: List[Union[Range,Expr]]
|
||||
|
||||
The region over which the inverse should be determined.
|
||||
Used for validating that the mapping is bijective over
|
||||
this range.
|
||||
analyzer : Optional[tvm.arith.Analyzer]
|
||||
The analyzer to use while deriving and validating the inverse.
|
||||
|
||||
Returns
|
||||
-------
|
||||
inverse : IndexMap
|
||||
|
||||
The inverse
|
||||
"""
|
||||
|
||||
shape = [dim if isinstance(dim, Range) else Range(0, dim) for dim in shape]
|
||||
return _ffi_api.IndexMapInverse(self, shape, analyzer)
|
||||
|
||||
def non_surjective_inverse(
|
||||
self, shape: list[Range | Expr], analyzer=None
|
||||
) -> tuple["IndexMap", Expr]:
|
||||
"""Return the inverse of the map
|
||||
|
||||
Can be applied to transformations that introduce padding.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
shape: List[Union[Range,Expr]]
|
||||
|
||||
The region over which the inverse should be determined.
|
||||
Used for determining the predicate.
|
||||
analyzer : Optional[tvm.arith.Analyzer]
|
||||
The analyzer to use while deriving the inverse and padding predicate.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : Tuple[IndexMap, Expr]
|
||||
|
||||
The inverse, and a predicate for which the inverse maps to
|
||||
a valid index in the input range.
|
||||
|
||||
Examples
|
||||
--------
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
index_map = IndexMap.from_func(lambda i: [i//4, i%4])
|
||||
inverse_map, predicate = index_map.non_surjective_inverse([14])
|
||||
assert inverse_map.is_equivalent_to(IndexMap.from_func(lambda j,k: [4*j + k])
|
||||
print(predicate) # Prints "(axis0==3) && (axis2 >= 2)"
|
||||
"""
|
||||
|
||||
shape = [dim if isinstance(dim, Range) else Range(0, dim) for dim in shape]
|
||||
return _ffi_api.IndexMapNonSurjectiveInverse(self, shape, analyzer)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,16 @@
|
||||
# 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.
|
||||
@@ -0,0 +1,8 @@
|
||||
# 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.
|
||||
"""Compatibility redirect for CUDA allocation pool helpers."""
|
||||
|
||||
from tvm.backend.cuda.lang.alloc_pool import * # noqa: F403 # pylint: disable=wildcard-import
|
||||
@@ -0,0 +1,8 @@
|
||||
# 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.
|
||||
"""Compatibility redirect for CUDA pipeline helpers."""
|
||||
|
||||
from tvm.backend.cuda.lang.pipeline import * # noqa: F403 # pylint: disable=wildcard-import
|
||||
@@ -0,0 +1,8 @@
|
||||
# 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.
|
||||
"""Compatibility redirect for CUDA shared-memory descriptors."""
|
||||
|
||||
from tvm.backend.cuda.lang.smem_desc import * # noqa: F403 # pylint: disable=wildcard-import
|
||||
@@ -0,0 +1,8 @@
|
||||
# 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.
|
||||
"""Compatibility redirect for CUDA tile schedulers."""
|
||||
|
||||
from tvm.backend.cuda.lang.tile_scheduler import * # noqa: F403 # pylint: disable=wildcard-import
|
||||
@@ -0,0 +1,8 @@
|
||||
# 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.
|
||||
"""Compatibility redirect for CUDA warp role helpers."""
|
||||
|
||||
from tvm.backend.cuda.lang.warp_role import * # noqa: F403 # pylint: disable=wildcard-import
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,41 @@
|
||||
# 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.
|
||||
|
||||
# `tile_primitive` defines Python Op classes (`Zero(UnaryOp)`, etc.) whose
|
||||
# class bodies call `Op.get("tirx.<name>")` at class-definition time, which
|
||||
# requires the compiler-side FFI. Load it lazily so that
|
||||
# `tvm.tirx.operator.intrinsics._common` (pure data) and other runtime-safe
|
||||
# submodules can be imported under `TVM_USE_RUNTIME_LIB=1`, matching apache's
|
||||
# discipline for `tvm.tirx`.
|
||||
def __getattr__(name):
|
||||
# `from . import tile_primitive` here would recurse: Python's import
|
||||
# machinery does `getattr(self, 'tile_primitive')` to see if the submodule
|
||||
# is already loaded, which goes back through this __getattr__. Use
|
||||
# importlib.import_module to bypass attribute lookup; it sets the attribute
|
||||
# on the parent package as a side effect, so subsequent lookups go through
|
||||
# the normal attribute path, not this __getattr__.
|
||||
import sys # pylint: disable=import-outside-toplevel
|
||||
from importlib import import_module # pylint: disable=import-outside-toplevel
|
||||
|
||||
tp_qualname = f"{__name__}.tile_primitive"
|
||||
tile_primitive = sys.modules.get(tp_qualname) or import_module(tp_qualname)
|
||||
if hasattr(tile_primitive, name):
|
||||
return getattr(tile_primitive, name)
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
__all__ = ["get_tirx_op"]
|
||||
@@ -0,0 +1,62 @@
|
||||
# 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.
|
||||
"""Shared enum / value tables for PTX intrinsic schemas and user wrappers.
|
||||
|
||||
Single source of truth. Both ``tvm.tirx.op`` (user wrappers that validate
|
||||
arguments via ``_choice``) and ``tvm.tirx.cuda.operator.intrinsics.*``
|
||||
(schema declarations using ``Choice(choices=...)`` / ``IntAttr(choices=...)``)
|
||||
import from here.
|
||||
|
||||
Adding a new modifier value requires changing exactly one place.
|
||||
"""
|
||||
|
||||
# Memory ordering / scope -----------------------------------------------------
|
||||
FENCE_SEM = ("sc", "acq_rel")
|
||||
FENCE_SCOPE = ("cta", "cluster", "gpu", "sys")
|
||||
FENCE_PROXY_ASYNC_SPACE = ("", "global", "shared::cta", "shared::cluster")
|
||||
CLUSTER_BARRIER_SEM = ("", "release", "relaxed")
|
||||
|
||||
# CTA group (used by tcgen05 and TMA) -----------------------------------------
|
||||
TCGEN05_CTA_GROUP = (1, 2)
|
||||
|
||||
# NVSHMEM ---------------------------------------------------------------------
|
||||
NVSHMEM_CMP = ("eq", "ne", "gt", "ge", "lt", "le")
|
||||
NVSHMEM_SIG_OP = ("set", "add")
|
||||
|
||||
# Floating-point rounding -----------------------------------------------------
|
||||
F32X2_ROUND = ("rz", "rn", "rm", "rp")
|
||||
|
||||
# cp.async (non-bulk) ---------------------------------------------------------
|
||||
CP_ASYNC_CACHE_HINT = ("", "evict_last", "evict_first", "evict_normal")
|
||||
CP_ASYNC_PREFETCH_SIZE = (-1, 64, 128, 256)
|
||||
CP_ASYNC_FILL_MODE = ("", "zero")
|
||||
|
||||
# cp.async.bulk (TMA) ---------------------------------------------------------
|
||||
CP_ASYNC_BULK_CACHE_HINT = ("", "evict_last", "evict_first", "evict_normal", "evict_last_use")
|
||||
CP_ASYNC_BULK_RED_OP = ("add", "min", "max", "inc", "dec", "and", "or", "xor")
|
||||
|
||||
# ldmatrix / stmatrix ---------------------------------------------------------
|
||||
LDMATRIX_DTYPE = (".b16", ".b8")
|
||||
LDMATRIX_NUM = (1, 2, 4)
|
||||
|
||||
# tcgen05.cp ------------------------------------------------------------------
|
||||
TCGEN05_CP_SHAPES = ("32x128b", "4x256b", "128x128b", "128x256b", "64x128b")
|
||||
TCGEN05_CP_MULTICAST = ("", "warpx4", "warpx2::02_13", "warpx2::01_23")
|
||||
TCGEN05_CP_DECOMPRESS = ("", "b8x16.b4x16_p64", "b8x16.b6x16_p32")
|
||||
|
||||
# tcgen05.ld / tcgen05.st -----------------------------------------------------
|
||||
TCGEN05_LDST_SHAPES = ("16x32bx2", "16x64b", "16x128b", "16x256b", "32x32b")
|
||||
@@ -0,0 +1,30 @@
|
||||
# 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.
|
||||
|
||||
# ruff: noqa: I001
|
||||
|
||||
# Op class declarations (Add, Sub, Gemm, ...) — must run first so their
|
||||
# `op = Op.get("tirx.<name>")` registrations execute before any dispatch
|
||||
# code refers to the same ops.
|
||||
from .ops import *
|
||||
|
||||
# Dispatch infrastructure. Per-backend schedule registrations are loaded via
|
||||
# ``tvm.backend.load(<name>)``.
|
||||
from .dispatcher import fail, list_registered_schedules, predicate, register_dispatch
|
||||
from .registry import DispatchContext
|
||||
|
||||
__all__ = ["DispatchContext", "fail", "list_registered_schedules", "predicate", "register_dispatch"]
|
||||
@@ -0,0 +1,45 @@
|
||||
# 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.
|
||||
"""TIRx operator dispatch common utilities."""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class MapOpType(Enum):
|
||||
"""Enumeration of common unary and binary operator types."""
|
||||
|
||||
ADD = 0
|
||||
SUB = 1
|
||||
MUL = 2
|
||||
FDIV = 3
|
||||
ZERO = 4
|
||||
SQRT = 5
|
||||
RECIPROCAL = 6
|
||||
FILL = 7
|
||||
MAX = 8
|
||||
MIN = 9
|
||||
EXP = 10
|
||||
EXP2 = 11
|
||||
SILU = 12
|
||||
|
||||
|
||||
class ReduceOpType(Enum):
|
||||
"""Enumeration of common reduce operator types."""
|
||||
|
||||
SUM = 0
|
||||
MAX = 1
|
||||
MIN = 2
|
||||
@@ -0,0 +1,209 @@
|
||||
# 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.
|
||||
"""TIRx operator dispatch context."""
|
||||
|
||||
from tvm_ffi import register_object
|
||||
|
||||
from tvm.ir import Range
|
||||
from tvm.runtime import Object, Scriptable
|
||||
from tvm.target import Target
|
||||
from tvm.tirx import Buffer, IterVar, Stmt, Var, _ffi_api
|
||||
from tvm.tirx.exec_scope import ExecScope
|
||||
|
||||
|
||||
@register_object("tirx.DispatchContext")
|
||||
class DispatchContext(Object, Scriptable):
|
||||
"""DispatchContext node.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
target : Target
|
||||
The target of the dispatch context.
|
||||
|
||||
exec_scope : ExecScope
|
||||
The execution scope of the dispatch context.
|
||||
|
||||
launch_params : Dict[str, Expr]
|
||||
The launch parameters of the dispatch context.
|
||||
|
||||
var_range_map : Dict[Var, Range]
|
||||
A map from loop variables to their ranges.
|
||||
|
||||
callbacks : Dict[str, Object]
|
||||
The callbacks of the dispatch context.
|
||||
|
||||
shared_state : Dict[str, Object]
|
||||
Shared state persisting across dispatch calls within a single lowering pass.
|
||||
"""
|
||||
|
||||
target: Target
|
||||
exec_scope: ExecScope
|
||||
launch_params: dict[str, IterVar]
|
||||
var_range_map: dict[Var, Range]
|
||||
alloc_only: bool
|
||||
callbacks: dict[str, Object]
|
||||
shared_state: dict[str, Object]
|
||||
inter: dict[str, list]
|
||||
intra: dict[str, list]
|
||||
scope_kind: str
|
||||
|
||||
kPrivateAlloc = "private_alloc"
|
||||
kDeviceInitStmt = "device_init_stmt"
|
||||
kHostInitStmt = "host_init_stmt"
|
||||
kPostBufferDefStmt = "post_buffer_def_stmt"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
target: Target,
|
||||
exec_scope: ExecScope,
|
||||
launch_params: dict[str, IterVar],
|
||||
var_range_map: dict[Var, Range],
|
||||
alloc_only: bool = False,
|
||||
callbacks: dict[str, Object] = {},
|
||||
shared_state: dict[str, Object] = {},
|
||||
inter: dict[str, list] | None = None,
|
||||
intra: dict[str, list] | None = None,
|
||||
scope_kind: str = "",
|
||||
) -> None:
|
||||
self.__init_handle_by_constructor__(
|
||||
_ffi_api.DispatchContext, # pylint: disable=no-member
|
||||
target,
|
||||
exec_scope,
|
||||
launch_params,
|
||||
var_range_map,
|
||||
alloc_only,
|
||||
callbacks,
|
||||
shared_state,
|
||||
inter or {},
|
||||
intra or {},
|
||||
scope_kind,
|
||||
)
|
||||
|
||||
def add_alloc_buffer(self, buffer: Buffer) -> None:
|
||||
"""Add an allocated buffer to the dispatch context.
|
||||
Can be called only if alloc_only is True.
|
||||
The buffer will be added to the workspace of operator (the key in the workspace is the buffer name).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
buffer : Buffer
|
||||
The buffer to be added.
|
||||
""" # noqa: E501
|
||||
_ffi_api.DispatchContextAddAllocBuffer(self, buffer) # pylint: disable=no-member
|
||||
|
||||
def add_init_stmt(self, stmt: Stmt, host: bool = False) -> None:
|
||||
"""Add an initialization statement to the dispatch context.
|
||||
Device initialization statements is only allowed if alloc_only is True.
|
||||
Host initialization statements will be ignored if alloc_only is True.
|
||||
The statements will be added to the beginning of the kernel.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
stmt : Stmt
|
||||
The initialization statement to be added.
|
||||
host : bool
|
||||
Whether the statement is a host statement.
|
||||
If True, the statement will be added to the host code (before the kernel).
|
||||
If False, the statement will be added to the kernel body (at the beginning of the kernel).
|
||||
""" # noqa: E501
|
||||
_ffi_api.DispatchContextAddInitStmt(self, stmt, host) # pylint: disable=no-member
|
||||
|
||||
def add_post_buffer_def_stmt(self, buffer: Buffer, stmt: Stmt) -> None:
|
||||
"""Add a statement to be inserted after a buffer's definition (DeclBuffer/AllocBuffer).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
buffer : Buffer
|
||||
The buffer whose definition scope the statement should appear in.
|
||||
stmt : Stmt
|
||||
The statement to be inserted.
|
||||
"""
|
||||
_ffi_api.DispatchContextAddPostBufferDefStmt(self, buffer, stmt) # pylint: disable=no-member
|
||||
|
||||
def cache_get(self, key: str) -> Object | None:
|
||||
"""Look up a cached value by key.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
key : str
|
||||
Cache key (built by the caller from construction parameters).
|
||||
|
||||
Returns
|
||||
-------
|
||||
Optional[Object]
|
||||
The cached value, or None on miss.
|
||||
"""
|
||||
return _ffi_api.DispatchContextSharedStateGet(self, key)
|
||||
|
||||
def cache_set(self, key: str, value: Object) -> None:
|
||||
"""Store a value in the cross-dispatch cache.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
key : str
|
||||
Cache key (built by the caller from construction parameters).
|
||||
value : Object
|
||||
The object to cache (e.g. a Buffer or Var).
|
||||
"""
|
||||
_ffi_api.DispatchContextSharedStateSet(self, key, value)
|
||||
|
||||
def is_cuda(self) -> bool:
|
||||
"""Check if the target is CUDA."""
|
||||
return self.target.kind.name == "cuda"
|
||||
|
||||
def is_trn(self) -> bool:
|
||||
"""Check if the target is Trainium."""
|
||||
return self.target.kind.name == "trn"
|
||||
|
||||
def is_target(self, name: str) -> bool:
|
||||
"""Check if the target kind matches ``name``."""
|
||||
return self.target.kind.name == name
|
||||
|
||||
# -- scope predicates ----------------------------------------------------
|
||||
#
|
||||
# Each ``is_<scope>`` returns True iff the op site is at that scope kind.
|
||||
# Backed by ``self.scope_kind``, which 1-1 maps to a canonical intra
|
||||
# TileLayout shape:
|
||||
# thread -> {}
|
||||
# warp -> {laneid}
|
||||
# warpgroup -> {laneid, wid_in_wg}
|
||||
# cta -> {laneid, warpid}
|
||||
# cluster -> {laneid, warpid, cta_id}
|
||||
#
|
||||
# Prefer these predicates over raw ``self.scope_kind == "..."`` comparisons
|
||||
# so dispatchers that later need stricter intra/inter shape checks can
|
||||
# tighten the predicate body without touching every call site.
|
||||
|
||||
@property
|
||||
def is_thread(self) -> bool:
|
||||
return self.scope_kind == "thread"
|
||||
|
||||
@property
|
||||
def is_warp(self) -> bool:
|
||||
return self.scope_kind == "warp"
|
||||
|
||||
@property
|
||||
def is_warpgroup(self) -> bool:
|
||||
return self.scope_kind == "warpgroup"
|
||||
|
||||
@property
|
||||
def is_cta(self) -> bool:
|
||||
return self.scope_kind == "cta"
|
||||
|
||||
@property
|
||||
def is_cluster(self) -> bool:
|
||||
return self.scope_kind == "cluster"
|
||||
@@ -0,0 +1,329 @@
|
||||
# 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.
|
||||
"""Rich dispatcher for TIRx operator dispatchs.
|
||||
|
||||
This module adds a structured dispatch table with predicates and
|
||||
deterministic failure reporting via exceptions.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import traceback
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from tvm.ir import Op
|
||||
from tvm.tirx import PrimFunc
|
||||
from tvm.tirx.operator import get_tirx_op
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from .dispatch_context import DispatchContext
|
||||
|
||||
|
||||
class DispatchFail(RuntimeError):
|
||||
"""Raised by variants or predicates to provide a reasoned failure."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class Predicate:
|
||||
"""A named predicate. The callable can return:
|
||||
|
||||
- bool
|
||||
- (bool, str) where the second element is an optional reason on failure
|
||||
- raise DispatchFail(reason)
|
||||
"""
|
||||
|
||||
name: str
|
||||
fn: Callable[[TilePrimitiveCall, DispatchContext], Any]
|
||||
kwargs: dict[str, Any]
|
||||
|
||||
def evaluate(
|
||||
self, op_call: TilePrimitiveCall, sctx: DispatchContext
|
||||
) -> tuple[bool, str | None]:
|
||||
try:
|
||||
out = self.fn(op_call, sctx, **self.kwargs)
|
||||
if isinstance(out, tuple):
|
||||
ok, reason = out
|
||||
return bool(ok), (str(reason) if not ok and reason is not None else None)
|
||||
return bool(out), None
|
||||
except DispatchFail as e: # surface explicit failure reasons
|
||||
return False, str(e)
|
||||
except Exception as e: # unexpected predicate exception
|
||||
return False, f"predicate exception: {type(e).__name__}: {e}"
|
||||
|
||||
|
||||
def predicate(
|
||||
name: str, fn: Callable[[TilePrimitiveCall, DispatchContext], Any], **kwargs
|
||||
) -> Predicate:
|
||||
"""Wrap a callable into a named predicate."""
|
||||
|
||||
return Predicate(name=name, fn=fn, kwargs=kwargs)
|
||||
|
||||
|
||||
def fail(reason: str) -> None:
|
||||
"""Helper for schedule variants to explain why they decline to handle the op."""
|
||||
|
||||
raise DispatchFail(reason)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DispatchCase:
|
||||
variant: str
|
||||
priority: int
|
||||
preds: list[Predicate]
|
||||
# Impl must either return a PrimFunc or raise DispatchFail
|
||||
impl: Callable[[TilePrimitiveCall, DispatchContext], PrimFunc]
|
||||
|
||||
|
||||
# Keyed by (Op, target_kind)
|
||||
_DISPATCH_TABLE: dict[tuple[Op, str], list[DispatchCase]] = {}
|
||||
|
||||
|
||||
def _target_kind_name(sctx: DispatchContext) -> str:
|
||||
"""Normalize target kind to a stable dispatch key."""
|
||||
|
||||
kind = getattr(getattr(sctx, "target", None), "kind", None)
|
||||
return getattr(kind, "name", str(kind))
|
||||
|
||||
|
||||
def register_dispatch(
|
||||
op_name: str,
|
||||
target_kind: str,
|
||||
*,
|
||||
variant: str,
|
||||
priority: int = 0,
|
||||
when: list[Predicate] | None = None,
|
||||
):
|
||||
"""Decorator to add a dispatch case for an op/target pair.
|
||||
|
||||
Cases with higher priority run earlier. When list predicates must all pass.
|
||||
The impl must return a PrimFunc on success, and must NOT return None.
|
||||
To decline handling, raise `fail("reason")` (or `DispatchFail`).
|
||||
"""
|
||||
|
||||
op = get_tirx_op(op_name)
|
||||
|
||||
def decorator(impl: Callable[[TilePrimitiveCall, DispatchContext], Any]):
|
||||
# Wrap impl to forbid returning None; require raise-or-PrimFunc
|
||||
def wrapped_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
|
||||
res = impl(op_call, sctx)
|
||||
if res is None:
|
||||
# Enforce raise-or-PrimFunc contract for schedule implementations
|
||||
raise DispatchFail(
|
||||
"impl returned None; schedule must return PrimFunc or raise fail()"
|
||||
)
|
||||
return res # type: ignore[return-value]
|
||||
|
||||
cases = _DISPATCH_TABLE.setdefault((op, target_kind), [])
|
||||
cases.append(
|
||||
DispatchCase(variant=variant, priority=priority, preds=when or [], impl=wrapped_impl)
|
||||
)
|
||||
return impl
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def list_registered_schedules() -> dict[str, dict[str, list[str]]]:
|
||||
"""Return a mapping: op_name -> target_kind -> [variant names]."""
|
||||
|
||||
out: dict[str, dict[str, list[str]]] = {}
|
||||
for (op, tgt), cases in _DISPATCH_TABLE.items():
|
||||
name = op.name
|
||||
out.setdefault(name, {}).setdefault(tgt, [])
|
||||
# keep insertion order by default; sort by priority desc for readability
|
||||
for c in sorted(cases, key=lambda x: (-x.priority, x.variant)):
|
||||
out[name][tgt].append(c.variant)
|
||||
return out
|
||||
|
||||
|
||||
def _format_opcall(op_call: TilePrimitiveCall) -> str:
|
||||
"""Return a readable representation of the failing opcall."""
|
||||
# Prefer TVMScript or IR text printer if available on this object
|
||||
try:
|
||||
script_method = getattr(op_call, "script", None)
|
||||
if callable(script_method):
|
||||
try:
|
||||
return str(script_method())
|
||||
except TypeError:
|
||||
# Some versions may require keyword args; fall back safely
|
||||
return str(script_method())
|
||||
astext_method = getattr(op_call, "astext", None)
|
||||
if callable(astext_method):
|
||||
return str(astext_method())
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
s = str(op_call)
|
||||
# constrain extremely long single-line prints from repr
|
||||
return s
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
args_len = len(getattr(op_call, "args", []))
|
||||
except Exception:
|
||||
args_len = -1
|
||||
try:
|
||||
op_name = op_call.op.name # type: ignore[attr-defined]
|
||||
except Exception:
|
||||
op_name = "<unknown-op>"
|
||||
return f"op={op_name}, args={args_len}"
|
||||
|
||||
|
||||
def _format_failure_table(header: str, rows: list[tuple[str, list[str]]]) -> str:
|
||||
"""Format failures into a readable ASCII table.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
header : str
|
||||
The header line describing the op/target
|
||||
rows : List[Tuple[str, str, Optional[str]]]
|
||||
Each row is (variant_label, error_summary, traceback_str)
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
The formatted report string
|
||||
"""
|
||||
# Compute column widths
|
||||
variant_header = "Variant"
|
||||
error_header = "Error"
|
||||
variant_col_w = (
|
||||
max(len(variant_header), *(len(v) for (v, _) in rows)) if rows else len(variant_header)
|
||||
)
|
||||
# Error column width needs to consider multi-line cells
|
||||
if rows:
|
||||
error_col_w = max(
|
||||
len(error_header), *(max(len(line) for line in errs) for (_, errs) in rows)
|
||||
)
|
||||
else:
|
||||
error_col_w = len(error_header)
|
||||
|
||||
def hline(sep: str = "+") -> str:
|
||||
return f"{sep}{'-' * (variant_col_w + 2)}{sep}{'-' * (error_col_w + 2)}{sep}"
|
||||
|
||||
lines: list[str] = [header]
|
||||
if not rows:
|
||||
# No rows; keep the header only
|
||||
return "\n".join(lines)
|
||||
|
||||
# Table header
|
||||
lines.append(hline("+"))
|
||||
lines.append(f"| {variant_header.ljust(variant_col_w)} | {error_header.ljust(error_col_w)} |")
|
||||
lines.append(hline("+"))
|
||||
|
||||
# Rows (support multi-line Error column)
|
||||
for variant, errs in rows:
|
||||
if not errs:
|
||||
errs = [""]
|
||||
for i, err_line in enumerate(errs):
|
||||
v_text = variant if i == 0 else ""
|
||||
lines.append(f"| {v_text.ljust(variant_col_w)} | {err_line.ljust(error_col_w)} |")
|
||||
lines.append(hline("+"))
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def run_dispatch(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
|
||||
"""Run structured dispatch.
|
||||
|
||||
Returns a PrimFunc on success. Otherwise, raises RuntimeError with
|
||||
an aggregated reason report.
|
||||
"""
|
||||
|
||||
target_kind = _target_kind_name(sctx)
|
||||
key = (op_call.op, target_kind)
|
||||
cases = _DISPATCH_TABLE.get(key)
|
||||
if not cases:
|
||||
header = f"TIRx schedule dispatch failed: op={op_call.op.name} target={target_kind}"
|
||||
report = _format_failure_table(header, [])
|
||||
# Append a simple reason when there are no variants at all
|
||||
report = "\n".join([report, "no registered variants for this op/target"])
|
||||
raise RuntimeError(report)
|
||||
|
||||
# Collect structured failure rows: (variant_label, error_lines)
|
||||
# error_lines: [summary, traceback lines...]
|
||||
failure_rows: list[tuple[str, list[str]]] = []
|
||||
last_exception: BaseException | None = None
|
||||
|
||||
# If explicit dispatch is set, filter to that variant only
|
||||
forced_variant = getattr(op_call, "dispatch", None)
|
||||
if forced_variant is not None:
|
||||
cases = [c for c in cases if c.variant == forced_variant]
|
||||
if not cases:
|
||||
msg_header = f"TIRx schedule dispatch failed: op={op_call.op.name} target={target_kind}"
|
||||
table = _format_failure_table(msg_header, [])
|
||||
msg = "\n".join([table, f"no variant named '{forced_variant}' is registered"])
|
||||
raise RuntimeError(msg)
|
||||
|
||||
for case in sorted(cases, key=lambda c: (-c.priority, c.variant)):
|
||||
# evaluate predicates
|
||||
pred_ok = True
|
||||
pred_msgs: list[str] = []
|
||||
for pred in case.preds:
|
||||
ok, reason = pred.evaluate(op_call, sctx)
|
||||
if not ok:
|
||||
pred_ok = False
|
||||
msg = f"rejected: {pred.name}"
|
||||
if reason:
|
||||
msg += f" — {reason}"
|
||||
pred_msgs.append(msg)
|
||||
if not pred_ok:
|
||||
# Include the offending TilePrimitiveCall IR in the error cell
|
||||
op_str = _format_opcall(op_call)
|
||||
op_lines = [line.rstrip("\n") for line in str(op_str).splitlines()] if op_str else []
|
||||
failure_rows.append(
|
||||
(
|
||||
f"{case.variant} (prio={case.priority})",
|
||||
["; ".join(pred_msgs), "opcall:", *op_lines],
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
# run impl
|
||||
try:
|
||||
res = case.impl(op_call, sctx)
|
||||
# Defensive check in case a legacy impl bypassed the wrapper
|
||||
if res is None: # pragma: no cover - legacy guard
|
||||
raise DispatchFail("impl returned None (legacy behavior not allowed)")
|
||||
return res
|
||||
except DispatchFail as e:
|
||||
op_str = _format_opcall(op_call)
|
||||
op_lines = [line.rstrip("\n") for line in str(op_str).splitlines()] if op_str else []
|
||||
failure_rows.append(
|
||||
(
|
||||
f"{case.variant} (prio={case.priority})",
|
||||
[f"declined — {e!s}", "opcall:", *op_lines],
|
||||
)
|
||||
)
|
||||
except Exception as e: # keep searching other variants
|
||||
exc_summary = f"exception — {type(e).__name__}: {e}"
|
||||
tb_str = "".join(traceback.format_exception(type(e), e, e.__traceback__))
|
||||
# Expand traceback into lines
|
||||
tb_lines = [line.rstrip("\n") for line in tb_str.splitlines()]
|
||||
op_str = _format_opcall(op_call)
|
||||
op_lines = [line.rstrip("\n") for line in str(op_str).splitlines()] if op_str else []
|
||||
error_lines = [exc_summary, "opcall:", *op_lines, *tb_lines]
|
||||
failure_rows.append((f"{case.variant} (prio={case.priority})", error_lines))
|
||||
last_exception = e
|
||||
|
||||
# no success
|
||||
header = f"TIRx schedule dispatch failed: op={op_call.op.name} target={target_kind}"
|
||||
report = _format_failure_table(header, failure_rows)
|
||||
if last_exception is not None:
|
||||
raise RuntimeError(report) from last_exception
|
||||
raise RuntimeError(report)
|
||||
@@ -0,0 +1,567 @@
|
||||
# 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.
|
||||
|
||||
"""Implementation of TIR operator."""
|
||||
|
||||
from tvm.ir import Op
|
||||
from tvm.tirx import Expr
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
|
||||
def get_tirx_op(op_name: str):
|
||||
assert isinstance(op_name, str)
|
||||
return Op.get("tirx.tile." + op_name)
|
||||
|
||||
|
||||
class ArgProperty:
|
||||
def __init__(self, index):
|
||||
self.index = index
|
||||
|
||||
def __get__(self, obj, objtype=None):
|
||||
assert obj is not None, "TilePrimitiveCall cannot be None"
|
||||
return obj.args[self.index]
|
||||
|
||||
|
||||
### Base Operator Classes ###
|
||||
class UnaryOp(TilePrimitiveCall):
|
||||
"""Base class for unary operators: unary(output, input).
|
||||
|
||||
Unary operators take a single input tensor and produce a single output tensor.
|
||||
"""
|
||||
|
||||
scalar_input = False
|
||||
output = ArgProperty(0)
|
||||
input = ArgProperty(1)
|
||||
|
||||
@property
|
||||
def srcs(self) -> list[Expr]:
|
||||
"""Get the source expression (input) of the operator."""
|
||||
return [self.input]
|
||||
|
||||
@property
|
||||
def dsts(self) -> list[Expr]:
|
||||
"""Get the destination expression (output) of the operator."""
|
||||
return [self.output]
|
||||
|
||||
|
||||
class UnaryOpWithBiasScale(UnaryOp):
|
||||
"""Extended unary operator with bias and scale parameters: unary_with_bias_scale(output, input, bias, scale).
|
||||
|
||||
These operators support additional bias and scale parameters for more complex operations (only on trn).
|
||||
output = unary(input * scale + bias)
|
||||
""" # noqa: E501
|
||||
|
||||
bias = ArgProperty(2)
|
||||
scale = ArgProperty(3)
|
||||
|
||||
@property
|
||||
def srcs(self) -> list[Expr]:
|
||||
"""Get the source expressions (inputs) of the operator."""
|
||||
return [self.input, self.bias, self.scale]
|
||||
|
||||
|
||||
class BinaryOp(TilePrimitiveCall):
|
||||
"""Base class for binary operators: binary(output, input0, input1).
|
||||
|
||||
Binary operators take two input tensors and produce a single output tensor.
|
||||
"""
|
||||
|
||||
lhs = ArgProperty(1)
|
||||
rhs = ArgProperty(2)
|
||||
output = ArgProperty(0)
|
||||
|
||||
@property
|
||||
def srcs(self) -> list[Expr]:
|
||||
"""Get the source expressions (inputs) of the operator."""
|
||||
return [self.lhs, self.rhs]
|
||||
|
||||
@property
|
||||
def dsts(self) -> list[Expr]:
|
||||
"""Get the destination expression (output) of the operator."""
|
||||
return [self.output]
|
||||
|
||||
|
||||
class ReduceOp(TilePrimitiveCall):
|
||||
"""Base class for reduction operators: reduce(output, input, reduce_axes, accum).
|
||||
|
||||
Reduction operators reduce one or more dimensions of the input tensor.
|
||||
"""
|
||||
|
||||
input = ArgProperty(1)
|
||||
output = ArgProperty(0)
|
||||
reduce_axes = ArgProperty(2)
|
||||
accum = ArgProperty(3)
|
||||
|
||||
@property
|
||||
def srcs(self) -> list[Expr]:
|
||||
"""Get the source expression (input) of the operator."""
|
||||
return [self.input]
|
||||
|
||||
@property
|
||||
def dsts(self) -> list[Expr]:
|
||||
"""Get the destination expression (output) of the operator."""
|
||||
return [self.output]
|
||||
|
||||
|
||||
### Schedule Operators ###
|
||||
class Zero(UnaryOp):
|
||||
"""Zero out all elements in src and store to dst."""
|
||||
|
||||
op = get_tirx_op("zero")
|
||||
|
||||
|
||||
class Sqrt(UnaryOpWithBiasScale):
|
||||
"""Compute square root of all elements in src and store to dst.
|
||||
|
||||
If bias and scale are provided: dst = sqrt(src * scale + bias)
|
||||
"""
|
||||
|
||||
op = get_tirx_op("sqrt")
|
||||
|
||||
|
||||
class Fill(UnaryOp):
|
||||
"""Fill dst with a scalar value."""
|
||||
|
||||
op = get_tirx_op("fill")
|
||||
scalar_input = True
|
||||
|
||||
|
||||
class Add(BinaryOp):
|
||||
"""Add src1 and src2 element-wise and store to dst."""
|
||||
|
||||
op = get_tirx_op("add")
|
||||
|
||||
|
||||
class Sub(BinaryOp):
|
||||
"""Subtract src2 from src1 element-wise and store to dst."""
|
||||
|
||||
op = get_tirx_op("sub")
|
||||
|
||||
|
||||
class Mul(BinaryOp):
|
||||
"""Multiply src1 and src2 element-wise and store to dst."""
|
||||
|
||||
op = get_tirx_op("mul")
|
||||
|
||||
|
||||
class FDiv(BinaryOp):
|
||||
"""Divide src1 by src2 element-wise using floating point division and store to dst."""
|
||||
|
||||
op = get_tirx_op("fdiv")
|
||||
|
||||
|
||||
class FMA(TilePrimitiveCall):
|
||||
"""Fused multiply-add: output = input * scale + bias.
|
||||
|
||||
fma(output, input, scale, bias)
|
||||
|
||||
scale and bias can each be either a BufferRegion or a Expr scalar.
|
||||
"""
|
||||
|
||||
op = get_tirx_op("fma")
|
||||
|
||||
output = ArgProperty(0)
|
||||
input = ArgProperty(1)
|
||||
scale = ArgProperty(2)
|
||||
bias = ArgProperty(3)
|
||||
|
||||
@property
|
||||
def srcs(self) -> list[Expr]:
|
||||
"""Get the source expressions (inputs) of the operator."""
|
||||
return [self.input, self.scale, self.bias]
|
||||
|
||||
@property
|
||||
def dsts(self) -> list[Expr]:
|
||||
"""Get the destination expression (output) of the operator."""
|
||||
return [self.output]
|
||||
|
||||
|
||||
class Cast(UnaryOp):
|
||||
"""Cast src to dst."""
|
||||
|
||||
op = get_tirx_op("cast")
|
||||
|
||||
|
||||
class Copy(TilePrimitiveCall):
|
||||
"""Copy all elements from src to dst.
|
||||
|
||||
Args:
|
||||
dst: Destination buffer region
|
||||
src: Source buffer region
|
||||
"""
|
||||
|
||||
op = get_tirx_op("copy")
|
||||
|
||||
dst = ArgProperty(0)
|
||||
src = ArgProperty(1)
|
||||
|
||||
@property
|
||||
def srcs(self) -> list[Expr]:
|
||||
"""Get the source expressions (inputs) of the operator."""
|
||||
return [self.src]
|
||||
|
||||
@property
|
||||
def dsts(self) -> list[Expr]:
|
||||
"""Get the destination expressions (outputs) of the operator."""
|
||||
return [self.dst]
|
||||
|
||||
|
||||
class CopyAsync(TilePrimitiveCall):
|
||||
"""Copy all elements from src to dst asynchronously.
|
||||
|
||||
Args:
|
||||
dst: Destination buffer region
|
||||
src: Source buffer region
|
||||
"""
|
||||
|
||||
op = get_tirx_op("copy_async")
|
||||
|
||||
dst = ArgProperty(0)
|
||||
src = ArgProperty(1)
|
||||
|
||||
@property
|
||||
def srcs(self) -> list[Expr]:
|
||||
"""Get the source expressions (inputs) of the operator."""
|
||||
return [self.src]
|
||||
|
||||
@property
|
||||
def dsts(self) -> list[Expr]:
|
||||
"""Get the destination expressions (outputs) of the operator."""
|
||||
return [self.dst]
|
||||
|
||||
|
||||
class Gemm(TilePrimitiveCall):
|
||||
"""General matrix multiplication: D = A * B * alpha + C * beta.
|
||||
|
||||
Args:
|
||||
D: Output matrix
|
||||
A: First input matrix
|
||||
B: Second input matrix
|
||||
C: Third input matrix (for bias)
|
||||
transpose_A: Whether to transpose A
|
||||
transpose_B: Whether to transpose B
|
||||
alpha: Scalar multiplier for A*B
|
||||
beta: Scalar multiplier for C
|
||||
"""
|
||||
|
||||
op = get_tirx_op("gemm")
|
||||
output = ArgProperty(0)
|
||||
lhs = ArgProperty(1)
|
||||
rhs = ArgProperty(2)
|
||||
bias = ArgProperty(3)
|
||||
transpose_A = ArgProperty(4)
|
||||
transpose_B = ArgProperty(5)
|
||||
alpha = ArgProperty(6)
|
||||
beta = ArgProperty(7)
|
||||
|
||||
@property
|
||||
def srcs(self) -> list[Expr]:
|
||||
"""Get the source matrices."""
|
||||
return [self.lhs, self.rhs, self.bias]
|
||||
|
||||
@property
|
||||
def dsts(self) -> list[Expr]:
|
||||
"""Get the destination matrix."""
|
||||
return [self.output]
|
||||
|
||||
|
||||
class GemmAsync(TilePrimitiveCall):
|
||||
"""General matrix multiplication asynchronously.
|
||||
|
||||
Supports two arg layouts:
|
||||
- Regular (6 args): C, A, B, transA, transB, accum
|
||||
- Block-scaled (8 args): C, A, B, SFA, SFB, transA, transB, accum
|
||||
"""
|
||||
|
||||
op = get_tirx_op("gemm_async")
|
||||
output = ArgProperty(0)
|
||||
lhs = ArgProperty(1)
|
||||
rhs = ArgProperty(2)
|
||||
|
||||
@property
|
||||
def is_block_scaled(self) -> bool:
|
||||
"""Whether this is a block-scaled MMA operation."""
|
||||
return len(self.args) == 8
|
||||
|
||||
@property
|
||||
def sfa(self):
|
||||
"""Get the scale factor buffer for A (None for regular MMA)."""
|
||||
return self.args[3] if self.is_block_scaled else None
|
||||
|
||||
@property
|
||||
def sfb(self):
|
||||
"""Get the scale factor buffer for B (None for regular MMA)."""
|
||||
return self.args[4] if self.is_block_scaled else None
|
||||
|
||||
@property
|
||||
def transA(self):
|
||||
return self.args[5] if self.is_block_scaled else self.args[3]
|
||||
|
||||
@property
|
||||
def transB(self):
|
||||
return self.args[6] if self.is_block_scaled else self.args[4]
|
||||
|
||||
@property
|
||||
def accum(self):
|
||||
return self.args[7] if self.is_block_scaled else self.args[5]
|
||||
|
||||
@property
|
||||
def srcs(self) -> list[Expr]:
|
||||
"""Get the source matrices (including scale factors if block-scaled)."""
|
||||
srcs = [self.lhs, self.rhs]
|
||||
if self.is_block_scaled:
|
||||
srcs.extend([self.sfa, self.sfb])
|
||||
return srcs
|
||||
|
||||
@property
|
||||
def dsts(self) -> list[Expr]:
|
||||
"""Get the destination matrix."""
|
||||
return [self.output]
|
||||
|
||||
|
||||
class Sum(ReduceOp):
|
||||
"""Sum elements in src along specified axes and store in dst."""
|
||||
|
||||
op = get_tirx_op("sum")
|
||||
|
||||
|
||||
class Max(ReduceOp):
|
||||
"""Compute maximum value in src along specified axes and store in dst."""
|
||||
|
||||
op = get_tirx_op("max")
|
||||
|
||||
|
||||
class Min(ReduceOp):
|
||||
"""Compute minimum value in src along specified axes and store in dst."""
|
||||
|
||||
op = get_tirx_op("min")
|
||||
|
||||
|
||||
class Reciprocal(UnaryOp):
|
||||
"""Compute reciprocal (1/x) for all elements in src and store to dst."""
|
||||
|
||||
op = get_tirx_op("reciprocal")
|
||||
|
||||
|
||||
class SiLU(UnaryOp):
|
||||
"""Compute SiLU (x * sigmoid(x)) for all elements in src and store to dst."""
|
||||
|
||||
op = get_tirx_op("silu")
|
||||
|
||||
|
||||
class Memset(UnaryOp):
|
||||
"""Set all elements in dst to a specified value."""
|
||||
|
||||
op = get_tirx_op("memset")
|
||||
scalar_input = True
|
||||
|
||||
|
||||
class Maximum(BinaryOp):
|
||||
"""Compute element-wise maximum of src1 and src2 and store to dst."""
|
||||
|
||||
op = get_tirx_op("maximum")
|
||||
|
||||
|
||||
class Minimum(BinaryOp):
|
||||
"""Compute element-wise minimum of src1 and src2 and store to dst."""
|
||||
|
||||
op = get_tirx_op("minimum")
|
||||
|
||||
|
||||
class Exp(UnaryOpWithBiasScale):
|
||||
"""Compute exponential (e^x) of all elements in src and store to dst.
|
||||
|
||||
If bias and scale are provided: dst = exp(src * scale + bias)
|
||||
"""
|
||||
|
||||
op = get_tirx_op("exp")
|
||||
|
||||
|
||||
class Exp2(UnaryOpWithBiasScale):
|
||||
"""Compute base-2 exponential (2^x) of all elements in src and store to dst.
|
||||
|
||||
If bias and scale are provided: dst = exp2(src * scale + bias)
|
||||
"""
|
||||
|
||||
op = get_tirx_op("exp2")
|
||||
|
||||
|
||||
class Select(BinaryOp):
|
||||
"""Select elements from src1 or src2 based on the predicate.
|
||||
|
||||
select(dst, src1, src2, predicate)
|
||||
"""
|
||||
|
||||
op = get_tirx_op("select")
|
||||
predicate = ArgProperty(3)
|
||||
|
||||
|
||||
### Compose Ops ###
|
||||
class BinaryReduce(TilePrimitiveCall):
|
||||
"""Combine a binary operation with a reduction operation.
|
||||
|
||||
binary_reduce(binary_output, reduce_output, binary_input1, binary_input2, binary_op, reduce_op, reduce_axes, )
|
||||
""" # noqa: E501
|
||||
|
||||
op = get_tirx_op("binary_reduce")
|
||||
|
||||
binary_output = ArgProperty(0)
|
||||
reduce_output = ArgProperty(1)
|
||||
binary_input1 = ArgProperty(2)
|
||||
binary_input2 = ArgProperty(3)
|
||||
binary_op = ArgProperty(4)
|
||||
reduce_op = ArgProperty(5)
|
||||
reduce_axes = ArgProperty(6)
|
||||
|
||||
@property
|
||||
def srcs(self) -> list[Expr]:
|
||||
"""Get the source expressions (inputs) of the operator."""
|
||||
return [self.binary_input1, self.binary_input2]
|
||||
|
||||
@property
|
||||
def dsts(self) -> list[Expr]:
|
||||
"""Get the destination expressions (outputs) of the operator."""
|
||||
return [self.binary_output, self.reduce_output]
|
||||
|
||||
|
||||
class UnaryReduce(TilePrimitiveCall):
|
||||
"""Combine a unary operation with a reduction operation.
|
||||
|
||||
unary_reduce(unary_output, reduce_output, unary_input, unary_op, reduce_op, bias, scale, reduce_axes)
|
||||
""" # noqa: E501
|
||||
|
||||
op = get_tirx_op("unary_reduce")
|
||||
|
||||
unary_output = ArgProperty(0)
|
||||
reduce_output = ArgProperty(1)
|
||||
unary_input = ArgProperty(2)
|
||||
unary_op = ArgProperty(3)
|
||||
reduce_op = ArgProperty(4)
|
||||
bias = ArgProperty(5)
|
||||
scale = ArgProperty(6)
|
||||
reduce_axes = ArgProperty(7)
|
||||
|
||||
@property
|
||||
def srcs(self) -> list[Expr]:
|
||||
"""Get the source expressions (inputs) of the operator."""
|
||||
return [self.unary_input, self.bias, self.scale]
|
||||
|
||||
@property
|
||||
def dsts(self) -> list[Expr]:
|
||||
"""Get the destination expressions (outputs) of the operator."""
|
||||
return [self.unary_output, self.reduce_output]
|
||||
|
||||
|
||||
class BinaryChain(TilePrimitiveCall):
|
||||
"""Chain multiple binary operations together.
|
||||
|
||||
binary_chain(output, data, operand0, operand1, op0, op1, reverse1)
|
||||
|
||||
if not reverse1:
|
||||
output = (operand0 op0 data) op1 operand1
|
||||
else:
|
||||
output = operand1 op1 (operand0 op0 data)
|
||||
"""
|
||||
|
||||
op = get_tirx_op("binary_chain")
|
||||
|
||||
output = ArgProperty(0)
|
||||
data = ArgProperty(1)
|
||||
operand0 = ArgProperty(2)
|
||||
operand1 = ArgProperty(3)
|
||||
op0 = ArgProperty(4)
|
||||
op1 = ArgProperty(5)
|
||||
reverse1 = ArgProperty(6)
|
||||
|
||||
@property
|
||||
def srcs(self) -> list[Expr]:
|
||||
"""Get the source expressions (inputs) of the operator."""
|
||||
return [self.data, self.operand0, self.operand1]
|
||||
|
||||
@property
|
||||
def dsts(self) -> list[Expr]:
|
||||
"""Get the destination expressions (outputs) of the operator."""
|
||||
return [self.output]
|
||||
|
||||
|
||||
class ReduceNegate(ReduceOp):
|
||||
"""
|
||||
Negate the result of a reduction operation.
|
||||
|
||||
reduce_negate(output, input, reduce_axes, accum, reduce_op)
|
||||
"""
|
||||
|
||||
op = get_tirx_op("reduce_negate")
|
||||
|
||||
reduce_op = ArgProperty(4)
|
||||
|
||||
|
||||
class ComposeOp(TilePrimitiveCall):
|
||||
"""Generic operator for composition of multiple operations.
|
||||
|
||||
Must be lowered to specific compose operations before operator-level passes.
|
||||
"""
|
||||
|
||||
# TODO: add a pass to lower generic compose_op to specific compose ops
|
||||
|
||||
op = get_tirx_op("compose_op")
|
||||
|
||||
@property
|
||||
def srcs(self) -> list[Expr]:
|
||||
"""Get the source expressions (inputs) of the operator."""
|
||||
raise NotImplementedError(
|
||||
"Generic compose_op must be lowered to specific compose ops before operator-level passes" # noqa: E501
|
||||
)
|
||||
|
||||
@property
|
||||
def dsts(self) -> list[Expr]:
|
||||
"""Get the destination expressions (outputs) of the operator."""
|
||||
raise NotImplementedError(
|
||||
"Generic compose_op must be lowered to specific compose ops before operator-level passes" # noqa: E501
|
||||
)
|
||||
|
||||
|
||||
class PermuteLayout(TilePrimitiveCall):
|
||||
"""Move data so the buffer's bytes are arranged under a different layout.
|
||||
|
||||
Logical shape is preserved; only the byte placement changes. ``dst`` and
|
||||
``src`` carry their own TileLayouts; on lowering, the dispatcher reads
|
||||
those layouts and emits a register-staged warp transpose, optionally
|
||||
inserting a bank-conflict-avoiding XOR-swizzle on the per-lane register
|
||||
slots.
|
||||
|
||||
Args: ``permute_layout(dst_region, src_region)``.
|
||||
``dst`` and ``src`` may alias the same underlying SMEM (in-place).
|
||||
"""
|
||||
|
||||
op = get_tirx_op("permute_layout")
|
||||
|
||||
@property
|
||||
def dst(self) -> Expr:
|
||||
return self.args[0]
|
||||
|
||||
@property
|
||||
def src(self) -> Expr:
|
||||
return self.args[1]
|
||||
|
||||
@property
|
||||
def srcs(self) -> list[Expr]:
|
||||
return [self.src]
|
||||
|
||||
@property
|
||||
def dsts(self) -> list[Expr]:
|
||||
return [self.dst]
|
||||
@@ -0,0 +1,66 @@
|
||||
# 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.
|
||||
"""TIRx operator dispatch registry.
|
||||
|
||||
All operator dispatch is handled by the rich dispatcher. This module exposes
|
||||
the global entry `tirx.f_op_dispatcher` used by the C++ lowering pass to query a
|
||||
dispatch result.
|
||||
"""
|
||||
|
||||
from tvm_ffi import register_global_func
|
||||
|
||||
from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
# Note: legacy `register_schedule` is intentionally removed.
|
||||
|
||||
|
||||
@register_global_func("tirx.f_op_dispatcher")
|
||||
def f_op_dispatcher(op_call: TilePrimitiveCall, sctx: DispatchContext):
|
||||
"""Find and return a schedule for the operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
op_call : TilePrimitiveCall
|
||||
The operator to be scheduled
|
||||
sctx : DispatchContext
|
||||
The dispatch context
|
||||
|
||||
Returns
|
||||
-------
|
||||
Optional[PrimFunc]
|
||||
The result of the operator implementation
|
||||
"""
|
||||
assert sctx.target is not None, "Target not found"
|
||||
(op_call.op, str(sctx.target.kind))
|
||||
|
||||
# Use rich dispatcher for all dispatching
|
||||
try:
|
||||
from .dispatcher import run_dispatch # local import to avoid cycles
|
||||
except Exception: # pragma: no cover - fallback if import fails
|
||||
run_dispatch = None # type: ignore
|
||||
|
||||
if run_dispatch is not None:
|
||||
try:
|
||||
res = run_dispatch(op_call, sctx)
|
||||
except Exception:
|
||||
# propagate exceptions from dispatcher
|
||||
raise
|
||||
if res is not None:
|
||||
return res
|
||||
# Dispatcher reports errors on failure; unreachable on success
|
||||
return None
|
||||
@@ -0,0 +1,45 @@
|
||||
# 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=no-member
|
||||
"""Async structures for TIRX"""
|
||||
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
|
||||
from tvm_ffi import register_object
|
||||
|
||||
from tvm.runtime import Object
|
||||
from tvm.tirx import Expr, Var
|
||||
|
||||
from . import _ffi_api
|
||||
|
||||
|
||||
@register_object("tirx.Predicate")
|
||||
class Predicate(Object):
|
||||
"""A predicate object for TIRX"""
|
||||
|
||||
vars: list[Var]
|
||||
pred: Expr
|
||||
|
||||
def __init__(self, f_pred: Callable[..., Expr]):
|
||||
vars = [Var(name, "int32") for name in inspect.signature(f_pred).parameters]
|
||||
pred = f_pred(*vars)
|
||||
self.__init_handle_by_constructor__(_ffi_api.Predicate, vars, pred)
|
||||
|
||||
def apply(self, indices: list[Expr]) -> Expr:
|
||||
"""Apply the predicate to the given indices"""
|
||||
return _ffi_api.PredicateApply(self, indices)
|
||||
@@ -0,0 +1,37 @@
|
||||
# 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.
|
||||
"""TIRX-layer TVMScript pieces (parser, builder).
|
||||
|
||||
After the per-dialect TVMScript restructure, the TIRX layer owns its own
|
||||
``script/{parser,builder}`` subpackages. ``tvm.script.tirx`` resolves to
|
||||
this module via the dialect registry, so the public parser surface
|
||||
(``prim_func``, ``Buffer``, ``Ptr``, etc.) is re-exported here.
|
||||
"""
|
||||
|
||||
# pylint: disable=redefined-builtin,wildcard-import,unused-wildcard-import
|
||||
from .parser import *
|
||||
from .parser import Buffer, Ptr, prim_func
|
||||
|
||||
try:
|
||||
from .parser import macro
|
||||
except ImportError:
|
||||
macro = None
|
||||
from tvm.tirx.lang.alloc_pool import SMEMPool, TMEMPool, TMEMStages
|
||||
|
||||
from . import tile
|
||||
from .builder.ir import TensorMap, meta_class
|
||||
from .tile import cluster, cta, thread, warp, warpgroup, wg
|
||||
@@ -0,0 +1,26 @@
|
||||
# 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.
|
||||
"""Package tvm.script.ir_builder.tirx"""
|
||||
|
||||
from .ir import * # pylint: disable=wildcard-import,redefined-builtin
|
||||
from .ir import boolean as bool # pylint: disable=redefined-builtin
|
||||
from .ir import buffer as Buffer
|
||||
from .utils import buffer_proxy, frame_scope, seq_scope
|
||||
from tvm.tirx.lang.alloc_pool import SMEMPool, TMEMPool, TMEMStages
|
||||
from . import tirx as tile
|
||||
from .tirx import cluster, cta, thread, warp, warpgroup, wg
|
||||
@@ -0,0 +1,21 @@
|
||||
# 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.
|
||||
"""FFI APIs"""
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
tvm_ffi.init_ffi_api("script.ir_builder.tirx", __name__) # pylint: disable=protected-access
|
||||
@@ -0,0 +1,245 @@
|
||||
# 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.
|
||||
# ruff: noqa: E722
|
||||
"""External kernel integration fro TIR"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
from tvm import __version__ as tvm_version
|
||||
from tvm import tirx
|
||||
from tvm.ir import Expr, PointerType, is_prim_expr
|
||||
from tvm.runtime import Module, const
|
||||
from tvm.support import nvcc
|
||||
|
||||
|
||||
class BaseKernel: # pylint: disable=too-few-public-methods
|
||||
"""Base class for external kernels."""
|
||||
|
||||
def compile_to_device_module(
|
||||
self, launch_args, *args, **kwargs
|
||||
) -> tuple[str, Module, list[Any]]:
|
||||
"""Compile the kernel to a device module."""
|
||||
raise NotImplementedError()
|
||||
|
||||
def _format_tvm_module_metadata(self, kernel_name, arg_types, launch_param_tags):
|
||||
"""Format the TVM module metadata."""
|
||||
tvm_metadata = """{{
|
||||
"tvm_version": "{version}",
|
||||
"func_info": {{
|
||||
"{kernel_name}": {{
|
||||
"name": "",
|
||||
"arg_types": {arg_types},
|
||||
"launch_param_tags": {launch_param_tags}
|
||||
}}
|
||||
}}
|
||||
}}""".format_map(
|
||||
{
|
||||
"version": tvm_version,
|
||||
"kernel_name": kernel_name,
|
||||
"arg_types": json.dumps(arg_types),
|
||||
"launch_param_tags": json.dumps(launch_param_tags),
|
||||
}
|
||||
)
|
||||
return tvm_metadata
|
||||
|
||||
def _create_cuda_module(
|
||||
self, binary_data, kernel_arg_types, launch_param_tags, kernel_name, fmt="ptx"
|
||||
):
|
||||
"""
|
||||
Create a CUDA module from compiled binary (PTX or cubin) and metadata.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
binary_data : str or bytes
|
||||
The compiled binary data (PTX as str, cubin as bytes).
|
||||
|
||||
kernel_arg_types : List[str]
|
||||
The types of the kernel arguments.
|
||||
|
||||
launch_param_tags : List[str]
|
||||
The tags of the launch parameters.
|
||||
|
||||
kernel_name : str
|
||||
The name of the kernel.
|
||||
|
||||
fmt : str
|
||||
The format of the binary data: "ptx" or "cubin".
|
||||
|
||||
Returns
|
||||
-------
|
||||
kernel_module : Module
|
||||
The CUDA module.
|
||||
"""
|
||||
tvm_metadata = self._format_tvm_module_metadata(
|
||||
kernel_name, kernel_arg_types, launch_param_tags
|
||||
)
|
||||
# Build the FunctionInfo map in-memory from the JSON metadata, then
|
||||
# construct the CUDA module via the FFI registry without going to
|
||||
# disk. Avoids the load_from_file dispatch path entirely.
|
||||
if isinstance(binary_data, str):
|
||||
binary_bytes = binary_data.encode("utf-8")
|
||||
else:
|
||||
binary_bytes = bytes(binary_data)
|
||||
load_meta = tvm_ffi.get_global_func("runtime.LoadMetaDataFromJSON")
|
||||
fmap = load_meta(tvm_metadata)
|
||||
create_cuda = tvm_ffi.get_global_func("ffi.Module.create.cuda")
|
||||
kernel_module = create_cuda(binary_bytes, fmt, fmap, {})
|
||||
return kernel_module
|
||||
|
||||
|
||||
class SourceKernel(BaseKernel): # pylint: disable=too-few-public-methods
|
||||
"""A kernel from source code."""
|
||||
|
||||
def __init__(self, source_code: str):
|
||||
self.source_code = source_code
|
||||
|
||||
def compile_to_device_module( # pylint: disable=arguments-differ
|
||||
self,
|
||||
grid: list[list[int | tirx.Expr]],
|
||||
*args: list[Any],
|
||||
**kwargs: dict[str, Any],
|
||||
) -> tuple[str, Module, list[Any]]:
|
||||
"""Compile the kernel to a device module."""
|
||||
from tvm.relax.frontend.nn import ( # pylint: disable=import-outside-toplevel
|
||||
SourceModule,
|
||||
)
|
||||
|
||||
kernel_name = kwargs["kernel_name"]
|
||||
assert len(grid) == 2, (
|
||||
"grid should be two list of integers, representing the dimension of "
|
||||
"['blockIdx.x', 'blockIdx.y', 'blockIdx.z'] and "
|
||||
"['threadIdx.x', 'threadIdx.y', 'threadIdx.z']"
|
||||
)
|
||||
assert isinstance(grid[0], list | tuple) and isinstance(grid[1], list | tuple)
|
||||
launch_param_tags = ["blockIdx.x", "blockIdx.y", "blockIdx.z"][: len(grid[0])] + [
|
||||
"threadIdx.x",
|
||||
"threadIdx.y",
|
||||
"threadIdx.z",
|
||||
][: len(grid[1])]
|
||||
runtime_args = [arg if isinstance(arg, Expr) else const(arg) for arg in args]
|
||||
kernel_arg_types = []
|
||||
for arg in runtime_args:
|
||||
if isinstance(arg.ty, PointerType):
|
||||
kernel_arg_types.append("handle")
|
||||
else:
|
||||
assert is_prim_expr(arg)
|
||||
kernel_arg_types.append(str(arg.ty.dtype))
|
||||
runtime_args = runtime_args + list(grid[0]) + list(grid[1])
|
||||
|
||||
# Reuse compilation path from SourceModule
|
||||
compile_options = SourceModule.get_compile_options("cu")
|
||||
source_code = self.source_code
|
||||
try:
|
||||
source_path = Path(source_code)
|
||||
if source_path.is_file():
|
||||
with open(source_path) as f:
|
||||
source_code = f.read()
|
||||
except: # pylint: disable=bare-except
|
||||
pass
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
# Check if NVSHMEM is used - requires cubin output for device library linking
|
||||
use_nvshmem = (
|
||||
"#include <nvshmem.h>" in source_code or "#include <nvshmemx.h>" in source_code
|
||||
)
|
||||
target_format = "cubin" if use_nvshmem else "ptx"
|
||||
output_path = f"{temp_dir}/{kernel_name}.{target_format}"
|
||||
|
||||
compiler = os.environ.get("TVM_CUDA_COMPILE_MODE", "nvrtc")
|
||||
nvcc.compile_cuda(
|
||||
source_code,
|
||||
target_format=target_format,
|
||||
options=compile_options,
|
||||
path_target=output_path,
|
||||
compiler=compiler,
|
||||
)
|
||||
|
||||
if target_format == "ptx":
|
||||
with open(output_path) as f:
|
||||
binary_data = f.read()
|
||||
else:
|
||||
with open(output_path, "rb") as f:
|
||||
binary_data = f.read()
|
||||
|
||||
kernel_module = self._create_cuda_module(
|
||||
binary_data, kernel_arg_types, launch_param_tags, kernel_name, fmt=target_format
|
||||
)
|
||||
|
||||
return kernel_name, kernel_module, runtime_args
|
||||
|
||||
|
||||
def call_kernel(
|
||||
kernel,
|
||||
launch_args: list[int | tirx.Expr | list[int | tirx.Expr]],
|
||||
*args: list[Any],
|
||||
**kwargs: dict[str, Any],
|
||||
):
|
||||
"""
|
||||
Call an external kernel.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
kernel : Any
|
||||
The external kernel to call.
|
||||
|
||||
launch_args : List[Union[int, tirx.Expr, List[Union[int, tirx.Expr]]]]
|
||||
The launch arguments. A list of integers for grid size, block size, and shared memory size.
|
||||
The actual requirements depend on the kernel.
|
||||
|
||||
args : List[tirx.Expr]
|
||||
The arguments to pass to the kernel.
|
||||
|
||||
kwargs : Dict[str, Any]
|
||||
Additional keyword arguments to pass to the kernel or compilation.
|
||||
"""
|
||||
from tvm.script.ir_builder.ir import ( # pylint: disable=import-outside-toplevel
|
||||
module_get_attr,
|
||||
module_set_attr,
|
||||
)
|
||||
|
||||
from .ir import call_packed # pylint: disable=import-outside-toplevel
|
||||
|
||||
kernel_type = f"{type(kernel).__module__}.{type(kernel).__qualname__}"
|
||||
if kernel_type == "triton.runtime.jit.JITFunction":
|
||||
from .triton import TritonKernel # pylint: disable=import-outside-toplevel
|
||||
|
||||
kernel = TritonKernel(kernel)
|
||||
elif kernel_type == "builtins.str":
|
||||
kernel = SourceKernel(kernel)
|
||||
else:
|
||||
raise ValueError(f"Unsupported kernel type {kernel_type}")
|
||||
|
||||
kernel_name, kernel_module, runtime_args = kernel.compile_to_device_module(
|
||||
launch_args, *args, **kwargs
|
||||
)
|
||||
|
||||
# Attach the kernel module to the current IRModule
|
||||
external_mods: list[Module] = module_get_attr("external_mods") or []
|
||||
kernel_exists = any([mod.implements_function(kernel_name) for mod in external_mods])
|
||||
if kernel_exists:
|
||||
logging.debug("Kernel %s already exists in the IRModule", kernel_name)
|
||||
else:
|
||||
external_mods.append(kernel_module)
|
||||
module_set_attr("external_mods", external_mods, True)
|
||||
return call_packed(kernel_name, *runtime_args)
|
||||
@@ -0,0 +1,110 @@
|
||||
# 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.
|
||||
"""IRBuilder for TIR"""
|
||||
|
||||
from tvm_ffi import register_object as _register_object
|
||||
|
||||
from tvm.script.ir_builder.base import IRBuilderFrame
|
||||
from tvm.tirx import Buffer, Var
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.TIRFrame")
|
||||
class TIRFrame(IRBuilderFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.PrimFuncFrame")
|
||||
class PrimFuncFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.SSBlockFrame")
|
||||
class SBlockFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.SBlockInitFrame")
|
||||
class BlockInitFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.ForFrame")
|
||||
class ForFrame(TIRFrame):
|
||||
def __enter__(self) -> Var | list[Var]: # type: ignore[override]
|
||||
super().__enter__()
|
||||
return self.vars if len(self.vars) > 1 else self.vars[0]
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.AssertFrame")
|
||||
class AssertFrame(TIRFrame): ...
|
||||
|
||||
|
||||
class LetFrame(TIRFrame):
|
||||
def __enter__(self) -> Var:
|
||||
super().__enter__()
|
||||
return self.var
|
||||
|
||||
|
||||
class AllocateFrame(TIRFrame):
|
||||
def __enter__(self) -> Buffer:
|
||||
super().__enter__()
|
||||
return self.buffer_var
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.AttrFrame")
|
||||
class AttrFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.WhileFrame")
|
||||
class WhileFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.IfFrame")
|
||||
class IfFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.ThenFrame")
|
||||
class ThenFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.ElseFrame")
|
||||
class ElseFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.DeclBufferFrame")
|
||||
class DeclBufferFrame(TIRFrame):
|
||||
def __enter__(self) -> Buffer:
|
||||
super().__enter__()
|
||||
return self.buffer
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.LaunchThreadFrame")
|
||||
class LaunchThreadFrame(TIRFrame):
|
||||
def __enter__(self) -> Var:
|
||||
super().__enter__()
|
||||
return self.iter_var.var
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.ComposeOpFrame")
|
||||
class ComposeOpFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.AllocBufferFrame")
|
||||
class AllocBufferFrame(TIRFrame):
|
||||
def __enter__(self) -> Buffer:
|
||||
super().__enter__()
|
||||
return self.buffer
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.HintFrame")
|
||||
class HintFrame(TIRFrame): ...
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,19 @@
|
||||
# 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.
|
||||
"""Re-export from canonical location."""
|
||||
|
||||
from tvm.tirx.lang.alloc_pool import TMEMPool, TMEMStages # noqa: F401
|
||||
@@ -0,0 +1,137 @@
|
||||
# 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.
|
||||
# ruff: noqa: RUF005
|
||||
"""Triton kernel integration with TIR"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import triton
|
||||
from packaging import version
|
||||
from triton.runtime.jit import type_canonicalisation_dict
|
||||
|
||||
from tvm import tirx
|
||||
from tvm.ir import PointerType, PrimType, is_prim_expr
|
||||
from tvm.runtime import Module
|
||||
from tvm.topi.utils import get_const_int
|
||||
|
||||
from .external_kernel import BaseKernel
|
||||
|
||||
if version.parse(triton.__version__) < version.parse("3.3.0"):
|
||||
raise ImportError(
|
||||
f"TIR Triton integration requires Triton >= 3.3.0, but found Triton {triton.__version__}"
|
||||
)
|
||||
|
||||
|
||||
class TritonKernel(BaseKernel):
|
||||
"""A kernel from Triton JIT function.
|
||||
|
||||
This class bridges the Triton kernel with TVM runtime. The compilation includes the following
|
||||
steps:
|
||||
- Deduce the kernel signature and generate the Triton kernel
|
||||
- Embed the compiled kernel into the current IRModule as an external module
|
||||
- Generate a call to the Triton kernel following its calling convention via call_packed.
|
||||
"""
|
||||
|
||||
def __init__(self, func):
|
||||
self.func = func
|
||||
|
||||
def compile_to_device_module(
|
||||
self,
|
||||
launch_args: list[int | tirx.Expr],
|
||||
*args: list[Any],
|
||||
**kwargs: dict[str, Any],
|
||||
) -> tuple[str, Module, list[Any]]:
|
||||
"""Compile the kernel to a device module.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
launch_args : List[int]
|
||||
The grid size of the kernel. A list of one to three expressions, representing the number
|
||||
of
|
||||
"blockIdx.x", "blockIdx.y", and "blockIdx.z" respectively.
|
||||
|
||||
args : List[Any]
|
||||
Arguments to the kernel function.
|
||||
|
||||
kwargs : Dict[str, Any]
|
||||
Additional options for the kernel compilation.
|
||||
"""
|
||||
triton_kernel, kernel_args = self._generate_triton_kernel(self.func, *args, **kwargs)
|
||||
kernel_metadata = triton_kernel.metadata
|
||||
ptx = triton_kernel.asm["ptx"]
|
||||
assert kernel_metadata.num_ctas == 1, "Cluster is not supported"
|
||||
num_warps = kernel_metadata.num_warps
|
||||
grid = launch_args
|
||||
launch_param_tags = ["threadIdx.x"] + ["blockIdx.x", "blockIdx.y", "blockIdx.z"][
|
||||
: len(grid)
|
||||
]
|
||||
launch_args = [num_warps * 32] + list(grid)
|
||||
kernel_arg_types = []
|
||||
for arg in kernel_args:
|
||||
if isinstance(arg, int):
|
||||
kernel_arg_types.append("int64")
|
||||
elif isinstance(arg.ty, PointerType):
|
||||
kernel_arg_types.append("handle")
|
||||
else:
|
||||
assert is_prim_expr(arg)
|
||||
kernel_arg_types.append(str(arg.ty.dtype))
|
||||
if triton_kernel.metadata.shared > 0:
|
||||
# Add shared memory size to the launch arguments
|
||||
launch_param_tags.append("tirx.use_dyn_shared_memory")
|
||||
launch_args.append(triton_kernel.metadata.shared)
|
||||
|
||||
kernel_module = self._create_cuda_module(
|
||||
ptx, kernel_arg_types, launch_param_tags, triton_kernel.name
|
||||
)
|
||||
|
||||
return triton_kernel.name, kernel_module, kernel_args + launch_args
|
||||
|
||||
def _generate_triton_kernel(
|
||||
self, func, *args, **kwargs
|
||||
) -> tuple["triton.compiler.CompiledKernel", list[tirx.Expr]]:
|
||||
"""Deduce the kernel signature and generate the Triton kernel"""
|
||||
|
||||
kernel_params = func.params
|
||||
assert len(kernel_params) == len(args), (
|
||||
f"Number of arguments does not match, expected {len(kernel_params)}, got {len(args)}"
|
||||
)
|
||||
|
||||
signature = {}
|
||||
constants = {}
|
||||
kernel_args = [] # Arguments to invoke the kernel
|
||||
for i, arg in enumerate(args):
|
||||
if kernel_params[i].is_constexpr:
|
||||
constants[kernel_params[i].name] = get_const_int(arg)
|
||||
signature[kernel_params[i].name] = "constexpr"
|
||||
kernel_args.append(arg)
|
||||
continue
|
||||
if isinstance(arg.ty, PointerType):
|
||||
assert isinstance(arg, tirx.Var)
|
||||
assert isinstance(arg.ty.element_type, PrimType)
|
||||
elem_type = arg.ty.element_type.dtype
|
||||
pointer_type = "*" + type_canonicalisation_dict[elem_type]
|
||||
signature[kernel_params[i].name] = pointer_type
|
||||
else:
|
||||
assert is_prim_expr(arg)
|
||||
signature[kernel_params[i].name] = type_canonicalisation_dict[arg.ty.dtype]
|
||||
kernel_args.append(arg)
|
||||
|
||||
# TODO: Support default argument in the kernel
|
||||
# TODO: Add specialization for aligned buffer pointers
|
||||
source = triton.compiler.ASTSource(fn=func, signature=signature, constexprs=constants)
|
||||
compiled = triton.compiler.compile(source, options=kwargs)
|
||||
return compiled, kernel_args
|
||||
@@ -0,0 +1,226 @@
|
||||
# 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.
|
||||
"""Utility helpers for TIR IRBuilder."""
|
||||
|
||||
import contextlib
|
||||
|
||||
from tvm import tirx
|
||||
from tvm.tirx import Buffer
|
||||
|
||||
from . import frame
|
||||
from . import ir as T
|
||||
|
||||
|
||||
class _FrameScope:
|
||||
"""Context manager to enter multiple IRBuilder frames without deep nesting.
|
||||
|
||||
This class allows entering multiple frames in a single `with` statement,
|
||||
avoiding the pyramid of nested context managers.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
frames : List[IRBuilderFrame]
|
||||
The list of frames to enter.
|
||||
"""
|
||||
|
||||
def __init__(self, frames):
|
||||
self.frames = frames if isinstance(frames, list | tuple) else [frames]
|
||||
self._stack = None
|
||||
|
||||
def __enter__(self):
|
||||
self._stack = contextlib.ExitStack()
|
||||
self._stack.__enter__()
|
||||
results = [self._stack.enter_context(f) for f in self.frames]
|
||||
return tuple(results) if len(results) > 1 else results[0]
|
||||
|
||||
def __exit__(self, *args):
|
||||
return self._stack.__exit__(*args)
|
||||
|
||||
|
||||
def frame_scope(frames: list[frame.TIRFrame]) -> _FrameScope:
|
||||
"""Enter multiple IRBuilder frames without deep nesting.
|
||||
|
||||
This function provides a way to enter multiple frames in a single `with`
|
||||
statement, which is particularly useful when migrating from cases where
|
||||
allocations don't require nested scopes.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
frames : List[frame.TIRFrame]
|
||||
The list of frames to enter. Each frame's `__enter__` return value
|
||||
will be collected and returned as a tuple.
|
||||
|
||||
Returns
|
||||
-------
|
||||
_FrameScope
|
||||
A context manager that enters all frames and returns their values.
|
||||
"""
|
||||
return _FrameScope(frames)
|
||||
|
||||
|
||||
def seq_scope():
|
||||
"""Create a scope that allows multiple consecutive statements.
|
||||
|
||||
The IRBuilder requires a parent frame when having multiple consecutive
|
||||
top-level statements (e.g., multiple loops). This function creates a
|
||||
dummy attr frame that serves as a parent scope.
|
||||
|
||||
Returns
|
||||
-------
|
||||
frame.AttrFrame
|
||||
A dummy attribute frame that wraps multiple statements.
|
||||
|
||||
Examples
|
||||
--------
|
||||
Without seq_scope, multiple consecutive loops fail:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
with IRBuilder() as ib:
|
||||
with T.serial(0, 10) as i:
|
||||
T.evaluate(i)
|
||||
with T.serial(0, 5) as j: # This would fail!
|
||||
T.evaluate(j)
|
||||
|
||||
With seq_scope, multiple consecutive statements work:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
with IRBuilder() as ib:
|
||||
with seq_scope():
|
||||
with T.serial(0, 10) as i:
|
||||
T.evaluate(i)
|
||||
with T.serial(0, 5) as j:
|
||||
T.evaluate(j)
|
||||
result = ib.get()
|
||||
"""
|
||||
return T.attr(tirx.const(0, "int32"), "pragma_scope", tirx.StringImm("seq"))
|
||||
|
||||
|
||||
def _unravel_index(index, shape):
|
||||
"""Convert a flat index to multi-dimensional indices.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
index : Expr
|
||||
The flat index.
|
||||
shape : Tuple
|
||||
The shape of the buffer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[Expr]
|
||||
The multi-dimensional indices.
|
||||
"""
|
||||
indices = []
|
||||
for i, dim in enumerate(reversed(shape)):
|
||||
if i == len(shape) - 1:
|
||||
# Outermost dimension: use remaining quotient directly (no modulo)
|
||||
indices.append(index)
|
||||
else:
|
||||
indices.append(index % dim)
|
||||
index = index // dim
|
||||
return list(reversed(indices))
|
||||
|
||||
|
||||
class _BufferProxy:
|
||||
"""Proxy for flat indexing on multi-dimensional buffers.
|
||||
|
||||
This class wraps a TIR Buffer and provides flat indexing that gets
|
||||
automatically converted to multi-dimensional indices. It also supports
|
||||
assignment syntax via __setitem__.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
buf : Buffer
|
||||
The TIR buffer to wrap.
|
||||
|
||||
Examples
|
||||
--------
|
||||
.. code-block:: python
|
||||
|
||||
buf = tvm.tirx.decl_buffer([2, 3], "float32")
|
||||
ptr = buffer_proxy(buf)
|
||||
|
||||
# Read with flat index (converted to [0, 1])
|
||||
val = ptr[1]
|
||||
|
||||
# Write with flat index
|
||||
ptr[1] = 42.0
|
||||
|
||||
# Multi-dimensional access still works
|
||||
val = ptr[0, 2]
|
||||
"""
|
||||
|
||||
def __init__(self, buf):
|
||||
self._buffer = buf
|
||||
self.dtype = buf.dtype
|
||||
self.shape = buf.shape
|
||||
self.name = buf.name
|
||||
self.data = buf.data
|
||||
|
||||
def _normalize_index(self, index):
|
||||
"""Convert flat index to multi-dimensional indices if needed."""
|
||||
try:
|
||||
index = [*index]
|
||||
except TypeError:
|
||||
index = [index]
|
||||
if len(index) == 1 and len(self._buffer.shape) != 1:
|
||||
index = _unravel_index(index[0], self._buffer.shape)
|
||||
return index
|
||||
|
||||
def __getitem__(self, index):
|
||||
index = self._normalize_index(index)
|
||||
return tirx.BufferLoad(self._buffer, index)
|
||||
|
||||
def __setitem__(self, index, value):
|
||||
index = self._normalize_index(index)
|
||||
T.buffer_store(self._buffer, value, index)
|
||||
|
||||
|
||||
def buffer_proxy(buf: Buffer) -> _BufferProxy:
|
||||
"""Create a buffer proxy for flat indexing on multi-dimensional buffers.
|
||||
|
||||
This provides flat indexing that gets converted to multi-dimensional indices.
|
||||
It also supports assignment syntax via __setitem__.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
buf : Buffer
|
||||
The TIR buffer to wrap.
|
||||
|
||||
Returns
|
||||
-------
|
||||
_BufferProxy
|
||||
A proxy object that supports flat indexing and assignment.
|
||||
|
||||
Examples
|
||||
--------
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.tirx.script.builder.utils import buffer_proxy
|
||||
|
||||
buf = tvm.tirx.decl_buffer([2, 3], "float32")
|
||||
ptr = buffer_proxy(buf)
|
||||
|
||||
# Flat indexing (index 1 -> indices [0, 1])
|
||||
val = ptr[1]
|
||||
|
||||
# Assignment syntax
|
||||
ptr[1] = 42.0
|
||||
"""
|
||||
return _BufferProxy(buf)
|
||||
@@ -0,0 +1,50 @@
|
||||
# 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.
|
||||
# ruff: noqa: RUF005
|
||||
"""The tirx parser"""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from tvm.tirx.script.builder import * # pylint: disable=redefined-builtin
|
||||
from tvm.tirx.script.builder import ir as _tir
|
||||
|
||||
from . import operation as _operation
|
||||
from . import parser as _parser
|
||||
from .entry import Buffer, Ptr, constexpr
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# pylint: disable=invalid-name
|
||||
# Define prim_func and make it type check as static method
|
||||
# so most tvmscript won't trigger pylint error here.
|
||||
prim_func = staticmethod
|
||||
jit = staticmethod
|
||||
else:
|
||||
from .entry import inline, jit, macro, prim_func
|
||||
|
||||
__all__ = _tir.__all__ + [
|
||||
"Buffer",
|
||||
"Ptr",
|
||||
"SMEMPool",
|
||||
"TMEMPool",
|
||||
"TMEMStages",
|
||||
"bool",
|
||||
"constexpr",
|
||||
"inline",
|
||||
"jit",
|
||||
"macro",
|
||||
"prim_func",
|
||||
]
|
||||
@@ -0,0 +1,520 @@
|
||||
# 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.
|
||||
"""The entry point of TVM parser for tirx."""
|
||||
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from tvm.ir.base import deprecated
|
||||
from tvm.script.parser._core import parse, scan_macro, utils
|
||||
from tvm.script.parser.core.parser import Parser, ScriptMacro, VarTable
|
||||
from tvm.tirx import Buffer, PrimFunc
|
||||
from tvm.tirx.script.builder import block_name_suffix_context, buffer, ptr
|
||||
|
||||
|
||||
def prim_func(
|
||||
func: Callable | None = None,
|
||||
private: bool = False,
|
||||
check_well_formed=True,
|
||||
s_tir: bool = False,
|
||||
persistent: bool = False,
|
||||
) -> PrimFunc | Callable:
|
||||
"""The parsing method for tirx prim func, by using `@prim_func` as decorator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func : Callable
|
||||
The function to be parsed as prim func.
|
||||
(Listed as optional to allow the decorator to be used
|
||||
without arguments, like `@prim_func`,
|
||||
or with an argument, `@prim_func(private=True)`)
|
||||
|
||||
private : bool, optional
|
||||
Whether the function should be treated as private.
|
||||
A private function has no global symbol attribute;
|
||||
if the function is not private, it will have a global symbol
|
||||
matching the function name.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Union[PrimFunc, Callable]
|
||||
The parsed tirx prim func.
|
||||
"""
|
||||
# pylint: disable=unused-argument
|
||||
# (private will be used in the parser, but not immediately)
|
||||
|
||||
# need to capture this var outside the wrapper because the wrapper
|
||||
# adds to the stack
|
||||
outer_stack = inspect.stack()
|
||||
|
||||
def decorator_wrapper(func):
|
||||
if not inspect.isfunction(func):
|
||||
raise TypeError(f"Expect a function, but got: {func}")
|
||||
if utils.is_defined_in_class(outer_stack, func):
|
||||
return func
|
||||
extra_vars = utils.inspect_function_capture(func)
|
||||
utils.resolve_closure_vars(func, extra_vars, outer_stack)
|
||||
f = parse(func, extra_vars, check_well_formed=check_well_formed, s_tir=s_tir)
|
||||
setattr(f, "__name__", func.__name__)
|
||||
return f
|
||||
|
||||
if func is not None:
|
||||
# no optional args given => use wrapper directly
|
||||
return decorator_wrapper(func)
|
||||
else:
|
||||
# if there is an optional arg given, return a new decorator
|
||||
# that will then be invoked
|
||||
setattr(decorator_wrapper, "dispatch_token", "tirx")
|
||||
return decorator_wrapper
|
||||
|
||||
|
||||
setattr(prim_func, "dispatch_token", "tirx")
|
||||
|
||||
|
||||
class TIRInline(ScriptMacro):
|
||||
"""Specialization of ScriptMacro for TIR with Python LEGB scoping.
|
||||
|
||||
Two definition paths:
|
||||
1. Outside @T.prim_func (standalone @T.inline): definition_depth is None,
|
||||
closure_vars captured at definition time are used (module globals are
|
||||
effectively late-bound since they don't change during parsing).
|
||||
2. Inside @T.prim_func (inline def in parsed body): definition_depth is set
|
||||
to the VarTable frame depth at definition time, and defining_var_table
|
||||
stores a reference to the VarTable that was active. At call time,
|
||||
defining_var_table.get_at_depth(definition_depth) reads current values
|
||||
from the lexically enclosing frames.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
definition_depth : Optional[int]
|
||||
VarTable frame depth at definition time, or None for outside-prim_func.
|
||||
defining_var_table : Optional[VarTable]
|
||||
Reference to the VarTable that was active at definition time.
|
||||
call_count : int
|
||||
Counter for unique block name suffixes.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source,
|
||||
closure_vars: dict[str, Any],
|
||||
func: Callable,
|
||||
definition_depth: int | None = None,
|
||||
defining_var_table: VarTable | None = None,
|
||||
) -> None:
|
||||
# hygienic=True for the base class (field kept for compat but not used in dispatch)
|
||||
super().__init__(source, closure_vars, func, hygienic=True)
|
||||
self.definition_depth = definition_depth
|
||||
self.defining_var_table = defining_var_table
|
||||
self.call_count = 0
|
||||
|
||||
def parse_macro(self, parser: Parser) -> None:
|
||||
macro_def = self.get_macro_def()
|
||||
suffix = f"_{self.call_count}" if self.call_count > 0 else ""
|
||||
self.call_count += 1
|
||||
with block_name_suffix_context(suffix):
|
||||
parser.visit_body(macro_def.body)
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
param_binding = inspect.signature(self.func).bind(*args, **kwargs)
|
||||
param_binding.apply_defaults()
|
||||
local_vars = param_binding.arguments
|
||||
parser = self._find_parser_def()
|
||||
|
||||
with parser.with_diag_source(self.source):
|
||||
if self.defining_var_table is not None:
|
||||
# Inside-prim_func path: LEGB late binding from the defining scope
|
||||
enclosing_vars = self.defining_var_table.get_at_depth(self.definition_depth)
|
||||
else:
|
||||
# Outside-prim_func path: use captured closure vars
|
||||
enclosing_vars = self.closure_vars
|
||||
|
||||
saved_var_table = parser.var_table
|
||||
parser.var_table = VarTable()
|
||||
|
||||
with parser.var_table.with_frame():
|
||||
for k, v in enclosing_vars.items():
|
||||
parser.var_table.add(k, v)
|
||||
with parser.var_table.with_frame():
|
||||
for k, v in local_vars.items():
|
||||
parser.var_table.add(k, v)
|
||||
|
||||
parse_result = self.parse_macro(parser)
|
||||
|
||||
parser.var_table = saved_var_table
|
||||
|
||||
return parse_result
|
||||
|
||||
|
||||
def inline(*args, definition_depth: int | None = None, defining_var_table=None) -> Callable:
|
||||
"""Decorator for inline function definitions with Python LEGB scoping.
|
||||
|
||||
@T.inline follows Python's lexical scoping with late binding:
|
||||
- At definition time, record which scopes are visible.
|
||||
- At call time, read current values from those scopes.
|
||||
|
||||
Example::
|
||||
|
||||
import tvm
|
||||
from tvm.script import tirx as T
|
||||
|
||||
x_value = 128
|
||||
|
||||
@T.inline
|
||||
def capture(A, B):
|
||||
B[()] = A[x_value] # x_value resolved from enclosing scope
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def use(A: T.Buffer((1024,), "int32"), B: T.Buffer((), "int32")) -> None:
|
||||
capture(A, B) # Produces B[()] = A[128]
|
||||
"""
|
||||
|
||||
def _decorator(func: Callable) -> Callable:
|
||||
source, closure_vars = scan_macro(func, utils.inspect_function_capture(func))
|
||||
obj = TIRInline(
|
||||
source,
|
||||
closure_vars,
|
||||
func,
|
||||
definition_depth=definition_depth,
|
||||
defining_var_table=defining_var_table,
|
||||
)
|
||||
|
||||
def wrapper(*args, **kwargs):
|
||||
return obj(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
if len(args) == 0:
|
||||
setattr(_decorator, "dispatch_token", "tir.inline")
|
||||
return _decorator
|
||||
if len(args) == 1 and inspect.isfunction(args[0]):
|
||||
return _decorator(args[0])
|
||||
|
||||
raise ValueError("Invalid use of T.inline. Usage: @T.inline or @T.inline()")
|
||||
|
||||
|
||||
setattr(inline, "dispatch_token", "tir.inline")
|
||||
|
||||
|
||||
class TIRJit:
|
||||
"""Top-level kernel decorator with constexpr params + ``.specialize()``.
|
||||
|
||||
Parses the function body lazily: parsing is deferred until ``.specialize()``
|
||||
supplies concrete values for the params annotated as ``T.constexpr``. The
|
||||
return type of ``.specialize()`` is a ``tvm.tirx.PrimFunc``, identical in
|
||||
type to what ``@T.prim_func`` produces today.
|
||||
|
||||
Constexpr params are removed from the resulting PrimFunc's parameter list;
|
||||
their values are baked into the IR (e.g. into ``T.Buffer((M, K), ...)``
|
||||
shape annotations and into the body).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
func: Callable,
|
||||
check_well_formed: bool = True,
|
||||
is_stir: bool = False,
|
||||
persistent: bool = False,
|
||||
private: bool = False,
|
||||
) -> None:
|
||||
self.func = func
|
||||
self.check_well_formed = check_well_formed
|
||||
self.is_stir = is_stir
|
||||
self.persistent = persistent # pylint: disable=unused-private-member
|
||||
self.private = private # pylint: disable=unused-private-member
|
||||
# Resolved closure vars (computed once; the function itself is the
|
||||
# capture point, so this never changes between specializations).
|
||||
self._closure_vars: dict[str, Any] = utils.inspect_function_capture(func)
|
||||
# Detect which params are marked T.constexpr. With PEP 563
|
||||
# (``from __future__ import annotations``), each annotation is a
|
||||
# string; we eval them one-by-one so a constexpr probe is not
|
||||
# blocked by sibling annotations that reference yet-undefined names
|
||||
# (e.g. ``A: T.Buffer((N,), ...)`` referencing constexpr ``N``).
|
||||
raw_anns = getattr(func, "__annotations__", {}) or {}
|
||||
eval_globals = {**func.__globals__, **self._closure_vars}
|
||||
sig = inspect.signature(func)
|
||||
constexpr_names: set[str] = set()
|
||||
constexpr_defaults: dict[str, Any] = {}
|
||||
for name, param in sig.parameters.items():
|
||||
ann = raw_anns.get(name)
|
||||
if isinstance(ann, str):
|
||||
try:
|
||||
ann = eval(ann, eval_globals) # pylint: disable=eval-used
|
||||
except Exception: # pylint: disable=broad-except
|
||||
ann = None
|
||||
if ann is constexpr:
|
||||
constexpr_names.add(name)
|
||||
if param.default is not inspect.Parameter.empty:
|
||||
constexpr_defaults[name] = param.default
|
||||
self.constexpr_names: frozenset[str] = frozenset(constexpr_names)
|
||||
self.constexpr_defaults: dict[str, Any] = constexpr_defaults
|
||||
self._cache: dict[tuple, PrimFunc] = {}
|
||||
|
||||
def specialize(self, **constexpr_kwargs) -> PrimFunc:
|
||||
"""Build a concrete PrimFunc by binding the constexpr params.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
**constexpr_kwargs
|
||||
One value per ``T.constexpr``-annotated parameter. All such
|
||||
parameters must be supplied; passing names that are not
|
||||
constexpr-annotated is an error.
|
||||
|
||||
Returns
|
||||
-------
|
||||
PrimFunc
|
||||
A concrete TIRx PrimFunc, identical in type to the output of
|
||||
``@T.prim_func``.
|
||||
"""
|
||||
extra = constexpr_kwargs.keys() - self.constexpr_names
|
||||
if extra:
|
||||
raise TypeError(
|
||||
f"{self.func.__name__}.specialize() got unexpected arg(s): "
|
||||
f"{sorted(extra)} (constexpr params are: {sorted(self.constexpr_names)})"
|
||||
)
|
||||
effective = {**self.constexpr_defaults, **constexpr_kwargs}
|
||||
missing = self.constexpr_names - effective.keys()
|
||||
if missing:
|
||||
raise TypeError(
|
||||
f"{self.func.__name__}.specialize() missing constexpr arg(s) "
|
||||
f"(no default provided): {sorted(missing)}"
|
||||
)
|
||||
|
||||
try:
|
||||
cache_key = tuple(sorted(effective.items()))
|
||||
cached = self._cache.get(cache_key)
|
||||
except TypeError as err:
|
||||
raise TypeError(
|
||||
f"{self.func.__name__}.specialize(): all constexpr values must "
|
||||
f"be hashable (got: {effective!r})"
|
||||
) from err
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
extra_vars = {**self._closure_vars, **effective}
|
||||
prim_func = parse(
|
||||
self.func,
|
||||
extra_vars,
|
||||
check_well_formed=self.check_well_formed,
|
||||
s_tir=self.is_stir,
|
||||
)
|
||||
setattr(prim_func, "__name__", self.func.__name__)
|
||||
self._cache[cache_key] = prim_func
|
||||
return prim_func
|
||||
|
||||
|
||||
def jit(
|
||||
func: Callable | None = None,
|
||||
private: bool = False,
|
||||
check_well_formed: bool = True,
|
||||
is_stir: bool = False,
|
||||
persistent: bool = False,
|
||||
) -> "TIRJit | Callable":
|
||||
"""Decorator: capture the kernel and defer parsing until ``.specialize()``.
|
||||
|
||||
Use ``@T.jit`` (instead of ``@T.prim_func``) when the kernel takes
|
||||
compile-time parameters annotated with ``T.constexpr``. The resulting
|
||||
object exposes ``.specialize(**constexpr_kwargs)``, which returns a
|
||||
``tvm.tirx.PrimFunc``.
|
||||
|
||||
Example::
|
||||
|
||||
from tvm.script import tirx as T
|
||||
|
||||
@T.jit
|
||||
def add(
|
||||
A: T.Buffer((N,), "float32"),
|
||||
B: T.Buffer((N,), "float32"),
|
||||
*,
|
||||
N: T.constexpr,
|
||||
):
|
||||
...
|
||||
|
||||
kernel = add.specialize(N=1024) # returns a PrimFunc
|
||||
"""
|
||||
|
||||
def decorator_wrapper(func: Callable) -> TIRJit:
|
||||
if not inspect.isfunction(func):
|
||||
raise TypeError(f"Expect a function, but got: {func}")
|
||||
return TIRJit(
|
||||
func,
|
||||
check_well_formed=check_well_formed,
|
||||
is_stir=is_stir,
|
||||
persistent=persistent,
|
||||
private=private,
|
||||
)
|
||||
|
||||
if func is not None:
|
||||
return decorator_wrapper(func)
|
||||
setattr(decorator_wrapper, "dispatch_token", "tirx")
|
||||
return decorator_wrapper
|
||||
|
||||
|
||||
setattr(jit, "dispatch_token", "tirx")
|
||||
|
||||
|
||||
class TIRMacro(ScriptMacro):
|
||||
"""Specialization of the ScriptMacro class for TIR.
|
||||
|
||||
Apache-compatible hygienic macro. Distinct from ``TIRInline`` (which
|
||||
uses Python LEGB late binding) so upstream code that relies on
|
||||
capture-at-definition-time semantics keeps working.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
call_count : int
|
||||
Counter for the number of times this macro has been invoked.
|
||||
Used to generate unique block name suffixes.
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.call_count = 0
|
||||
|
||||
def parse_macro(self, parser: Parser) -> None:
|
||||
macro_def = self.get_macro_def()
|
||||
suffix = f"_{self.call_count}" if self.call_count > 0 else ""
|
||||
self.call_count += 1
|
||||
with block_name_suffix_context(suffix):
|
||||
parser.visit_body(macro_def.body)
|
||||
|
||||
|
||||
def macro(*args, hygienic: bool = True) -> Callable:
|
||||
"""Decorator for macro definitions with hygienic capture.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
hygienic: bool
|
||||
Specifies whether the macro is hygienic or not. A hygienic macro
|
||||
resolves symbols at definition time; a non-hygienic macro at use
|
||||
time. Defaults to ``True``.
|
||||
"""
|
||||
|
||||
def _decorator(func: Callable) -> TIRMacro:
|
||||
source, closure_vars = scan_macro(func, utils.inspect_function_capture(func))
|
||||
obj = TIRMacro(source, closure_vars, func, hygienic)
|
||||
|
||||
def wrapper(*args, **kwargs):
|
||||
return obj(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
if len(args) == 0:
|
||||
return _decorator
|
||||
if len(args) == 1 and inspect.isfunction(args[0]):
|
||||
return _decorator(args[0])
|
||||
|
||||
raise ValueError("Invalid use of T.macro. Usage: @T.macro or @T.macro()")
|
||||
|
||||
|
||||
setattr(macro, "dispatch_token", "tir.macro")
|
||||
|
||||
|
||||
class BufferProxy:
|
||||
"""Buffer proxy class for constructing tirx buffer."""
|
||||
|
||||
def __or__(self, other):
|
||||
"""Support ``T.Buffer | None`` union syntax in annotations."""
|
||||
return self
|
||||
|
||||
def __ror__(self, other):
|
||||
"""Support ``None | T.Buffer`` union syntax in annotations."""
|
||||
return self
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
shape,
|
||||
dtype="float32",
|
||||
data=None,
|
||||
strides=None,
|
||||
elem_offset=None,
|
||||
byte_offset=None,
|
||||
scope="global",
|
||||
align=0,
|
||||
offset_factor=0,
|
||||
buffer_type="",
|
||||
axis_separators=None,
|
||||
layout="default",
|
||||
) -> Buffer:
|
||||
return buffer(
|
||||
shape,
|
||||
dtype=dtype,
|
||||
data=data,
|
||||
strides=strides,
|
||||
elem_offset=elem_offset,
|
||||
byte_offset=byte_offset,
|
||||
scope=scope,
|
||||
align=align,
|
||||
offset_factor=offset_factor,
|
||||
buffer_type=buffer_type,
|
||||
axis_separators=axis_separators,
|
||||
layout=layout,
|
||||
)
|
||||
|
||||
@deprecated("T.Buffer[...]", "T.Buffer(...)")
|
||||
def __getitem__(self, keys) -> Buffer:
|
||||
if not isinstance(keys, tuple):
|
||||
return self(keys)
|
||||
if len(keys) >= 2 and not isinstance(keys[1], str):
|
||||
return self(keys)
|
||||
return self(*keys) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
class PtrProxy:
|
||||
"""Ptr proxy class for constructing tirx pointer."""
|
||||
|
||||
def __or__(self, other):
|
||||
"""Support union syntax in annotations."""
|
||||
return self
|
||||
|
||||
def __ror__(self, other):
|
||||
"""Support union syntax in annotations."""
|
||||
return self
|
||||
|
||||
@deprecated("T.Ptr(...)", "T.handle(...)")
|
||||
def __call__(self, dtype, storage_scope="global"):
|
||||
if callable(dtype):
|
||||
dtype = dtype().ty.dtype
|
||||
return ptr(dtype, storage_scope) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
@deprecated("T.Ptr[...]", "T.handle(...)")
|
||||
def __getitem__(self, keys):
|
||||
if not isinstance(keys, tuple):
|
||||
return self(keys)
|
||||
return self(*keys)
|
||||
|
||||
|
||||
class _ConstexprProxy:
|
||||
"""Sentinel marker for compile-time (specialization-time) parameters.
|
||||
|
||||
Used as a parameter annotation in ``@T.jit`` decorated functions to mark
|
||||
a parameter as constexpr — its value is supplied to ``.specialize(**kwargs)``
|
||||
rather than at call time, and it is removed from the generated PrimFunc's
|
||||
runtime parameter list.
|
||||
"""
|
||||
|
||||
def __or__(self, other):
|
||||
return self
|
||||
|
||||
def __ror__(self, other):
|
||||
return self
|
||||
|
||||
|
||||
Buffer = BufferProxy() # pylint: disable=invalid-name
|
||||
Ptr = PtrProxy() # pylint: disable=invalid-name
|
||||
constexpr = _ConstexprProxy() # pylint: disable=invalid-name
|
||||
@@ -0,0 +1,167 @@
|
||||
# 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.
|
||||
"""The tirx expression operation registration"""
|
||||
|
||||
import tvm
|
||||
from tvm import tirx
|
||||
from tvm.ir import PrimType
|
||||
from tvm.runtime import DataTypeCode
|
||||
from tvm.script.parser._core import OpMethod, doc, register_op
|
||||
from tvm.tirx import IntImm
|
||||
from tvm.tirx.expr import FloatImm
|
||||
|
||||
|
||||
def _register_expr_op(ty: type): # pylint: disable=invalid-name
|
||||
ty._dispatch_type = ty # pylint: disable=protected-access
|
||||
|
||||
def _expr_ty(expr):
|
||||
ty = expr.ty if tvm.ir.is_prim_expr(expr) else None
|
||||
if not isinstance(ty, PrimType):
|
||||
ty = expr.expr_ty()
|
||||
if not isinstance(ty, PrimType):
|
||||
raise TypeError(f"Expected a PrimType expression, but got {ty}")
|
||||
return ty
|
||||
|
||||
def _and(a, b):
|
||||
if isinstance(a, bool):
|
||||
a = IntImm("bool", a)
|
||||
if isinstance(b, bool):
|
||||
b = IntImm("bool", b)
|
||||
if not _expr_ty(a).is_scalar() or not _expr_ty(b).is_scalar():
|
||||
return a & b
|
||||
else:
|
||||
return tirx.And(a, b)
|
||||
|
||||
def _or(a, b):
|
||||
if isinstance(a, bool):
|
||||
a = IntImm("bool", a)
|
||||
if isinstance(b, bool):
|
||||
b = IntImm("bool", b)
|
||||
if not _expr_ty(a).is_scalar() or not _expr_ty(b).is_scalar():
|
||||
return a | b
|
||||
else:
|
||||
return tirx.Or(a, b)
|
||||
|
||||
def _get_type_str(ty: PrimType):
|
||||
dtype_str = str(ty.dtype)
|
||||
if ty.is_scalar():
|
||||
return dtype_str
|
||||
index = dtype_str.find("x")
|
||||
return dtype_str[0:index]
|
||||
|
||||
def _auto_broadcast(a, b, op):
|
||||
if isinstance(a, int):
|
||||
if tvm.ir.is_prim_expr(b) or hasattr(b, "expr_ty"):
|
||||
b_ty = _expr_ty(b)
|
||||
if b_ty.matches_code(DataTypeCode.INT, DataTypeCode.UINT, DataTypeCode.BOOL):
|
||||
a = IntImm(_get_type_str(b_ty), a)
|
||||
elif b_ty.matches_code(DataTypeCode.FLOAT):
|
||||
a = FloatImm(_get_type_str(b_ty), a)
|
||||
elif isinstance(b, float):
|
||||
a = FloatImm("float32", a)
|
||||
else:
|
||||
a = IntImm("int32", a)
|
||||
elif isinstance(a, float):
|
||||
b_ty = _expr_ty(b)
|
||||
if b_ty.matches_code(DataTypeCode.FLOAT):
|
||||
a = FloatImm(_get_type_str(b_ty), a)
|
||||
else:
|
||||
a = FloatImm("float32", a)
|
||||
|
||||
assert tvm.ir.is_prim_expr(a), "Operand should be a Expr."
|
||||
if isinstance(b, int):
|
||||
a_ty = _expr_ty(a)
|
||||
if a_ty.matches_code(DataTypeCode.INT, DataTypeCode.UINT, DataTypeCode.BOOL):
|
||||
b = IntImm(_get_type_str(a_ty), b)
|
||||
elif a_ty.matches_code(DataTypeCode.FLOAT):
|
||||
b = FloatImm(_get_type_str(a_ty), b)
|
||||
elif isinstance(b, float):
|
||||
b = FloatImm(_get_type_str(_expr_ty(a)), b)
|
||||
|
||||
a_ty = _expr_ty(a)
|
||||
b_ty = _expr_ty(b)
|
||||
if a_ty.dtype.lanes == b_ty.dtype.lanes:
|
||||
return op(a, b)
|
||||
elif a_ty.is_scalar() and a_ty.dtype.lanes != b_ty.dtype.lanes:
|
||||
broadcast_a = tirx.Broadcast(a, b_ty.dtype.lanes)
|
||||
return op(broadcast_a, b)
|
||||
elif b_ty.is_scalar() and a_ty.dtype.lanes != b_ty.dtype.lanes:
|
||||
broadcast_b = tirx.Broadcast(b, a_ty.dtype.lanes)
|
||||
return op(a, broadcast_b)
|
||||
else:
|
||||
raise TypeError("do not know how to deal with it.")
|
||||
|
||||
def _eq(a, b):
|
||||
return _auto_broadcast(a, b, tirx.EQ)
|
||||
|
||||
def _ne(a, b):
|
||||
return _auto_broadcast(a, b, tirx.NE)
|
||||
|
||||
def _lt(a, b):
|
||||
return _auto_broadcast(a, b, tirx.LT)
|
||||
|
||||
def _le(a, b):
|
||||
return _auto_broadcast(a, b, tirx.LE)
|
||||
|
||||
def _gt(a, b):
|
||||
return _auto_broadcast(a, b, tirx.GT)
|
||||
|
||||
def _ge(a, b):
|
||||
return _auto_broadcast(a, b, tirx.GE)
|
||||
|
||||
def r(op: type, i: int, m: OpMethod): # pylint: disable=invalid-name
|
||||
register_op(ty, op, i)(m)
|
||||
|
||||
for i in [0, 1]:
|
||||
# Case 1. binop
|
||||
# doc.Add <-- is overloaded
|
||||
# doc.Sub <-- is overloaded
|
||||
# doc.Mult <-- is overloaded
|
||||
# doc.Div <-- is overloaded
|
||||
# doc.FloorDiv <-- is overloaded
|
||||
# doc.Mod <-- is overloaded
|
||||
# doc.LShift <-- is overloaded
|
||||
# doc.RShift <-- is overloaded
|
||||
# doc.BitOr <-- is overloaded
|
||||
# doc.BitXor <-- is overloaded
|
||||
# doc.BitAnd <-- is overloaded
|
||||
# doc.MatMult <-- not implemented
|
||||
# doc.Pow <-- not implemented
|
||||
# Case 2. cmpop
|
||||
r(doc.Eq, i, _eq)
|
||||
r(doc.NotEq, i, _ne)
|
||||
r(doc.Lt, i, _lt)
|
||||
r(doc.LtE, i, _le)
|
||||
r(doc.Gt, i, _gt)
|
||||
r(doc.GtE, i, _ge)
|
||||
# doc.Is <-- not implemented
|
||||
# doc.IsNot <-- not implemented
|
||||
# doc.In <-- not implemented
|
||||
# doc.NotIn <-- not implemented
|
||||
# Case 3. boolop
|
||||
r(doc.And, i, _and)
|
||||
r(doc.Or, i, _or)
|
||||
for i in [0]:
|
||||
# Case 4. unaryop
|
||||
# doc.Invert <-- is overloaded
|
||||
r(doc.Not, i, tirx.Not)
|
||||
# doc.UAdd <-- is overloaded
|
||||
# doc.USub <-- is overloaded
|
||||
|
||||
|
||||
_register_expr_op(tirx.Expr)
|
||||
_register_expr_op(tirx.IterVar)
|
||||
@@ -0,0 +1,914 @@
|
||||
# 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.
|
||||
"""The base parser for tirx"""
|
||||
|
||||
import ast
|
||||
import contextlib
|
||||
from copy import deepcopy
|
||||
from functools import partial
|
||||
from typing import Any
|
||||
|
||||
import tvm
|
||||
from tvm.ir import Expr, GlobalVar, PrimType
|
||||
from tvm.script.ir_builder import ir as I
|
||||
from tvm.script.ir_builder.base import IRBuilder
|
||||
from tvm.script.ir_builder.base import IRBuilderFrame as Frame
|
||||
from tvm.script.parser._core import Parser, dispatch, doc
|
||||
from tvm.script.parser.core.doc import from_doc
|
||||
from tvm.tirx import Buffer, IterVar, Layout, Var
|
||||
from tvm.tirx.script import builder as T
|
||||
from tvm.tirx.script.builder.ir import name_meta_class_value
|
||||
from tvm.tirx.stmt import BufferRegion
|
||||
|
||||
from .entry import constexpr as _constexpr_sentinel
|
||||
from .entry import inline
|
||||
|
||||
|
||||
def slice_buffer_from_region(br: BufferRegion) -> Buffer:
|
||||
"""Create a matched DeclBuffer from a BufferRegion.
|
||||
|
||||
Slices the layout (if present) or computes elem_offset for the sub-region,
|
||||
producing a DeclBuffer that views the same underlying data.
|
||||
"""
|
||||
import functools # pylint: disable=import-outside-toplevel
|
||||
|
||||
buf = br.buffer
|
||||
region = br.region
|
||||
new_shape = [r.extent for r in region]
|
||||
sliced_layout = None
|
||||
if buf.layout is not None:
|
||||
range_pairs = [(r.min, r.min + r.extent) for r in region]
|
||||
sliced_layout = buf.layout.slice(list(buf.shape), range_pairs)
|
||||
if sliced_layout is not None:
|
||||
return T.decl_buffer(
|
||||
new_shape,
|
||||
buf.dtype,
|
||||
buf.data,
|
||||
buf.strides,
|
||||
buf.elem_offset,
|
||||
None,
|
||||
buf.scope(),
|
||||
buf.data_alignment,
|
||||
buf.offset_factor,
|
||||
"",
|
||||
buf.axis_separators,
|
||||
sliced_layout,
|
||||
)
|
||||
# Fallback: compute elem_offset for default/no layout
|
||||
strides = []
|
||||
for i in range(len(buf.shape)):
|
||||
stride = functools.reduce(
|
||||
lambda x, y: x * y, buf.shape[i + 1 :], tvm.tirx.const(1, "int32")
|
||||
)
|
||||
strides.append(stride)
|
||||
offset = tvm.tirx.const(0, "int32")
|
||||
for i, r in enumerate(region):
|
||||
offset = offset + r.min * strides[i]
|
||||
new_elem_offset = buf.elem_offset + offset
|
||||
return T.decl_buffer(
|
||||
new_shape,
|
||||
buf.dtype,
|
||||
buf.data,
|
||||
buf.strides,
|
||||
new_elem_offset,
|
||||
None,
|
||||
buf.scope(),
|
||||
buf.data_alignment,
|
||||
buf.offset_factor,
|
||||
"",
|
||||
buf.axis_separators,
|
||||
buf.layout,
|
||||
)
|
||||
|
||||
|
||||
def bind_with_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
|
||||
"""Value binding methods when parsing with statement.
|
||||
e.g. binding i, j, k with T.grid(128, 128, 128), when parsing
|
||||
with T.grid(128, 128, 18) as i, j, k.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The current parser.
|
||||
|
||||
node : doc.expr
|
||||
The doc AST expression node for error reporting.
|
||||
|
||||
var_name : str
|
||||
The variable name.
|
||||
|
||||
value : Any
|
||||
The value to be bound with.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The bound value.
|
||||
"""
|
||||
if isinstance(value, list | tuple):
|
||||
for i, v in enumerate(value):
|
||||
bind_with_value(self, node, f"{var_name}_{i}", v)
|
||||
return value
|
||||
elif isinstance(value, Buffer | Var):
|
||||
IRBuilder.name(var_name, value)
|
||||
return value
|
||||
else:
|
||||
self.report_error(node, f"Do not know how to bind type: {type(value)} in with statement")
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def bind_for_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
|
||||
"""Value binding methods when parsing for statement.
|
||||
e.g. binding i, j, k with T.grid(128, 128, 128), when parsing
|
||||
for i, j, k in T.grid(128, 128, 128).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The current parser.
|
||||
|
||||
node : doc.expr
|
||||
The doc AST expression node for error reporting.
|
||||
|
||||
var_name : str
|
||||
The variable name.
|
||||
|
||||
value : Any
|
||||
The value to be bound with.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The bound value.
|
||||
"""
|
||||
if isinstance(value, list | tuple | tvm.ir.Array):
|
||||
for i, v in enumerate(value):
|
||||
bind_for_value(self, node, f"{var_name}_{i}", v)
|
||||
return value
|
||||
elif isinstance(value, Var):
|
||||
IRBuilder.name(var_name, value)
|
||||
return value
|
||||
else:
|
||||
self.report_error(node, f"Do not know how to bind type: {type(value)} in for statement")
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
|
||||
"""Value binding methods when parsing assign statement.
|
||||
e.g. binding vi, vj, vk with T.axis.remap("SSR", [i, j, k]), when parsing
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k]).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The current parser.
|
||||
|
||||
node : doc.expr
|
||||
The doc AST expression node for error reporting.
|
||||
|
||||
var_name : str
|
||||
The variable name.
|
||||
|
||||
value : Any
|
||||
The value to be bound with.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The bound value.
|
||||
"""
|
||||
if isinstance(value, T.scalar_wrapper): # pylint: disable=protected-access
|
||||
# special case for scalar, name the buffer, but the var is used as BufferLoad
|
||||
assert isinstance(value.scalar, T.BufferLoad)
|
||||
IRBuilder.name(var_name, value.scalar.buffer)
|
||||
return value.scalar
|
||||
if isinstance(value, T.meta_var):
|
||||
return value.value
|
||||
elif getattr(type(value), "_is_meta_class", False):
|
||||
name_meta_class_value(var_name, value)
|
||||
return value
|
||||
elif isinstance(value, list | tuple):
|
||||
# Tuple-unpacking with a starred target (e.g. ``vi, *vs = T.axis.remap(...)``)
|
||||
# collects multiple elements into a single list bound here. Recurse so each
|
||||
# element gets a per-index name; this matches apache's behavior.
|
||||
for i, v in enumerate(value):
|
||||
bind_assign_value(self, node, f"{var_name}_{i}", v)
|
||||
return value
|
||||
elif isinstance(value, BufferRegion):
|
||||
return value
|
||||
elif isinstance(value, Frame):
|
||||
value.add_callback(partial(value.__exit__, None, None, None))
|
||||
res = value.__enter__()
|
||||
IRBuilder.name(var_name, res)
|
||||
return res
|
||||
elif isinstance(value, Buffer | IterVar | Layout) or (
|
||||
isinstance(value, Var) and not self.var_table.exist(value)
|
||||
):
|
||||
IRBuilder.name(var_name, value)
|
||||
return value
|
||||
else:
|
||||
if not tvm.ir.is_prim_expr(value):
|
||||
value = tvm.tirx.const(value)
|
||||
if not isinstance(value, tvm.tirx.StringImm):
|
||||
# x = expr -> scalar (auto-typed from value)
|
||||
scalar = T.local_scalar(dtype=str(value.ty.dtype))
|
||||
IRBuilder.name(var_name, scalar.scalar.buffer)
|
||||
T.buffer_store(scalar.scalar.buffer, value, [0])
|
||||
return scalar.scalar
|
||||
else:
|
||||
# StringImm: x = expr -> immutable Bind var
|
||||
ann_var = tvm.tirx.Var(var_name, value.ty)
|
||||
IRBuilder.name(var_name, ann_var)
|
||||
T.Bind(value, var=ann_var)
|
||||
return ann_var
|
||||
|
||||
|
||||
def find_decorator_annotation(node: doc.FunctionDef, annotation: str, default: bool = True) -> bool:
|
||||
"""
|
||||
Check the value of given annotation (argument name) in the prim_func decorator.
|
||||
Returns the value of the annotation if present, otherwise giving the default value.
|
||||
"""
|
||||
# look for the named argument in the prim_func / jit decorator
|
||||
for dec in node.decorator_list:
|
||||
if not isinstance(dec, doc.Call) or dec.func.attr not in ("prim_func", "jit"):
|
||||
continue
|
||||
for keyword in dec.keywords:
|
||||
if keyword.arg == annotation:
|
||||
return keyword.value.value
|
||||
return default
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="For")
|
||||
def visit_for(self: Parser, node: doc.For) -> None:
|
||||
"""The for visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.For
|
||||
The doc AST for node.
|
||||
"""
|
||||
# Intercept range() at AST level so it works with both Python ints and PrimExprs.
|
||||
# In other contexts (e.g. list comprehensions), range remains Python's builtin.
|
||||
if (
|
||||
isinstance(node.iter, doc.Call)
|
||||
and isinstance(node.iter.func, doc.Name)
|
||||
and node.iter.func.id == "range"
|
||||
):
|
||||
args = [self.eval_expr(a) for a in node.iter.args]
|
||||
kwargs = {kw.arg: self.eval_expr(kw.value) for kw in node.iter.keywords}
|
||||
if len(args) == 1:
|
||||
for_frame = T.serial(0, args[0], **kwargs)
|
||||
elif len(args) == 2:
|
||||
for_frame = T.serial(args[0], args[1], **kwargs)
|
||||
elif len(args) == 3:
|
||||
for_frame = T.serial(args[0], args[1], step=args[2], **kwargs)
|
||||
else:
|
||||
self.report_error(node.iter, "range() takes 1 to 3 arguments")
|
||||
else:
|
||||
for_frame = self.eval_expr(node.iter)
|
||||
if not isinstance(for_frame, T.frame.ForFrame):
|
||||
self.report_error(
|
||||
node.iter,
|
||||
"Expect the for loop to be one of the following: "
|
||||
"range, T.serial, T.grid, T.parallel, T.vectorized, T.unroll, T.thread_binding",
|
||||
)
|
||||
with self.var_table.with_frame():
|
||||
with for_frame as iters:
|
||||
self.eval_assign(target=node.target, source=iters, bind_value=bind_for_value)
|
||||
self.visit_body(node.body)
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="While")
|
||||
def visit_while(self: Parser, node: doc.While) -> None:
|
||||
"""The while visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.While
|
||||
The doc AST while node.
|
||||
"""
|
||||
with self.var_table.with_frame():
|
||||
cond = self.eval_expr(node.test)
|
||||
with T.While(cond):
|
||||
self.visit_body(node.body)
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="Break")
|
||||
def visit_break(self: Parser, node: doc.Break) -> None:
|
||||
"""The break visiting method for tir.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.Break
|
||||
The doc AST break node.
|
||||
"""
|
||||
T.evaluate(T.break_loop())
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="Continue")
|
||||
def visit_continue(self: Parser, node: doc.Continue) -> None:
|
||||
"""The continue visiting method for tir.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.Continue
|
||||
The doc AST continue node.
|
||||
"""
|
||||
T.evaluate(T.continue_loop())
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="Assign")
|
||||
def visit_assign(self: Parser, node: doc.Assign) -> None:
|
||||
"""The assign visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.Assign
|
||||
The doc AST assign node.
|
||||
"""
|
||||
if len(node.targets) != 1:
|
||||
self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
|
||||
lhs = node.targets[0]
|
||||
|
||||
if isinstance(node.value, doc.Subscript):
|
||||
check_slices = []
|
||||
if isinstance(node.value.slice, doc.Slice):
|
||||
check_slices = [node.value.slice]
|
||||
elif isinstance(node.value.slice, doc.Tuple):
|
||||
for p in node.value.slice.elts:
|
||||
if isinstance(p, doc.Slice):
|
||||
check_slices.append(p)
|
||||
for s in check_slices:
|
||||
if not s.step and s.upper and s.lower:
|
||||
s.step = doc.Constant(
|
||||
1,
|
||||
None,
|
||||
s.upper.lineno,
|
||||
s.upper.end_col_offset + 1,
|
||||
s.upper.lineno,
|
||||
s.upper.end_col_offset + 2,
|
||||
)
|
||||
|
||||
rhs = self.eval_expr(node.value)
|
||||
if isinstance(lhs, doc.Subscript):
|
||||
if isinstance(lhs.slice, doc.Tuple):
|
||||
indices = []
|
||||
for index in lhs.slice.elts:
|
||||
if isinstance(index, doc.Starred):
|
||||
# x[*y]
|
||||
indices.extend(self.eval_expr(index.value))
|
||||
else:
|
||||
indices.append(self.eval_expr(index))
|
||||
else:
|
||||
indices = self.eval_expr(lhs.slice)
|
||||
T.buffer_store(self.eval_expr(lhs.value), rhs, indices)
|
||||
else:
|
||||
# special case for scalar buffers
|
||||
# scalar = xxx <=> scalar.buffer[()] = xxx
|
||||
# or for a normal 1-dim buffer with shape (1,)
|
||||
# buffer = xxx <=> buffer[()] = xxx
|
||||
# Try to resolve lhs as a buffer/scalar variable. eval_expr may raise
|
||||
# if the name is not yet defined (i.e. this is a new variable binding),
|
||||
# which is the expected fallthrough case.
|
||||
lhs_value = None
|
||||
try:
|
||||
lhs_copy = deepcopy(lhs)
|
||||
if hasattr(lhs_copy, "ctx"):
|
||||
lhs_copy.ctx = doc.Load()
|
||||
lhs_value = self.eval_expr(lhs_copy)
|
||||
except Exception: # pylint: disable=broad-except
|
||||
pass
|
||||
# Buffer check and store are intentionally outside the try/except so
|
||||
# that genuine errors (e.g. wrong shape, bad store) are not swallowed.
|
||||
# Only TypeError from FFI type mismatch (e.g. rhs is a meta_var, not
|
||||
# a Expr or auto-convertible scalar) triggers fallthrough.
|
||||
if isinstance(lhs_value, T.scalar_wrapper | T.BufferLoad | tvm.tirx.Buffer):
|
||||
if isinstance(lhs_value, T.scalar_wrapper):
|
||||
buffer = lhs_value.scalar.buffer
|
||||
else:
|
||||
buffer = lhs_value.buffer if isinstance(lhs_value, T.BufferLoad) else lhs_value
|
||||
if len(buffer.shape) == 1 and bool(buffer.shape[0] == 1):
|
||||
# only 1-dim buffer with shape (1,) can be assigned directly
|
||||
# Note that shape can be a Expr, so we only judge by
|
||||
# bool(shape[0] == 1) rather than int(shape[0]) == 1.
|
||||
try:
|
||||
T.buffer_store(buffer, rhs, [0])
|
||||
return
|
||||
except TypeError:
|
||||
pass # rhs not compatible with buffer_store, fall through
|
||||
# otherwise
|
||||
self.eval_assign(target=lhs, source=rhs, bind_value=bind_assign_value)
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="AugAssign")
|
||||
def visit_aug_assign(self: Parser, node: doc.AugAssign) -> None:
|
||||
"""The augmented assign visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.AugAssign
|
||||
The doc AST augmented assign node.
|
||||
"""
|
||||
lhs_pos = (
|
||||
node.target.lineno,
|
||||
node.target.col_offset,
|
||||
node.target.end_lineno,
|
||||
node.target.end_col_offset,
|
||||
)
|
||||
rhs_pos = (
|
||||
node.value.lineno,
|
||||
node.value.col_offset,
|
||||
node.value.end_lineno,
|
||||
node.value.end_col_offset,
|
||||
)
|
||||
node.target.ctx = doc.Load()
|
||||
with self.var_table.with_frame():
|
||||
lhs_name = "__tvm_tmp_value_aug_assign_lhs"
|
||||
rhs_name = "__tvm_tmp_value_aug_assign_rhs"
|
||||
lhs_expr = self.eval_expr(node.target)
|
||||
rhs_expr = self.eval_expr(node.value)
|
||||
self.var_table.add(lhs_name, lhs_expr)
|
||||
self.var_table.add(rhs_name, rhs_expr)
|
||||
op = doc.BinOp(
|
||||
doc.Name(lhs_name, doc.Load(), *lhs_pos),
|
||||
node.op,
|
||||
doc.Name(rhs_name, doc.Load(), *rhs_pos),
|
||||
*lhs_pos,
|
||||
)
|
||||
rhs = self.eval_expr(op)
|
||||
lhs = node.target
|
||||
lhs.ctx = doc.Store()
|
||||
if isinstance(lhs, doc.Subscript):
|
||||
if isinstance(lhs.slice, doc.Tuple):
|
||||
indices = []
|
||||
for index in lhs.slice.elts:
|
||||
if isinstance(index, doc.Starred):
|
||||
# x[*y]
|
||||
indices.extend(self.eval_expr(index.value))
|
||||
else:
|
||||
indices.append(self.eval_expr(index))
|
||||
else:
|
||||
indices = [self.eval_expr(lhs.slice)]
|
||||
T.buffer_store(self.eval_expr(lhs.value), rhs, indices)
|
||||
else:
|
||||
lhs_value = None
|
||||
try:
|
||||
lhs_copy = deepcopy(lhs)
|
||||
if hasattr(lhs_copy, "ctx"):
|
||||
lhs_copy.ctx = doc.Load()
|
||||
lhs_value = self.eval_expr(lhs_copy)
|
||||
except Exception: # pylint: disable=broad-except
|
||||
pass
|
||||
if isinstance(lhs_value, T.scalar_wrapper | T.BufferLoad | tvm.tirx.Buffer):
|
||||
if isinstance(lhs_value, T.scalar_wrapper):
|
||||
buffer = lhs_value.scalar.buffer
|
||||
else:
|
||||
buffer = lhs_value.buffer if isinstance(lhs_value, T.BufferLoad) else lhs_value
|
||||
if len(buffer.shape) == 1 and bool(buffer.shape[0] == 1):
|
||||
try:
|
||||
T.buffer_store(buffer, rhs, [0])
|
||||
return
|
||||
except TypeError:
|
||||
pass
|
||||
self.eval_assign(target=lhs, source=rhs, bind_value=bind_assign_value)
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="AnnAssign")
|
||||
def visit_ann_assign(self: Parser, node: doc.AnnAssign) -> None:
|
||||
"""The annotated assign visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.AnnAssign
|
||||
The doc AST annotated assign node.
|
||||
"""
|
||||
lhs = node.target
|
||||
rhs = self.eval_expr(node.value) if node.value is not None else None
|
||||
raw_ann = self.eval_expr(node.annotation)
|
||||
|
||||
if isinstance(raw_ann, T.LocalVectorAnnotation):
|
||||
# x: T.float32[N] or x: T.f32[M, N] -> local buffer allocation
|
||||
if rhs is not None:
|
||||
self.report_error(node, "Vector annotation does not support initial value")
|
||||
buf = T.alloc_local(shape=raw_ann.shape, dtype=raw_ann.dtype)
|
||||
self.eval_assign(target=lhs, source=buf, bind_value=bind_assign_value)
|
||||
elif isinstance(raw_ann, T.LetAnnotation):
|
||||
# T.let or T.let[type] -> immutable Bind var
|
||||
if rhs is None:
|
||||
self.report_error(node, "T.let annotation requires a value")
|
||||
if not isinstance(rhs, Expr):
|
||||
if isinstance(rhs, str):
|
||||
rhs = tvm.tirx.StringImm(rhs)
|
||||
else:
|
||||
rhs = tvm.tirx.const(rhs)
|
||||
if raw_ann.type_spec is not None:
|
||||
ann_var = raw_ann.as_var()
|
||||
else:
|
||||
ann_var = raw_ann.as_var(rhs_dtype=rhs.ty)
|
||||
if not isinstance(ann_var, Var):
|
||||
self.report_error(node.annotation, "Annotation should resolve to Var")
|
||||
self.eval_assign(target=lhs, source=ann_var, bind_value=bind_assign_value)
|
||||
T.Bind(rhs, var=ann_var)
|
||||
else:
|
||||
ann_var = raw_ann() if callable(raw_ann) else raw_ann
|
||||
if not isinstance(ann_var, Var):
|
||||
self.report_error(node.annotation, "Annotation should resolve to Var")
|
||||
if not isinstance(ann_var.ty, PrimType):
|
||||
self.report_error(
|
||||
node.annotation,
|
||||
"Use T.let[...] for non-PrimType annotations (e.g. PointerType, handle)",
|
||||
)
|
||||
if str(ann_var.ty) == "handle":
|
||||
self.report_error(
|
||||
node.annotation,
|
||||
"handle type cannot be used as scalar annotation; use T.let[T.handle] instead",
|
||||
)
|
||||
# x: T.int32 = expr -> scalar (mutable scalar buffer)
|
||||
scalar = T.local_scalar(dtype=str(ann_var.ty))
|
||||
self.eval_assign(target=lhs, source=scalar, bind_value=bind_assign_value)
|
||||
if rhs is not None:
|
||||
T.buffer_store(scalar.scalar.buffer, rhs, [0])
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="With")
|
||||
def visit_with(self: Parser, node: doc.With) -> None:
|
||||
"""The with visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.With
|
||||
The doc AST with node.
|
||||
"""
|
||||
with contextlib.ExitStack() as stack:
|
||||
stack.enter_context(self.var_table.with_frame())
|
||||
for item in node.items:
|
||||
frame = self.eval_expr(item.context_expr)
|
||||
if not isinstance(frame, Frame) and not (
|
||||
hasattr(frame, "__enter__") and hasattr(frame, "__exit__")
|
||||
):
|
||||
self.report_error(
|
||||
item.context_expr,
|
||||
"Invalid context expression in the with-statement.",
|
||||
)
|
||||
rhs = stack.enter_context(frame)
|
||||
if item.optional_vars is not None:
|
||||
self.eval_assign(target=item.optional_vars, source=rhs, bind_value=bind_with_value)
|
||||
self.visit_body(node.body)
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="FunctionDef")
|
||||
def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
|
||||
"""The function definition visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.FunctionDef
|
||||
The doc AST function definition node.
|
||||
"""
|
||||
supplied_annotation = self.function_annotations
|
||||
func_annotation = supplied_annotation.get(node.name, {})
|
||||
privacy = find_decorator_annotation(node, "private", default=False)
|
||||
s_tir = find_decorator_annotation(node, "s_tir", default=False)
|
||||
persistent = find_decorator_annotation(node, "persistent", default=False)
|
||||
self.function_annotations = None
|
||||
with self.var_table.with_frame():
|
||||
prim_func_ctx = T.prim_func(is_private=privacy, s_tir=s_tir, persistent=persistent)
|
||||
with prim_func_ctx:
|
||||
T.func_name(node.name)
|
||||
if node.returns is not None:
|
||||
ret_type = self.eval_expr(node.returns)
|
||||
if callable(ret_type):
|
||||
ret_type = ret_type().ty
|
||||
T.func_ret(ret_type)
|
||||
with self.with_dispatch_token("tirx"):
|
||||
# TODO: handle different types of arguments:
|
||||
# - vararg: arg | None
|
||||
# - kwonlyargs: list[arg]
|
||||
# - kw_defaults: list[expr | None]
|
||||
# - kwarg: arg | None
|
||||
# - defaults: list[expr]
|
||||
# - posonlyargs: list[arg]
|
||||
for arg in node.args.args:
|
||||
if arg.annotation is None:
|
||||
self.report_error(arg, "Type annotation required for function parameters.")
|
||||
try:
|
||||
ann = self.eval_expr(arg.annotation)
|
||||
if callable(ann) and ann is not _constexpr_sentinel:
|
||||
ann = ann()
|
||||
except Exception: # pylint: disable=broad-except
|
||||
ann = func_annotation.get(arg.arg, None)
|
||||
if ann is None:
|
||||
raise
|
||||
if ann is _constexpr_sentinel:
|
||||
# T.constexpr param: value was bound in extra_vars by
|
||||
# TIRJit.specialize() and lives in an outer var_table
|
||||
# frame; do not register a runtime PrimFunc param.
|
||||
continue
|
||||
param = T.arg(arg.arg, ann)
|
||||
self.var_table.add(arg.arg, param)
|
||||
self.visit_body(node.body)
|
||||
self.function_annotations = supplied_annotation
|
||||
|
||||
|
||||
@dispatch.register(token="tir.inline", type_name="FunctionDef")
|
||||
def visit_inline_function_def(self: Parser, node: doc.FunctionDef) -> None:
|
||||
"""The function definition visiting method for inline functions in tir.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.FunctionDef
|
||||
The doc AST function definition node.
|
||||
"""
|
||||
# remove the inline decorator
|
||||
node.decorator_list.pop()
|
||||
# adjust the node location to the source code location
|
||||
node.lineno += self.diag.source.start_line - 1
|
||||
node.col_offset += self.diag.source.start_column + 1
|
||||
node.end_lineno += self.diag.source.start_line - 1
|
||||
node.end_col_offset += self.diag.source.start_column + 1
|
||||
|
||||
# Record definition depth for LEGB late binding
|
||||
definition_depth = len(self.var_table.frames)
|
||||
|
||||
def get_func():
|
||||
func_ast = from_doc(node)
|
||||
module_ast = ast.Module(body=[func_ast], type_ignores=[])
|
||||
ast.fix_missing_locations(module_ast)
|
||||
# set the filename to the source name, so that the error message can be reported correctly
|
||||
code_obj = compile(module_ast, filename=self.diag.source.source_name, mode="exec")
|
||||
namespace = self.var_table.get()
|
||||
exec(code_obj, namespace) # pylint: disable=exec-used
|
||||
func_name = func_ast.name
|
||||
func = namespace[func_name]
|
||||
return func, func_name
|
||||
|
||||
func, func_name = get_func()
|
||||
wrapper = inline(func, definition_depth=definition_depth, defining_var_table=self.var_table)
|
||||
|
||||
self.var_table.add(func_name, wrapper, allow_shadowing=False)
|
||||
return None
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="tvm_annotation")
|
||||
def visit_tvm_annotation(self: Parser, node: doc.expr):
|
||||
"""The TVM annotation visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.expr
|
||||
The doc AST expr node.
|
||||
"""
|
||||
annotation = self.eval_expr(node)
|
||||
if callable(annotation):
|
||||
annotation = annotation()
|
||||
return annotation
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="Expr")
|
||||
def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
|
||||
"""The expr statement visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.Expr
|
||||
The doc AST Expr node.
|
||||
"""
|
||||
|
||||
res = self.eval_expr(node.value)
|
||||
if res is None:
|
||||
pass
|
||||
elif isinstance(res, Frame):
|
||||
res.add_callback(partial(res.__exit__, None, None, None))
|
||||
res.__enter__()
|
||||
elif hasattr(res, "frames") and hasattr(res, "__enter__"):
|
||||
# _FrameScope from T.attr({...}) — enter each inner frame for concise scoping
|
||||
for f in res.frames:
|
||||
f.add_callback(partial(f.__exit__, None, None, None))
|
||||
f.__enter__()
|
||||
elif isinstance(res, Var):
|
||||
# Standalone Var expression (e.g. from T.bind(value, var=v)) --
|
||||
# the Bind statement was already emitted to the parent frame by the FFI call,
|
||||
# so just discard the returned Var.
|
||||
pass
|
||||
elif tvm.ir.is_prim_expr(res):
|
||||
T.evaluate(res)
|
||||
elif isinstance(res, int | bool):
|
||||
T.evaluate(tvm.tirx.const(res))
|
||||
elif isinstance(res, tvm.ir.Call) and not tvm.ir.is_prim_expr(res):
|
||||
if isinstance(res.op, tvm.ir.GlobalVar) and res.ty.is_missing():
|
||||
# GlobalVar calls with a missing return type are ambiguous, as each IR has a
|
||||
# different function Call representation. Convert to the TIR representation.
|
||||
T.evaluate(tvm.tirx.call_tir(res.op, *res.args))
|
||||
else:
|
||||
# Pointer-valued TIR calls are general Expr rather than PrimExpr,
|
||||
# but are still valid standalone Evaluate statements.
|
||||
T.evaluate(res)
|
||||
elif isinstance(res, str):
|
||||
# Ignore docstrings
|
||||
pass
|
||||
elif isinstance(res, tvm.tirx.stmt.BufferStore):
|
||||
T.buffer_store(res.buffer, res.value, res.indices, res.predicate)
|
||||
elif isinstance(res, tvm.tirx.Buffer):
|
||||
# ``T.match_buffer(...)`` used as a bare statement (no LHS) — the
|
||||
# buffer object is discarded; the underlying side effect (the
|
||||
# match_buffer node) has already been emitted into the frame.
|
||||
pass
|
||||
else:
|
||||
self.report_error(node, f"Parsing resulted in unexpected type {type(res)}")
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="If")
|
||||
def visit_if(self: Parser, node: doc.If) -> None:
|
||||
"""The if visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.If
|
||||
The doc AST if node.
|
||||
"""
|
||||
with self.var_table.with_frame():
|
||||
predicate = self.eval_expr(node.test)
|
||||
if tvm.ir.is_prim_expr(predicate) or isinstance(predicate, tvm.tirx.expr.ExprOp):
|
||||
with T.If(self.eval_expr(node.test)):
|
||||
with T.Then():
|
||||
with self.var_table.with_frame():
|
||||
self.visit_body(node.body)
|
||||
if node.orelse:
|
||||
with T.Else():
|
||||
with self.var_table.with_frame():
|
||||
self.visit_body(node.orelse)
|
||||
elif isinstance(predicate, bool):
|
||||
if predicate:
|
||||
with self.var_table.with_frame():
|
||||
self.visit_body(node.body)
|
||||
elif node.orelse:
|
||||
with self.var_table.with_frame():
|
||||
self.visit_body(node.orelse)
|
||||
else:
|
||||
self.report_error(
|
||||
node.test,
|
||||
f"If condition must be a boolean expression, but got {predicate}",
|
||||
)
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="Assert")
|
||||
def visit_assert(self: Parser, node: doc.Assert) -> None:
|
||||
"""The assert visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.Assert
|
||||
The doc AST assert node.
|
||||
|
||||
The assert message can be either:
|
||||
- A plain string: ``assert cond, "message"``
|
||||
- A tuple of (kind, [parts...]): ``assert cond, ("ValueError", ["part0", "part1"])``
|
||||
"""
|
||||
cond = self.eval_expr(node.test)
|
||||
msg = self.eval_expr(node.msg)
|
||||
|
||||
kind = "RuntimeError"
|
||||
message = msg
|
||||
|
||||
if isinstance(msg, tuple):
|
||||
if len(msg) != 2:
|
||||
self.report_error(
|
||||
node,
|
||||
f"Assert message tuple must have exactly 2 elements (kind, [parts...]), "
|
||||
f"got {len(msg)} elements",
|
||||
)
|
||||
kind_str, parts = msg
|
||||
if isinstance(kind_str, tvm.tirx.StringImm):
|
||||
kind_str = kind_str.value
|
||||
if not isinstance(kind_str, str):
|
||||
self.report_error(
|
||||
node,
|
||||
f"Assert message tuple first element must be a string (error kind like "
|
||||
f'"ValueError"), got {type(kind_str).__name__}',
|
||||
)
|
||||
kind = kind_str
|
||||
message = parts
|
||||
|
||||
if isinstance(message, list | tuple):
|
||||
message = [p.value if isinstance(p, tvm.tirx.StringImm) else str(p) for p in message]
|
||||
|
||||
frame = T.Assert(cond, message, error_kind=kind)
|
||||
frame.add_callback(partial(frame.__exit__, None, None, None))
|
||||
frame.__enter__()
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="Return")
|
||||
def visit_return(self: Parser, node: doc.Return) -> None:
|
||||
"""The return visiting method for tirx.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.Return
|
||||
The doc AST return node.
|
||||
"""
|
||||
value = self.eval_expr(node.value)
|
||||
if value is None:
|
||||
self.report_error(node, "Expression to be returned must be a Expr")
|
||||
T.evaluate(tvm.tirx.ret(value))
|
||||
|
||||
|
||||
@dispatch.register(token="tirx", type_name="tvm_declare_function")
|
||||
def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
|
||||
"""The function declaration step for tirx
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.Return
|
||||
The doc AST return node.
|
||||
"""
|
||||
|
||||
supplied_annotation = self.function_annotations
|
||||
func_annotation = supplied_annotation.get(node.name, {})
|
||||
|
||||
ret_type = None
|
||||
with self.var_table.with_frame():
|
||||
if node.returns is not None:
|
||||
ret_type = self.eval_expr(node.returns)
|
||||
if callable(ret_type):
|
||||
ret_type = ret_type().ty
|
||||
|
||||
arg_annotations = []
|
||||
for arg in node.args.args:
|
||||
if arg.annotation is None:
|
||||
self.report_error(arg, "Type annotation required for function parameters.")
|
||||
try:
|
||||
ann = self.eval_expr(arg.annotation)
|
||||
if callable(ann):
|
||||
ann = ann()
|
||||
except Exception: # pylint: disable=broad-except
|
||||
ann = func_annotation.get(arg.arg, None)
|
||||
if ann is None:
|
||||
raise
|
||||
|
||||
IRBuilder.name(arg.arg, ann)
|
||||
arg_annotations.append(ann)
|
||||
|
||||
func_signature = tvm.tirx.PrimFunc(arg_annotations, None, ret_type=ret_type)
|
||||
return I.decl_function(node.name, func_signature)
|
||||
@@ -0,0 +1,119 @@
|
||||
# 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.
|
||||
"""Tile primitive shorthand namespace for TIRx TVMScript."""
|
||||
|
||||
import functools
|
||||
|
||||
from tvm.tirx import Buffer, BufferRegion
|
||||
|
||||
from .builder import tirx as _builder
|
||||
|
||||
_TILE_ARG_TYPES = (Buffer, BufferRegion)
|
||||
|
||||
|
||||
def _get_arg(args, kwargs, index, name):
|
||||
if len(args) > index:
|
||||
return args[index]
|
||||
return kwargs.get(name)
|
||||
|
||||
|
||||
def _require_buffer_arg(op_name, arg_name, value):
|
||||
if not isinstance(value, _TILE_ARG_TYPES):
|
||||
raise TypeError(
|
||||
f"Tx.{op_name} is tile-only and expects `{arg_name}` to be a Buffer "
|
||||
f"or BufferRegion; use T.{op_name} for expression/builtin calls"
|
||||
)
|
||||
|
||||
|
||||
def _validate_tile_call(op_name, args, kwargs):
|
||||
dst = _get_arg(args, kwargs, 0, "dst")
|
||||
_require_buffer_arg(op_name, "dst", dst)
|
||||
|
||||
if op_name in {"cast", "max", "min", "permute_layout", "silu"}:
|
||||
src = _get_arg(args, kwargs, 1, "src")
|
||||
_require_buffer_arg(op_name, "src", src)
|
||||
elif op_name in {"sqrt", "exp", "exp2", "reciprocal"}:
|
||||
src = _get_arg(args, kwargs, 1, "src")
|
||||
if src is not None:
|
||||
_require_buffer_arg(op_name, "src", src)
|
||||
|
||||
|
||||
def _tile_scoped_op(op_name):
|
||||
scoped_op = getattr(_builder, op_name)
|
||||
|
||||
@functools.wraps(scoped_op._fn) # pylint: disable=protected-access
|
||||
def wrapper(*args, scope=None, **kwargs):
|
||||
_validate_tile_call(op_name, args, kwargs)
|
||||
return scoped_op._fn(*args, scope=scope, **kwargs) # pylint: disable=protected-access
|
||||
|
||||
return _builder.ScopedOp(wrapper)
|
||||
|
||||
|
||||
_SCOPED_TILE_OP_NAMES = [
|
||||
"add",
|
||||
"binary_chain",
|
||||
"binary_reduce",
|
||||
"cast",
|
||||
"copy",
|
||||
"copy_async",
|
||||
"exp",
|
||||
"exp2",
|
||||
"fdiv",
|
||||
"fill",
|
||||
"fma",
|
||||
"gemm",
|
||||
"gemm_async",
|
||||
"max",
|
||||
"maximum",
|
||||
"memset",
|
||||
"min",
|
||||
"minimum",
|
||||
"mul",
|
||||
"permute_layout",
|
||||
"reciprocal",
|
||||
"reduce_negate",
|
||||
"select",
|
||||
"silu",
|
||||
"sqrt",
|
||||
"sub",
|
||||
"sum",
|
||||
"unary_reduce",
|
||||
"zero",
|
||||
]
|
||||
|
||||
for _op_name in _SCOPED_TILE_OP_NAMES:
|
||||
globals()[_op_name] = _tile_scoped_op(_op_name)
|
||||
|
||||
cluster = _builder.ScopeNamespace("cluster", "cluster")
|
||||
cta = _builder.ScopeNamespace("cta", "cta")
|
||||
wg = _builder.ScopeNamespace("warpgroup", "wg")
|
||||
warpgroup = _builder.ScopeNamespace("warpgroup", "warpgroup")
|
||||
warp = _builder.ScopeNamespace("warp", "warp")
|
||||
thread = _builder.ScopeNamespace("thread", "thread")
|
||||
|
||||
compose_op = _builder.compose_op
|
||||
|
||||
__all__ = [
|
||||
*_SCOPED_TILE_OP_NAMES,
|
||||
"cluster",
|
||||
"compose_op",
|
||||
"cta",
|
||||
"thread",
|
||||
"warp",
|
||||
"warpgroup",
|
||||
"wg",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,22 @@
|
||||
# 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.
|
||||
"""Namespace of all TIR transformations"""
|
||||
# pylint: disable=wildcard-import, invalid-name
|
||||
|
||||
from .function_pass import prim_func_pass, PrimFuncPass
|
||||
from .transform import *
|
||||
@@ -0,0 +1,21 @@
|
||||
# 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.
|
||||
"""FFI APIs for tvm.tirx.transform"""
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
tvm_ffi.init_ffi_api("tirx.transform", __name__)
|
||||
@@ -0,0 +1,191 @@
|
||||
# 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.
|
||||
|
||||
|
||||
from tvm.ir import Call, Op, is_prim_expr
|
||||
from tvm.tirx import (
|
||||
AllocBuffer,
|
||||
BufferLoad,
|
||||
BufferRegion,
|
||||
BufferStore,
|
||||
DeclBuffer,
|
||||
Evaluate,
|
||||
Expr,
|
||||
Stmt,
|
||||
TilePrimitiveCall,
|
||||
Var,
|
||||
decl_buffer,
|
||||
)
|
||||
from tvm.tirx.buffer import Buffer
|
||||
from tvm.tirx.layout import Iter, TileLayout
|
||||
from tvm.tirx.stmt_functor import StmtExprMutator, StmtMutator
|
||||
|
||||
|
||||
class BufferReplacer(StmtExprMutator):
|
||||
"""
|
||||
Replace buffer with another buffer.
|
||||
Also replace the data of the buffer with another var.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, buffer_map: dict[Buffer, Buffer] | None = None, var_map: dict[Var, Var] | None = None
|
||||
):
|
||||
super().__init__()
|
||||
self.buffer_map = buffer_map if buffer_map is not None else {}
|
||||
self.var_map = var_map if var_map is not None else {}
|
||||
self.buffer_attr_var_mutated = False
|
||||
for old_buffer, new_buffer in self.buffer_map.items():
|
||||
self.var_map[old_buffer.data] = new_buffer.data
|
||||
|
||||
def mutate_buffer(self, buffer: Buffer):
|
||||
if buffer in self.buffer_map:
|
||||
return self.buffer_map[buffer]
|
||||
|
||||
# Track mutations for this specific buffer only. Without this reset,
|
||||
# unrelated buffers can be spuriously cloned and introduce alias buffers.
|
||||
prev_mutated = self.buffer_attr_var_mutated
|
||||
self.buffer_attr_var_mutated = False
|
||||
new_data = self.visit_expr(buffer.data)
|
||||
new_shape = [self.visit_expr(expr) for expr in buffer.shape]
|
||||
new_strides = [self.visit_expr(expr) for expr in buffer.strides]
|
||||
new_elem_offset = (
|
||||
self.visit_expr(buffer.elem_offset) if buffer.elem_offset is not None else None
|
||||
)
|
||||
if isinstance(buffer.layout, TileLayout):
|
||||
new_shard = []
|
||||
new_replicate = []
|
||||
for iter in buffer.layout.shard:
|
||||
new_iter = Iter(
|
||||
self.visit_expr(iter.extent), self.visit_expr(iter.stride), iter.axis
|
||||
)
|
||||
new_shard.append(new_iter)
|
||||
for iter in buffer.layout.replica:
|
||||
new_iter = Iter(
|
||||
self.visit_expr(iter.extent), self.visit_expr(iter.stride), iter.axis
|
||||
)
|
||||
new_replicate.append(new_iter)
|
||||
new_layout = TileLayout.from_iters(
|
||||
new_shard, new_replicate, offset=buffer.layout.offset
|
||||
)
|
||||
else:
|
||||
new_layout = buffer.layout
|
||||
buffer_attr_mutated = self.buffer_attr_var_mutated
|
||||
self.buffer_attr_var_mutated = prev_mutated or buffer_attr_mutated
|
||||
if not buffer_attr_mutated:
|
||||
return None
|
||||
new_buffer = decl_buffer(
|
||||
new_shape,
|
||||
buffer.dtype,
|
||||
buffer.name,
|
||||
new_data,
|
||||
new_strides,
|
||||
new_elem_offset,
|
||||
buffer.scope(),
|
||||
buffer.data_alignment,
|
||||
buffer.offset_factor,
|
||||
layout=new_layout,
|
||||
)
|
||||
self.buffer_map[buffer] = new_buffer
|
||||
return new_buffer
|
||||
|
||||
def visit_var_(self, op: Var):
|
||||
op = super().visit_var_(op)
|
||||
if op in self.var_map:
|
||||
self.buffer_attr_var_mutated = True
|
||||
return self.var_map[op]
|
||||
return op
|
||||
|
||||
def visit_buffer_load_(self, op: BufferLoad):
|
||||
new_buffer = self.mutate_buffer(op.buffer)
|
||||
op = super().visit_buffer_load_(op)
|
||||
if new_buffer is not None:
|
||||
return BufferLoad(new_buffer, op.indices)
|
||||
return op
|
||||
|
||||
def visit_buffer_store_(self, op: BufferStore):
|
||||
new_buffer = self.mutate_buffer(op.buffer)
|
||||
op = super().visit_buffer_store_(op)
|
||||
if new_buffer is not None:
|
||||
return BufferStore(new_buffer, op.value, op.indices)
|
||||
return op
|
||||
|
||||
def visit_buffer_region_(self, op: BufferRegion):
|
||||
new_buffer = self.mutate_buffer(op.buffer)
|
||||
op = super().visit_buffer_region_(op)
|
||||
if new_buffer is not None:
|
||||
return BufferRegion(new_buffer, op.region)
|
||||
return op
|
||||
|
||||
def visit_decl_buffer_(self, op: DeclBuffer):
|
||||
new_buffer = self.mutate_buffer(op.buffer)
|
||||
op = super().visit_decl_buffer_(op)
|
||||
if new_buffer is not None:
|
||||
return DeclBuffer(new_buffer, op.span)
|
||||
return op
|
||||
|
||||
def visit_array_prim_expr_(self, op: list[Expr]):
|
||||
return [self.visit_expr(expr) for expr in op]
|
||||
|
||||
def visit_alloc_buffer_(self, op: AllocBuffer):
|
||||
op = super().visit_alloc_buffer_(op)
|
||||
if op.buffer in self.buffer_map:
|
||||
return AllocBuffer(self.buffer_map[op.buffer], op.annotations, op.span)
|
||||
return op
|
||||
|
||||
def visit_op_call_(self, op):
|
||||
op = super().visit_op_call_(op)
|
||||
new_workspace = {}
|
||||
for key, value in op.workspace.items():
|
||||
new_buffer = self.mutate_buffer(value)
|
||||
if new_buffer is not None:
|
||||
new_workspace[key] = new_buffer
|
||||
else:
|
||||
new_workspace[key] = value
|
||||
new_config = {}
|
||||
for key, value in op.config.items():
|
||||
if is_prim_expr(value):
|
||||
new_config[key] = self.visit_expr(value)
|
||||
else:
|
||||
new_config[key] = value
|
||||
args = list()
|
||||
for arg in op.args:
|
||||
args.append(arg)
|
||||
return TilePrimitiveCall(
|
||||
*args,
|
||||
op=op.op,
|
||||
workspace=new_workspace,
|
||||
config=new_config,
|
||||
dispatch=op.dispatch,
|
||||
scope=op.scope,
|
||||
)
|
||||
|
||||
|
||||
class KernelReplacePointSearcher(StmtMutator):
|
||||
def __init__(self, body: Stmt):
|
||||
super().__init__()
|
||||
self.body = body
|
||||
|
||||
def visit_evaluate_(self, op: Evaluate):
|
||||
value = op.value
|
||||
if isinstance(value, Call) and value.op.same_as(Op.get("tirx.tvm_kernel_replace_point")):
|
||||
return self.body
|
||||
return super().visit_evaluate_(op)
|
||||
|
||||
|
||||
def seek_kernel_replace_point(stmt: Stmt, body: Stmt) -> Stmt:
|
||||
"""replace kernel replace point in stmt with body"""
|
||||
return KernelReplacePointSearcher(body)(stmt)
|
||||
@@ -0,0 +1,163 @@
|
||||
# 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.
|
||||
"""TIR specific function pass support."""
|
||||
|
||||
import functools
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
from tvm.ir.transform import Pass, PassInfo
|
||||
|
||||
from . import _ffi_api
|
||||
|
||||
|
||||
@tvm_ffi.register_object("tirx.PrimFuncPass")
|
||||
class PrimFuncPass(Pass):
|
||||
"""A pass that works on each :py:func:`tvm.tirx.PrimFunc` in a module. A function
|
||||
pass class should be created through py:func:`tvm.tirx.transform.function_pass`.
|
||||
"""
|
||||
|
||||
|
||||
def _wrap_class_function_pass(pass_cls, pass_info):
|
||||
"""Wrap a python class as function pass"""
|
||||
|
||||
class PyFunctionPass(PrimFuncPass):
|
||||
"""Internal wrapper class to create a class instance."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
inst = pass_cls(*args, **kwargs)
|
||||
|
||||
# it is important not to capture self to
|
||||
# avoid a cyclic dependency
|
||||
def _pass_func(func, mod, ctx):
|
||||
return inst.transform_function(func, mod, ctx)
|
||||
|
||||
self.__init_handle_by_constructor__(
|
||||
_ffi_api.CreatePrimFuncPass,
|
||||
_pass_func,
|
||||
pass_info, # type: ignore
|
||||
)
|
||||
|
||||
self._inst = inst
|
||||
|
||||
def __getattr__(self, name):
|
||||
# fall back to instance attribute if there is not any
|
||||
return self._inst.__getattribute__(name)
|
||||
|
||||
functools.update_wrapper(PyFunctionPass.__init__, pass_cls.__init__)
|
||||
PyFunctionPass.__name__ = pass_cls.__name__
|
||||
PyFunctionPass.__doc__ = pass_cls.__doc__
|
||||
PyFunctionPass.__module__ = pass_cls.__module__
|
||||
return PyFunctionPass
|
||||
|
||||
|
||||
def prim_func_pass(
|
||||
pass_func=None,
|
||||
opt_level: int | None = None,
|
||||
name: str | None = None,
|
||||
required: list[str] | None = None,
|
||||
traceable=False,
|
||||
) -> Callable | PrimFuncPass:
|
||||
"""Decorate a function pass.
|
||||
|
||||
This function returns a callback when pass_func
|
||||
is provided. Otherwise, it returns the created function pass using the
|
||||
given optimization function.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
pass_func : Optional[Callable[(tvm.tirx.PrimFunc, IRModule, PassContext) -> tvm.tirx.PrimFunc]]
|
||||
The transformation function or class.
|
||||
|
||||
opt_level : int
|
||||
The optimization level of this module pass.
|
||||
|
||||
name : Optional[str]
|
||||
The name of the function pass. The name could be empty. In this case, the
|
||||
name of the optimization function will be used as the pass name.
|
||||
|
||||
required : Optional[List[str]]
|
||||
The list of passes that the function pass is dependent on.
|
||||
|
||||
Returns
|
||||
-------
|
||||
create_function_pass : Union[Callable, FunctionPass]
|
||||
|
||||
A decorator will be returned if pass_func is not provided,
|
||||
otherwise return the decorated result.
|
||||
The returned decorator has two behaviors depending on the input:
|
||||
A new FunctionPass will be returned when we decorate a pass function.
|
||||
A new FunctionPass class will be returned when we decorate a class type.
|
||||
|
||||
Examples
|
||||
--------
|
||||
The following code block decorates a function pass class.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@tvm.tirx.transform.prim_func_pass(opt_level=1)
|
||||
class TestReplaceFunc:
|
||||
def __init__(self, new_func):
|
||||
self.new_func = new_func
|
||||
|
||||
def transform_function(self, func, mod, ctx):
|
||||
# just for demo purposes
|
||||
# transform func to new_func
|
||||
return self.new_func
|
||||
|
||||
The following code creates a function pass by decorating
|
||||
a user defined transform function.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@tvm.tirx.transform.prim_func_pass(opt_level=2)
|
||||
def transform(func, mod, ctx):
|
||||
# my transformations here.
|
||||
return func
|
||||
|
||||
function_pass = transform
|
||||
assert isinstance(function_pass, transform.FunctionPass)
|
||||
assert function_pass.info.opt_level == 2
|
||||
|
||||
# Given a module m, the optimization could be invoked as the following:
|
||||
updated_mod = function_pass(m)
|
||||
# Now constant folding should have been applied to every function in
|
||||
# the provided module m. And the updated module will be returned.
|
||||
"""
|
||||
|
||||
if opt_level is None:
|
||||
raise ValueError("Please provide opt_level for the function pass.")
|
||||
|
||||
required = required if required else []
|
||||
if not isinstance(required, list | tuple):
|
||||
raise TypeError("Required is expected to be the type of " + "list/tuple.")
|
||||
|
||||
def create_function_pass(pass_arg):
|
||||
"""Internal function that creates a function pass"""
|
||||
fname = name if name else pass_arg.__name__
|
||||
info = PassInfo(opt_level, fname, required, traceable)
|
||||
if inspect.isclass(pass_arg):
|
||||
return _wrap_class_function_pass(pass_arg, info)
|
||||
if not callable(pass_arg):
|
||||
raise TypeError("pass_func must be a callable for Module pass")
|
||||
return _ffi_api.CreatePrimFuncPass(pass_arg, info) # type: ignore
|
||||
|
||||
if pass_func:
|
||||
return create_function_pass(pass_func)
|
||||
return create_function_pass
|
||||
@@ -0,0 +1,528 @@
|
||||
# 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.
|
||||
"""Wrapping existing transformations."""
|
||||
# pylint: disable=invalid-name, unsupported-binary-operation
|
||||
|
||||
import enum
|
||||
from collections.abc import Callable
|
||||
|
||||
import tvm_ffi as _ffi
|
||||
|
||||
from . import _ffi_api
|
||||
from . import function_pass as _fpass
|
||||
|
||||
|
||||
def Apply(ftransform):
|
||||
"""Apply ftransform to each function in the Module.
|
||||
|
||||
This function is a thin wrapper around tvm.tirx.transform.prim_func_pass
|
||||
|
||||
Parameters
|
||||
----------
|
||||
ftransform: tvm.tirx.PrimFunc -> tvm.tirx.PrimFunc
|
||||
The transformation pass.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
|
||||
# pylint: disable=unused-argument
|
||||
def _transform(func, mod, ctx):
|
||||
return ftransform(func)
|
||||
|
||||
return _fpass.prim_func_pass(_transform, opt_level=0, name="Apply") # type: ignore
|
||||
|
||||
|
||||
def VectorizeLoop(enable_vectorize: bool = True):
|
||||
"""Lower vectorization loops.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
enable_vectorize : bool
|
||||
Whether vectorization is enabled.
|
||||
Will lower to scalar loop when it is turned off.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.VectorizeLoop(enable_vectorize) # type: ignore
|
||||
|
||||
|
||||
def StorageRewrite():
|
||||
"""Rewrite storage allocation pattern.
|
||||
|
||||
Moves the allocation to outer most possible scope.
|
||||
Trying to share space between allocations to make
|
||||
a static allocation plan when possible.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.StorageRewrite() # type: ignore
|
||||
|
||||
|
||||
def InlinePrivateFunctions():
|
||||
"""Inline calls to private functions
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.InlinePrivateFunctions() # type: ignore
|
||||
|
||||
|
||||
def PointerValueTypeRewrite():
|
||||
"""
|
||||
Rewrite the pointer content type of arguments, as well as Alloc internal to the function to use
|
||||
the most frequently accessed type for load/store to avoid pointer casting in backend when
|
||||
possible.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.PointerValueTypeRewrite() # type: ignore
|
||||
|
||||
|
||||
@_ffi.register_object("tirx.transform.UnrollLoopConfig")
|
||||
class UnrollLoopConfig(_ffi.Object):
|
||||
"""Config for unroll loop pass"""
|
||||
|
||||
|
||||
def UnrollLoop():
|
||||
"""Unroll the constant loop marked by unroll.
|
||||
|
||||
This pass also automatically attach pragma unroll tag to loops which meets the standard.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.UnrollLoop() # type: ignore
|
||||
|
||||
|
||||
@_ffi.register_object("tirx.transform.RemoveNoOpConfig")
|
||||
class RemoveNoOpConfig(_ffi.Object):
|
||||
"""Config for remove no op pass"""
|
||||
|
||||
|
||||
def RemoveNoOp():
|
||||
"""Remove No Op from the Stmt.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.RemoveNoOp() # type: ignore
|
||||
|
||||
|
||||
def RemoveAssume():
|
||||
"""Remove all instances of builtin::assume
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.RemoveAssume() # type: ignore
|
||||
|
||||
|
||||
def BF16ComputeLegalize():
|
||||
"""Legalize bf16 compute Ops.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.BF16ComputeLegalize() # type: ignore
|
||||
|
||||
|
||||
def FP8ComputeLegalize(promote_dtype: str = "float32"):
|
||||
"""Legalize fp8 compute Ops.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
promote_dtype : str
|
||||
The data type we promote fp8 to, options: float16/float32.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.FP8ComputeLegalize(promote_dtype) # type: ignore
|
||||
|
||||
|
||||
def BF16StorageLegalize():
|
||||
"""Legalize bf16 storage types to u16.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.BF16StorageLegalize() # type: ignore
|
||||
|
||||
|
||||
def FP8StorageLegalize():
|
||||
"""Legalize fp8 storage types to u8.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.FP8StorageLegalize() # type: ignore
|
||||
|
||||
|
||||
def CommonSubexprElim():
|
||||
"""Replace redundant computations by new variables.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.CommonSubexprElim() # type: ignore
|
||||
|
||||
|
||||
@_ffi.register_object("tirx.transform.StmtSimplifyConfig")
|
||||
class StmtSimplifyConfig(_ffi.Object):
|
||||
"""Config for stmt simplify pass"""
|
||||
|
||||
|
||||
def StmtSimplify():
|
||||
"""Run statement-level arithmetic simplifications on the TIR PrimFunc.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.StmtSimplify() # type: ignore
|
||||
|
||||
|
||||
def ConvertSSA():
|
||||
"""Convert an IRModule to be SSA form.
|
||||
|
||||
This pass handles cases where the same `tirx.Var` appears in
|
||||
multiple functions within the same module. For example, after
|
||||
extracting a fragment from one function into another, where the
|
||||
same `tirx.Var` may be defined both as within the body of the
|
||||
original function, and as a parameter within the hoisted function.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
|
||||
"""
|
||||
return _ffi_api.ConvertSSA() # type: ignore
|
||||
|
||||
|
||||
def MakePackedAPI():
|
||||
"""Transform the PrimFuncs in the module to a packed func API.
|
||||
|
||||
Prior to this pass, the PrimFunc may have Buffer arguments defined
|
||||
in the `PrimFuncNode::buffer_map`. This pass consumes the
|
||||
`buffer_map`, using it to generate arguments that implement
|
||||
the packed based TVM FFI API.
|
||||
|
||||
For static shapes, the `BufferNode::shape`, `BufferNode::strides`,
|
||||
and `BufferNode::elem_offset` member variables are used to
|
||||
generate runtime checks on the corresponding member variables in
|
||||
the user-provided `DLTensor*` or `tvm.runtime.tensor` argument. (e.g. A
|
||||
PrimFunc that accepts a buffer of shape `[16,32]` validates that
|
||||
the `DLTensor::shape` array is `[16,32]`.)
|
||||
|
||||
For dynamic Buffers, in which one or more of these `BufferNode` member
|
||||
variables use `tirx.Var` that are not defined by other PrimFunc
|
||||
parameters, these are instead used to define the variables based on
|
||||
the corresponding `DLTensor` members. (e.g. A PrimFunc that accepts a
|
||||
buffer of shape `[tirx.Var("n"), tirx.Var("m")]`, when passed a
|
||||
`DLTensor` of shape `[16,32]`, will define `n = 16` and `n=32`, based
|
||||
on the argument's shape.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.MakePackedAPI() # type: ignore
|
||||
|
||||
|
||||
def SplitHostDevice():
|
||||
"""Annotate, split, and lower host/device functions.
|
||||
|
||||
This pass first annotates device regions within host functions,
|
||||
then splits them into host and device-side PrimFuncs, and finally
|
||||
lowers host-to-device calls into the device kernel launch ABI.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.SplitHostDevice() # type: ignore
|
||||
|
||||
|
||||
def SkipAssert():
|
||||
"""Skip assert stmt.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.SkipAssert() # type: ignore
|
||||
|
||||
|
||||
def LowerWarpMemory():
|
||||
"""Lower warp memory access to low-level device related function calls.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.LowerWarpMemory() # type: ignore
|
||||
|
||||
|
||||
def LowerTVMBuiltin():
|
||||
"""Lower tvm builtin intrinsics.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.LowerTVMBuiltin() # type: ignore
|
||||
|
||||
|
||||
def LowerIntrin():
|
||||
"""Lower target specific intrinsic calls.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.LowerIntrin() # type: ignore
|
||||
|
||||
|
||||
def NarrowDataType(target_bits: int):
|
||||
"""Narrow down Expr datatype in stmt to target_bits.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
target_bits : int
|
||||
The target bit configuration.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
|
||||
Note
|
||||
----
|
||||
Run this pass after FlattenBuffer.
|
||||
"""
|
||||
return _ffi_api.NarrowDataType(target_bits) # type: ignore
|
||||
|
||||
|
||||
def ForceNarrowIndexToInt32():
|
||||
"""Force narrow down indexing expressions and integer buffers to int32 dtype.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
|
||||
Note
|
||||
----
|
||||
This pass should not be used in default cases.
|
||||
"""
|
||||
return _ffi_api.ForceNarrowIndexToInt32() # type: ignore
|
||||
|
||||
|
||||
def VerifyMemory():
|
||||
"""Verify if func contains illegal host side direct memory access.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.VerifyMemory() # type: ignore
|
||||
|
||||
|
||||
@_ffi.register_object("s_tir.transform.HoistIfThenElseConfig")
|
||||
class HoistIfThenElseConfig(_ffi.Object):
|
||||
"""Config for hoist if then else pass"""
|
||||
|
||||
|
||||
class HoistedConditionals(enum.Flag):
|
||||
"""Flags for use in HoistExpressionConfig.conditional_types
|
||||
|
||||
Each bitflag represents a type of expression that should be
|
||||
hoisted to the outermost loop possible.
|
||||
"""
|
||||
|
||||
Never = 0
|
||||
""" No hoisting of conditionals """
|
||||
|
||||
IfElseStmt = 1
|
||||
""" If set, look for hoist candidates in IfElseStmt """
|
||||
|
||||
IfElseExpr = 2
|
||||
""" If set, look for hoist candidates in tirx.if_then_else """
|
||||
|
||||
BooleanExpression = 4
|
||||
""" If set, look for hoist candidates in all boolean expressions """
|
||||
|
||||
UsingBlockVar = 8
|
||||
""" If set, allow hoisting of conditionals that use a block variable (e.g. threadIdx.x) """
|
||||
|
||||
All = IfElseStmt | IfElseExpr | BooleanExpression | UsingBlockVar
|
||||
""" Enable all hoisting of conditionals"""
|
||||
|
||||
|
||||
class HoistedLetBindings(enum.Flag):
|
||||
"""Flags for use in HoistExpressionConfig.let_binding_types
|
||||
|
||||
Each bitflag represents a type of let binding expression that should be
|
||||
hoisted to the outermost loop possible.
|
||||
"""
|
||||
|
||||
Never = 0
|
||||
""" No hoisting of let bindings """
|
||||
|
||||
RequiredByConditional = 1
|
||||
""" Bindings that are used by a hoisted conditional """
|
||||
|
||||
Bind = 2
|
||||
""" Bindings occurring in Bind nodes """
|
||||
|
||||
LetExpr = 4
|
||||
""" Bindings occurring in Let expressions """
|
||||
|
||||
All = RequiredByConditional | Bind | LetExpr
|
||||
""" Enable all hoisting of let bindings """
|
||||
|
||||
|
||||
@_ffi.register_object("s_tir.transform.HoistExpressionConfig")
|
||||
class HoistExpressionConfig(_ffi.Object):
|
||||
"""Config for hoist expression pass"""
|
||||
|
||||
|
||||
def FlattenBuffer():
|
||||
"""Flatten the multi-dimensional BufferLoad and BufferStore to single dimensional
|
||||
BufferLoad/BufferStore for the TIR not contains opaque block.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.FlattenBuffer() # type: ignore
|
||||
|
||||
|
||||
def BindTarget(target):
|
||||
"""Annotate a PrimFunc with a given target.
|
||||
Parameters
|
||||
-------
|
||||
target : tvm.target.Target
|
||||
target
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.BindTarget(target) # type: ignore
|
||||
|
||||
|
||||
def AnnotateEntryFunc():
|
||||
"""Set a PrimFunc as the entry point if it is only function in IRModule.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.AnnotateEntryFunc() # type: ignore
|
||||
|
||||
|
||||
def Filter(fcond: Callable):
|
||||
"""Filter out PrimFuncs that does not satisfy the given condition.
|
||||
`fcond` should be a function that takes a primfunc and returns boolean.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.Filter(fcond) # type: ignore
|
||||
|
||||
|
||||
def TilePrimitiveDispatch():
|
||||
"""Lower TIRx tile primitive calls through the active backend dispatch table.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.TilePrimitiveDispatch() # type: ignore
|
||||
|
||||
|
||||
def LowerTIRx():
|
||||
"""Lower TIR to a lower-level IR.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.LowerTIRx() # type: ignore
|
||||
|
||||
|
||||
def LowerTIRxOpaque():
|
||||
"""Lower opaque constructs in TIRX programs.
|
||||
|
||||
Handles AllocBuffer lowering, For(thread_binding) to AttrStmt(thread_extent)
|
||||
conversion, unit loop elimination, and pragma annotation handling.
|
||||
This is the tirx-specific counterpart of s_tir.LowerOpaqueBlock,
|
||||
without any SBlock/SBlockRealize handling.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.LowerTIRxOpaque() # type: ignore
|
||||
Reference in New Issue
Block a user