wehub-resource-sync
1384 lines
35 KiB
Python
1384 lines
35 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.
|
|
# pylint: disable=redefined-builtin
|
|
"""TIR expression nodes.
|
|
|
|
Each expression node have subfields that can be visited from python side.
|
|
For example, you can use addexp.a to get the left operand of an Add node.
|
|
|
|
.. code-block:: python
|
|
|
|
x = tvm.tirx.Var("n", "int32")
|
|
y = x + 2
|
|
assert(isinstance(y, tvm.tirx.Add))
|
|
assert(y.a == x)
|
|
"""
|
|
|
|
import tvm_ffi
|
|
|
|
import tvm.ir._ffi_api
|
|
import tvm.ir._overload_prim_expr as _overload_prim_expr
|
|
from tvm import ir
|
|
from tvm.ir import Expr
|
|
from tvm.ir.base import Span
|
|
from tvm.runtime import DataTypeCode, Object, ObjectConvertible, Scriptable, const
|
|
|
|
from . import _ffi_api
|
|
from .buffer import Buffer, DataProducer
|
|
|
|
|
|
def convert(expr) -> Expr:
|
|
return _ffi_api.convert(expr)
|
|
|
|
|
|
def div_ambiguity_error() -> RuntimeError:
|
|
return RuntimeError(
|
|
"TVM supports multiple types of integer divisions, "
|
|
+ "please call div, indexdiv/indexmod, floordiv/floormod "
|
|
+ " or truncdiv/truncmod directly to avoid ambiguity in the code."
|
|
)
|
|
|
|
|
|
def _dtype_is_int(value):
|
|
if isinstance(value, int):
|
|
return True
|
|
if isinstance(value, ExprOp):
|
|
return value.expr_ty().matches_code(DataTypeCode.INT)
|
|
if ir.is_prim_expr(value):
|
|
return value.ty.matches_code(DataTypeCode.INT)
|
|
return False
|
|
|
|
|
|
def _dtype_is_float(value):
|
|
if isinstance(value, float):
|
|
return True
|
|
if isinstance(value, ExprOp):
|
|
return value.expr_ty().matches_code(DataTypeCode.FLOAT)
|
|
if ir.is_prim_expr(value):
|
|
return value.ty.matches_code(DataTypeCode.FLOAT)
|
|
return False
|
|
|
|
|
|
def _is_scalar_operand(value):
|
|
if isinstance(value, ExprOp | int | float) or ir.is_prim_expr(value):
|
|
return True
|
|
|
|
# BufferRegion is a C++ PrimExprConvertible, but its Python wrapper is not an ExprOp.
|
|
from .stmt import BufferRegion # pylint: disable=import-outside-toplevel
|
|
|
|
return isinstance(value, BufferRegion)
|
|
|
|
|
|
class ExprOp:
|
|
"""Operator overloading for Expr like expressions."""
|
|
|
|
# TODO(tkonolige): use inspect to add source information to these objects
|
|
|
|
def expr_ty(self) -> ir.PrimType:
|
|
"""Return the compile-time primitive type for expression operators."""
|
|
ty = getattr(self, "ty", None)
|
|
if isinstance(ty, ir.PrimType):
|
|
return ty
|
|
raise TypeError(f"Cannot determine PrimType for {type(self).__name__}")
|
|
|
|
def __add__(self, other: Expr) -> Expr:
|
|
if not _is_scalar_operand(other):
|
|
return NotImplemented
|
|
return _ffi_api._OpAdd(self, other, None) # type: ignore
|
|
|
|
def __radd__(self, other: Expr) -> Expr:
|
|
if not _is_scalar_operand(other):
|
|
return NotImplemented
|
|
return _ffi_api._OpAdd(other, self, None) # type: ignore
|
|
|
|
def __sub__(self, other: Expr) -> Expr:
|
|
if not _is_scalar_operand(other):
|
|
return NotImplemented
|
|
return _ffi_api._OpSub(self, other, None) # type: ignore
|
|
|
|
def __rsub__(self, other: Expr) -> Expr:
|
|
if not _is_scalar_operand(other):
|
|
return NotImplemented
|
|
return _ffi_api._OpSub(other, self, None) # type: ignore
|
|
|
|
def __mul__(self, other: Expr) -> Expr:
|
|
if not _is_scalar_operand(other):
|
|
return NotImplemented
|
|
return _ffi_api._OpMul(self, other, None) # type: ignore
|
|
|
|
def __rmul__(self, other: Expr) -> Expr:
|
|
if not _is_scalar_operand(other):
|
|
return NotImplemented
|
|
return _ffi_api._OpMul(other, self, None) # type: ignore
|
|
|
|
def __div__(self, other: Expr) -> Expr:
|
|
if not _is_scalar_operand(other):
|
|
return NotImplemented
|
|
if _dtype_is_int(self) and _dtype_is_int(other):
|
|
raise div_ambiguity_error()
|
|
return _ffi_api._OpDiv(self, other, None) # type: ignore
|
|
|
|
def __rdiv__(self, other: Expr) -> Expr:
|
|
if not _is_scalar_operand(other):
|
|
return NotImplemented
|
|
if _dtype_is_int(self) and _dtype_is_int(other):
|
|
raise div_ambiguity_error()
|
|
return _ffi_api._OpDiv(other, self, None) # type: ignore
|
|
|
|
def __truediv__(self, other: Expr) -> Expr:
|
|
if not _is_scalar_operand(other):
|
|
return NotImplemented
|
|
if _dtype_is_int(self) and _dtype_is_int(other):
|
|
raise div_ambiguity_error()
|
|
return _ffi_api._OpDiv(self, other, None) # type: ignore
|
|
|
|
def __rtruediv__(self, other: Expr) -> Expr:
|
|
if not _is_scalar_operand(other):
|
|
return NotImplemented
|
|
if _dtype_is_int(self) and _dtype_is_int(other):
|
|
raise div_ambiguity_error()
|
|
return _ffi_api._OpDiv(other, self, None) # type: ignore
|
|
|
|
def __floordiv__(self, other: Expr) -> Expr:
|
|
return _ffi_api._OpFloorDiv(self, other, None) # type: ignore
|
|
|
|
def __rfloordiv__(self, other: Expr) -> Expr:
|
|
return _ffi_api._OpFloorDiv(other, self, None) # type: ignore
|
|
|
|
def __mod__(self, other: Expr) -> Expr:
|
|
return _ffi_api._OpFloorMod(self, other, None) # type: ignore
|
|
|
|
def __rmod__(self, other: Expr) -> Expr:
|
|
return _ffi_api._OpFloorMod(other, self, None) # type: ignore
|
|
|
|
def __neg__(self) -> Expr:
|
|
neg_one = const(-1, self.expr_ty().dtype)
|
|
return self.__mul__(neg_one)
|
|
|
|
def __lshift__(self, other: Expr) -> Expr:
|
|
return _ffi_api.left_shift(self, other, None) # type: ignore
|
|
|
|
def __rlshift__(self, other: Expr) -> Expr:
|
|
return _ffi_api.left_shift(other, self, None) # type: ignore
|
|
|
|
def __rshift__(self, other: Expr) -> Expr:
|
|
return _ffi_api.right_shift(self, other, None) # type: ignore
|
|
|
|
def __rrshift__(self, other: Expr) -> Expr:
|
|
return _ffi_api.right_shift(other, self, None) # type: ignore
|
|
|
|
def __and__(self, other: Expr) -> Expr:
|
|
return _ffi_api.bitwise_and(self, other, None) # type: ignore
|
|
|
|
def __rand__(self, other: Expr) -> Expr:
|
|
return _ffi_api.bitwise_and(other, self, None) # type: ignore
|
|
|
|
def __or__(self, other: Expr) -> Expr:
|
|
return _ffi_api.bitwise_or(self, other, None) # type: ignore
|
|
|
|
def __ror__(self, other: Expr) -> Expr:
|
|
return _ffi_api.bitwise_or(other, self, None) # type: ignore
|
|
|
|
def __xor__(self, other: Expr) -> Expr:
|
|
return _ffi_api.bitwise_xor(self, other, None) # type: ignore
|
|
|
|
def __rxor__(self, other: Expr) -> Expr:
|
|
return _ffi_api.bitwise_xor(other, self, None) # type: ignore
|
|
|
|
def __invert__(self) -> Expr:
|
|
if _dtype_is_float(self):
|
|
raise RuntimeError("Cannot use ~ operator on float type Expr.")
|
|
return _ffi_api.bitwise_not(self, None) # type: ignore
|
|
|
|
def __lt__(self, other: Expr) -> Expr:
|
|
return _ffi_api._OpLT(self, other, None) # type: ignore
|
|
|
|
def __le__(self, other: Expr) -> Expr:
|
|
return _ffi_api._OpLE(self, other, None) # type: ignore
|
|
|
|
def __eq__(self, other: Expr) -> Expr:
|
|
return EqualOp(self, other)
|
|
|
|
def __ne__(self, other: Expr) -> Expr:
|
|
return NotEqualOp(self, other)
|
|
|
|
def __gt__(self, other: Expr) -> Expr:
|
|
return _ffi_api._OpGT(self, other, None) # type: ignore
|
|
|
|
def __ge__(self, other: Expr) -> Expr:
|
|
return _ffi_api._OpGE(self, other, None) # type: ignore
|
|
|
|
def __nonzero__(self):
|
|
raise ValueError(
|
|
"Cannot use and / or / not operator to Expr, hint: use tvm.tirx.all / "
|
|
"tvm.tirx.any, if it is None checking, use node is not None"
|
|
)
|
|
|
|
def __bool__(self) -> bool:
|
|
return self.__nonzero__()
|
|
|
|
def equal(self, other: Expr, span: Span | None = None) -> bool:
|
|
"""Build an equal check expression with other expr.
|
|
|
|
Parameters
|
|
----------
|
|
other : Expr
|
|
The other expression
|
|
|
|
span : Optional[Span]
|
|
The location of the cast in the source.
|
|
|
|
Returns
|
|
-------
|
|
ret : Expr
|
|
The equality expression.
|
|
"""
|
|
return _ffi_api._OpEQ(self, other, span) # type: ignore
|
|
|
|
def astype(self, dtype: str | ir.PrimType, span: Span | None = None) -> Expr:
|
|
"""Cast the expression to other type.
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The type of new expression
|
|
|
|
span : Optional[Span]
|
|
The location of the cast in the source.
|
|
|
|
Returns
|
|
-------
|
|
expr : Expr
|
|
Expression with new type
|
|
"""
|
|
return _ffi_api._cast(dtype, self, span) # type: ignore
|
|
|
|
|
|
_overload_prim_expr.__add__ = ExprOp.__add__
|
|
_overload_prim_expr.__radd__ = ExprOp.__radd__
|
|
_overload_prim_expr.__sub__ = ExprOp.__sub__
|
|
_overload_prim_expr.__rsub__ = ExprOp.__rsub__
|
|
_overload_prim_expr.__mul__ = ExprOp.__mul__
|
|
_overload_prim_expr.__rmul__ = ExprOp.__rmul__
|
|
_overload_prim_expr.__div__ = ExprOp.__div__
|
|
_overload_prim_expr.__rdiv__ = ExprOp.__rdiv__
|
|
_overload_prim_expr.__truediv__ = ExprOp.__truediv__
|
|
_overload_prim_expr.__rtruediv__ = ExprOp.__rtruediv__
|
|
_overload_prim_expr.__floordiv__ = ExprOp.__floordiv__
|
|
_overload_prim_expr.__rfloordiv__ = ExprOp.__rfloordiv__
|
|
_overload_prim_expr.__mod__ = ExprOp.__mod__
|
|
_overload_prim_expr.__rmod__ = ExprOp.__rmod__
|
|
_overload_prim_expr.__neg__ = ExprOp.__neg__
|
|
_overload_prim_expr.__lshift__ = ExprOp.__lshift__
|
|
_overload_prim_expr.__rlshift__ = ExprOp.__rlshift__
|
|
_overload_prim_expr.__rshift__ = ExprOp.__rshift__
|
|
_overload_prim_expr.__rrshift__ = ExprOp.__rrshift__
|
|
_overload_prim_expr.__and__ = ExprOp.__and__
|
|
_overload_prim_expr.__rand__ = ExprOp.__rand__
|
|
_overload_prim_expr.__or__ = ExprOp.__or__
|
|
_overload_prim_expr.__ror__ = ExprOp.__ror__
|
|
_overload_prim_expr.__xor__ = ExprOp.__xor__
|
|
_overload_prim_expr.__rxor__ = ExprOp.__rxor__
|
|
_overload_prim_expr.__invert__ = ExprOp.__invert__
|
|
_overload_prim_expr.__lt__ = ExprOp.__lt__
|
|
_overload_prim_expr.__le__ = ExprOp.__le__
|
|
_overload_prim_expr.__eq__ = ExprOp.__eq__
|
|
_overload_prim_expr.__ne__ = ExprOp.__ne__
|
|
_overload_prim_expr.__gt__ = ExprOp.__gt__
|
|
_overload_prim_expr.__ge__ = ExprOp.__ge__
|
|
_overload_prim_expr.equal = ExprOp.equal
|
|
_overload_prim_expr.astype = ExprOp.astype
|
|
|
|
|
|
class EqualOp(ObjectConvertible, ExprOp):
|
|
"""Deferred equal operator.
|
|
|
|
This is used to support sugar that a == b can either
|
|
mean Object.same_as or Object.equal.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
Left operand.
|
|
|
|
b : Expr
|
|
Right operand.
|
|
|
|
span : Optional[Span]
|
|
The location of the cast in the source.
|
|
"""
|
|
|
|
# This class is not manipulated by C++. So use python's identity check function is sufficient
|
|
same_as = object.__eq__
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None):
|
|
self.a = a
|
|
self.b = b
|
|
self.span = span
|
|
|
|
def __nonzero__(self) -> bool:
|
|
return self.a.same_as(self.b)
|
|
|
|
def __bool__(self) -> bool:
|
|
return self.__nonzero__()
|
|
|
|
def asobject(self) -> Expr:
|
|
"""Convert object."""
|
|
return _ffi_api._OpEQ(self.a, self.b, self.span) # type: ignore
|
|
|
|
def expr_ty(self) -> ir.PrimType:
|
|
"""Compile-time type of the equality result."""
|
|
return ir.PrimType("bool")
|
|
|
|
def __repr__(self) -> str:
|
|
return f"EqualOp({self.a!r}, {self.b!r})"
|
|
|
|
|
|
class NotEqualOp(ObjectConvertible, ExprOp):
|
|
"""Deferred NE operator.
|
|
|
|
This is used to support sugar that a != b can either
|
|
mean not Object.same_as or make.NE.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
Left operand.
|
|
|
|
b : Expr
|
|
Right operand.
|
|
|
|
span : Optional[Span]
|
|
The location of the cast in the source.
|
|
"""
|
|
|
|
# This class is not manipulated by C++. So use python's identity check function is sufficient
|
|
same_as = object.__eq__
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.a = a
|
|
self.b = b
|
|
self.span = span
|
|
|
|
def __nonzero__(self) -> bool:
|
|
return not self.a.same_as(self.b)
|
|
|
|
def __bool__(self) -> bool:
|
|
return self.__nonzero__()
|
|
|
|
def asobject(self) -> Expr:
|
|
"""Convert object."""
|
|
return _ffi_api._OpNE(self.a, self.b, self.span) # type: ignore
|
|
|
|
def expr_ty(self) -> ir.PrimType:
|
|
"""Compile-time type of the inequality result."""
|
|
return ir.PrimType("bool")
|
|
|
|
def __repr__(self) -> str:
|
|
return f"NotEqualOp({self.a!r}, {self.b!r})"
|
|
|
|
|
|
class IntImmEnum(ObjectConvertible):
|
|
"""Lazily evaluate an IntImm in case
|
|
the constructor is not available in runtime.
|
|
|
|
Parameters
|
|
----------
|
|
value : int
|
|
The enum value
|
|
|
|
span : Optional[Span]
|
|
The location of the cast in the source.
|
|
"""
|
|
|
|
def __init__(self, value: int, span: Span | None = None) -> None:
|
|
self.value = value
|
|
self.span = span
|
|
|
|
def asobject(self) -> "IntImm":
|
|
"""Convert object."""
|
|
return IntImm("int32", self.value, self.span) # type: ignore
|
|
|
|
|
|
class ExprWithOp(ExprOp, Expr, Scriptable):
|
|
"""Helper base class to inherit from Expr."""
|
|
|
|
# In Python3, We have to explicitly tell interpreter to retain __hash__ if we overide __eq__
|
|
# https://docs.python.org/3.1/reference/datamodel.html#object.__hash__
|
|
__hash__ = Expr.__hash__
|
|
|
|
|
|
class ConstExpr(ExprWithOp):
|
|
pass
|
|
|
|
|
|
class BinaryOpExpr(ExprWithOp):
|
|
a: Expr
|
|
b: Expr
|
|
|
|
|
|
class CmpExpr(ExprWithOp):
|
|
a: Expr
|
|
b: Expr
|
|
|
|
|
|
class LogicalExpr(ExprWithOp):
|
|
pass
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Var")
|
|
class Var(ExprWithOp):
|
|
"""Symbolic variable.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name
|
|
|
|
dtype : Union[str, ir.Type]
|
|
The data type
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
name_hint: str
|
|
|
|
def __init__(self, name: str, dtype: str | ir.Type, span: Span | None = None) -> None:
|
|
if isinstance(dtype, str) and dtype == "handle":
|
|
dtype = ir.PointerType(ir.PrimType("void"))
|
|
self.__init_handle_by_constructor__(_ffi_api.Var, name, dtype, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.IterVar")
|
|
class IterVar(ExprOp, Object, Scriptable):
|
|
"""Represent iteration variable.
|
|
|
|
IterVar represents axis iterations in the computation.
|
|
|
|
Parameters
|
|
----------
|
|
dom : Range
|
|
The domain of the iteration.
|
|
|
|
var : Union[Var, str]
|
|
The internal variable that is used for iteration.
|
|
|
|
iter_type : int
|
|
The iteration type.
|
|
|
|
thread_tag : str
|
|
The thread type tag.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
|
|
See Also
|
|
--------
|
|
te.thread_axis: Create thread axis IterVar.
|
|
te.reduce_axis: Create reduce axis IterVar.
|
|
"""
|
|
|
|
DataPar = 0
|
|
ThreadIndex = 1
|
|
CommReduce = 2
|
|
Ordered = 3
|
|
Opaque = 4
|
|
Unrolled = 5
|
|
Vectorized = 6
|
|
Parallelized = 7
|
|
Tensorized = 8
|
|
|
|
dom: ir.Range
|
|
var: Var
|
|
iter_type: int
|
|
thread_tag: str
|
|
|
|
def __init__(
|
|
self,
|
|
dom: ir.Range,
|
|
var: Var | str,
|
|
iter_type: int,
|
|
thread_tag: str = "",
|
|
span: Span | None = None,
|
|
) -> None:
|
|
if dom is not None:
|
|
if isinstance(dom, list | tuple):
|
|
if len(dom) != 2:
|
|
raise TypeError("need to be list of ranges")
|
|
dom = tvm.ir.Range(dom[0], dom[1])
|
|
|
|
if not isinstance(dom, tvm.ir.Range):
|
|
raise TypeError("dom need to be Range")
|
|
|
|
name = var if var is not None else "iter"
|
|
dtype = "int32" if dom is None else dom.extent.ty
|
|
var = Var(name, dtype=dtype, span=span) if not isinstance(var, Var) else var
|
|
if dom is not None:
|
|
assert var.ty == dom.extent.ty, "IterVar's Var type must match its domain's extent type"
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.IterVar,
|
|
dom,
|
|
var,
|
|
iter_type,
|
|
thread_tag,
|
|
span, # type: ignore
|
|
)
|
|
|
|
def expr_ty(self) -> ir.PrimType:
|
|
"""Compile-time type of the iteration variable."""
|
|
return self.var.ty
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.CommReducer")
|
|
class CommReducer(Object, Scriptable):
|
|
"""Commutative reduce operator
|
|
|
|
Parameters
|
|
----------
|
|
lhs : List[Var]
|
|
The left arguments of the reducer.
|
|
|
|
rhs : List[Var]
|
|
The right arguments of the reducer.
|
|
|
|
result : List[Expr]
|
|
The reduction results.
|
|
|
|
identity_element : List[Expr]
|
|
The identity elements.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
lhs: list[Var]
|
|
rhs: list[Var]
|
|
result: list[Expr]
|
|
identity_element: list[Expr]
|
|
|
|
def __init__(
|
|
self,
|
|
lhs: list[Var],
|
|
rhs: list[Var],
|
|
result: list[Expr],
|
|
identity_element: list[Expr],
|
|
span: Span | None = None,
|
|
) -> None:
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.CommReducer,
|
|
lhs,
|
|
rhs,
|
|
result,
|
|
identity_element,
|
|
span, # type: ignore
|
|
)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Reduce")
|
|
class Reduce(ExprWithOp):
|
|
"""Reduce node.
|
|
|
|
Parameters
|
|
----------
|
|
combiner : CommReducer
|
|
The combiner.
|
|
|
|
src : list of Expr
|
|
The source expression.
|
|
|
|
rdom : list of IterVar
|
|
The iteration domain
|
|
|
|
condition : Expr
|
|
The reduce condition.
|
|
|
|
value_index : int
|
|
The value index.
|
|
|
|
init : list of Expr
|
|
The initial value for output. This can be an int, float or ProducerLoad
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
combiner: CommReducer
|
|
source: list[Expr]
|
|
init: list[Expr]
|
|
axis: list[IterVar]
|
|
condition: Expr
|
|
value_index: int
|
|
|
|
def __init__(
|
|
self,
|
|
combiner: CommReducer,
|
|
src: list[Expr],
|
|
rdom: list[IterVar],
|
|
condition: Expr,
|
|
value_index: int,
|
|
init: list[Expr] | None = None,
|
|
span: Span | None = None,
|
|
) -> None:
|
|
init = [] if init is None else init
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.Reduce,
|
|
combiner,
|
|
src,
|
|
rdom,
|
|
condition,
|
|
value_index,
|
|
init,
|
|
span, # type: ignore
|
|
)
|
|
|
|
|
|
@tvm_ffi.register_object("ir.FloatImm")
|
|
class FloatImm(ConstExpr):
|
|
"""Float constant.
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The data type
|
|
|
|
value : float
|
|
The constant value.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
value: float
|
|
|
|
def __init__(self, dtype: str | ir.PrimType, value: float, span: Span | None = None) -> None:
|
|
if isinstance(dtype, ir.PrimType):
|
|
dtype = dtype.dtype
|
|
self.__init_handle_by_constructor__(
|
|
tvm.ir._ffi_api.FloatImm,
|
|
dtype,
|
|
value,
|
|
span, # type: ignore
|
|
)
|
|
|
|
def __float__(self) -> float:
|
|
return self.value
|
|
|
|
|
|
@tvm_ffi.register_object("ir.IntImm")
|
|
class IntImm(ConstExpr):
|
|
"""Int constant.
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The data type
|
|
|
|
value : int
|
|
The constant value.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
value: int
|
|
|
|
def __init__(self, dtype: str | ir.PrimType, value: int, span: Span | None = None) -> None:
|
|
if isinstance(dtype, ir.PrimType):
|
|
dtype = dtype.dtype
|
|
self.__init_handle_by_constructor__(
|
|
tvm.ir._ffi_api.IntImm,
|
|
dtype,
|
|
value,
|
|
span, # type: ignore
|
|
)
|
|
|
|
def __hash__(self) -> int:
|
|
return self.value
|
|
|
|
def __int__(self) -> int:
|
|
return self.value
|
|
|
|
def __nonzero__(self) -> bool:
|
|
return self.value != 0
|
|
|
|
def __eq__(self, other: Expr) -> Expr:
|
|
return _ffi_api._OpEQ(self, other, None) # type: ignore
|
|
|
|
def __ne__(self, other: Expr) -> Expr:
|
|
return _ffi_api._OpNE(self, other, None) # type: ignore
|
|
|
|
def __bool__(self) -> bool:
|
|
return self.__nonzero__()
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.StringImm") # type: ignore
|
|
class StringImm(ConstExpr):
|
|
"""String constant.
|
|
|
|
Parameters
|
|
----------
|
|
value : str
|
|
The value of the function.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
value: str
|
|
|
|
def __init__(self, value: str, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.StringImm, value, span) # type: ignore
|
|
|
|
def __eq__(self, other: Expr) -> bool:
|
|
if isinstance(other, ConstExpr):
|
|
return self.value == other.value
|
|
return self.value == other
|
|
|
|
def __ne__(self, other: Expr) -> bool:
|
|
if isinstance(other, ConstExpr):
|
|
return self.value != other.value
|
|
return self.value != other
|
|
|
|
def __hash__(self) -> int:
|
|
return Expr.__hash__(self)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Cast")
|
|
class Cast(ExprWithOp):
|
|
"""Cast expression.
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The data type
|
|
|
|
value : Expr
|
|
The value of the function.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
value: Expr
|
|
|
|
def __init__(self, dtype: str | ir.PrimType, value, span: Span | None = None) -> None:
|
|
if isinstance(dtype, ir.PrimType):
|
|
dtype = dtype.dtype
|
|
self.__init_handle_by_constructor__(_ffi_api.Cast, dtype, value, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Add")
|
|
class Add(BinaryOpExpr):
|
|
"""Add node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.Add, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Sub")
|
|
class Sub(BinaryOpExpr):
|
|
"""Sub node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.Sub, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Mul")
|
|
class Mul(BinaryOpExpr):
|
|
"""Mul node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.Mul, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Div")
|
|
class Div(BinaryOpExpr):
|
|
"""Div node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.Div, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Mod")
|
|
class Mod(BinaryOpExpr):
|
|
"""Mod node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.Mod, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.FloorDiv")
|
|
class FloorDiv(BinaryOpExpr):
|
|
"""FloorDiv node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.FloorDiv, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.FloorMod")
|
|
class FloorMod(BinaryOpExpr):
|
|
"""FloorMod node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.FloorMod, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Min")
|
|
class Min(BinaryOpExpr):
|
|
"""Min node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.Min, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Max")
|
|
class Max(BinaryOpExpr):
|
|
"""Max node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.Max, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.EQ")
|
|
class EQ(CmpExpr):
|
|
"""EQ node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.EQ, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.NE")
|
|
class NE(CmpExpr):
|
|
"""NE node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.NE, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.LT")
|
|
class LT(CmpExpr):
|
|
"""LT node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.LT, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.LE")
|
|
class LE(CmpExpr):
|
|
"""LE node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.LE, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.GT")
|
|
class GT(CmpExpr):
|
|
"""GT node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.GT, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.GE")
|
|
class GE(CmpExpr):
|
|
"""GE node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.GE, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.And")
|
|
class And(LogicalExpr):
|
|
"""And node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.And, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Or")
|
|
class Or(LogicalExpr):
|
|
"""Or node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand.
|
|
|
|
b : Expr
|
|
The right hand operand.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
a: Expr
|
|
b: Expr
|
|
|
|
def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.Or, a, b, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Not")
|
|
class Not(LogicalExpr):
|
|
"""Not node.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The input value
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
a: Expr
|
|
|
|
def __init__(self, a: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.Not, a, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Select")
|
|
class Select(ExprWithOp):
|
|
"""Select node.
|
|
|
|
Note
|
|
----
|
|
Select may compute both true_value and false_value.
|
|
Use :py:class:`tvm.tirx.if_then_else` instead if you want to
|
|
get a conditional expression that only evaluates
|
|
the correct branch.
|
|
|
|
Parameters
|
|
----------
|
|
condition : Expr
|
|
The condition expression.
|
|
|
|
true_value : Expr
|
|
The value to take when condition is true.
|
|
|
|
false_value : Expr
|
|
The value to take when condition is false.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
condition: Expr
|
|
true_value: Expr
|
|
false_value: Expr
|
|
|
|
def __init__(
|
|
self,
|
|
condition: Expr,
|
|
true_value: Expr,
|
|
false_value: Expr,
|
|
span: Span | None = None,
|
|
) -> None:
|
|
if isinstance(condition, bool):
|
|
condition = IntImm("bool", condition)
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.Select,
|
|
condition,
|
|
true_value,
|
|
false_value,
|
|
span, # type: ignore
|
|
)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.BufferLoad")
|
|
class BufferLoad(ExprWithOp):
|
|
"""Buffer load node.
|
|
|
|
Parameters
|
|
----------
|
|
buffer : Buffer
|
|
The buffer to be loaded.
|
|
|
|
indices : List[Expr]
|
|
The buffer indices to load values from.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
|
|
predicate : Optional[Expr]
|
|
A vector mask of boolean values indicating which lanes of a vector are to be
|
|
loaded. The number lanes of the mask must be equal to the number of lanes being loaded.
|
|
"""
|
|
|
|
buffer: Buffer
|
|
indices: list[Expr]
|
|
|
|
def __init__(
|
|
self,
|
|
buffer: Buffer,
|
|
indices: list[Expr],
|
|
predicate: Expr | None = None,
|
|
span: Span | None = None,
|
|
) -> None:
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.BufferLoad,
|
|
buffer,
|
|
indices,
|
|
predicate,
|
|
span, # type: ignore
|
|
)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.ProducerLoad")
|
|
class ProducerLoad(ExprWithOp):
|
|
"""Producer load node.
|
|
|
|
Parameters
|
|
----------
|
|
producer : DataProducer
|
|
The buffer to be loaded.
|
|
|
|
indices : List[Expr]
|
|
The buffer indices.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
producer: DataProducer
|
|
indices: list[Expr]
|
|
|
|
def __init__(
|
|
self, producer: DataProducer, indices: list[Expr], span: Span | None = None
|
|
) -> None:
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.ProducerLoad,
|
|
producer,
|
|
indices,
|
|
span, # type: ignore
|
|
)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Ramp")
|
|
class Ramp(ExprWithOp):
|
|
"""Ramp node.
|
|
|
|
Parameters
|
|
----------
|
|
base : Expr
|
|
The base expression.
|
|
|
|
stride : Expr
|
|
The stride of the ramp.
|
|
|
|
lanes : Expr
|
|
The lanes of the expression.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
base: Expr
|
|
stride: Expr
|
|
lanes: Expr
|
|
|
|
def __init__(self, base: Expr, stride: Expr, lanes: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.Ramp,
|
|
base,
|
|
stride,
|
|
lanes,
|
|
span, # type: ignore
|
|
)
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Broadcast")
|
|
class Broadcast(ExprWithOp):
|
|
"""Broadcast node.
|
|
|
|
Parameters
|
|
----------
|
|
value : Expr
|
|
The value of the expression.
|
|
|
|
lanes : Expr
|
|
The lanes of the expression.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
value: Expr
|
|
lanes: Expr
|
|
|
|
def __init__(self, value: Expr, lanes: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.Broadcast, value, lanes, span) # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Shuffle")
|
|
class Shuffle(ExprWithOp):
|
|
"""Shuffle node.
|
|
|
|
Parameters
|
|
----------
|
|
vectors : List[Expr]
|
|
The vectors
|
|
|
|
indices : List[Expr]
|
|
The indices
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
vectors: list[Expr]
|
|
indices: list[Expr]
|
|
|
|
def __init__(self, vectors: list[Expr], indices: list[Expr], span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.Shuffle,
|
|
vectors,
|
|
indices,
|
|
span, # type: ignore
|
|
)
|
|
|
|
|
|
class CallEffectKind:
|
|
"""Possible kinds of Call effects."""
|
|
|
|
# only expose up to opaque
|
|
ExprAnnotation = IntImmEnum(0)
|
|
Pure = IntImmEnum(1)
|
|
ReadState = IntImmEnum(2)
|
|
UpdateState = IntImmEnum(3)
|
|
Opaque = UpdateState
|
|
|
|
|
|
@tvm_ffi.register_object("tirx.Let")
|
|
class Let(ExprWithOp):
|
|
"""Let node.
|
|
|
|
Parameters
|
|
----------
|
|
var : Var
|
|
The variable in the binding.
|
|
|
|
value : Expr
|
|
The value in to be bound.
|
|
|
|
body : Expr
|
|
The body expression.
|
|
|
|
span : Optional[Span]
|
|
The location of this expression in the source code.
|
|
"""
|
|
|
|
var: Var
|
|
value: Expr
|
|
body: Expr
|
|
|
|
def __init__(self, var: Var, value: Expr, body: Expr, span: Span | None = None) -> None:
|
|
self.__init_handle_by_constructor__(_ffi_api.Let, var, value, body, span) # type: ignore
|