wehub-resource-sync
1094 lines
31 KiB
Python
1094 lines
31 KiB
Python
# 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.
|
|
"""Statement AST Node in TVM.
|
|
|
|
Each statement node have subfields that can be visited from python side.
|
|
|
|
.. code-block:: python
|
|
|
|
x = tvm.tirx.Var("n", "int32")
|
|
buffer = tvm.tirx.decl_buffer((16,), "float32")
|
|
st = tvm.tirx.stmt.BufferStore(buffer, 1, (x,))
|
|
assert isinstance(st, tvm.tirx.stmt.BufferStore)
|
|
assert(st.buffer == buffer)
|
|
"""
|
|
|
|
from collections.abc import Mapping
|
|
from enum import IntEnum
|
|
from typing import TYPE_CHECKING, Any, ClassVar
|
|
|
|
import tvm_ffi
|
|
|
|
from tvm.ir import Expr, Op, Range, Span, is_prim_expr
|
|
from tvm.runtime import Object, Scriptable, const
|
|
from tvm.tirx import FloatImm, IntImm
|
|
|
|
from . import _ffi_api
|
|
from .buffer import Buffer
|
|
from .exec_scope import ExecScope, ScopeIdDef
|
|
from .expr import IterVar, StringImm, Var
|
|
|
|
if TYPE_CHECKING:
|
|
from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Stmt")
|
|
class Stmt(Object, Scriptable):
|
|
"""Base class of all the statements."""
|
|
|
|
|
|
def _normalize_legacy_stmt(stmt: Stmt | None) -> Stmt | None:
|
|
"""Expand legacy body-carrying leaf stmt wrappers into SeqStmt form.
|
|
|
|
Legacy python compatibility may attach a `body` attribute to leaf statements
|
|
(Bind/DeclBuffer/AllocBuffer). This helper converts such wrappers to the new
|
|
leaf + SeqStmt representation when embedding inside another statement node.
|
|
"""
|
|
|
|
if stmt is None:
|
|
return None
|
|
|
|
prefix: list[Stmt] = []
|
|
cur = stmt
|
|
while True:
|
|
if isinstance(cur, DeclBuffer) and hasattr(cur, "body"):
|
|
prefix.append(DeclBuffer(cur.buffer, cur.span))
|
|
cur = cur.body
|
|
continue
|
|
if isinstance(cur, AllocBuffer) and hasattr(cur, "body"):
|
|
prefix.append(AllocBuffer(cur.buffer, cur.annotations, cur.span))
|
|
cur = cur.body
|
|
continue
|
|
break
|
|
|
|
if not prefix:
|
|
return stmt
|
|
|
|
normalized_tail = _normalize_legacy_stmt(cur)
|
|
if normalized_tail is not None:
|
|
prefix.append(normalized_tail)
|
|
if len(prefix) == 1:
|
|
return prefix[0]
|
|
return SeqStmt(prefix)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Bind")
|
|
class Bind(Stmt):
|
|
"""Bind node.
|
|
|
|
Bind a variable to a value in the enclosing scope.
|
|
Bind has no body field.
|
|
The bound variable is visible in all subsequent statements
|
|
within the same enclosing scope (SeqStmt, ForNode.body, etc.).
|
|
|
|
Parameters
|
|
----------
|
|
var : Var
|
|
The variable in the binding.
|
|
|
|
value : Expr
|
|
The value to be bound.
|
|
|
|
span : Optional[Span]
|
|
The location of the stmt in the source code.
|
|
"""
|
|
|
|
var: Var
|
|
value: Expr
|
|
span: Span | None
|
|
|
|
def __init__(self, var: Var, value: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.Bind,
|
|
var,
|
|
value,
|
|
span, # type: ignore
|
|
)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.AssertStmt")
|
|
class AssertStmt(Stmt):
|
|
"""AssertStmt node.
|
|
|
|
Parameters
|
|
----------
|
|
kind : StringImm
|
|
The error kind, e.g. "RuntimeError", "TypeError", "ValueError".
|
|
|
|
condition : Expr
|
|
The assert condition.
|
|
|
|
message_parts : list[StringImm]
|
|
Error message fragments, concatenated at runtime when assertion fails.
|
|
|
|
span : Span | None
|
|
The location of the stmt in the source code.
|
|
"""
|
|
|
|
kind: StringImm
|
|
condition: Expr
|
|
message_parts: list
|
|
span: Span | None
|
|
|
|
def __init__(
|
|
self,
|
|
kind: StringImm,
|
|
condition: Expr,
|
|
message_parts: list | None = None,
|
|
span: Span | None = None,
|
|
) -> None:
|
|
if message_parts is None:
|
|
message_parts = []
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.AssertStmt,
|
|
kind,
|
|
condition,
|
|
message_parts,
|
|
span, # type: ignore
|
|
)
|
|
|
|
|
|
class ForKind(IntEnum):
|
|
"""The kind of the for loop.
|
|
|
|
note
|
|
----
|
|
ForKind can change the control flow semantics
|
|
of the loop and need to be considered in all TIR passes.
|
|
"""
|
|
|
|
SERIAL = 0
|
|
PARALLEL = 1
|
|
VECTORIZED = 2
|
|
UNROLLED = 3
|
|
THREAD_BINDING = 4 # pylint: disable=invalid-name
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.For")
|
|
class For(Stmt):
|
|
"""For node.
|
|
|
|
Parameters
|
|
----------
|
|
loop_var : Var
|
|
The loop variable.
|
|
|
|
min : Expr
|
|
The beginning value.
|
|
|
|
extent : Expr
|
|
The length of the loop.
|
|
|
|
kind : ForKind
|
|
The type of the for.
|
|
|
|
body : Stmt
|
|
The body statement.
|
|
|
|
thread_binding: Optional[tirx.IterVar]
|
|
The thread this loop binds to. Only valid
|
|
if kind is ThreadBinding
|
|
|
|
step : Expr
|
|
The loop step. Default to none which
|
|
represent one.
|
|
|
|
annotations: Optional[Mapping[str, Object]]
|
|
Additional annotation hints.
|
|
|
|
span : Optional[Span]
|
|
The location of the stmt in the source code.
|
|
"""
|
|
|
|
loop_var: Var
|
|
min: Expr
|
|
extent: Expr
|
|
kind: ForKind
|
|
body: Stmt
|
|
thread_binding: IterVar | None
|
|
annotations: Mapping[str, Object]
|
|
step: Expr | None
|
|
span: Span | None
|
|
|
|
def __init__(
|
|
self,
|
|
loop_var: Var,
|
|
min: Expr, # pylint: disable=redefined-builtin
|
|
extent: Expr,
|
|
kind: ForKind,
|
|
body: Stmt,
|
|
thread_binding: IterVar | None = None,
|
|
annotations: Mapping[str, Object] | None = None,
|
|
step: Expr | None = None,
|
|
span: Span | None = None,
|
|
) -> None:
|
|
body = _normalize_legacy_stmt(body)
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.For, # type: ignore
|
|
loop_var,
|
|
min,
|
|
extent,
|
|
kind,
|
|
body,
|
|
thread_binding,
|
|
annotations,
|
|
step,
|
|
span,
|
|
)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.While")
|
|
class While(Stmt):
|
|
"""While node.
|
|
|
|
Parameters
|
|
----------
|
|
condition : Expr
|
|
The termination condition.
|
|
|
|
body : Stmt
|
|
The body statement.
|
|
|
|
span : Optional[Span]
|
|
The location of the stmt in the source code.
|
|
"""
|
|
|
|
condition: Expr
|
|
body: Stmt
|
|
span: Span | None
|
|
|
|
def __init__(self, condition: Expr, body: Stmt, span: Span | None = None) -> None:
|
|
body = _normalize_legacy_stmt(body)
|
|
self.__init_handle_by_constructor__(_ffi_api.While, condition, body, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.BufferStore")
|
|
class BufferStore(Stmt):
|
|
"""Buffer store node.
|
|
|
|
Parameters
|
|
----------
|
|
buffer : Buffer
|
|
The buffer.
|
|
|
|
value : Expr
|
|
The value we to be stored.
|
|
|
|
indices : List[Expr]
|
|
The indices location to be stored.
|
|
|
|
predicate : Optional[Expr]
|
|
A vector mask of boolean values indicating which lanes of a vector are to be
|
|
stored. The number lanes of the mask must be equal to the number of lanes in
|
|
value.
|
|
|
|
span : Optional[Span]
|
|
The location of the stmt in the source code.
|
|
"""
|
|
|
|
buffer: Buffer
|
|
value: Expr
|
|
indices: list[Expr]
|
|
predicate: Expr | None
|
|
span: Span | None
|
|
|
|
def __init__(
|
|
self,
|
|
buffer: Buffer,
|
|
value: Expr,
|
|
indices: list[Expr],
|
|
predicate: Expr | None = None,
|
|
span: Span | None = None,
|
|
) -> None:
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.BufferStore,
|
|
buffer,
|
|
value,
|
|
indices,
|
|
predicate,
|
|
span, # type: ignore
|
|
)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.AllocBuffer")
|
|
class AllocBuffer(Stmt):
|
|
"""AllocBuffer node.
|
|
|
|
Allocates a buffer and declares it in scope.
|
|
|
|
Parameters
|
|
----------
|
|
buffer: Buffer
|
|
The buffer being allocated and declared.
|
|
|
|
annotations: Optional[dict]
|
|
Additional annotations about the allocation.
|
|
|
|
span: Optional[Span]
|
|
The location of this AllocBuffer in the source code.
|
|
"""
|
|
|
|
buffer: Buffer
|
|
span: Span | None
|
|
|
|
def __init__(self, buffer: Buffer, *args, **kwargs) -> None:
|
|
body: Stmt | None = None
|
|
annotations: dict | None = None
|
|
span: Span | None = None
|
|
|
|
idx = 0
|
|
argc = len(args)
|
|
|
|
# Legacy form: AllocBuffer(buffer, body[, annotations][, span])
|
|
if idx < argc and isinstance(args[idx], Stmt):
|
|
body = args[idx]
|
|
idx += 1
|
|
|
|
if idx < argc:
|
|
arg = args[idx]
|
|
if isinstance(arg, Mapping):
|
|
annotations = dict(arg)
|
|
idx += 1
|
|
elif arg is None:
|
|
annotations = None
|
|
idx += 1
|
|
elif isinstance(arg, Span):
|
|
span = arg
|
|
idx += 1
|
|
else:
|
|
raise TypeError(
|
|
"AllocBuffer expects (buffer[, annotations][, span]) or "
|
|
"legacy (buffer, body[, annotations][, span])"
|
|
)
|
|
|
|
if idx < argc:
|
|
arg = args[idx]
|
|
if arg is None or isinstance(arg, Span):
|
|
span = arg
|
|
idx += 1
|
|
else:
|
|
raise TypeError("AllocBuffer span must be a Span or None")
|
|
|
|
if idx != argc:
|
|
raise TypeError(
|
|
"AllocBuffer expects (buffer[, annotations][, span]) or "
|
|
"legacy (buffer, body[, annotations][, span])"
|
|
)
|
|
|
|
if kwargs:
|
|
invalid_keys = set(kwargs.keys()) - {"body", "annotations", "span"}
|
|
if invalid_keys:
|
|
raise TypeError(f"Unexpected keyword arguments for AllocBuffer: {invalid_keys}")
|
|
if "body" in kwargs:
|
|
kw_body = kwargs["body"]
|
|
if kw_body is not None and not isinstance(kw_body, Stmt):
|
|
raise TypeError("AllocBuffer body must be a Stmt or None")
|
|
if body is not None and kw_body is not None and body is not kw_body:
|
|
raise TypeError("AllocBuffer body specified by both args and kwargs")
|
|
body = kw_body if kw_body is not None else body
|
|
if "annotations" in kwargs:
|
|
kw_ann = kwargs["annotations"]
|
|
if kw_ann is not None and not isinstance(kw_ann, Mapping):
|
|
raise TypeError("AllocBuffer annotations must be Mapping or None")
|
|
if annotations is not None and kw_ann is not None and annotations != dict(kw_ann):
|
|
raise TypeError("AllocBuffer annotations specified by both args and kwargs")
|
|
annotations = dict(kw_ann) if kw_ann is not None else annotations
|
|
if "span" in kwargs:
|
|
kw_span = kwargs["span"]
|
|
if kw_span is not None and not isinstance(kw_span, Span):
|
|
raise TypeError("AllocBuffer span must be a Span or None")
|
|
if span is not None and kw_span is not None and span is not kw_span:
|
|
raise TypeError("AllocBuffer span specified by both args and kwargs")
|
|
span = kw_span if kw_span is not None else span
|
|
|
|
self.__init_handle_by_constructor__(_ffi_api.AllocBuffer, buffer, annotations, span)
|
|
# Legacy compatibility. Body is carried on python side only.
|
|
if body is not None:
|
|
self.body = body
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.DeclBuffer")
|
|
class DeclBuffer(Stmt):
|
|
"""DeclBuffer node.
|
|
|
|
Parameters
|
|
----------
|
|
buffer: Buffer
|
|
The buffer being declared.
|
|
|
|
span: Optional[Span]
|
|
The location of this DeclBuffer in the source code.
|
|
"""
|
|
|
|
buffer: Buffer
|
|
span: Span | None
|
|
|
|
def __init__(self, buffer: Buffer, *args, **kwargs) -> None:
|
|
body: Stmt | None = None
|
|
span: Span | None = None
|
|
|
|
if len(args) == 1:
|
|
arg0 = args[0]
|
|
if isinstance(arg0, Stmt):
|
|
body = arg0
|
|
elif arg0 is None or isinstance(arg0, Span):
|
|
span = arg0
|
|
else:
|
|
raise TypeError(
|
|
"DeclBuffer expects (buffer[, span]) or legacy (buffer, body[, span])"
|
|
)
|
|
elif len(args) == 2:
|
|
body, span = args
|
|
if body is not None and not isinstance(body, Stmt):
|
|
raise TypeError("Legacy DeclBuffer body must be a Stmt or None")
|
|
if span is not None and not isinstance(span, Span):
|
|
raise TypeError("DeclBuffer span must be a Span or None")
|
|
elif len(args) > 2:
|
|
raise TypeError("DeclBuffer expects (buffer[, span]) or legacy (buffer, body[, span])")
|
|
|
|
if kwargs:
|
|
invalid_keys = set(kwargs.keys()) - {"body", "span"}
|
|
if invalid_keys:
|
|
raise TypeError(f"Unexpected keyword arguments for DeclBuffer: {invalid_keys}")
|
|
if "body" in kwargs:
|
|
kw_body = kwargs["body"]
|
|
if kw_body is not None and not isinstance(kw_body, Stmt):
|
|
raise TypeError("DeclBuffer body must be a Stmt or None")
|
|
if body is not None and kw_body is not None and body is not kw_body:
|
|
raise TypeError("DeclBuffer body specified by both args and kwargs")
|
|
body = kw_body if kw_body is not None else body
|
|
if "span" in kwargs:
|
|
kw_span = kwargs["span"]
|
|
if kw_span is not None and not isinstance(kw_span, Span):
|
|
raise TypeError("DeclBuffer span must be a Span or None")
|
|
if span is not None and kw_span is not None and span is not kw_span:
|
|
raise TypeError("DeclBuffer span specified by both args and kwargs")
|
|
span = kw_span if kw_span is not None else span
|
|
|
|
self.__init_handle_by_constructor__(_ffi_api.DeclBuffer, buffer, span)
|
|
# Legacy compatibility. Body is carried on python side only.
|
|
if body is not None:
|
|
self.body = body
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.AttrStmt")
|
|
class AttrStmt(Stmt):
|
|
"""AttrStmt node.
|
|
|
|
Parameters
|
|
----------
|
|
node : Object
|
|
The node to annotate the attribute
|
|
|
|
attr_key : str
|
|
Attribute type key.
|
|
|
|
value : Expr
|
|
The value of the attribute
|
|
|
|
body : Stmt
|
|
The body statement.
|
|
|
|
span : Optional[Span]
|
|
The location of the stmt in the source code.
|
|
"""
|
|
|
|
node: Object
|
|
attr_key: str
|
|
value: Expr
|
|
body: Stmt
|
|
span: Span | None
|
|
|
|
def __init__(
|
|
self, node: Object, attr_key: str, value: Expr, body: Stmt, span: Span | None = None
|
|
) -> None:
|
|
body = _normalize_legacy_stmt(body)
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.AttrStmt,
|
|
node,
|
|
attr_key,
|
|
value,
|
|
body,
|
|
span, # type: ignore
|
|
)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.SeqStmt")
|
|
class SeqStmt(Stmt):
|
|
"""Sequence of statements.
|
|
|
|
Parameters
|
|
----------
|
|
seq : List[Stmt]
|
|
The statements
|
|
|
|
span : Optional[Span]
|
|
The location of the stmt in the source code.
|
|
"""
|
|
|
|
seq: list[Stmt]
|
|
span: Span | None
|
|
|
|
def __init__(self, seq: list[Stmt], span: Span | None = None) -> None:
|
|
seq = [_normalize_legacy_stmt(s) for s in seq]
|
|
self.__init_handle_by_constructor__(_ffi_api.SeqStmt, seq, span) # type: ignore
|
|
|
|
def __getitem__(self, i: int):
|
|
return self.seq[i]
|
|
|
|
def __len__(self):
|
|
return len(self.seq)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.IfThenElse")
|
|
class IfThenElse(Stmt):
|
|
"""IfThenElse node.
|
|
|
|
Parameters
|
|
----------
|
|
condition : Expr
|
|
The expression
|
|
|
|
then_case : Stmt
|
|
The statement to execute if condition is true.
|
|
|
|
else_case : Optional[Stmt]
|
|
The statement to execute if condition is false.
|
|
|
|
span : Optional[Span]
|
|
The location of the stmt in the source code.
|
|
"""
|
|
|
|
condition: Expr
|
|
then_case: Stmt
|
|
else_case: Stmt | None
|
|
|
|
def __init__(
|
|
self, condition: Expr, then_case: Stmt, else_case: Stmt | None, span: Span | None = None
|
|
) -> None:
|
|
then_case = _normalize_legacy_stmt(then_case)
|
|
else_case = _normalize_legacy_stmt(else_case)
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.IfThenElse,
|
|
condition,
|
|
then_case,
|
|
else_case,
|
|
span, # type: ignore
|
|
)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Evaluate")
|
|
class Evaluate(Stmt):
|
|
"""Evaluate node.
|
|
|
|
Parameters
|
|
----------
|
|
value : Expr
|
|
The expression to be evaluated.
|
|
|
|
span : Optional[Span]
|
|
The location of the stmt in the source code.
|
|
"""
|
|
|
|
value: Expr
|
|
span: Span | None
|
|
|
|
def __init__(self, value: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.Evaluate, value, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.BufferRegion")
|
|
class BufferRegion(Object, Scriptable):
|
|
"""BufferRegion node.
|
|
|
|
Parameters
|
|
----------
|
|
buffer : Buffer
|
|
The buffer of the buffer region
|
|
|
|
region : List[Range]
|
|
The region array of the buffer region
|
|
"""
|
|
|
|
buffer: Buffer
|
|
region: list[Range]
|
|
|
|
def __init__(self, buffer: Buffer, region: list[Range]) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.BufferRegion, buffer, region) # type: ignore
|
|
|
|
def __getitem__(self, indices):
|
|
from ..arith import Analyzer
|
|
|
|
if not isinstance(indices, tuple | list):
|
|
indices = [indices]
|
|
|
|
has_step = any(
|
|
isinstance(i, slice) and (i.step is not None and i.step != 1) for i in indices
|
|
)
|
|
if has_step:
|
|
raise ValueError("BufferRegion slicing does not support steps")
|
|
|
|
analyzer = Analyzer()
|
|
new_region = []
|
|
for i, index in enumerate(indices):
|
|
old_range = self.region[i]
|
|
if isinstance(index, slice):
|
|
start = 0 if index.start is None else index.start
|
|
stop = old_range.extent if index.stop is None else index.stop
|
|
new_min = old_range.min + start
|
|
new_extent = analyzer.simplify(stop - start)
|
|
new_region.append(Range.from_min_extent(new_min, new_extent))
|
|
else:
|
|
new_min = old_range.min + index
|
|
new_region.append(
|
|
Range.from_min_extent(
|
|
new_min, IntImm(index.ty, 1) if is_prim_expr(index) else 1
|
|
)
|
|
)
|
|
# Fill remaining dimensions with their original ranges
|
|
for i in range(len(indices), len(self.region)):
|
|
new_region.append(self.region[i])
|
|
return BufferRegion(self.buffer, new_region)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.MatchBufferRegion")
|
|
class MatchBufferRegion(Object, Scriptable):
|
|
"""MatchBufferRegion node.
|
|
|
|
Parameters
|
|
----------
|
|
buffer : Buffer
|
|
The target buffer
|
|
|
|
source : BufferRegion
|
|
The region of source buffer
|
|
"""
|
|
|
|
buffer: Buffer
|
|
source: BufferRegion
|
|
|
|
def __init__(self, buffer: Buffer, source: BufferRegion) -> None:
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.MatchBufferRegion,
|
|
buffer,
|
|
source, # type: ignore
|
|
)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.SBlock")
|
|
class SBlock(Stmt):
|
|
"""SBlock node.
|
|
|
|
Parameters
|
|
----------
|
|
iter_vars : List[IterVar]
|
|
The block Variable.
|
|
|
|
reads : List[BufferRegion]
|
|
The read buffer regions of the block.
|
|
|
|
writes: List[BufferRegion]
|
|
The write buffer regions of the block.
|
|
|
|
name_hint: str
|
|
the name_hint of the block.
|
|
|
|
body: Stmt
|
|
The body of the block.
|
|
|
|
init: Optional[Stmt]
|
|
The init block of the reduction block
|
|
|
|
alloc_buffers: Optional[list[Buffer]]
|
|
The buffer allocations
|
|
|
|
match_buffers: Optional[List[MatchBufferRegion]]
|
|
The subregion buffer match
|
|
|
|
annotations: Optional[Mapping[str, Object]]
|
|
Additional annotation hints.
|
|
|
|
span : Optional[Span]
|
|
The location of this block in the source code.
|
|
"""
|
|
|
|
iter_vars: list[IterVar]
|
|
reads: list[BufferRegion]
|
|
writes: list[BufferRegion]
|
|
name_hint: str
|
|
body: Stmt
|
|
init: Stmt | None
|
|
alloc_buffers: list[Buffer]
|
|
match_buffers: list[MatchBufferRegion]
|
|
annotations: Mapping[str, Object]
|
|
span: Span | None
|
|
|
|
def __init__(
|
|
self,
|
|
iter_vars: list[IterVar],
|
|
reads: list[BufferRegion],
|
|
writes: list[BufferRegion],
|
|
name_hint: str,
|
|
body: Stmt,
|
|
init: Stmt | None = None,
|
|
alloc_buffers: list[Buffer] | None = None,
|
|
match_buffers: list[MatchBufferRegion] | None = None,
|
|
annotations: Mapping[str, Object] | None = None,
|
|
span: Span | None = None,
|
|
) -> None:
|
|
if alloc_buffers is None:
|
|
alloc_buffers = []
|
|
if match_buffers is None:
|
|
match_buffers = []
|
|
if annotations is None:
|
|
annotations = {}
|
|
body = _normalize_legacy_stmt(body)
|
|
init = _normalize_legacy_stmt(init)
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.SBlock, # type: ignore
|
|
iter_vars,
|
|
reads,
|
|
writes,
|
|
name_hint,
|
|
body,
|
|
init,
|
|
alloc_buffers,
|
|
match_buffers,
|
|
annotations,
|
|
span,
|
|
) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.SBlockRealize")
|
|
class SBlockRealize(Stmt):
|
|
"""SBlockRealize node.
|
|
|
|
Parameters
|
|
----------
|
|
iter_values : List[Expr]
|
|
The binding values of the block var.
|
|
|
|
predicate : Union[Expr, bool]
|
|
The predicate of the block.
|
|
|
|
block : SBlock
|
|
The block to realize
|
|
|
|
span : Optional[Span]
|
|
The location of this block_realize in the source code.
|
|
"""
|
|
|
|
iter_values: list[Expr]
|
|
predicate: Expr
|
|
block: SBlock
|
|
span: Span | None
|
|
|
|
def __init__(
|
|
self,
|
|
iter_values: list[Expr],
|
|
predicate: Expr | bool,
|
|
block: SBlock,
|
|
span: Span | None = None,
|
|
) -> None:
|
|
if isinstance(predicate, bool):
|
|
predicate = const(predicate, "bool")
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.SBlockRealize, # type: ignore
|
|
iter_values,
|
|
predicate,
|
|
block,
|
|
span,
|
|
) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.ScopeIdDefStmt")
|
|
class ScopeIdDefStmt(Stmt):
|
|
"""ScopeIdDefStmt node.
|
|
|
|
Leaf statement that introduces scope-identifier vars
|
|
(``wg_id = Tx.warpgroup_id([N])``, ``warp_id = Tx.warp_id_in_wg([4])``,
|
|
``lane_id = Tx.lane_id([32])``, …) at the kernel-body top level. The
|
|
underlying ``ScopeIdDef`` carries the def vars, their extents, and
|
|
the parent/child scope binding.
|
|
|
|
Note: the C++ field is named ``def`` (a Python keyword). Access it
|
|
via ``getattr(stmt, "def")`` or ``stmt.__getattribute__("def")`` —
|
|
the type-annotation alias here is purely for documentation.
|
|
|
|
Parameters
|
|
----------
|
|
def_ : ScopeIdDef
|
|
The scope-id definition (def vars, extents, scope binding).
|
|
|
|
span : Optional[Span]
|
|
The location of this statement in the source code.
|
|
"""
|
|
|
|
span: Span | None
|
|
|
|
def __init__(self, def_: ScopeIdDef, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.ScopeIdDefStmt, # type: ignore
|
|
def_,
|
|
span,
|
|
) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Break")
|
|
class Break(Stmt):
|
|
"""Break node.
|
|
|
|
Parameters
|
|
----------
|
|
"""
|
|
|
|
def __init__(self, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.Break, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Continue")
|
|
class Continue(Stmt):
|
|
"""Continue node.
|
|
|
|
Parameters
|
|
----------
|
|
"""
|
|
|
|
def __init__(self, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.Continue, span) # type: ignore
|
|
|
|
|
|
def stmt_seq(*args: Expr | Stmt) -> SeqStmt:
|
|
"""Make sequence of statements
|
|
|
|
Parameters
|
|
----------
|
|
*args : Union[Expr, Stmt]
|
|
List of statements to be combined as sequence.
|
|
|
|
Returns
|
|
-------
|
|
stmt : Stmt
|
|
The combined statement.
|
|
"""
|
|
ret = []
|
|
for value in args:
|
|
if not isinstance(value, Stmt):
|
|
value = Evaluate(value)
|
|
ret.append(value)
|
|
if len(ret) == 1:
|
|
return ret[0]
|
|
return SeqStmt(ret)
|
|
|
|
|
|
def stmt_list(stmt: Stmt) -> list[Stmt]:
|
|
"""Make list of stmt from blocks.
|
|
|
|
Parameters
|
|
----------
|
|
stmt : Stmt
|
|
The input statement.
|
|
|
|
Returns
|
|
-------
|
|
stmt_list : List[Stmt]
|
|
The unpacked list of statements
|
|
"""
|
|
if isinstance(stmt, SeqStmt):
|
|
res = []
|
|
for x in stmt:
|
|
res += stmt_list(x)
|
|
return res
|
|
return [stmt]
|
|
|
|
|
|
def normalize_const_arg(arg) -> Expr:
|
|
if isinstance(arg, float):
|
|
return FloatImm("float32", arg)
|
|
return arg
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.TilePrimitiveCall")
|
|
class TilePrimitiveCall(Stmt):
|
|
"""TilePrimitiveCall node.
|
|
|
|
Parameters
|
|
----------
|
|
op : Op
|
|
The operator.
|
|
|
|
args : List[Expr]
|
|
The arguments.
|
|
|
|
workspace : Map[str, Buffer]
|
|
The workspace.
|
|
|
|
config : Map[str, ObjectRef]
|
|
The scheduler/config dictionary.
|
|
|
|
dispatch : Optional[str]
|
|
The explicit variant name to dispatch to.
|
|
|
|
scope : ExecScope
|
|
The cooperation scope of this call. Defaults to ``thread`` (an unscoped call).
|
|
"""
|
|
|
|
args: list[Expr]
|
|
workspace: dict[str, Buffer]
|
|
config: dict[str, Any]
|
|
dispatch: str | None
|
|
scope: ExecScope
|
|
_registry: ClassVar[dict[Op, type["TilePrimitiveCall"]]] = {}
|
|
|
|
def __init__(
|
|
self,
|
|
*args: list[Expr],
|
|
op: Op | None = None,
|
|
workspace: dict[str, Buffer] | None = None,
|
|
config: dict[str, Any] | None = None,
|
|
dispatch: str | None = None,
|
|
scope: ExecScope | None = None,
|
|
) -> None:
|
|
if workspace is None:
|
|
workspace = {}
|
|
if config is None:
|
|
config = {}
|
|
if scope is None:
|
|
scope = ExecScope("thread")
|
|
if op is None:
|
|
assert self.__class__ != TilePrimitiveCall, (
|
|
"Directly instantiating TilePrimitiveCall needs to specify the op"
|
|
)
|
|
op = self.__class__.op
|
|
args = list(map(normalize_const_arg, args))
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.TilePrimitiveCall,
|
|
op,
|
|
args,
|
|
workspace,
|
|
config,
|
|
dispatch,
|
|
scope, # pylint: disable=no-member
|
|
)
|
|
|
|
def __init_subclass__(cls, **kwargs):
|
|
super().__init_subclass__(**kwargs)
|
|
if hasattr(cls, "op"):
|
|
cls._registry[cls.op] = cls
|
|
|
|
@classmethod
|
|
def downcast(cls, instance: "TilePrimitiveCall") -> "TilePrimitiveCall":
|
|
subclass = cls._registry.get(instance.op)
|
|
if subclass is None:
|
|
return instance # Unknown op: return as-is
|
|
new_instance = subclass.__new__(subclass)
|
|
new_instance.__init_handle_by_constructor__(
|
|
_ffi_api.TilePrimitiveCallCopyHandle,
|
|
instance, # pylint: disable=no-member
|
|
)
|
|
return new_instance
|
|
|
|
def replace(self, **changes: Any) -> "TilePrimitiveCall":
|
|
"""Return a copy of this call with selected fields replaced.
|
|
|
|
Every field that is not overridden in ``changes`` is preserved from
|
|
``self`` (including ``scope``), so rebuilds never silently drop fields.
|
|
The returned node is downcast to the registered subclass for ``op``.
|
|
|
|
Parameters
|
|
----------
|
|
**changes : Any
|
|
Field overrides; any of ``op``, ``args``, ``workspace``, ``config``,
|
|
``dispatch``, ``scope``.
|
|
|
|
Returns
|
|
-------
|
|
new_call : TilePrimitiveCall
|
|
A new call with the requested fields replaced.
|
|
"""
|
|
unknown = set(changes) - {"op", "args", "workspace", "config", "dispatch", "scope"}
|
|
if unknown:
|
|
raise TypeError(f"Unknown field(s) for TilePrimitiveCall.replace: {sorted(unknown)}")
|
|
new_call = TilePrimitiveCall(
|
|
*changes.get("args", self.args),
|
|
op=changes.get("op", self.op),
|
|
workspace=changes.get("workspace", self.workspace),
|
|
config=changes.get("config", self.config),
|
|
dispatch=changes.get("dispatch", self.dispatch),
|
|
scope=changes.get("scope", self.scope),
|
|
)
|
|
return TilePrimitiveCall.downcast(new_call)
|
|
|
|
def with_workspace(self, workspace: dict[str, Buffer]) -> "TilePrimitiveCall":
|
|
"""Return a copy with ``workspace`` replaced, preserving all other fields."""
|
|
return self.replace(workspace=workspace)
|
|
|
|
@property
|
|
def srcs(self) -> list[Expr]:
|
|
raise NotImplementedError("Subclass must implement this method")
|
|
|
|
@property
|
|
def dsts(self) -> list[Expr]:
|
|
raise NotImplementedError("Subclass must implement this method")
|
|
|
|
def get_private_buffers(
|
|
self, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: "DispatchContext"
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Create private (intermediate) buffers needed in this operator.
|
|
|
|
Parameters
|
|
----------
|
|
buffer_dict: Dict[Any, Tuple[Buffer, Optional[Stmt]]]
|
|
A dictionary containing private buffers (and their init stmts) in other operators.
|
|
Key can be anything to reference the buffer.
|
|
This is used to reuse private buffers in other operators (like identity tensor etc.).
|
|
If the buffer is not found in the buffer_dict, it will be created and added to
|
|
the buffer_dict.
|
|
If the buffer is found in the buffer_dict but smaller than required, it will be
|
|
enlarged and updated.
|
|
|
|
sctx: DispatchContext
|
|
The dispatch context.
|
|
This is used to get the target and reuse op dispatch implementations.
|
|
|
|
Returns:
|
|
private_buffer_refs: Dict[str, Any]
|
|
The references to private buffers created in this operator.
|
|
Key will be the name to add into workspace.
|
|
private buffer can be accessed by buffer_dict[private_buffer_refs[name]]
|
|
"""
|
|
if sctx.target.kind.name == "trn":
|
|
return self.get_private_buffers_trn(buffer_dict, sctx)
|
|
elif sctx.target.kind.name == "cuda":
|
|
return self.get_private_buffers_cuda(buffer_dict, sctx)
|
|
else:
|
|
raise ValueError(f"Unsupported target: {sctx.target.kind.name}")
|
|
|
|
def get_private_buffers_trn(
|
|
self, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: "DispatchContext"
|
|
) -> dict[str, Any]:
|
|
return {}
|
|
|
|
def get_private_buffers_cuda(
|
|
self, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: "DispatchContext"
|
|
) -> dict[str, Any]:
|
|
return {}
|
|
|
|
def validate(self) -> None:
|
|
pass
|