wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions 6 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
apache--tvm/python/tvm/s_tir/schedule/schedule.py
T
wehub-resource-sync 26446540fa
Lint / lint (push) Waiting to run
Details
CI / MacOS (push) Waiting to run
Details
CI / Windows (push) Waiting to run
Details
chore: import upstream snapshot with attribution
2026-07-13 13:36:25 +08:00

4109 lines
152 KiB
Python
Raw Permalink Blame History

# 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 TensorIR schedule class"""
import inspect
from collections.abc import Callable
from typing import Literal
from tvm_ffi import register_object as _register_object
from tvm.error import register_error
from tvm.ir import Expr, GlobalVar, IRModule, is_prim_expr
from tvm.runtime import DataTypeCode, Object
from tvm.tirx import Buffer, FloatImm, For, IntImm, PrimFunc, SBlock
from tvm.tirx.function import IndexMap
from . import _ffi_api
from ._type_checker import type_checked
from .state import ScheduleState, StmtSRef, _parse_debug_mask, _parse_mod
from .trace import Trace
@register_error
class ScheduleError(RuntimeError):
"""Error that happens during TensorIR scheduling."""
@_register_object("s_tir.LoopRV")
class LoopRV(Object):
"""A random variable that refers to a loop"""
def __init__(self) -> None:
"""Construct a new LoopRV."""
self.__init_handle_by_constructor__(
_ffi_api.LoopRV # type: ignore # pylint: disable=no-member
)
@_register_object("s_tir.SBlockRV")
class SBlockRV(Object):
"""A random variable that refers to a block"""
def __init__(self) -> None:
"""Construct a new SBlockRV."""
self.__init_handle_by_constructor__(
_ffi_api.SBlockRV # type: ignore # pylint: disable=no-member
)
# It is a workaround for mypy: https://github.com/python/mypy/issues/7866#issuecomment-549454370
# This feature is not supported until python 3.10:
# https://docs.python.org/3.10/whatsnew/3.10.html#pep-613-typealias
# A random variable that evaluates to an integer
ExprRV = Expr # pylint: disable=invalid-name
RAND_VAR_TYPE = ExprRV | SBlockRV | LoopRV # pylint: disable=invalid-name
_ERROR_RENDER_LEVEL: dict[Literal["detail", "fast", "none"], int] = {
"detail": 0,
"fast": 1,
"none": 2,
}
def _parse_error_render_level(error_render_level: str) -> int:
if error_render_level not in _ERROR_RENDER_LEVEL:
raise ValueError(
'error_render_level can be "detail", "fast", or "none", but got: '
+ f"{error_render_level}"
)
return _ERROR_RENDER_LEVEL.get(error_render_level)
def _parse_enable_checks(enable_checks: bool) -> bool:
if not isinstance(enable_checks, bool):
raise TypeError(f"enable_checks only accepts bool value, got {type(enable_checks)} instead")
return enable_checks
def _parse_seed(seed: int | None) -> int:
if seed is None:
return -1
if not isinstance(seed, int):
raise TypeError(f"Expected `seed` to be int or None, but gets: {seed}")
if seed < 1 or seed > 2147483647:
raise ValueError(f"seed must be in the range [1, 2147483647], but gets: {seed}")
return seed
def _get_sblock_default_dtype(block: SBlock) -> str:
for i in block.iter_vars:
return str(i.var.ty)
for buffer_region in list(block.reads) + list(block.writes):
for dom in buffer_region.region:
return str(dom.min.ty)
return "int64"
@_register_object("s_tir.Schedule")
class Schedule(Object):
"""The user-facing schedule class
A schedule is a set of transformations that change the order of computation but
preserve the semantics of computation. Some example of schedules:
1) Split a loop into two;
2) Reorder two loops;
3) Inline the computation of a specific buffer into its consumer
The schedule class stores auxiliary information to schedule correctly and efficiently.
Link to tutorial: https://tvm.apache.org/docs/tutorials/language/schedule_primitives.html
"""
@type_checked
def __init__(
self,
mod: PrimFunc | IRModule,
*,
seed: int | None = None,
debug_mask: str | int = "none",
error_render_level: str = "detail",
enable_check: bool = True,
) -> None:
"""Construct a TensorIR schedule class from an IRModule
Parameters
----------
mod : PrimFunc | IRModule
The IRModule or PrimFunc to be scheduled
seed: Optional[int]
The seed value for schedule's random state
Note that None and -1 means use device random, otherwise only integer between 1 and
2147483647 is allowed.
debug_mask : str | int
Do extra correctness checking after the class creation and each time
after calling the Replace method.
Possible choices of `debug_mask`:
1) "all" - Turn on all the checks
2) "none" - Turn off all the checks
3) An integer - Turn on checks according to the bitmasks provided in ScheduleDebugMask
error_render_level : str = "detail"
The level of error rendering. Choices: "detail", "fast", "none".
- "detail": Render a detailed error message, with the TIR and error locations printed
- "fast: Show a simple error message without rendering or string manipulation
- "none": Do not show any error message.
enable_check : bool = True
The default schedule checks are too strict and might prevent us performing some valid
schedules. `enable_check` is an argument to control whether we enable prerequisite
checks for some schedule primitives or not:
- true: perform prerequisite check before applying some schedules.
- false: do not perform some check before applying schedules, but still raise error
if schedule fails.
It's user duty to guarantee schedule correctness if `enable_check` is set to `False`.
Note
----
The checks performed includes:
1) VerifySRefTree
2) VerifyCachedFlags
"""
# call the constructor
self.__init_handle_by_constructor__(
_ffi_api.TracedSchedule, # type: ignore # pylint: disable=no-member
_parse_mod(mod),
_parse_seed(seed),
_parse_debug_mask(debug_mask),
_parse_error_render_level(error_render_level),
_parse_enable_checks(enable_check),
)
@staticmethod
def _create_non_traced(
mod: PrimFunc | IRModule,
*,
seed: int | None = None,
debug_mask: str | int = "none",
error_render_level: str = "detail",
enable_check: bool = True,
) -> "Schedule":
"""Construct a non-traced TensorIR schedule class from an IRModule."""
return _ffi_api.ConcreteSchedule( # type: ignore # pylint: disable=no-member
_parse_mod(mod),
_parse_seed(seed),
_parse_debug_mask(debug_mask),
_parse_error_render_level(error_render_level),
_parse_enable_checks(enable_check),
)
########## Utilities ##########
@property
def mod(self) -> IRModule:
"""Returns the AST of the module being scheduled"""
return _ffi_api.ScheduleGetMod(self) # type: ignore # pylint: disable=no-member
@property
def state(self) -> ScheduleState:
"""Returns the ScheduleState in the current schedule class"""
return _ffi_api.ScheduleGetState(self) # type: ignore # pylint: disable=no-member
@property
def trace(self) -> Trace | None:
"""Returns the internally maintained trace of scheduling program execution"""
return _ffi_api.ScheduleGetTrace(self) # type: ignore # pylint: disable=no-member
@property
def func_working_on(self) -> GlobalVar | None:
"""Returns the GlobalVar of the func that the schedule is currently working on"""
return _ffi_api.ScheduleGetFuncWorkingOn(self) # type: ignore # pylint: disable=no-member
def work_on(self, func_name: str) -> None:
"""Instruct the schedule to work on a function in the IRModule.
By default, the schedule works on the function with the name "main", or the only function in
the IRModule if there is only one. If there is multiple functions in the IRModule, and none
of their names are "main", users will have to call this method to explicitly specify which
function to work on.
This sugar function will guide the `GetSBlock` method if its `func_name` is not specified.
Parameters
----------
func_name : str
The name of the function to work on.
"""
_ffi_api.ScheduleWorkOn(self, func_name) # type: ignore # pylint: disable=no-member
def copy(self) -> "Schedule":
"""Returns a copy of the schedule, including both the state and the symbol table,
* guaranteeing that
* 1) SRef tree is completely reconstructed;
* 2) The IRModule being scheduled is untouched;
* 3) All the random variables are valid in the copy, pointing to the corresponding sref
* reconstructed
Returns
-------
copy : Schedule
A new copy of the schedule
"""
return _ffi_api.ScheduleCopy(self) # type: ignore # pylint: disable=no-member
@type_checked
def seed(self, seed: int) -> None:
"""Seed the randomness
Parameters
----------
seed : int
The new random seed, -1 if use device random, otherwise non-negative
"""
return _ffi_api.ScheduleSeed(self, seed) # type: ignore # pylint: disable=no-member
def fork_seed(self) -> int:
"""Returns a forked random state as seed for new schedules
Returns
-------
seed : int
The forked random state, not the same as the current random state
"""
return _ffi_api.ScheduleForkSeed(self) # type: ignore # pylint: disable=no-member
def show(self, *args, **kwargs) -> None:
"""A sugar for print highlighted TVM script.
All parameters are forwarded to the underlying `Module.show`
and `Trace.show` methods.
"""
mod = self.mod
if mod is not None:
mod.show(*args, **kwargs)
trace = self.trace
if trace is not None:
# Trace.show only supports the style and black_format arguments
param_binding = inspect.signature(mod.show).bind(*args, **kwargs)
param_binding.apply_defaults()
bound_args = param_binding.arguments
trace.show(style=bound_args["style"], black_format=bound_args["black_format"])
########## Lookup ##########
@type_checked
def get(self, rand_var_or_sref: RAND_VAR_TYPE | StmtSRef) -> int | SBlock | For | None:
"""Returns:
- the corresponding SBlock that a SBlockRV evaluates to;
- the corresponding For that a LoopRV evaluates to;
- the corresponding integer that a ExprRV evaluates to;
- the corresponding SBlock that a SBlock sref points to;
- the corresponding For that a loop sref points to;
Parameters
----------
rand_var_or_sref : ExprRV | SBlockRV | LoopRV | StmtSRef
The random variable / sref to be evaluated
Returns
-------
result : Optional[int | SBlock | For]
The corresponding result
"""
if isinstance(rand_var_or_sref, StmtSRef):
return rand_var_or_sref.stmt
# pylint: disable-next=no-member
result = _ffi_api.ScheduleGet(self, rand_var_or_sref) # type: ignore
if isinstance(result, IntImm):
result = result.value
return result
@type_checked
def get_sref(self, rand_var_or_stmt: SBlockRV | LoopRV | SBlock | For) -> StmtSRef | None:
"""Returns the corresponding sref to the given
1) LoopRV
2) SBlockRV
3) Block
4) For
Parameters
----------
rand_var_or_stmt : SBlockRV | LoopRV | SBlock | For
The random variable / sref to be evaluated
Returns
-------
result : Optional[StmtSRef]
The corresponding result
"""
return _ffi_api.ScheduleGetSRef( # type: ignore # pylint: disable=no-member
self, rand_var_or_stmt
)
@type_checked
def remove_rv(self, rand_var: RAND_VAR_TYPE) -> None:
"""Remove a random variable from the symbol table
Parameters
----------
rand_var : SBlockRV | LoopRV | ExprRV
The random variable to be removed
"""
return _ffi_api.ScheduleRemoveRV(self, rand_var) # type: ignore # pylint: disable=no-member
########## Schedule: Sampling ##########
@type_checked
def sample_categorical(
self, candidates: list[int], probs: list[float], decision: int | None = None
) -> ExprRV:
"""Sample an integer given the probability distribution
Parameters
----------
candidates : List[int]
The candidates to be sampled from
probs : List[float]
The probability of each candidate
decision : Optional[int]
The sampling decision, if any
Returns
-------
result : ExprRV
The random variable sampled from candidates
"""
return _ffi_api.ScheduleSampleCategorical( # type: ignore # pylint: disable=no-member
self, candidates, probs, decision
)
@type_checked
def sample_perfect_tile(
self,
loop: LoopRV,
n: int,
max_innermost_factor: int = 16,
decision: list[int] | None = None,
) -> list[ExprRV]:
"""Sample the factors to perfect tile a specific loop
Parameters
----------
loop : LoopRV
The loop to be tiled
n : int
The number of tiles to be sampled
max_innermost_factor : int
The maximum tile size allowed to be sampled in the innermost loop
decision: Optional[List[int]]
The sampling decision, if any
Returns
-------
result : List[ExprRV]
A list of length `n`, the random perfect tile sizes sampled
"""
return list(
_ffi_api.ScheduleSamplePerfectTile( # type: ignore # pylint: disable=no-member
self, loop, n, max_innermost_factor, decision
)
)
@type_checked
def sample_partitioned_tile(
self,
loop: LoopRV,
n: int,
partition_pos: int = 0,
innerpart_factor: int = 1,
decision: list[int] | None = None,
) -> list[ExprRV]:
"""Sample the factors to a partitioned tile for a specific loop
Parameters
----------
loop : LoopRV
The loop to be tiled
n : int
The number of tiles to be sampled
partition_pos : int
The position to partition tiles to two parts
innerpart_factor : int
The factor of the second part
decision: Optional[List[int]]
The sampling decision, if any
Returns
-------
result : List[ExprRV]
A list of length `n`, the random partitioned tile sizes sampled
"""
return list(
_ffi_api.ScheduleSamplePartitionedTile( # type: ignore # pylint: disable=no-member
self,
loop,
n,
partition_pos,
innerpart_factor,
decision,
)
)
@type_checked
def sample_compute_location(self, block: SBlockRV | str, decision: int | None = None) -> LoopRV:
"""Sample a compute-at location of the given block
Parameters
----------
block : SBlockRV | str
The block whose compute-at location is to be sampled
decision : Optional[int]
The sampling decision
Returns
-------
result : LoopRV
The sampled loop where the input block is to be computed at
"""
block = self._normalize_block_arg(block)
return _ffi_api.ScheduleSampleComputeLocation( # type: ignore # pylint: disable=no-member
self, block, decision
)
########## Schedule: Get blocks & loops ##########
@type_checked
def get_sblock(self, name: str, func_name: str | None = None) -> SBlockRV:
"""Retrieve a block in a specific function with its name
By default, if `func_name` is not specified, the schedule will search for the block in the
function that is currently being "worked on". To switch the function to be worked on, use
`work_on` before calling this method.
Parameters
----------
name : str
The name of the block
func_name : Optional[str] = None
The name of the function
Returns
-------
block : SBlockRV
The block retrieved
IndexError is raised if 0 or multiple blocks exist with the specific name.
"""
return _ffi_api.ScheduleGetSBlock( # type: ignore # pylint: disable=no-member
self, name, func_name
)
@type_checked
def get_loops(self, block: SBlockRV | str) -> list[LoopRV]:
"""Get the parent loops of the block in its scope, from outer to inner
Parameters
----------
block : SBlockRV | str
The query block
Returns
-------
loops : List[LoopRV]
A list of loops above the given block in its scope, from outer to inner
"""
block = self._normalize_block_arg(block)
# pylint: disable-next=no-member
return list(_ffi_api.ScheduleGetLoops(self, block)) # type: ignore
@type_checked
def get_child_blocks(self, block_or_loop: SBlockRV | LoopRV) -> list[SBlockRV]:
"""Get the leaf blocks of a specific block/loop
Parameters
----------
block_or_loop : SBlockRV | LoopRV
The query block/loop
Returns
-------
blocks : List[LoopRV]
A list of leaf blocks inside a specific block/loop
"""
# pylint: disable-next=no-member
return list(_ffi_api.ScheduleGetChildBlocks(self, block_or_loop)) # type: ignore
@type_checked
def get_producers(self, block: SBlockRV | str) -> list[SBlockRV]:
"""Get the producers of a specific block
Parameters
----------
block : SBlockRV | str
The block in the query
Returns
-------
producers : List[SBlockRV]
A list of producers of the given block
"""
block = self._normalize_block_arg(block)
# pylint: disable-next=no-member
return list(_ffi_api.ScheduleGetProducers(self, block)) # type: ignore
@type_checked
def get_consumers(self, block: SBlockRV | str) -> list[SBlockRV]:
"""Get the consumers of a specific block
Parameters
----------
block : SBlockRV | str
The block in the query
Returns
-------
consumers : List[SBlockRV]
A list of consumers of the given block
"""
block = self._normalize_block_arg(block)
# pylint: disable-next=no-member
return list(_ffi_api.ScheduleGetConsumers(self, block)) # type: ignore
@type_checked
def get_output_blocks(self, scope_block: SBlockRV | str) -> list[SBlockRV]:
"""Get the list of output blocks within the given scope
An output block is a block which has atleast one buffer being written
to, but is not allocated within the PrimFunc
Parameters
----------
scope_block : SBlockRV | str,
The scope block from which output blocks are collected
Returns
-------
output_blocks : List[SBlockRV]
A list of all blocks that write to some output buffer
"""
scope_block = self._normalize_block_arg(scope_block)
# pylint: disable-next=no-member
return list(_ffi_api.ScheduleGetOutputBlocks(self, scope_block)) # type: ignore
########## Schedule: Transform loops ##########
@type_checked
def merge(self, *loops: list[LoopRV]) -> LoopRV:
"""Merge a list of loops into one. The loops under their LCA requires:
1) Under the same scope.
2) Can't have annotations or thread bindings.
3) Start with 0 and have same extent and same nesting depth.
4) From target loop to their LCA, The inner loop must be the only child of the outer loop.
Parameters
----------
*loops : List[LoopRV]
The loops to be merged
Returns
-------
fused_loop : LoopRV
The new loop after merge
Examples
--------
Before applying merge, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_merge(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
C = T.match_buffer(c, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = A[vi, vj] * 2.0
Create the schedule and do fuse:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_fuse)
i1, _ = sch.get_loops(sch.get_sblock("B"))
i2, _ = sch.get_loops(sch.get_sblock("C"))
sch.merge(i1, i2)
print(sch.mod["main"].script())
After applying fuse, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_fuse(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
C = T.match_buffer(c, (128, 128))
# the 2 loops are merged into 1
for i_m in range(128):
for j in range(128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i_m, j])
T.reads(A[vi, vj])
T.writes(B[vi, vj])
B[vi, vj] = A[vi, vj] * T.float32(2)
for j in range(128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i_m, j])
T.reads(A[vi, vj])
T.writes(C[vi, vj])
C[vi, vj] = A[vi, vj] * T.float32(2)
"""
return _ffi_api.ScheduleMerge(self, loops) # type: ignore # pylint: disable=no-member
@type_checked
def fuse(self, *loops: list[LoopRV], preserve_unit_iters: bool = True) -> LoopRV:
"""Fuse a list of consecutive loops into one. It requires:
1) The loops can't have annotations or thread bindings.
2) The (i+1)-th loop must be the only child of the i-th loop.
3) All loops must start with 0.
4) The domain of a loop to be fused cannot depend on another loop to be fused.
Parameters
----------
*loops : List[LoopRV]
The loops to be fused
Returns
-------
fused_loop : LoopRV
The new loop after fusion
Examples
--------
Before applying fuse, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_fuse(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
Create the schedule and do fuse:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_fuse)
i, j = sch.get_loops(sch.get_sblock("B"))
sch.fuse(i, j)
print(sch.mod["main"].script())
After applying fuse, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_fuse(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
# the 2 loops are fused into 1
for i_j_fused in T.serial(0, 16384):
with T.sblock("B"):
vi = T.axis.S(128, T.floordiv(i_j_fused, 128))
vj = T.axis.S(128, T.floormod(i_j_fused, 128))
B[vi, vj] = A[vi, vj] * 2.0
"""
# pylint: disable-next=no-member
return _ffi_api.ScheduleFuse(self, loops, preserve_unit_iters) # type: ignore
@type_checked
def split(
self,
loop: LoopRV,
factors: list[int | ExprRV | None],
preserve_unit_iters: bool = True,
disable_predication: bool = False,
) -> list[LoopRV]:
"""Split a loop into a list of consecutive loops. It requires:
- The loop can't have annotation or thread binding.
- The loop must start with 0.
Predicates may be added to ensure the total loop numbers keeps unchanged.
In `factors`, at most one of the factors can be None,
which will be automatically inferred.
Parameters
----------
loop : LoopRV
The loop to be split
factors: List[int | ExprRV | None]
The splitting factors
Potential inputs are:
- None
- ExprRV
- Positive constant integers
preserve_unit_iters : bool
Whether or not to preserve unit iterators in block bindings
disable_predication : bool
If enabled, don't create a predicate for guarding the loop. This can
be useful when splitting with scalable factors that the schedule writer
knows are divisible by the loop bound.
Warning: enabling this feature may result in incorrect code generation
if not used carefully.
Returns
-------
split_loops : List[LoopRV]
The new loops after split
Examples
--------
Before split, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_split(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
Create the schedule and do split:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_split)
i, j = sch.get_loops(sch.get_sblock("B"))
sch.split(i, factors=[2, 64])
print(sch.mod["main"].script())
After applying split, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_split(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
# the original loop is split into 2 loops
for i0, i1, j in T.grid(2, 64, 128):
with T.sblock("B"):
vi = T.axis.S(128, i0 * 64 + i1)
vj = T.axis.S(128, j)
B[vi, vj] = A[vi, vj] * 2.0
"""
# it will be checked later in C++ implementation
# that there is at most one None in `factors`
return list(
_ffi_api.ScheduleSplit( # type: ignore # pylint: disable=no-member
self,
loop,
factors,
preserve_unit_iters,
disable_predication,
)
)
@type_checked
def loop_partition(
self,
loop: LoopRV,
factors: list[int | ExprRV | None],
preserve_unit_iters: bool = True,
) -> list[LoopRV]:
"""Partition a loop into a list of consecutive loops. It requires:
1) The loop can't have annotation or thread binding.
Predicates may be added to ensure the total loop numbers keeps unchanged.
In `factors`, at most one of the factors can be None,
which will be automatically inferred.
Parameters
----------
loop : LoopRV
The loop to be partition
factors: List[int | ExprRV | None]
The partitioning factors
Potential inputs are:
- None
- ExprRV
- Positive constant integers
preserve_unit_iters : bool
Whether or not to preserve unit iterators in block bindings
Returns
-------
partition_loops : List[LoopRV]
The new loops after partition
Examples
--------
Before partition, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_partition(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
Create the schedule and do partition:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_partition)
i, j = sch.get_loops(sch.get_sblock("B"))
sch.partition(i, factors=[2, 64])
print(sch.mod["main"].script())
After applying partition, the IR becomes:
.. code-block:: python
def after_partition(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
# the original loop is partition into 3 loops
with T.sblock("root"):
T.reads()
T.writes()
with T.sblock("B_i_common"):
T.reads()
T.writes()
with T.sblock("B_i0_partition"):
T.reads()
T.writes()
for i0, j in T.grid(2, 128):
with T.sblock("B_i0"):
vi, vj = T.axis.remap("SS", [i0, j])
T.reads(A[0:2, 0:128])
T.writes(B[0:2, 0:128])
B[vi, vj] = A[vi, vj] * T.float32(2)
with T.sblock("B_i1_partition"):
T.reads()
T.writes()
for i1 in range(2, 66):
for j in range(128):
with T.sblock("B_i1"):
vi, vj = T.axis.remap("SS", [i1, j])
T.reads(A[2:66, 0:128])
T.writes(B[2:66, 0:128])
B[vi, vj] = A[vi, vj] * T.float32(2)
with T.sblock("B_partition_2"):
T.reads()
T.writes()
for i2 in range(66, 128):
for j in range(128):
with T.sblock("B_i2"):
vi, vj = T.axis.remap("SS", [i2, j])
T.reads(A[66:128, 0:128])
T.writes(B[66:128, 0:128])
B[vi, vj] = A[vi, vj] * T.float32(2)
"""
return list(
_ffi_api.ScheduleLoopPartition( # type: ignore # pylint: disable=no-member
self, loop, factors, preserve_unit_iters
)
)
@type_checked
def reorder(self, *ordered_loops: list[LoopRV]) -> None:
"""
Reorder a list of loops. It doesn't require the loops to be consecutive.
It requires:
1) The loops are in the same chain. That means: the loops can be ordered to [l_1, l_2, ... ,
l_n] where l_i is an ancestor of l_{i+1} and there are only single-branch loops between
l_1 and l_n (which also indicates they are under the same scope).
2) After reordering, the domain of an outer loop cannot depend on any of the inner loops.
3) For every block under the loop nests, its block binding must be affine, and the block
variables must be either data parallel or reduction.
4) No duplicated loops are allowed in the arguments.
Parameters
----------
*ordered_loops : List[LoopRV]
The loops in the new order
Examples
--------
Before reorder, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_reorder(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
Create the schedule and do reorder:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_reorder)
i, j = sch.get_loops(sch.get_sblock("B"))
sch.reorder(j, i)
print(sch.mod["main"].script())
After applying reorder, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_reorder(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
# Here j and i are reordered
for j, i in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
"""
_ffi_api.ScheduleReorder(self, ordered_loops) # type: ignore # pylint: disable=no-member
@type_checked
def reorder_block_iter_var(self, block: SBlockRV, new_order: list[int]) -> None:
"""Reorder the itervars inside a given block.
Parameters
----------
block : SBlockRV
The block to be transformed.
new_order : List[int]
The new block itervar order.
Examples
--------
Before reorder_block_iter_var, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def matmul(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
C: T.Buffer((128, 128), "float32"),
) -> None:
for i, j, k in T.grid(128, 128, 128):
with T.sblock("C"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
with T.init():
C[vi, vj] = 0.0
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
Create the schedule and do reorder_block_iter_var:
.. code-block:: python
sch = tvm.s_tir.Schedule(matmul)
C = sch.get_sblock("C")
sch.reorder_block_iter_var(C, [2, 1, 0])
After applying reorder_block_iter_var, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def matmul_after_reorder_block_iter_var(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
C: T.Buffer((128, 128), "float32"),
):
for i, j, k in T.grid(128, 128, 128):
with T.sblock("C"):
vk, vj, vi = T.axis.remap("RSS", [k, j, i])
T.reads(A[vi, vk], B[vj, vk])
T.writes(C[vi, vj])
with T.init():
C[vi, vj] = T.float32(0)
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
See Also
--------
reorder
"""
# pylint: disable-next=no-member
_ffi_api.ScheduleReorderBlockIterVar(self, block, new_order) # type: ignore
@type_checked
def add_unit_loop(self, block_or_loop: LoopRV | SBlockRV) -> LoopRV:
"""Create a new unit loop on top of the specific block or loop.
Parameters
----------
block_or_loop : LoopRV | SBlockRV
The block above which the new loop is created
Returns
-------
new_loop : LoopRV
The new unit loop
Examples
--------
Before add_unit_loop, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_add_unit_loop(
A: T.Buffer((), "int32"),
B: T.Buffer((), "int32"),
C: T.Buffer((), "int32"),
) -> None:
with T.sblock("C"):
vi = T.axis.spatial(1, 0)
C[()] = A[()] + B[()]
Create the schedule and do add-unit-loop:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_add_unit_loop)
sch.add_unit_loop(sch.get_sblock("C"))
print(sch.mod["main"].script())
After applying add-unit-loop, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_add_unit_loop(
A: T.Buffer((), "int32"),
B: T.Buffer((), "int32"),
C: T.Buffer((), "int32"),
) -> None:
for u in T.serial(1):
with T.sblock("C"):
vi = T.axis.spatial(1, 0)
C[()] = A[()] + B[()]
"""
# pylint: disable-next=no-member
return _ffi_api.ScheduleAddUnitLoop(self, block_or_loop) # type: ignore
########## Schedule: Manipulate ForKind ##########
@type_checked
def parallel(self, loop: LoopRV) -> None:
"""Parallelize the input loop. It requires:
- The scope block that the loop is in should have stage-pipeline property.
- All the blocks under the loop are complete blocks or reduction blocks, and have affine
bindings.
- For each block under the loop, the loop can only be contained in data-parallel block
iters' bindings.
Parameters
----------
loop : LoopRV
The loop to be parallelized
Examples
--------
Before parallel, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_parallel(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
Create the schedule and do parallel:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_parallel)
i, j = sch.get_loops(sch.get_sblock("B"))
sch.parallel(i)
After applying parallel, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_parallel(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i in T.parallel(0, 128):
for j in T.serial(0, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
"""
_ffi_api.ScheduleParallel(self, loop) # type: ignore # pylint: disable=no-member
@type_checked
def vectorize(self, loop: LoopRV) -> None:
"""Vectorize the input loop. It requires:
- The scope block that the loop is in should have stage-pipeline property.
- All the blocks under the loop are complete blocks or reduction blocks, and have affine
bindings.
- For each block under the loop, the loop can only be contained in data-parallel block
iters' bindings.
Parameters
----------
loop : LoopRV
The loop to be vectorized
Examples
--------
Before vectorize, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_vectorize(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
Create the schedule and do vectorize:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_vectorize)
i, j = sch.get_loops(sch.get_sblock("B"))
sch.vectorize(j)
After applying vectorize, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_vectorize(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i in T.serial(0, 128):
for j in T.vectorized(0, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
"""
_ffi_api.ScheduleVectorize(self, loop) # type: ignore # pylint: disable=no-member
@type_checked
def bind(self, loop: LoopRV, thread_axis: str) -> None:
"""Bind the input loop to the given thread axis. It requires:
- The scope block that the loop is in should have stage-pipeline property.
- All the blocks under the loop are complete blocks or reduction blocks, and have affine
bindings.
- For each block under the loop, if the thread axis starts with ``threadIdx``, the loop can
only be contained in data-parallel block iter and reduction block iters' bindings.
Otherwise the loop can only be contained in data-parallel block iters' bindings.
Parameters
----------
loop : LoopRV
The loop to be bound to the thread axis
thread_axis : str
The thread axis to be bound to the loop. Possible candidates are ``blockIdx.x/y/z``,
``threadIdx.x/y/z``, ``vthread.x/y/z``, and ``vthread``. The ``vthread`` value is a
legacy behavior that will be deprecated. Please use ``vthread.x/y/z`` instead.
Examples
--------
Before bind, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_bind(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
Create the schedule and do bind:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_bind)
i, j = sch.get_loops(sch.get_sblock("B"))
sch.bind(i, "blockIdx.x")
sch.bind(j, "threadIdx.x")
After applying bind, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_bind(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i in T.thread_binding(0, 128, thread = "blockIdx.x"):
for j in T.thread_binding(0, 128, thread = "threadIdx.x"):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
"""
_ffi_api.ScheduleBind(self, loop, thread_axis) # type: ignore # pylint: disable=no-member
@type_checked
def unroll(self, loop: LoopRV) -> None:
"""Unroll the input loop. It requires nothing
Parameters
----------
loop : LoopRV
The loop to be unrolled
Examples
--------
Before unroll, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_unroll(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
Create the schedule and do unroll:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_unroll)
i, j = sch.get_loops(sch.get_sblock("B"))
sch.unroll(i)
After applying unroll, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_unroll(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i in T.unroll(0, 128):
for j in T.serial(0, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
"""
_ffi_api.ScheduleUnroll(self, loop) # type: ignore # pylint: disable=no-member
########## Schedule: Insert cache stages ##########
@type_checked
def cache_read(
self,
block: SBlockRV | str,
read_buffer_index: int | str | Buffer,
storage_scope: str,
consumer_blocks: list[SBlockRV | str] | None = None,
) -> SBlockRV:
"""Create a block that reads a buffer region into a read cache. It requires:
1) There is at most one block who write the buffer in the scope.
2) The scope block have stage-pipeline property.
Parameters
----------
block : SBlockRV | str
The consumer block of the target buffer.
buffer: int | str | Buffer
The index of the buffer in block's read region, the unique
name of a read buffer in the block, or a Buffer object
that is within the blocks read region.
storage_scope: str
The target storage scope.
consumer_blocks: Optional[List[SBlockRV | str]]
An optional list of consumers that should read from the cache. If not specified,
all consumers will use the cache.
Returns
-------
cached_block : SBlockRV
The block of the cache stage
Examples
--------
Before cache_read, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_cache_read(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
Create the schedule and cache_read:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_cache_read)
block_b = sch.get_sblock("B")
sch.cache_read(block_b, 0, "local")
print(sch.mod["main"].script())
After applying cache_read, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_cache_read(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
A_local = T.sblock_alloc_buffer((128, 128), scope="local")
for i, j in T.grid(128, 128):
with T.sblock("A_local"):
vi, vj = T.axis.remap("SS", [i, j])
A_local[vi, vj] = A[vi, vj]
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A_local[vi, vj] * 2.0
"""
if consumer_blocks is None:
consumer_blocks = []
# Convert any string SBlock names into SBlock RVs.
consumer_blocks = [self._normalize_block_arg(b) for b in consumer_blocks]
block = self._normalize_block_arg(block)
if not isinstance(read_buffer_index, int):
_, read_buffer_index, _ = self._normalize_buffer_arg(
block, read_buffer_index, required_buffer_type="read"
)
return _ffi_api.ScheduleCacheRead( # type: ignore # pylint: disable=no-member
self, block, read_buffer_index, storage_scope, consumer_blocks
)
@type_checked
def cache_write(
self,
block: SBlockRV | str,
write_buffer_index: int | str | Buffer,
storage_scope: str,
consumer_blocks: list[SBlockRV | str] | None = None,
) -> SBlockRV:
"""Create a block that reads a buffer region into a write cache. It requires:
1) There is only one block who write the buffer in the scope.
2) The scope block have stage-pipeline property.
Parameters
----------
block : SBlockRV | str
The producer block of the target buffer.
write_buffer_index: int
The index of the buffer in block's write region, the unique
name of a write buffer in the block, or a Buffer object
that is within the blocks write region.
storage_scope: str
The target storage scope.
consumer_blocks: Optional[List[SBlockRV | str]]
An optional list of consumers that should read directly from the cache.
If not specified, all consumers will read from the original buffer.
Returns
-------
cached_block : SBlockRV
The block of the cache stage
Examples
--------
Before cache_write, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_cache_write(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
Create the schedule and cache_write:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_cache_write)
block_b = sch.get_sblock("B")
sch.cache_write(block_b, 0, "local")
print(sch.mod["main"].script())
After applying cache_write, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_cache_write(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
B_local = T.sblock_alloc_buffer((128, 128), scope="local")
for i, j in T.grid(128, 128):
with T.sblock("A_local"):
vi, vj = T.axis.remap("SS", [i, j])
B_local[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = B_local[vi, vj]
"""
if consumer_blocks is None:
consumer_blocks = []
# Convert any string SBlock names into SBlock RVs.
consumer_blocks = [self._normalize_block_arg(b) for b in consumer_blocks]
block = self._normalize_block_arg(block)
if not isinstance(write_buffer_index, int):
_, write_buffer_index, _ = self._normalize_buffer_arg(
block, write_buffer_index, required_buffer_type="write"
)
return _ffi_api.ScheduleCacheWrite( # type: ignore # pylint: disable=no-member
self, block, write_buffer_index, storage_scope, consumer_blocks
)
@type_checked
def reindex_cache_read(
self,
block: SBlockRV | str,
read_buffer_index: int,
storage_scope: str,
index_map: IndexMap | Callable,
) -> SBlockRV:
"""Create a block that reads a buffer region into a read cache using customized
indices specified by index map. The read region of the buffer must be a single point.
The cache stage block follows the original order of loops and block itervars in the block.
If a block itervar does not appear in the buffer access region, it and its corresponding
loop variables will be omitted. User can then use `transform_block_layout` primitive to
reorder the block itervars and surrounding loops of the cache read/write block.
Unlike `cache_read`, `reindex_cache_read` only supports single consumer, please use
`cache_read` when there are multiple consumers.
Parameters
----------
block : SBlockRV
The consumer block of the target buffer.
read_buffer_index: int
The index of the buffer in block's read region.
storage_scope: str
The target storage scope.
index_map: IndexMap | Callable
User defined indices to access allocated cache buffer, maps from block iter vars.
Returns
-------
cached_block : SBlockRV
The block of the cache stage
Examples
--------
Before reindex_cache_read, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_reindex_cache_read(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
Create the schedule and reindex_cache_read:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_cache_read)
block_b = sch.get_sblock("B")
sch.reindex_cache_read(block_b, 0, "local", lambda vi, vj: (vj, vi))
print(sch.mod["main"].script())
After applying reindex_cache_read, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_reindex_cache_read(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
A_local = T.sblock_alloc_buffer((128, 128), scope="local")
for i, j in T.grid(128, 128):
with T.sblock("A_local"):
vi, vj = T.axis.remap("SS", [i, j])
A_local[vj, vi] = A[vi, vj]
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A_local[vj, vi] * 2.0
See Also
--------
reindex_cache_write
transform_block_layout
transform_layout
cache_read
reindex
"""
# Convert any string SBlock names into SBlock RVs.
block = self._normalize_block_arg(block)
if callable(index_map):
index_map = IndexMap.from_func(
index_map,
index_dtype=_get_sblock_default_dtype(self.get(block)),
)
return _ffi_api.ScheduleReindexCacheRead( # type: ignore # pylint: disable=no-member
self, block, read_buffer_index, storage_scope, index_map
)
@type_checked
def reindex_cache_write(
self,
block: SBlockRV | str,
write_buffer_index: int,
storage_scope: str,
index_map: Callable | IndexMap,
) -> SBlockRV:
r"""Create a block that reads a buffer region into a write cache using customized
indices specified by index map. The write region of the buffer must be a single point.
The cache stage block follows the original order of loops and block itervars in the block.
If a block itervar does not appear in the buffer access region, it and its corresponding
loop variables will be omitted. User can then use `transform_block_layout` primitive to
reorder the block itervars and surrounding loops of the cache read/write block.
Unlike `cache_write`, `reindex_cache_write` only supports single consumer, please use
`cache_write` when there are multiple consumers.
Parameters
----------
block : SBlockRV | str
The consumer block of the target buffer.
write_buffer_index: int
The index of the buffer in block's write region.
storage_scope: str
The target storage scope.
index_map: Callable | IndexMap
User defined indices to access allocated cache buffer, maps from block iter vars.
consumer_blocks: Optional[List[SBlockRV | str]]
An optional list of consumers that should read directly from the cache.
If not specified, all consumers will read from the original buffer.
Returns
-------
cached_block : SBlockRV
The block of the cache stage
Examples
--------
Before reindex_cache_write, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_reindex_cache_write(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
Create the schedule and reindex_cache_write:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_cache_write)
block_b = sch.get_sblock("B")
sch.reindex_cache_write(block_b, 0, "local", lambda vi, vj: (vi // 2, vi % 2, vj))
print(sch.mod["main"].script())
After applying reindex_cache_write, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_cache_write(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (64, 2, 128))
B_local = T.sblock_alloc_buffer((128, 128), scope="local")
for i, j in T.grid(128, 128):
with T.sblock("A_local"):
vi, vj = T.axis.remap("SS", [i, j])
B_local[vi % 2, vi // 2, vj] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = B_local[vi % 2, vi // 2, vj]
See Also
--------
reindex_cache_read
transform_block_layout
transform_layout
cache_write
reindex
"""
# Convert any string SBlock names into SBlock RVs.
block = self._normalize_block_arg(block)
if callable(index_map):
index_map = IndexMap.from_func(
index_map,
index_dtype=_get_sblock_default_dtype(self.get(block)),
)
return _ffi_api.ScheduleReindexCacheWrite( # type: ignore # pylint: disable=no-member
self, block, write_buffer_index, storage_scope, index_map
)
@type_checked
def cache_inplace(
self,
block: SBlockRV | str,
read_buffer_index: int | str | Buffer,
storage_scope: str,
) -> list[SBlockRV]:
"""Create blocks that reads & write a buffer region into a cache block.
It requires the target block both read & write the target buffer.
Mainly for inplace operation.
Parameters
----------
block : SBlockRV | str
The target block operates on the target buffer.
read_buffer_index: int
The index of the buffer in block's read region, the unique
name of a read buffer in the block, or a Buffer object
that is within the blocks read region.
storage_scope: str
The target storage scope.
Returns
-------
cached_blocks : List[SBlockRV]
The blocks of the cache stage, read cache first, write cache second
Examples
--------
Before cache_inplace, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_cache_inplace(data_io: T.Buffer((64), "int32")):
for i0 in T.serial(1):
with T.sblock("A"):
T.reads(data_io[:64])
T.writes(data_io[:64])
T.evaluate(T.call_extern("call_impl", data_io.data, dtype=""))
Create the schedule and cache_inplace:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_cache_inplace)
block_a = sch.get_sblock("A")
sch.cache_inplace(block_a, 0, "local")
print(sch.mod["main"].script())
After applying cache_inplace, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def cache_inplace(data_io: T.Buffer(64, "int32")) -> None:
data_io_local = T.sblock_alloc_buffer([64], dtype="int32", scope="local")
for i0 in T.serial(1):
for ax0 in T.serial(64):
with T.sblock("data_io_local"):
v0 = T.axis.spatial(64, ax0)
T.reads(data_io[v0])
T.writes(data_io_local[v0])
data_io_local[v0] = data_io[v0]
with T.sblock("A"):
T.reads(data_io_local[0 : 64])
T.writes(data_io_local[0 : 64])
T.evaluate(T.call_extern("call_impl", data_io_local.data, dtype=""))
for ax0 in T.serial(64):
with T.sblock("data_io_local"):
v0 = T.axis.spatial(64, ax0)
T.reads(data_io_local[v0])
T.writes(data_io[v0])
data_io[v0] = data_io_local[v0]
"""
block = self._normalize_block_arg(block)
if not isinstance(read_buffer_index, int):
_, read_buffer_index, _ = self._normalize_buffer_arg(
block, read_buffer_index, required_buffer_type="read"
)
return _ffi_api.ScheduleCacheInplace( # type: ignore # pylint: disable=no-member
self, block, read_buffer_index, storage_scope
)
@type_checked
def cache_index(
self, block: SBlockRV | str, storage_scope: str, cse_thresh: int = 0
) -> list[SBlockRV]:
"""Create a block to cache precomputed index for later use.
if there is no index computation, keep unchanged.
Parameters
----------
block : SBlockRV | str
The target block operates on the target buffer.
storage_scope: str
The storage scope of cached block.
cse_thresh: int
The repeat threshold that determines a common sub expr,
default 0 means cache all index computation.
Returns
-------
cached_blocks : List[SBlockRV]
The blocks of the stage writing the cache buffers
Examples
--------
Before cache_inplace, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def resize(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (1, 3, 40, 40))
B = T.match_buffer(b, (1, 3, 80, 80))
for i0, i1, i2, i3 in T.grid(1, 3, 80, 80):
with T.sblock("A"):
n, c, vi, vj = T.axis.remap("SSSS", [i0, i1, i2, i3])
B[n, c, vi, vj] = A[n, c, vi//4 + vj//4, vj//2]
Create the schedule and cache_index:
.. code-block:: python
sch = tvm.s_tir.Schedule(resize)
block_a = sch.get_sblock("A")
sch.cache_index(block_a, "global", 1)
print(sch.mod["main"].script())
After applying cache_index, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def resize_cache_index(
A: T.Buffer((1, 3, 40, 40), "float32"), B: T.Buffer((1, 3, 80, 80), "float32")
) -> None:
index_var_0 = T.sblock_alloc_buffer([80, 80], dtype="int32", strides=[1])
index_var_1 = T.sblock_alloc_buffer([80], dtype="int32", strides=[1])
for ax0, ax1 in T.grid(80, 80):
with T.sblock("index_0"):
v0 = T.axis.spatial(80, ax0)
v1 = T.axis.spatial(80, ax1)
T.reads()
T.writes(index_var_0[v0, v1])
index_var_0[v0, v1] = v0 // 4 + v1 // 4
for ax0 in T.serial(80):
with T.sblock("index_1"):
v0 = T.axis.spatial(80, ax0)
T.reads()
T.writes(index_var_1[v0])
index_var_1[v0] = v0 // 2
for i0, i1, i2, i3 in T.grid(1, 3, 80, 80):
with T.sblock("A"):
n, c, vi, vj = T.axis.remap("SSSS", [i0, i1, i2, i3])
T.reads(A[n, c, vi // 4 + vj // 4, vj // 2])
T.writes(B[n, c, vi, vj])
B[n, c, vi, vj] = A[n, c, index_var_0[vi, vj], index_var_1[vj]]
"""
block = self._normalize_block_arg(block)
return _ffi_api.ScheduleCacheIndex( # type: ignore # pylint: disable=no-member
self, block, storage_scope, cse_thresh
)
@type_checked
def reindex(self, block: SBlockRV | str, buffer: tuple[str, int] | str | Buffer) -> SBlockRV:
"""Create a block that read/write a buffer region into a read/write cache with reindexing.
The layout of the cache will be the same as by the iterators of the block that reads/writes
the buffer. It requires:
1) There is only one block who reads/writes the target buffer
2) There is only one buffer load/store of this buffer in the block
Parameters
----------
block : SBlockRV | str
The block that accesses the target buffer. If a string,
this must uniquely identify a block.
buffer: Union[Tuple[str,int], Buffer, str]
The buffer to be transformed, or a specification of how to
identify the buffer to be transformed.
If `buffer` if a tuple of ``(str,int)``, the first item
should be either "read" or "write", and the second item is
an index into the block's read or write regions.
If `buffer` is a string, it is the name of the buffer,
which must exist within the reads/writes of the block. In
addition, the reads/writes of the block may not contain
more than one buffer with this name.
If `buffer` is a Buffer object, it must exist within the
reads/writes of the block.
Returns
-------
reindex_block : SBlockRV
The block of the reindex stage
Examples
--------
Before reindex, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_reindex(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32")
) -> None:
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vj, vi] * 2.0
Create the schedule and do reindex:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_reindex)
block = sch.get_sblock("B")
sch.reindex(block, ("read", 0))
After applying reindex, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_reindex(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32")
) -> None:
A_reindex = T.sblock_alloc_buffer((128, 128), "float32")
for i, j in T.grid(128, 128):
with T.sblock("A_reindex"):
vi, vj = T.axis.remap("SS", [i, j])
A_reindex[vi, vj] = A[vj, vi]
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A_reindex[vi, vj] * 2.0
"""
block = self._normalize_block_arg(block)
buffer_index_type, buffer_index, _ = self._normalize_buffer_arg(block, buffer)
assert buffer_index_type in ["read", "write"], "Invalid buffer_index_type"
buffer_index_type_enum = 0 if buffer_index_type == "read" else 1
return _ffi_api.ScheduleReIndex( # type: ignore # pylint: disable=no-member
self, block, buffer_index, buffer_index_type_enum
)
########## Schedule: Data movement ##########
def read_at(
self, loop: LoopRV, block: SBlockRV, read_buffer_index: int, storage_scope: str
) -> SBlockRV:
return _ffi_api.ScheduleReadAt( # type: ignore # pylint: disable=no-member
self, loop, block, read_buffer_index, storage_scope
)
def write_at(
self, loop: LoopRV, block: SBlockRV, write_buffer_index: int, storage_scope: str
) -> SBlockRV:
return _ffi_api.ScheduleWriteAt( # type: ignore # pylint: disable=no-member
self, loop, block, write_buffer_index, storage_scope
)
########## Schedule: Compute location ##########
@type_checked
def compute_at(
self,
block: SBlockRV | str,
loop: LoopRV,
preserve_unit_loops: bool = False,
index: int = -1,
) -> None:
"""Compute-At. Move a producer block under the specific loop, and regenerate the
loops induced by the block so that the buffer region produced by the producer block could
cover those regions consumed by its consumer blocks under the given loop. It requires:
1) `block` and `loop` are under the same scope, `loop` is not the ancestor of `block`
2) The scope block has stage-pipeline property
3) The subtree of the scope block, where the given block is in, satisfies the compact
dataflow condition. i.e. all the blocks in the scope block's subtree must be either
complete block or reduction block
4) The block is not an output block with regard to the scope block, i.e. the buffers written
by the block are allocated under the scope block
5) All the consumers of the block are under the given loop
Parameters
----------
block : SBlockRV | str
The block to be moved
loop: LoopRV
The loop where the block to be moved under
preserve_unit_loops: bool
Whether to keep the trivial loops whose extents are 1
index: int
The block index of the loop body subtree blocks:
- `index = -1` means inserted into the last possible insertion point;
- `index = -2` means inserted into the first possible insertion point;
- Otherwise, `index` is a nonnegative number that indicates the insertion point
Examples
--------
Before compute-at, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_compute_at(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128), "float32")
B = T.sblock_alloc_buffer((128, 128), "float32")
C = T.match_buffer(c, (128, 128), "float32")
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + 1.0
Create the schedule and do compute-at:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_compute_at)
block = sch.get_sblock("B")
loop, _ = sch.get_loops(sch.get_sblock("C"))
sch.compute_at(block, loop, preserve_unit_loops=False)
print(sch.mod["main"].script())
After applying compute-at, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_compute_at(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128), "float32")
B = T.sblock_alloc_buffer((128, 128), "float32")
C = T.match_buffer(c, (128, 128), "float32")
for i in T.serial(0, 128):
for j in T.serial(0, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
for j in T.serial(0, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + 1.0
"""
block = self._normalize_block_arg(block)
_ffi_api.ScheduleComputeAt( # type: ignore # pylint: disable=no-member
self, block, loop, preserve_unit_loops, index
)
@type_checked
def reverse_compute_at(
self,
block: SBlockRV | str,
loop: LoopRV,
preserve_unit_loops: bool = False,
index: int = -1,
) -> None:
"""Reverse-Compute-At. Move a consumer block under the specific loop, and regenerate the
loops induced by the block so that the buffer region consumed by the consumer block could
cover those regions produced by its producer blocks under the given loop. It requires:
1) `block` and `loop` are under the same scope, `loop` is not the ancestor of `block`
2) The scope block has stage-pipeline property
3) The subtree of the scope block, where the given block is in, satisfies the compact
dataflow condition. i.e. all the blocks in the scope block's subtree must be either
complete block or reduction block
4) All the producers of the block are under the given loop
Parameters
----------
block : SBlockRV | str
The block to be moved
loop: LoopRV
The loop where the block to be moved under
preserve_unit_loops: bool
Whether to keep the trivial loops whose extents are 1
index: int
The block index of the loop body subtree blocks:
- `index = -1` means inserted into the last possible insertion point;
- `index = -2` means inserted into the first possible insertion point;
- Otherwise, `index` is a nonnegative number that indicates the insertion point
Examples
--------
Before reverse-compute-at, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_reverse_compute_at(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128), "float32")
B = T.sblock_alloc_buffer((128, 128), "float32")
C = T.match_buffer(c, (128, 128), "float32")
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + 1.0
Create the schedule and do reverse-compute-at:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_reverse_compute_at)
block = sch.get_sblock("C")
loop, _ = sch.get_loops(sch.get_sblock("B"))
sch.reverse_compute_at(block, loop, preserve_unit_loops=False)
print(sch.mod["main"].script())
After applying reverse-compute-at, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_reverse_compute_at(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128), "float32")
B = T.sblock_alloc_buffer((128, 128), "float32")
C = T.match_buffer(c, (128, 128), "float32")
for i in T.serial(0, 128):
for j in T.serial(0, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
for j in T.serial(0, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + 1.0
"""
block = self._normalize_block_arg(block)
_ffi_api.ScheduleReverseComputeAt( # type: ignore # pylint: disable=no-member
self, block, loop, preserve_unit_loops, index
)
@type_checked
def compute_inline(self, block: SBlockRV | str) -> None:
"""Inline a block into its consumer(s). It requires:
1) The block is a complete non-root block, which only produces one buffer
2) The block must not be the only leaf in the scope.
3) The body of the block must be a BufferStore statement in
the form of, ``A[i, j, k, ...] = ...`` where the indices of
the LHS are all distinct atomic variables, and no variables
other than those indexing variables are allowed in the
statement.
Parameters
----------
block : SBlockRV | str
The block to be inlined to its consumer(s)
Examples
--------
Before compute-inline, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_inline(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.sblock_alloc_buffer((128, 128))
C = T.match_buffer(c, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + 1.0
Create the schedule and do compute-inline:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_inline)
sch.compute_inline(sch.get_sblock("B"))
print(sch.mod["main"].script())
After applying compute-inline, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_inline(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
C = T.match_buffer(c, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = A[vi, vj] * 2.0 + 1.0
"""
block = self._normalize_block_arg(block)
_ffi_api.ScheduleComputeInline(self, block) # type: ignore # pylint: disable=no-member
@type_checked
def reverse_compute_inline(self, block: SBlockRV | str) -> None:
"""Inline a block into its only producer. It requires:
1) The block is a complete non-root block, which only produces and consumes one buffer
2) The block must not be the only leaf in the scope.
3) The only producer of the block is a read-after-write producer and a
complete non-root block
4) The body of the block must be a BufferStore statement in the form of,
``B[f(i, j, k, ...)] = g(i, j, k, A[i, j, k, ...] ...)`` where the
indices of each `BufferLoad` on the RHS are all distinct atomic
variables, and no variables other than those indexing variables are
allowed in the statement.
Parameters
----------
block : SBlockRV | str
The block to be inlined to its producer
Examples
--------
Before reverse-compute-inline, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_inline(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.sblock_alloc_buffer((128, 128))
C = T.match_buffer(c, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + 1.0
Create the schedule and do reverse-compute-inline:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_inline)
sch.reverse_compute_inline(sch.get_sblock("C"))
print(sch.mod["main"].script())
After applying reverse-compute-inline, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_inline(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
C = T.match_buffer(c, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = A[vi, vj] * 2.0 + 1.0
"""
block = self._normalize_block_arg(block)
# pylint: disable-next=no-member
_ffi_api.ScheduleReverseComputeInline(self, block) # type: ignore
@type_checked
def fuse_reduction_epilogue(
self,
reduction_block: SBlockRV | str,
epilogue_block: SBlockRV | str,
) -> None:
"""Fuse an epilogue block into a reduction block.
It requires:
1) The reduction block is a complete reduction block
2) The epilogue block only reads from the reduction block's output
3) The epilogue matches one of the supported patterns:
- Bias: ``output = reduction_result + bias``
- BiasReLU: ``output = max(reduction_result + bias, 0)``
- Clipping: ``output = min(max(reduction_result, lower), upper)``
or their commutative variants
.. warning::
**Semantic Change for Non-Linear Epilogues (BiasReLU, Clipping):**
For non-linear epilogues (BiasReLU and Clipping), fusion changes the
computation semantics from post-reduction application to per-iteration
application. This can lead to different numerical results.
**Example with Clipping to [-5, 5] and inputs [6, -2]:**
- **Post-reduction clipping** (original): ``clip(sum([6, -2])) = clip(4) = 4``
- **Per-iteration clipping** (fused): ``acc=0 → clip(0+6)=5 → clip(5+(-2))=3``
The fused version applies clipping at each reduction iteration, which
may be an intended optimization for some models but can cause unexpected
correctness issues if users are not aware of this behavior.
For linear epilogues (Bias), fusion preserves exact numerical equivalence.
Parameters
----------
reduction_block : SBlockRV | str
The reduction block (e.g., matmul)
epilogue_block : SBlockRV | str
The epilogue block to be fused (e.g., bias add, ReLU, clipping)
Examples
--------
See :py:func:`test_tir_schedule_fuse_reduction_epilogue` for examples.
"""
reduction_block = self._normalize_block_arg(reduction_block)
epilogue_block = self._normalize_block_arg(epilogue_block)
# pylint: disable-next=no-member
_ffi_api.ScheduleFuseReductionEpilogue(self, reduction_block, epilogue_block) # type: ignore
########## Schedule: Reduction ##########
@type_checked
def decompose_reduction(self, block: SBlockRV | str, loop: LoopRV) -> SBlockRV:
"""Decompose a reduction block into two separate blocks.
a) The init block, which is translated from the init statement of the reduction block;
b) The update block, which is the original block without init statement.
The init block is inserted right before the given loop.
The schedule primitive requires:
1) The input block is a reduction block.
2) The input loop is the ancestor of the block.
3) The input loop is not lower than all the loops related to reduce block var.
Parameters
----------
block : SBlockRV | str
The reduction block to be decomposed
loop : LoopRV
The loop above which the init block is inserted before.
Returns
-------
init_block : SBlockRV
The init block
Examples
--------
Before decompose-reduction, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_decompose(a: ty.handle, b: ty.handle, c: ty.handle) -> None:
A = tirx.match_buffer(a, [128, 128])
B = tirx.match_buffer(b, [128, 128])
C = tirx.match_buffer(c, [128, 128])
for i, j, k in tirx.grid(128, 128, 128):
with tirx.block([128, 128, tirx.reduce_axis(0, 128)], "C") as [vi, vj, vk]:
with tirx.init():
C[vi, vj] = 0.0
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
Create the schedule and do decompose-reduction with specified loop:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_decompose)
C = sch.get_sblock("C")
i, j, k = sch.get_loops(C)
sch.decompose_reduction(C, i)
print(sch.mod["main"].script())
After applying decompose-reduction, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_decompose(a: ty.handle, b: ty.handle, c: ty.handle) -> None:
A = tirx.match_buffer(a, [128, 128])
B = tirx.match_buffer(b, [128, 128])
C = tirx.match_buffer(c, [128, 128])
for i in tirx.serial(128):
for j in tirx.serial(128):
with tirx.block([128, 128]) as [vi, vj]:
C[vi, vj] = 0.0
for i, j, k in tirx.grid(128, 128, 128):
with tirx.block([128, 128, tirx.reduce_axis(0, 128)], "C") as [vi, vj, vk]:
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
"""
block = self._normalize_block_arg(block)
# pylint: disable-next=no-member
return _ffi_api.ScheduleDecomposeReduction(self, block, loop) # type: ignore
@type_checked
def rfactor(self, loop: LoopRV, factor_axis: int) -> SBlockRV:
"""Factorize an associative reduction block by the specified loop.
An associative reduction cannot be parallelized directly,
because it leads to potential race condition during accumulation.
Alternatively, the reduction could be factorized on a loop with the following steps:
- Step 1: evenly slice the reduction into `n` separate chunks, where `n` is the loop extent
- Step 2: compute the chunks separately and write the result into `n` intermediate buffers;
- Step 3: accumulate the `n` separate buffer into the result buffer.
Note that the Step 2 above introduces opportunities for parallelization.
RFactor is a schedule primitive that implements the transformation described above:
Given a block that writes to buffer `B`, it factorizes a loop of extent `n`.
For example, the pseudocode below accumulates `B[i] = sum(A[i, : , : ])`:
.. code-block:: python
for i in range(128): # loop i is a data parallel loop
for j in range(128): # loop j is a reduction loop
for k in range(128): # loop k is a reduction loop
B[i] = B[i] + A[i, j, k]
Suppose RFactor is applied on the innermost loop `k` and `factor_axis = 1`.
RFactor then creates an intermediate buffer and two blocks.
1. The intermediate buffer, or "rf-buffer" is a buffer of rank `ndim(B) + 1` and
size `size(B) * n`, whose shape expands from `shape(B)` by adding an axis of `n`
at the position specified by `factor_axis`. For example,
* shape(B) = [1, 2, 3], factor_axis = 0 => shape(B_rf) = [n, 1, 2, 3]
* shape(B) = [1, 2, 3], factor_axis = 1 => shape(B_rf) = [1, n, 2, 3]
* shape(B) = [1, 2, 3], factor_axis = 2 => shape(B_rf) = [1, 2, n, 3]
* shape(B) = [1, 2, 3], factor_axis = 3 => shape(B_rf) = [1, 2, 3, n]
2. The rfactor block, or "rf-block", is a block that writes to the `rf-buffer` without
accumulating over the loop `k`, i.e. the loop `k` is converted from a reduction loop
to a data parallel loop. In our example, the rf-block is:
.. code-block:: python
B_rf = np.zeros((128, 128)) # the rf-buffer
for k in range(128): # loop k is converted to a data parallel loop
for i in range(128): # loop i is a data parallel loop (unchanged)
for j in range(128): # loop j is a reduction loop (unchanged)
B_rf[i, k] = B_rf[i, k] + A[i, j, k]
3. The write-back block, or `wb-block`, is a block that accumulates the rf-buffer into
the result buffer. All the reduction loops are removed except the loop `k` for accumulation.
In our example, the wb-block is:
.. code-block:: python
for i in range(128): # loop i is a data parallel loop (unchanged)
# loop j is removed because it is a reduction loop
for k in range(128): # loop k is a reduction loop (unchanged)
B[i] = B[i] + B_rf[i, k]
Parameters
----------
loop : LoopRV
The loop outside block for which we want to do rfactor
factor_axis : int
The position where the new dimension is placed in the new introduced rfactor buffer
Returns
-------
rf_block : SBlockRV
The block which computes partial results over each slices (i.e., the first block
as described in the above illustration)
Examples
--------
Before rfactor, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_rfactor(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128, 128))
B = T.match_buffer(b, (128,))
for ii, i, j in T.grid(128, 128, 128):
with T.sblock("B"):
vii, vi, vj = T.axis.remap("SRR", [ii, i, j])
with T.init():
B[vii] = 0.0
B[vii] = B[vii] + A[vii, vi, vj]
Create the schedule and do rfactor:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_rfactor)
_, _, k = sch.get_loops(sch.get_sblock("B"))
sch.rfactor(k, 0)
print(sch.mod["main"].script())
After applying rfactor, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_rfactor(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, [128, 128, 128])
B = T.match_buffer(b, [128])
B_rf = T.sblock_alloc_buffer([128, 128])
for i2, ii, i in T.grid(128, 128, 128):
with T.sblock("B_rf"):
vi2, vii, vi = T.axis.remap("SSR", [i2, ii, i])
with T.init():
B_rf[vi2, vii] = 0.0
B_rf[vi2, vii] = (B_rf[vi2, vii] + A[vii, vi, vi2])
for ii, i2 in T.grid(128, 128):
with T.sblock("B"):
vii, vi2 = T.axis.remap("SR", [ii, i2])
with T.init():
B[vii] = 0.0
B[vii] = B[vii] + B_rf[vi2, vii]
Note
----
Rfactor requires:
1) `loop` has only one child block, and it is a reduction block;
2) `loop` is a reduction loop, i.e. the loop variable is bound to only reduction variables
in the block binding;
3) `loop` is not parallelized, vectorized, unrolled or bound to any thread axis;
4) The block scope that `loop` is in is a staged-pipeline;
5) The outermost loop outside the reduction block should has the reduction block as its
first child block;
6) The outermost reduction loop should have only one child block;
7) An unary extent loop that is not bound to any reduction or data parallel variables in
the block binding should not appear under some reduction loop;
8) The reduction block should write to only one buffer, and its init and body are both
simple `BufferStore`s, and the pattern is registered as an associative reducer.
The pre-defined patterns include: plus, multiplication, min and max;
9) Each of the loops on top of the block cannot be bound to a data parallel and a
reduction block binding at the same time;
10) `factor_axis` should be in range `[-ndim(B) - 1, ndim(B)]`,
where `B` is the buffer that the reduction block writes to.
Negative indexing is normalized according to numpy convention.
"""
# pylint: disable-next=no-member
return _ffi_api.ScheduleRFactor(self, loop, factor_axis) # type: ignore
######## Schedule: SBlock annotation ########
@type_checked
def storage_align( # pylint: disable=too-many-arguments
self, block: SBlockRV | str, buffer_index: int, axis: int, factor: int, offset: int
) -> None:
"""Set alignment requirement for specific dimension such that
stride[axis] == k * factor + offset for some k. This is useful to set memory layout for more
friendly memory access pattern. For example, we can set alignment to be factor=2, offset=1
to avoid bank conflict for thread access on higher dimension in GPU shared memory.
Parameters
----------
block : SBlockRV | str
The producer block of the buffer.
buffer_index : int
The index of the buffer in block's write region.
axis : int
The dimension to be specified for alignment.
factor : int
The factor multiple of alignment.
offset : int
The required offset factor.
Examples
--------
Before storage_align, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_storage_align(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.sblock_alloc_buffer((128, 128))
C = T.match_buffer(c, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + 1.0
Create the schedule and do storage_align:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_storage_align)
sch.storage_align(sch.get_sblock("B"), buffer_index=0, axis=0, factor=128, offset=1)
print(sch.mod["main"].script())
After applying storage_align, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_storage_align(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.sblock_alloc_buffer((128, 128))
C = T.match_buffer(c, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
T.sblock_attr({"buffer_dim_align": [[[0, 128, 1]]]})
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + 1.0
After lowering passes, buffer B will have strides as [129, 1].
Note
----
Storage_align requires the buffer to be an intermediate buffer defined via `alloc_buffer`.
"""
block = self._normalize_block_arg(block)
_ffi_api.ScheduleStorageAlign( # type: ignore # pylint: disable=no-member
self, block, buffer_index, axis, factor, offset
)
@type_checked
def set_scope(
self, block: SBlockRV | str, buffer_index: int | str | Buffer, storage_scope: str
) -> None:
"""Set the storage scope of a buffer, where the buffer is
specified by the a block and a write-index.
Parameters
----------
block : SBlockRV | str
The producer block of the buffer
buffer_index : int
The index of the buffer in block's write region
storage_scope : str
The storage scope to be set
Examples
--------
Before set_scope, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_set_scope(
A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
B = T.sblock_alloc_buffer((128, 128), dtype="float32")
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + 1.0
Create the schedule and do set_scope:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_set_scope)
sch.set_scope(sch.get_sblock("B"), buffer_index=0, storage_scope="shared")
print(sch.mod["main"].script())
After applying set_scope, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_set_scope(
A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
B_shared = T.sblock_alloc_buffer([128, 128], dtype="float32", scope="shared")
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B_shared[vi, vj] = A[vi, vj] * T.float32(2)
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B_shared[vi, vj] + T.float32(1)
Note
----
`set_scope` requires the buffer to be an intermediate buffer defined via `alloc_buffer`.
"""
block = self._normalize_block_arg(block)
if not isinstance(buffer_index, int):
_, buffer_index, _ = self._normalize_buffer_arg(
block, buffer_index, required_buffer_type="write"
)
_ffi_api.ScheduleSetScope( # type: ignore # pylint: disable=no-member
self, block, buffer_index, storage_scope
)
@type_checked
def unsafe_set_dtype(self, block: SBlockRV | str, buffer_index: int, dtype: str) -> None:
"""Set the data type of a buffer, where the buffer is
specified by the a block and write-index.
This schedule primitive is unsafe and may change the correctness of program because of
type conversion, please use with caution.
Parameters
----------
block : SBlockRV | str
The producer block of the buffer
buffer_index : int
The index of the buffer in block's write region
dtype : str
The data type to be set
Examples
--------
Before unsafe_set_dtype, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_set_dtype(
A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
B = T.sblock_alloc_buffer((128, 128), dtype="float32")
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j]
C[vi, vj] = B[vi, vj] + 1.0
Create the schedule and do unsafe_set_dtype:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_set_dtype)
sch.unsafe_set_dtype("B", buffer_index=0, dtype="float16")
print(sch.mod["main"].script())
After applying set_dtype, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_set_dtype(
A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
B = T.sblock_alloc_buffer((128, 128), dtype="float16")
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = T.cast(A[vi, vj] * 2.0, "float16")
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j]
C[vi, vj] = T.cast(B[vi, vj], "float32") + 1.0
Note
----
`unsafe_set_dtype` requires the buffer to be an intermediate buffer defined via
`alloc_buffer`.
"""
block = self._normalize_block_arg(block)
_ffi_api.ScheduleUnsafeSetDType( # type: ignore # pylint: disable=no-member
self, block, buffer_index, dtype
)
########## Schedule: Blockize & Tensorize ##########
@type_checked
def blockize(
self, target: LoopRV | list[SBlockRV], preserve_unit_iters: bool = True
) -> SBlockRV:
"""Convert multiple blocks or the subtree rooted at a specific loop into a block.
Parameters
----------
target : LoopRV or List[SBlockRV]
The root of the subtree or the specified blocks.
preserve_unit_iters : bool
Whether or not to preserve unit iterators in block bindings
Returns
-------
result : SBlockRV
The new block.
Examples
--------
Before blockize, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_blockize(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32")
) -> None:
for i_0, j_0, i_1, j_1 in T.grid(8, 8, 16, 16):
with T.sblock("B"):
vi = T.axis.spatial(128, i_0 * 16 + i_1)
vj = T.axis.spatial(128, j_0 * 16 + j_1)
T.reads(A[vi, vj])
T.writes(B[vi, vj])
B[vi, vj] = A[vi, vj] * T.float32(2)
Create the schedule and do set_scope:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_blockize)
B = sch.get_sblock("B")
_, _, i1, _ = sch.get_loops(B)
sch.blockize(i1)
print(sch.mod["main"].script())
After applying blockize, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_blockize(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32")
)-> None:
for i_0, j_0 in T.grid(8, 8):
with T.sblock("B_o"):
vio, vjo = T.axis.remap("SS", [i_0, j_0])
T.reads(A[vio * 16 : vio * 16 + 16, vjo * 16 : vjo * 16 + 16])
T.writes(B[vio * 16 : vio * 16 + 16, vjo * 16 : vjo * 16 + 16])
for i_1, j_1 in T.grid(16, 16):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i_1, j_1])
T.reads(A[vio * 16 + vi, vjo * 16 + vj])
T.writes(B[vio * 16 + vi, vjo * 16 + vj])
B[vio * 16 + vi, vjo * 16 + vj] = A[vio * 16 + vi, vjo * 16 + vj] \
* T.float32(2)
Note
----
blockize requires there is exactly one block under the given loop and the bindings of the
block are divisible by the subspace represented by the loops starting at the given loop.
"""
# pylint: disable-next=no-member
return _ffi_api.ScheduleBlockize(self, target, preserve_unit_iters) # type: ignore
@type_checked
def tensorize(
self,
block_or_loop: SBlockRV | LoopRV,
tensor_intrin: str,
preserve_unit_iters: bool = True,
) -> None:
"""Tensorize the computation enclosed by loop with the tensor intrinsic.
Parameters
----------
block_or_loop : SBlockRV | LoopRV
The loop to be tensorized.
tensor_intrin : str
The tensor intrin or the name of the tensor intrin.
preserve_unit_iters : bool
Whether or not to preserve unit iterators in block bindings
Examples
--------
Before tensorize, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_tensorize(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
C: T.Buffer((128, 128), "float32"),
) -> None:
# body
# with T.sblock("root")
for i_0, j_0, k_0, i_1, j_1, k_1 in T.grid(8, 8, 8, 16, 16, 16):
with T.sblock("update"):
vi = T.axis.spatial(128, i_0 * 16 + i_1)
vj = T.axis.spatial(128, j_0 * 16 + j_1)
vk = T.axis.reduce(128, k_0 * 16 + k_1)
T.reads(C[vi, vj], A[vi, vk], B[vj, vk])
T.writes(C[vi, vj])
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
Declare and register the tensor intrinsic:
.. code-block:: python
@T.prim_func(s_tir=True)
def mma_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (16, 16), align=128, offset_factor=1)
B = T.match_buffer(b, (16, 16), align=128, offset_factor=1)
C = T.match_buffer(c, (16, 16), align=128, offset_factor=1)
with T.sblock("root"):
T.reads(C[0 : 16, 0 : 16], A[0 : 16, 0 : 16], B[0 : 16, 0 : 16])
T.writes(C[0 : 16, 0 : 16])
for i, j, k in T.grid(16, 16, 16):
with T.sblock("update"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
@T.prim_func(s_tir=True)
def mma_intrin(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (16, 16), align=128, offset_factor=1)
B = T.match_buffer(b, (16, 16), align=128, offset_factor=1)
C = T.match_buffer(c, (16, 16), align=128, offset_factor=1)
with T.sblock("root"):
T.reads(C[0 : 16, 0 : 16], A[0 : 16, 0 : 16], B[0 : 16, 0 : 16])
T.writes(C[0 : 16, 0 : 16])
T.evaluate(
T.tvm_mma_sync(
C.data,
C.elem_offset // 256,
A.data,
A.elem_offset // 256,
B.data,
B.elem_offset // 256,
C.data,
C.elem_offset // 256,
dtype="void",
)
)
tirx.TensorIntrin.register("test_mma_intrin", mma_desc, mma_intrin)
Create the schedule and do tensorize:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_tensorize)
update = sch.get_sblock("update")
_, _, _, i1, _, _ = sch.get_loops(update)
sch.tensorize(i1, "test_mma_intrin")
print(sch.mod["main"].script())
After applying tensorize, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_tensorize(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
C: T.Buffer((128, 128), "float32"),
) -> None:
# body
# with T.sblock("root")
for i_0, j_0, k_0 in T.grid(8, 8, 8):
with T.sblock("update_o"):
vio, vjo, vko = T.axis.remap("SSR", [i_0, j_0, k_0])
T.reads(
C[vio * 16 : vio * 16 + 16, vjo * 16 : vjo * 16 + 16],
A[vio * 16 : vio * 16 + 16, vko * 16 : vko * 16 + 16],
B[vjo * 16 : vjo * 16 + 16, vko * 16 : vko * 16 + 16],
)
T.writes(C[vio * 16 : vio * 16 + 16, vjo * 16 : vjo * 16 + 16])
A_1 = T.match_buffer(
A[vio * 16 : vio * 16 + 16, vko * 16 : vko * 16 + 16],
[16, 16],
dtype="float32",
offset_factor=1,
)
B_1 = T.match_buffer(
B[vjo * 16 : vjo * 16 + 16, vko * 16 : vko * 16 + 16],
[16, 16],
dtype="float32",
offset_factor=1,
)
C_1 = T.match_buffer(
C[vio * 16 : vio * 16 + 16, vjo * 16 : vjo * 16 + 16],
[16, 16],
dtype="float32",
offset_factor=1,
)
T.evaluate(
T.tvm_mma_sync(
C_1.data,
C_1.elem_offset // 256,
A_1.data,
A_1.elem_offset // 256,
B_1.data,
B_1.elem_offset // 256,
C_1.data,
C_1.elem_offset // 256,
dtype="void",
)
)
"""
_ffi_api.ScheduleTensorize( # type: ignore # pylint: disable=no-member
self, block_or_loop, tensor_intrin, preserve_unit_iters
)
########## Schedule: Annotation ##########
PrimAnnotationValueT = str | int | float | ExprRV
AnnotationValueT = (
PrimAnnotationValueT
| list[PrimAnnotationValueT]
| dict[str, PrimAnnotationValueT | list[PrimAnnotationValueT]]
)
@type_checked
def annotate(
self, block_or_loop: SBlockRV | LoopRV, ann_key: str, ann_val: AnnotationValueT
) -> None:
"""Annotate a block/loop with a key value pair
Parameters
----------
block_or_loop: SBlockRV | LoopRV
The block/loop to be annotated
ann_key : str
The annotation key
ann_val : AnnotationValueT
The annotation value
Examples
--------
Before annotate, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_annotate(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
Create the schedule and do annotate:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_annotate)
sch.annotate(sch.get_sblock("B"), "ann_key", "ann_value")
print(sch.mod["main"].script())
After applying annotate, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_annotate(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
T.sblock_attr({"ann_key", "ann_value"})
B[vi, vj] = A[vi, vj] * 2.0
"""
_ffi_api.ScheduleAnnotate( # type: ignore # pylint: disable=no-member
self, block_or_loop, ann_key, ann_val
)
@type_checked
def unannotate(self, block_or_loop: SBlockRV | LoopRV, ann_key: str) -> None:
"""Unannotate a block/loop's annotation with key ann_key
Parameters
----------
block_or_loop: SBlockRV | LoopRV
The block/loop to be unannotated
ann_key : str
The annotation key
Examples
--------
Before unannotate, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_unannotate(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
T.sblock_attr({"ann_key", "ann_value"})
B[vi, vj] = A[vi, vj] * 2.0
Create the schedule and do annotate:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_unannotate)
sch.unannotate(sch.get_sblock("B"), "ann_key")
print(sch.mod["main"].script())
After applying unannotate, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_unannotate(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
"""
_ffi_api.ScheduleUnannotate( # type: ignore # pylint: disable=no-member
self, block_or_loop, ann_key
)
########## Schedule: Layout transformation ##########
def _normalize_block_arg(self, block: SBlockRV | str) -> SBlockRV:
if isinstance(block, str):
return self.get_sblock(block)
return block
def _normalize_buffer_arg(
self,
block: SBlockRV,
buffer: tuple[str, int] | int | str | Buffer,
required_buffer_type=None,
) -> tuple[str, int, Buffer]:
block_obj: SBlock = self.get(block)
block_name = block_obj.name_hint
def iter_buffers():
for i, read in enumerate(block_obj.reads):
yield "read", i, read.buffer
for i, write in enumerate(block_obj.writes):
yield "write", i, write.buffer
if isinstance(buffer, int):
buffer = (required_buffer_type, buffer)
if isinstance(buffer, str):
possible_buffers = {}
# String lookup requires ensuring that the name is unique
for buffer_index_type, buffer_index, buf in iter_buffers():
if buf.name == buffer:
possible_buffers[buf] = (buffer_index_type, buffer_index)
assert possible_buffers, f"Could not find buffer '{buffer}' in block '{block_name}'"
assert len(possible_buffers) == 1, (
f"Multiple buffers named '{buffer}' in block '{block_name}'"
)
buffer_obj, (buffer_index_type, buffer_index) = next(iter(possible_buffers.items()))
elif isinstance(buffer, Buffer):
# Buffer lookup has unique id, can break out early
found = False
for buffer_index_type, buffer_index, buffer_obj in iter_buffers():
if buffer_obj.same_as(buffer):
found = True
break
assert found, f"Could not find buffer '{buffer.name}' in block '{block_name}'"
elif isinstance(buffer, tuple):
buffer_index_type, buffer_index = buffer
assert buffer_index_type in ["read", "write"], (
f"Invalid buffer_index_type. "
f"Expected 'read' or 'write', "
f"but received {buffer_index_type}"
)
buffer_list = block_obj.reads if buffer_index_type == "read" else block_obj.writes
assert 0 <= buffer_index < len(buffer_list), (
f"Invalid buffer_index {buffer_index}. "
f"Block {block_name} has only "
f"{len(buffer_list)} {buffer_index_type} buffers."
)
buffer_obj = buffer_list[buffer_index].buffer
else:
raise TypeError(f"Invalid type for argument 'buffer': {type(buffer)}")
if required_buffer_type is not None:
assert buffer_index_type == required_buffer_type, (
f"Expected buffer to be read buffer, "
f"but {buffer_obj.name} was a {buffer_index_type} buffer "
f"in the specified block"
)
return (buffer_index_type, buffer_index, buffer_obj)
@type_checked
def transform_layout(
self,
block: SBlockRV | str,
buffer: tuple[str, int] | str | Buffer,
index_map: IndexMap | Callable,
pad_value: int | float | Expr | IndexMap | Callable | None = None,
*,
assume_injective_transform: bool = False,
) -> None:
"""Apply a transformation represented by IndexMap to buffer
Parameters
----------
block : SBlockRV | str
The block that accesses the target buffer. If a string,
this must uniquely identify a block.
buffer: Union[Tuple[str,int], Buffer, str]
The buffer to be transformed, or a specification of how to
identify the buffer to be transformed.
If `buffer` if a tuple of ``(str,int)``, the first item
should be either "read" or "write", and the second item is
an index into the block's read or write regions.
If `buffer` is a string, it is the name of the buffer,
which must exist within the reads/writes of the block. In
addition, the reads/writes of the block may not contain
more than one buffer with this name.
If `buffer` is a Buffer object, it must exist within the
reads/writes of the block.
index_map : IndexMap | Callable
The transformation to apply.
If `index_map` is a callable, and the returned list
contains IndexMap.AXIS_SEPARATOR, the SetAxisSeparators
primitive will be called in addition to the
TransformLayout primitive.
pad_value: Optional[int | float | Expr | IndexMap | Callable]
The value to be used for any padding introduced by the
transformation. If the schedule contains a producer block
for the specified buffer, the pad value will be written as
part of the producer block if possible, or after the producer
block otherwise. Otherwise, if the buffer is an input, will
insert an annotation block to state that the padding contains
the known value.
The pad value may not contain instances of BufferLoad,
except where it loads a value from the buffer being
transformed (e.g. to create a circular buffer with
padding that consists of repeated elements).
Note: If applied to an input buffer, the calling scope is
responsible for ensuring that the pad_value is present.
Algebraic symplifications, branch elimination, and other
optimizations may assume that this precondition is met, and
may result in incorrect results being returned.
If None, the transformation may not introduce padding.
If an int, float or Expr, the transformation is the
specific value to be present in the padding.
If an IndexMap or Callable, the transformation is the
value to be present in the padding in terms of the
transformed index.
assume_injective_transform : bool
If set to true, the schedule primitive will assume the index_map is injective and skip
checking overlapping of the mapped indices. This can be useful for complicated index_map
that the analysis does not cover. It is the callers' responsibility to ensure the
index map is injective, otherwise, the correctness of the schedule is not guaranteed.
Examples
--------
Before transform_layout, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_transform_layout(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128), "float32")
B = T.sblock_alloc_buffer((128, 128), "float32")
C = T.match_buffer(c, (128, 128), "float32")
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + 1.0
Create the schedule and do transform_layout:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_storage_align)
sch.transform_layout(sch.get_sblock("B"), buffer=("write",0),
index_map=lambda m, n: (m // 16, n // 16, m % 16, n % 16))
print(sch.mod["main"].script())
After applying transform_layout, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128), "float32")
B = T.sblock_alloc_buffer((8, 8, 16, 16), "float32")
C = T.match_buffer(c, (128, 128), "float32")
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi // 16, vj // 16, vi % 16, vj % 16] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi // 16, vj // 16, vi % 16, vj % 16] + 1.0
"""
block = self._normalize_block_arg(block)
buffer_index_type, buffer_index, buffer_obj = self._normalize_buffer_arg(block, buffer)
ndim = len(buffer_obj.shape)
if callable(index_map):
index_map, axis_separators = IndexMap.from_func_with_separators(
index_map,
ndim=ndim,
index_dtype=_get_sblock_default_dtype(self.get(block)),
)
else:
axis_separators = []
if pad_value is None:
pass
elif callable(pad_value):
pad_value = IndexMap.from_func(
pad_value,
ndim=len(index_map.final_indices),
index_dtype=_get_sblock_default_dtype(self.get(block)),
)
elif not isinstance(pad_value, IndexMap):
# Explicitly convert python int/float arguments to the
# buffer's type. If the default `tvm.runtime.convert`
# behavior is applied, these would be converted to
# int32/float32, which may not match the buffer's type.
if buffer_obj.dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT) and isinstance(
pad_value, int
):
pad_value = IntImm(buffer_obj.dtype.dtype, pad_value)
elif buffer_obj.dtype.matches_code(DataTypeCode.FLOAT, DataTypeCode.BFLOAT) and (
isinstance(pad_value, float)
):
pad_value = FloatImm(buffer_obj.dtype.dtype, pad_value)
pad_value = IndexMap.from_func(
lambda *indices: pad_value,
ndim=len(index_map.final_indices),
index_dtype=_get_sblock_default_dtype(self.get(block)),
)
buffer_index_type_enum = 0 if buffer_index_type == "read" else 1
_ffi_api.ScheduleTransformLayout( # type: ignore # pylint: disable=no-member
self,
block,
buffer_index,
buffer_index_type_enum,
index_map,
pad_value,
assume_injective_transform,
)
if axis_separators:
_ffi_api.ScheduleSetAxisSeparator( # type: ignore # pylint: disable=no-member
self, block, buffer_index, buffer_index_type_enum, axis_separators
)
@type_checked
def transform_block_layout(self, block: SBlockRV | str, index_map: IndexMap | Callable) -> None:
"""Apply a transformation represented by IndexMap to block
Parameters
----------
block : SBlockRV | str
The block to be transformed
index_map : IndexMap | Callable
The transformation to apply.
Examples
--------
Before transform_block_layout, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_transform_block_layout(
A: T.Buffer((16, 16), "float32"),
B: T.Buffer((16, 16), "float32")
) -> None:
for i, j in T.grid(16, 16):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
Create the schedule and do transform_block_layout:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_transform_block_layout)
sch.transform_block_layout(sch.get_sblock("B"), lambda i, j: (i * 16 + j,))
print(sch.mod["main"].script())
After applying transform_block_layout, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_transform_block_layout(
A: T.Buffer((16, 16), "float32"),
B: T.Buffer((16, 16), "float32")
) -> None:
for i in range(256):
with T.sblock("B"):
vi, = T.axis.remap("S", [i])
B[vi // 16, vi % 16] = A[vi // 16, vi % 16] * 2.0
"""
block = self._normalize_block_arg(block)
if callable(index_map):
index_map = IndexMap.from_func(
index_map,
index_dtype=_get_sblock_default_dtype(self.get(block)),
)
_ffi_api.ScheduleTransformBlockLayout( # type: ignore # pylint: disable=no-member
self, block, index_map
)
def set_axis_separator(
self,
block: SBlockRV | str,
buffer: tuple[str, int] | str | Buffer,
axis_separators: list[int] | None,
) -> None:
"""Set the axis separator of a buffer, where the buffer is specified by a block and a read
or write index.
Parameters
----------
block : SBlockRV | str
The block that accesses the target buffer. If a string,
this must uniquely identify a block.
buffer: Union[Tuple[str,int], Buffer, str]
The buffer to be transformed, or a specification of how to
identify the buffer to be transformed.
If `buffer` if a tuple of ``(str,int)``, the first item
should be either "read" or "write", and the second item is
an index into the block's read or write regions.
If `buffer` is a string, it is the name of the buffer,
which must exist within the reads/writes of the block. In
addition, the reads/writes of the block may not contain
more than one buffer with this name.
If `buffer` is a Buffer object, it must exist within the
reads/writes of the block.
axis_separators : Optional[List[int]]
The axis separators.
Examples
--------
Before set_axis_separator, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_set_axis_separator(
A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
B = T.sblock_alloc_buffer((128, 128), dtype="float32")
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + 1.0
Create the schedule and do set_axis_separator:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_set_axis_separator)
sch.set_axis_separators(sch.get_sblock("B"), buffer=("write", 0),
axis_separators=[1])
print(sch.mod["main"].script())
After applying set_axis_separator, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_set_axis_separators(
A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
B = T.sblock_alloc_buffer([128, 128], dtype="float32", axis_separators=[1])
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * T.float32(2)
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + T.float32(1)
"""
axis_separators = axis_separators or []
block = self._normalize_block_arg(block)
buffer_index_type, buffer_index, _ = self._normalize_buffer_arg(block, buffer)
buffer_index_type_enum = 0 if buffer_index_type == "read" else 1
_ffi_api.ScheduleSetAxisSeparator( # type: ignore # pylint: disable=no-member
self, block, buffer_index, buffer_index_type_enum, axis_separators
)
########## Schedule: Padding decomposition #########
@type_checked
def decompose_padding(self, block: SBlockRV | str, loop: LoopRV) -> SBlockRV:
"""Decompose a block of padding computation pattern into two separate blocks.
a) The block which fill const pad values into full write region;
b) The block which fill in-bound values into region where pad predicate is true.
The pad value filling block is inserted right before the given loop.
The schedule primitive requires:
1) The input block is a complete block.
2) The input loop is the ancestor of the block.
3) The input block is a block which match padding pattern.
Parameters
----------
block : SBlockRV | str
The padding block to be decomposed.
loop : LoopRV
The loop above which the pad value filling block is inserted before.
Returns
-------
pad_value_block : SBlockRV
The block filling const pad values.
Examples
--------
Before decompose-padding, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_decompose(x: T.Buffer(128, "int32"), y: T.Buffer(140, "int32")):
for i in range(140):
with T.sblock("block"):
vi = T.axis.remap("S", [i])
y[vi] = T.if_then_else(vi >= 6 and vi < 134, x[vi - 6], 0, dtype="int32")
Create the schedule and do decompose-padding with specified loop:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_decompose, debug_mask="all")
block = sch.get_sblock("block")
sch.decompose_padding(block, sch.get_loops(block)[0])
print(sch.mod["main].script())
After applying decompose-padding, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_decompose(x: T.Buffer(128, "int32"), y: T.Buffer(140, "int32")):
for i in T.serial(140):
with T.sblock("block_pad_const"):
vi = T.axis.spatial(140, i)
y[vi] = 0
for i in T.serial(128):
with T.sblock("block"):
vi = T.axis.spatial(128, i)
y[vi + 6] = x[vi]
"""
block = self._normalize_block_arg(block)
return _ffi_api.ScheduleDecomposePadding( # type: ignore # pylint: disable=no-member
self, block, loop
)
@type_checked
def can_decompose_padding(self, block: SBlockRV | str, loop: LoopRV) -> bool:
"""Check whether the block match padding pattern and can be decomposed."""
# pylint: disable-next=no-member
return _ffi_api.CanDecomposePadding(self, block, loop) # type: ignore
@type_checked
def pad_einsum(self, block: SBlockRV | str, padding: list[int]) -> None:
"""Pad the computation of Einsum.
On a block with trivial binding, this primitive pads the iteration domain of the block by
the given padding factors, for example, 127 -> 128, 132 -> 144 when padding factor is 16.
Extra producer and consumer padding blocks will be generated to avoid out-of-bound buffer
access.
Einsum pattern means all the indices on the buffer access are either by constants
(e.g. B[0]) or by variables (e.g. B[i]), but not by composite expressions (e.g. B[i + 1]).
Parameters
----------
block : SBlockRV | str
The block that matches the Einsum pattern.
padding : List[int]
The padding for each block iter.
Examples
--------
Before applying pad-einsum, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_pad_einsum(
A: T.Buffer((127, 127), "float32"),
B: T.Buffer((127, 127), "float32"),
C: T.Buffer((127, 127), "float32"),
) -> None:
for i0, i1, i2 in T.grid(127, 127, 127):
with T.sblock("C_shared"):
i, j, k = T.axis.remap("SSR", [i0, i1, i2])
with T.init():
C[i, j] = T.float32(0)
C[i, j] = C[i, j] + A[i, k] * B[k, j]
Create the schedule and do pad-einsum with specified block:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_pad_einsum, debug_mask="all")
block = sch.get_sblock("C_shared")
sch.pad_einsum(block, [32, 32, 32])
print(sch.mod["main"].script())
After applying decompose-padding, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def main(
A: T.Buffer((127, 127), "float32"),
B: T.Buffer((127, 127), "float32"),
C: T.Buffer((127, 127), "float32"),
):
# with T.sblock("root"):
A_pad = T.sblock_alloc_buffer((128, 128))
B_pad = T.sblock_alloc_buffer((128, 128))
C_pad = T.sblock_alloc_buffer((128, 128))
for i0, i1 in T.grid(128, 128):
with T.sblock("A_pad"):
v0, v1 = T.axis.remap("SS", [i0, i1])
A_pad[v0, v1] = T.if_then_else(
v0 < 127 and v1 < 127,
A[v0, v1],
T.float32(0),
)
for i0, i1 in T.grid(128, 128):
with T.sblock("B_pad"):
v0, v1 = T.axis.remap("SS", [i0, i1])
B_pad[v0, v1] = T.if_then_else(
v0 < 127 and v1 < 127,
B[v0, v1],
T.float32(0),
)
for i0, i1, i2 in T.grid(128, 128, 128):
with T.sblock("C_shared"):
i, j, k = T.axis.remap("SSR", [i0, i1, i2])
with T.init():
C_pad[i, j] = T.float32(0)
C_pad[i, j] = C_pad[i, j] + A_pad[i, k] * B_pad[k, j]
for i0, i1 in T.grid(127, 127):
with T.sblock("C_pad"):
v0, v1 = T.axis.remap("SS", [i0, i1])
C[v0, v1] = C_pad[v0, v1]
"""
block = self._normalize_block_arg(block)
return _ffi_api.SchedulePadEinsum( # type: ignore # pylint: disable=no-member
self, block, padding
)
######## Schedule: Buffer transformation ########
@type_checked
def rolling_buffer(self, block: SBlockRV | str, write_buffer_index: int) -> None:
"""Compute the target buffer via rolling buffering, select the outermost rollable
axis with a positive bound overlap that appears in the block's ancestor loops
as `rolling axis`, fold and circularize the buffer along the rolling dimension,
append block predicate to avoid recomputing overlapping elements. It requires:
1) The block is not an output block and has only RAW dependencies.
2) The buffer to be an intermediate buffer defined via `alloc_buffer`.
3) The LCA of the producer and consumer of the buffer is a for loop, typically,
the producer and consumer of the buffer are cascaded through compute_at.
4) The access region of the buffer has at least one dimension that contains
a positive bound overlap.
Parameters
----------
block : SBlockRV | str
The producer block of the buffer.
write_buffer_index : int
The index of the buffer in block's write region.
Examples
--------
Before rolling_buffer, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_rolling_buffer(
A: T.Buffer((12, 12), "int8"), C: T.Buffer((8, 8), "int8")
) -> None:
# body
# with T.sblock("root")
B = T.sblock_alloc_buffer([10, 10], dtype="int8")
for i0, i1 in T.grid(2, 2):
for ax0, ax1, ax2, ax3 in T.grid(6, 6, 3, 3):
with T.sblock("B"):
ax0_1 = T.axis.spatial(10, i0 * 4 + ax0)
ax1_1 = T.axis.spatial(10, i1 * 4 + ax1)
rv0, rv1 = T.axis.remap("RR", [ax2, ax3])
B[ax0_1, ax1_1] = T.max(
B[ax0_1, ax1_1], A[ax0_1 + rv0, ax1_1 + rv1]
)
for ax0, ax1, ax2, ax3 in T.grid(4, 4, 3, 3):
with T.sblock("C"):
ax0_1 = T.axis.spatial(8, i0 * 4 + ax0)
ax1_1 = T.axis.spatial(8, i1 * 4 + ax1)
rv0, rv1 = T.axis.remap("RR", [ax2, ax3])
C[ax0_1, ax1_1] = T.max(
C[ax0_1, ax1_1], B[ax0_1 + rv0, ax1_1 + rv1]
)
Create the schedule and do rolling_buffer:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_rolling_buffer)
sch.rolling_buffer(sch.get_sblock("B"), write_buffer_index=0)
print(sch.mod["main"].script())
After applying rolling_buffer, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_rolling_buffer(
A: T.Buffer((12, 12), "int8"),
C: T.Buffer((8, 8), "int8")
) -> None:
# body
# with T.sblock("root")
B = T.sblock_alloc_buffer([6, 10], dtype="int8")
for i0, i1 in T.grid(2, 2):
for ax0, ax1, ax2, ax3 in T.grid(6, 6, 3, 3):
with T.sblock("B"):
T.where((i0 < 1 or 2 <= ax0) and (i1 < 1 or 2 <= ax1))
ax0_1 = T.axis.spatial(10, i0 * 4 + ax0)
ax1_1 = T.axis.spatial(10, i1 * 4 + ax1)
rv0, rv1 = T.axis.remap("RR", [ax2, ax3])
B[ax0_1 % 6, ax1_1] = T.max(
B[ax0_1 % 6, ax1_1], A[ax0_1 + rv0, ax1_1 + rv1]
)
for ax0, ax1, ax2, ax3 in T.grid(4, 4, 3, 3):
with T.sblock("C"):
ax0_1 = T.axis.spatial(8, i0 * 4 + ax0)
ax1_1 = T.axis.spatial(8, i1 * 4 + ax1)
rv0, rv1 = T.axis.remap("RR", [ax2, ax3])
C[ax0_1, ax1_1] = T.max(
C[ax0_1, ax1_1], B[ax0_1 % 6 + rv0, ax1_1 + rv1]
)
Note
----
The region_cover property of the consumer block of the target buffer will become false.
"""
block = self._normalize_block_arg(block)
# pylint: disable-next=no-member
return _ffi_api.ScheduleRollingBuffer(self, block, write_buffer_index) # type: ignore
########## Schedule: Misc ##########
@type_checked
def enter_postproc(self) -> None:
"""A no-op that marks the start of postprocessing phase of scheduling"""
_ffi_api.ScheduleEnterPostproc(self) # type: ignore # pylint: disable=no-member
@type_checked
def unsafe_hide_buffer_access(
self, block: SBlockRV, buf_type: str, buf_index_array: list[int]
) -> None:
"""Hide some buffer access in a given block. This is an unsafe schedule primitive.
Parameters
----------
block : SBlockRV
The block where we hide read access.
buf_type : str
The buffer type: "read"/"write".
buf_index_array : List[int]
The array of buffer indices we hide access.
Note
----
This schedule primitive is unsafe, and may fail dependency analysis.
One use case of `unsafe_hide_buffer_access` is to hide the buffer access
to indices buffers (e.g. in sparse computation) so that we can further tensorize
the block (the indices buffers appeared in read/write regions may fail the pattern
matching in `tensorize` primitive, and hide the access to these buffers could address
the issue).
"""
_ffi_api.ScheduleUnsafeHideBufferAccess( # type: ignore # pylint: disable=no-member
self,
block,
buf_type,
buf_index_array,
)
@type_checked
def annotate_buffer_access(
self, block: SBlockRV, buffer_index: int, buf_type: str, gen_new_ranges: Callable
) -> None:
"""Annotate the read or write region of a block
Parameters
----------
block : SBlockRV
The block to be annotated
buffer_index : int
The index of the buffer in block's read or write region
buf_type : str
The buffer type: "read" or "write"
gen_new_ranges : Callable
A function that takes the block's iter_vars and returns a
Tuple[Union[Expr, Tuple[Expr, Expr]], ...]
which defines the new read or write region for the buffer.
Each element in the tuple can be:
- A single Expr representing the iter_var itself
- A tuple of two PrimExprs representing the range (begin, end)
Examples
--------
Annotate a 2D read region for a buffer.
Before annotate_buffer_access, in TensorIR, the IR is:
.. code-block:: python
@T.prim_func(s_tir=True)
def before_annotate_buffer_access(
A: T.Buffer((128, 128), "float32"),
C: T.Buffer((128, 128), "float32")
) -> None:
B = T.sblock_alloc_buffer((128, 128), "float32")
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + 1.0
Create the schedule and do annotate_buffer_access:
.. code-block:: python
sch = tvm.s_tir.Schedule(before_annotate_buffer_access)
block = sch.get_sblock("B")
sch.annotate_buffer_access(block, 0, "read",
lambda vi, vj: ((vi - 1, vi + 1), (vj - 1, vj + 1)))
print(sch.mod["main"].script())
After applying annotate_buffer_access, the IR becomes:
.. code-block:: python
@T.prim_func(s_tir=True)
def after_annotate_buffer_access(
A: T.Buffer((128, 128), "float32"),
C: T.Buffer((128, 128), "float32")
) -> None:
B = T.sblock_alloc_buffer((128, 128), "float32")
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
T.reads(A[vi - 1:vi + 1, vj - 1:vj + 1])
T.writes(B[vi, vj])
T.sblock_attr({"explicit_read_region": 0})
B[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + 1.0
This annotates the read region for buffer A (index 0) in block "B" to be
[vi-1:vi+1, vj-1:vj+1] for each (vi, vj) in the block's iteration domain.
Note
----
This function allows manual specification of read or write regions, which
can be useful in cases where the compiler cannot accurately infer the
access pattern, such as complex data-dependent accesses.
It overrides the automatically inferred region for the specified buffer.
The function adds an annotation to the block, indicating that an explicit
region has been provided for the buffer at the given index. This annotation
is used in the CompactBufferAllocation pass to respect the manually specified
region instead of relying on automatic inference.
Caution should be exercised when using this function, as incorrect annotations
may lead to incorrect code generation or runtime errors. It's crucial to
ensure that the specified region covers all actual reads or writes performed
by the block for the given buffer.
"""
block_obj = self.get(block)
iter_vars = [x.var for x in block_obj.iter_vars]
new_ranges_spec = gen_new_ranges(*iter_vars)
if len(iter_vars) != len(new_ranges_spec):
raise ValueError(
f"Number of iter_vars ({len(iter_vars)}) must match "
f"number of new_ranges_spec ({len(new_ranges_spec)})"
)
result = []
for rng in new_ranges_spec:
if isinstance(rng, tuple | list):
if len(rng) != 2:
raise ValueError(
"Tuple must have exactly 2 elements to represent (begin, end)."
)
result.extend(rng)
elif is_prim_expr(rng):
result.extend([rng, rng + 1]) # Single point represented as (rng, rng + 1)
else:
raise TypeError(f"Expected Expr or tuple of Expr, got {type(rng)}")
# Create index_map using IndexMap constructor
index_map = IndexMap(
initial_indices=iter_vars,
final_indices=result,
inverse_index_map=None,
)
if buf_type == "read":
buffer_index_type = 0
elif buf_type == "write":
buffer_index_type = 1
else:
raise ValueError(f"Invalid buf_type: {buf_type}. Expected 'read' or 'write'.")
return _ffi_api.ScheduleAnnotateBufferAccess(
self, block, buffer_index, buffer_index_type, index_map
)
Reference in New Issue View Git Blame Copy Permalink