This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 45e1b8233a Refactor Tensor arithmetic dispatch away from tirx.generic
(#19943)
45e1b8233a is described below
commit 45e1b8233a0f443311853d08e65a92c36b5c79bd
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Jul 5 08:18:43 2026 +0800
Refactor Tensor arithmetic dispatch away from tirx.generic (#19943)
## Summary
- Move whole-Tensor arithmetic and cast dispatch onto `te.Tensor` while
scalar TIRx smart constructors decline whole-Tensor operands.
- Remove the legacy `tirx.generic` module, TOPI import-time mutation
bridge, and obsolete aliases.
- Migrate scan and cast callers while preserving identity-gated Thrust
sum selection.
Whole-Tensor behavior now lives with TE, leaving scalar TIRx
construction independent of TOPI initialization.
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 3 +-
python/tvm/relax/transform/legalize_ops/qdq.py | 4 +-
python/tvm/te/__init__.py | 2 -
python/tvm/te/_te_tensor_overload.py | 61 +++++++
python/tvm/te/tensor.py | 202 +++++++++++++++++++++++-
python/tvm/tirx/__init__.py | 1 -
python/tvm/tirx/expr.py | 61 +++++--
python/tvm/tirx/generic.py | 145 -----------------
python/tvm/tirx/script/builder/ir.py | 9 +-
python/tvm/topi/__init__.py | 2 +-
python/tvm/topi/_te_tensor_overload.py | 64 ++++++++
python/tvm/topi/generic_op_impl.py | 105 ------------
python/tvm/topi/gpu/scan.py | 56 ++++---
python/tvm/topi/gpu/sort.py | 14 +-
python/tvm/topi/scan.py | 7 +-
15 files changed, 418 insertions(+), 318 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index ef08c3f690..f460418315 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -61,7 +61,6 @@ from tvm import relax, tirx, topi
from tvm.ir import IRModule
from tvm.ir.supply import UniqueNameSupply
from tvm.runtime import DataType, DataTypeCode
-from tvm.tirx.generic import cast
from tvm.topi.utils import get_const_tuple
from ..common import autopad
@@ -3257,7 +3256,7 @@ class Resize(OnnxOpConverter):
sizes = []
for i, dim in enumerate(x.ty.shape):
- sizes.append(cast(scales[i] * dim, "int64"))
+ sizes.append((scales[i] * dim).astype("int64"))
sizes = sizes[2:]
else:
if isinstance(sizes, relax.Constant):
diff --git a/python/tvm/relax/transform/legalize_ops/qdq.py
b/python/tvm/relax/transform/legalize_ops/qdq.py
index 0f85c6bbdd..9cea923146 100644
--- a/python/tvm/relax/transform/legalize_ops/qdq.py
+++ b/python/tvm/relax/transform/legalize_ops/qdq.py
@@ -147,8 +147,8 @@ def _dequantize(bb: BlockBuilder, call: Call) -> Expr:
if data.dtype.matches_code(DataTypeCode.FLOAT,
DataTypeCode.BFLOAT)
else "int32"
)
- sub = te.subtract(data[indices].astype(dtype), zp_value)
- out = te.multiply(sub, scale_value.astype("float32"))
+ sub = data[indices].astype(dtype) - zp_value
+ out = sub * scale_value.astype("float32")
if out_dtype == "float32":
return out
return clip_cast(out, out_dtype)
diff --git a/python/tvm/te/__init__.py b/python/tvm/te/__init__.py
index c923efde00..4773b8bcb9 100644
--- a/python/tvm/te/__init__.py
+++ b/python/tvm/te/__init__.py
@@ -27,8 +27,6 @@ from tvm.tirx import trunc, abs, round, nearbyint, power,
popcount, fmod, if_the
from tvm.tirx import isnan, isfinite, isinf
from tvm.tirx import div, indexdiv, indexmod, truncdiv, truncmod, floordiv,
floormod, logaddexp
from tvm.tirx import comm_reducer, min, max, sum
-from tvm.tirx import add, subtract, multiply
-
from .tensor import TensorSlice, Tensor
from .tag import tag_scope
from .operation import placeholder, compute, scan, extern, var, const
diff --git a/python/tvm/te/_te_tensor_overload.py
b/python/tvm/te/_te_tensor_overload.py
new file mode 100644
index 0000000000..d52ec703f1
--- /dev/null
+++ b/python/tvm/te/_te_tensor_overload.py
@@ -0,0 +1,61 @@
+# 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.
+"""Tensor overload hooks for TE tensors and tensor slices."""
+
+
+def __add__(_lhs, _rhs):
+ return NotImplemented
+
+
+def __radd__(_lhs, _rhs):
+ return NotImplemented
+
+
+def __sub__(_lhs, _rhs):
+ return NotImplemented
+
+
+def __rsub__(_lhs, _rhs):
+ return NotImplemented
+
+
+def __mul__(_lhs, _rhs):
+ return NotImplemented
+
+
+def __rmul__(_lhs, _rhs):
+ return NotImplemented
+
+
+def __div__(_lhs, _rhs):
+ return NotImplemented
+
+
+def __rdiv__(_lhs, _rhs):
+ return NotImplemented
+
+
+def __truediv__(_lhs, _rhs):
+ return NotImplemented
+
+
+def __rtruediv__(_lhs, _rhs):
+ return NotImplemented
+
+
+def astype(_value, _dtype, _span=None):
+ return NotImplemented
diff --git a/python/tvm/te/tensor.py b/python/tvm/te/tensor.py
index 29f9b185ad..ae1fad55a4 100644
--- a/python/tvm/te/tensor.py
+++ b/python/tvm/te/tensor.py
@@ -19,14 +19,18 @@
# pylint: disable=invalid-name
import tvm_ffi
-from tvm.runtime import Object, ObjectConvertible
+from tvm.runtime import Object, ObjectConvertible, const
from tvm.tirx import DataProducer
from tvm.tirx import expr as _expr
-from . import _ffi_api
+from . import _ffi_api, _te_tensor_overload
-class TensorSlice(ObjectConvertible, _expr.ExprOp):
+def _as_scalar_operand(value):
+ return value.asobject() if isinstance(value, TensorSlice) else value
+
+
+class TensorSlice(ObjectConvertible):
"""Auxiliary data structure for enable slicing syntax from tensor."""
def __init__(self, tensor, indices):
@@ -53,9 +57,199 @@ class TensorSlice(ObjectConvertible, _expr.ExprOp):
"""Compile-time element type of the tensor."""
return self.tensor.expr_ty()
+ def __add__(self, other):
+ result = _te_tensor_overload.__add__(self, other)
+ if result is not NotImplemented:
+ return result
+ return _expr.ExprOp.__add__(self.asobject(), _as_scalar_operand(other))
+
+ def __radd__(self, other):
+ result = _te_tensor_overload.__radd__(self, other)
+ if result is not NotImplemented:
+ return result
+ return _expr.ExprOp.__radd__(self.asobject(),
_as_scalar_operand(other))
+
+ def __sub__(self, other):
+ result = _te_tensor_overload.__sub__(self, other)
+ if result is not NotImplemented:
+ return result
+ return _expr.ExprOp.__sub__(self.asobject(), _as_scalar_operand(other))
+
+ def __rsub__(self, other):
+ result = _te_tensor_overload.__rsub__(self, other)
+ if result is not NotImplemented:
+ return result
+ return _expr.ExprOp.__rsub__(self.asobject(),
_as_scalar_operand(other))
+
+ def __mul__(self, other):
+ result = _te_tensor_overload.__mul__(self, other)
+ if result is not NotImplemented:
+ return result
+ return _expr.ExprOp.__mul__(self.asobject(), _as_scalar_operand(other))
+
+ def __rmul__(self, other):
+ result = _te_tensor_overload.__rmul__(self, other)
+ if result is not NotImplemented:
+ return result
+ return _expr.ExprOp.__rmul__(self.asobject(),
_as_scalar_operand(other))
+
+ def __div__(self, other):
+ result = _te_tensor_overload.__div__(self, other)
+ if result is not NotImplemented:
+ return result
+ return _expr.ExprOp.__div__(self.asobject(), _as_scalar_operand(other))
+
+ def __rdiv__(self, other):
+ result = _te_tensor_overload.__rdiv__(self, other)
+ if result is not NotImplemented:
+ return result
+ return _expr.ExprOp.__rdiv__(self.asobject(),
_as_scalar_operand(other))
+
+ def __truediv__(self, other):
+ result = _te_tensor_overload.__truediv__(self, other)
+ if result is not NotImplemented:
+ return result
+ return _expr.ExprOp.__truediv__(self.asobject(),
_as_scalar_operand(other))
+
+ def __rtruediv__(self, other):
+ result = _te_tensor_overload.__rtruediv__(self, other)
+ if result is not NotImplemented:
+ return result
+ return _expr.ExprOp.__rtruediv__(self.asobject(),
_as_scalar_operand(other))
+
+ def __floordiv__(self, other):
+ return _expr.ExprOp.__floordiv__(self.asobject(),
_as_scalar_operand(other))
+
+ def __rfloordiv__(self, other):
+ return _expr.ExprOp.__rfloordiv__(self.asobject(),
_as_scalar_operand(other))
+
+ def __mod__(self, other):
+ return _expr.ExprOp.__mod__(self.asobject(), _as_scalar_operand(other))
+
+ def __rmod__(self, other):
+ return _expr.ExprOp.__rmod__(self.asobject(),
_as_scalar_operand(other))
+
+ def __neg__(self):
+ return _expr.ExprOp.__neg__(self.asobject())
+
+ def __lshift__(self, other):
+ return _expr.ExprOp.__lshift__(self.asobject(),
_as_scalar_operand(other))
+
+ def __rlshift__(self, other):
+ return _expr.ExprOp.__rlshift__(self.asobject(),
_as_scalar_operand(other))
+
+ def __rshift__(self, other):
+ return _expr.ExprOp.__rshift__(self.asobject(),
_as_scalar_operand(other))
+
+ def __rrshift__(self, other):
+ return _expr.ExprOp.__rrshift__(self.asobject(),
_as_scalar_operand(other))
+
+ def __and__(self, other):
+ return _expr.ExprOp.__and__(self.asobject(), _as_scalar_operand(other))
+
+ def __rand__(self, other):
+ return _expr.ExprOp.__rand__(self.asobject(),
_as_scalar_operand(other))
+
+ def __or__(self, other):
+ return _expr.ExprOp.__or__(self.asobject(), _as_scalar_operand(other))
+
+ def __ror__(self, other):
+ return _expr.ExprOp.__ror__(self.asobject(), _as_scalar_operand(other))
+
+ def __xor__(self, other):
+ return _expr.ExprOp.__xor__(self.asobject(), _as_scalar_operand(other))
+
+ def __rxor__(self, other):
+ return _expr.ExprOp.__rxor__(self.asobject(),
_as_scalar_operand(other))
+
+ def __invert__(self):
+ return _expr.ExprOp.__invert__(self.asobject())
+
+ def __lt__(self, other):
+ return _expr.ExprOp.__lt__(self.asobject(), _as_scalar_operand(other))
+
+ def __le__(self, other):
+ return _expr.ExprOp.__le__(self.asobject(), _as_scalar_operand(other))
+
+ def __eq__(self, other):
+ return _expr.ExprOp.__eq__(self.asobject(), _as_scalar_operand(other))
+
+ def __ne__(self, other):
+ return _expr.ExprOp.__ne__(self.asobject(), _as_scalar_operand(other))
+
+ def __gt__(self, other):
+ return _expr.ExprOp.__gt__(self.asobject(), _as_scalar_operand(other))
+
+ def __ge__(self, other):
+ return _expr.ExprOp.__ge__(self.asobject(), _as_scalar_operand(other))
+
+ def __nonzero__(self):
+ return _expr.ExprOp.__nonzero__(self.asobject())
+
+ def __bool__(self):
+ return self.__nonzero__()
+
+ def equal(self, other, span=None):
+ return _expr.ExprOp.equal(self.asobject(), _as_scalar_operand(other),
span)
+
+ def astype(self, dtype, span=None):
+ return _expr.ExprOp.astype(self.asobject(), dtype, span)
+
+
+class TensorOpBase:
+ """Operator overloads for whole TE Tensor values."""
+
+ def __add__(self, other):
+ return _te_tensor_overload.__add__(self, other)
+
+ def __radd__(self, other):
+ return _te_tensor_overload.__radd__(self, other)
+
+ def __sub__(self, other):
+ return _te_tensor_overload.__sub__(self, other)
+
+ def __rsub__(self, other):
+ return _te_tensor_overload.__rsub__(self, other)
+
+ def __mul__(self, other):
+ return _te_tensor_overload.__mul__(self, other)
+
+ def __rmul__(self, other):
+ return _te_tensor_overload.__rmul__(self, other)
+
+ def __div__(self, other):
+ return _te_tensor_overload.__div__(self, other)
+
+ def __rdiv__(self, other):
+ return _te_tensor_overload.__rdiv__(self, other)
+
+ def __truediv__(self, other):
+ return _te_tensor_overload.__truediv__(self, other)
+
+ def __rtruediv__(self, other):
+ return _te_tensor_overload.__rtruediv__(self, other)
+
+ def __neg__(self):
+ return self.__mul__(const(-1, self.expr_ty()))
+
+ def __nonzero__(self):
+ return _expr.ExprOp.__nonzero__(self)
+
+ def __bool__(self):
+ return self.__nonzero__()
+
+ def equal(self, other, span=None):
+ return _expr.ExprOp.equal(self, other, span)
+
+ def astype(self, dtype, span=None):
+ result = _te_tensor_overload.astype(self, dtype, span)
+ if result is NotImplemented:
+ raise TypeError("TE Tensor overload astype is not registered")
+ return result
+
@tvm_ffi.register_object("te.Tensor")
-class Tensor(DataProducer, _expr.ExprOp):
+class Tensor(DataProducer, TensorOpBase):
"""Tensor object, to construct, see function.Tensor"""
def __call__(self, *indices):
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 5a7c308bb9..efec3e35d6 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -82,7 +82,6 @@ from .op import start_profile_intrinsic, end_profile_intrinsic
from .op import vscale, get_active_lane_mask, get_vscale_expr
from .op import dp4a
from .op import ignore_loop_partition
-from .generic import add, subtract, multiply
# TIRX-specific imports (must come before subpackage imports to avoid circular
imports)
from .exec_scope import ExecScope, ScopeIdDef
diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py
index f45539d3b8..2e64b9316a 100644
--- a/python/tvm/tirx/expr.py
+++ b/python/tvm/tirx/expr.py
@@ -38,7 +38,6 @@ from tvm.ir.base import Span
from tvm.runtime import DataTypeCode, Object, ObjectConvertible, Scriptable,
const
from . import _ffi_api
-from . import generic as _generic
from .buffer import Buffer, DataProducer
@@ -59,6 +58,8 @@ def _dtype_is_int(value):
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
@@ -67,9 +68,21 @@ def _dtype_is_float(value):
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."""
@@ -83,48 +96,68 @@ class ExprOp:
raise TypeError(f"Cannot determine PrimType for {type(self).__name__}")
def __add__(self, other: Expr) -> Expr:
- return _generic.add(self, other)
+ if not _is_scalar_operand(other):
+ return NotImplemented
+ return _ffi_api._OpAdd(self, other, None) # type: ignore
def __radd__(self, other: Expr) -> Expr:
- return _generic.add(other, self)
+ if not _is_scalar_operand(other):
+ return NotImplemented
+ return _ffi_api._OpAdd(other, self, None) # type: ignore
def __sub__(self, other: Expr) -> Expr:
- return _generic.subtract(self, other)
+ if not _is_scalar_operand(other):
+ return NotImplemented
+ return _ffi_api._OpSub(self, other, None) # type: ignore
def __rsub__(self, other: Expr) -> Expr:
- return _generic.subtract(other, self)
+ if not _is_scalar_operand(other):
+ return NotImplemented
+ return _ffi_api._OpSub(other, self, None) # type: ignore
def __mul__(self, other: Expr) -> Expr:
- return _generic.multiply(self, other)
+ if not _is_scalar_operand(other):
+ return NotImplemented
+ return _ffi_api._OpMul(self, other, None) # type: ignore
def __rmul__(self, other: Expr) -> Expr:
- return _generic.multiply(other, self)
+ 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 _generic.divide(self, other)
+ 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 _generic.divide(other, self)
+ 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 _generic.divide(self, other)
+ 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 _generic.divide(other, self)
+ return _ffi_api._OpDiv(other, self, None) # type: ignore
def __floordiv__(self, other: Expr) -> Expr:
- return _generic.floordiv(self, other)
+ return _ffi_api._OpFloorDiv(self, other, None) # type: ignore
def __rfloordiv__(self, other: Expr) -> Expr:
- return _generic.floordiv(other, self, None)
+ return _ffi_api._OpFloorDiv(other, self, None) # type: ignore
def __mod__(self, other: Expr) -> Expr:
return _ffi_api._OpFloorMod(self, other, None) # type: ignore
@@ -232,7 +265,7 @@ class ExprOp:
expr : Expr
Expression with new type
"""
- return _generic.cast(self, dtype, span)
+ return _ffi_api._cast(dtype, self, span) # type: ignore
_overload_prim_expr.__add__ = ExprOp.__add__
diff --git a/python/tvm/tirx/generic.py b/python/tvm/tirx/generic.py
deleted file mode 100644
index 132a09e302..0000000000
--- a/python/tvm/tirx/generic.py
+++ /dev/null
@@ -1,145 +0,0 @@
-# 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.
-"""Generic opertors in TVM.
-We follow the numpy naming convention for this interface
-(e.g., tvm.tirx.generic.multitply ~ numpy.multiply).
-The default implementation is used by tvm.ExprOp.
-"""
-
-# pylint: disable=unused-argument
-from . import _ffi_api
-
-# Operator precedence used when overloading.
-__op_priority__ = 0
-
-
-def add(lhs, rhs, span=None):
- """Generic add operator.
-
- Parameters
- ----------
- lhs : object
- The left operand.
- rhs : object
- The right operand.
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- op : tvm.Expr
- The result Expr of add operaton.
- """
- return _ffi_api._OpAdd(lhs, rhs, span) # type: ignore
-
-
-def subtract(lhs, rhs, span=None):
- """Generic subtract operator.
-
- Parameters
- ----------
- lhs : object
- The left operand.
- rhs : object
- The right operand.
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- op : tvm.Expr
- The result Expr of subtract operaton.
- """
- return _ffi_api._OpSub(lhs, rhs, span) # type: ignore
-
-
-def multiply(lhs, rhs, span=None):
- """Generic multiply operator.
-
- Parameters
- ----------
- lhs : object
- The left operand.
- rhs : object
- The right operand.
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- op : tvm.Expr
- The result Expr of multiply operaton.
- """
- return _ffi_api._OpMul(lhs, rhs, span) # type: ignore
-
-
-def divide(lhs, rhs, span=None):
- """Generic divide operator.
-
- Parameters
- ----------
- lhs : object
- The left operand.
- rhs : object
- The right operand.
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- op : tvm.Expr
- The result Expr of divide operaton.
- """
- return _ffi_api._OpDiv(lhs, rhs, span) # type: ignore
-
-
-def floordiv(lhs, rhs, span=None):
- """Generic floordiv operator.
-
- Parameters
- ----------
- lhs : object
- The left operand.
- rhs : object
- The right operand.
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- op : tvm.Expr
- The result Expr of floordiv operaton.
- """
- return _ffi_api._OpFloorDiv(lhs, rhs, span) # type: ignore
-
-
-def cast(src, dtype, span=None):
- """Generic cast operator.
-
- Parameters
- ----------
- src : object
- The source operand.
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- op : tvm.Expr
- The result Expr of cast operaton.
- """
- return _ffi_api._cast(dtype, src, span) # type: ignore
diff --git a/python/tvm/tirx/script/builder/ir.py
b/python/tvm/tirx/script/builder/ir.py
index 4a16cbffb2..4e80e6cd07 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/tirx/script/builder/ir.py
@@ -44,6 +44,7 @@ from tvm.target import Target
# pylint: disable=unused-import
from tvm.target.codegen import llvm_lookup_intrinsic_id
from tvm.tirx import Buffer, BufferRegion, Expr, IndexMap, type_annotation
+from tvm.tirx import _ffi_api as _tirx_ffi_api
from tvm.tirx import op as _tir_op
from tvm.tirx.exec_scope import ExecScope, ScopeIdDef, Var
@@ -82,7 +83,6 @@ from tvm.tirx.expr import (
StringImm,
Sub,
)
-from tvm.tirx.generic import cast
from tvm.tirx.layout import (
ComposeLayout,
Iter,
@@ -100,6 +100,11 @@ from .external_kernel import call_kernel
# pylint: enable=unused-import
+def cast(value, dtype, span=None):
+ """Cast an expression to the requested data type."""
+ return _tirx_ffi_api._cast(dtype, value, span) # type:
ignore[attr-defined]
+
+
def _current_s_tir() -> bool:
"""Return True if the innermost enclosing PrimFuncFrame has ``s_tir=True``.
@@ -2355,7 +2360,7 @@ def evaluate(value: Expr) -> None:
if isinstance(value, str):
value = StringImm(value)
if isinstance(value, bool):
- value = cast(value, "bool")
+ value = IntImm("bool", value)
return _ffi_api.Evaluate(value) # type: ignore[attr-defined] # pylint:
disable=no-member
diff --git a/python/tvm/topi/__init__.py b/python/tvm/topi/__init__.py
index f0db4488d3..f0db21581c 100644
--- a/python/tvm/topi/__init__.py
+++ b/python/tvm/topi/__init__.py
@@ -34,11 +34,11 @@ from . import cpp
from .math import *
from .tensor import *
-from .generic_op_impl import *
from .index_put import *
from .reduction import *
from .transform import *
from .broadcast import *
+from . import _te_tensor_overload
from .sort import *
from .scatter import *
from .scatter_elements import *
diff --git a/python/tvm/topi/_te_tensor_overload.py
b/python/tvm/topi/_te_tensor_overload.py
new file mode 100644
index 0000000000..67a5347438
--- /dev/null
+++ b/python/tvm/topi/_te_tensor_overload.py
@@ -0,0 +1,64 @@
+# 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.
+"""Register TOPI implementations for TE tensor overload hooks."""
+
+from tvm import te
+from tvm.te import _te_tensor_overload as _overload
+from tvm.tirx import expr as _expr
+
+from . import broadcast as _broadcast
+from . import math as _math
+
+
+def _is_integer(value):
+ if isinstance(value, te.Tensor | te.TensorSlice):
+ return value.dtype.matches_code(_expr.DataTypeCode.INT)
+ return _expr._dtype_is_int(value)
+
+
+def _binary(op, reflected=False, check_integer=False):
+ def implementation(lhs, rhs):
+ if not isinstance(lhs, te.Tensor) and not isinstance(rhs, te.Tensor):
+ return NotImplemented
+ if reflected:
+ lhs, rhs = rhs, lhs
+ if check_integer and _is_integer(lhs) and _is_integer(rhs):
+ raise _expr.div_ambiguity_error()
+ return op(lhs, rhs)
+
+ return implementation
+
+
+_overload.__add__ = _binary(_broadcast.add)
+_overload.__radd__ = _binary(_broadcast.add, reflected=True)
+_overload.__sub__ = _binary(_broadcast.subtract)
+_overload.__rsub__ = _binary(_broadcast.subtract, reflected=True)
+_overload.__mul__ = _binary(_broadcast.multiply)
+_overload.__rmul__ = _binary(_broadcast.multiply, reflected=True)
+_overload.__div__ = _binary(_broadcast.divide, check_integer=True)
+_overload.__rdiv__ = _binary(_broadcast.divide, reflected=True,
check_integer=True)
+_overload.__truediv__ = _overload.__div__
+_overload.__rtruediv__ = _overload.__rdiv__
+
+
+def _astype(value, dtype, span=None):
+ if not isinstance(value, te.Tensor):
+ return NotImplemented
+ return _math.cast(value, dtype, span)
+
+
+_overload.astype = _astype
diff --git a/python/tvm/topi/generic_op_impl.py
b/python/tvm/topi/generic_op_impl.py
deleted file mode 100644
index a26e435df9..0000000000
--- a/python/tvm/topi/generic_op_impl.py
+++ /dev/null
@@ -1,105 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""Implementation of generic operators in the presence of Tensor"""
-
-# pylint: disable=invalid-name, too-many-arguments
-import tvm
-from tvm import te
-
-from . import broadcast as _broadcast
-from . import math as _math
-
-
-def _make_bop(broadcast_bop, orig_bop):
- """Make a specific overloaded binary operator of Tensor when applicable;
- apply the original operator if it is not supposed to be overloaded.
-
- Consider the following scenario:
- OP : + | - | * | /
- R0 : int | float | Expr | TensorSlice | Tensor (rank zero)
- R1 : Tensor (positive rank)
-
- In terms of (LHS OP RHS), we apply the following overloading rules:
- (1) We use broadcast_OP(LHS, RHS), when both LHS and RHS are R1.
- (2) We perform element-wise operation of Tensor and scalar,
- when one of LHS and RHS is R1 and another is R0.
- (3) We do not overload OP (i.e. stick to orig_bop) otherwise.
-
- Parameters
- ----------
- broadcast_bop : operator function
- Operator for broadcast tensor-tensor operation, for rule (1).
-
- orig_bop: operator function
- Operator before overloading, for rule (3).
-
- Returns
- -------
- ret : operator function
- The overloaded operator function if applicable or orig_bop otherwise.
- """
-
- name = orig_bop.__name__
-
- def _tensor_bop_impl(lhs, rhs):
- """Overloaded {op} operator.
-
- If both operands are non-zero-rank Tensors, it performs
- tensor-tensor {op} operation, and broadcasts inputs when necessary.
-
- If one operand is non-zero-rank Tensor, while the other operand is
- scalar like type (e.g., numeric types, Expr, or TensorSlice),
- it performs tensor-scalar {op} operation on an element-wise basis.
-
- Otherwise, it performs default generic.{op} operation, as defined
- in tvm.tirx.generic module.
-
- Parameters
- ----------
- lhs : object
- Left operand.
- rhs : object
- Right operand.
-
- Returns
- -------
- ret : tvm.te.Tensor (if at least one operand is non-zero-rank Tensor)
- tvm.Expr (otherwise)
- The result of {op} operation.
- """
- if not isinstance(lhs, te.tensor.Tensor) and not isinstance(rhs,
te.tensor.Tensor):
- return orig_bop(lhs, rhs)
- return broadcast_bop(lhs, rhs)
-
- _tensor_bop_impl.__doc__ = _tensor_bop_impl.__doc__.format(op=name)
- return _tensor_bop_impl
-
-
-def _bind_generic_ops():
- """Bind generic operators for Tensor."""
- # Check __op_priority__ to make sure the binding happens only once.
- __op_priority__ = 1
- if __op_priority__ > tvm.tirx.generic.__op_priority__:
- tvm.tirx.generic.__op_priority__ = __op_priority__
- tvm.tirx.generic.add = _make_bop(_broadcast.add, tvm.tirx.generic.add)
- tvm.tirx.generic.subtract = _make_bop(_broadcast.subtract,
tvm.tirx.generic.subtract)
- tvm.tirx.generic.multiply = _make_bop(_broadcast.multiply,
tvm.tirx.generic.multiply)
- tvm.tirx.generic.divide = _make_bop(_broadcast.divide,
tvm.tirx.generic.divide)
- tvm.tirx.generic.cast = _math.cast
-
-
-_bind_generic_ops()
diff --git a/python/tvm/topi/gpu/scan.py b/python/tvm/topi/gpu/scan.py
index 0235c8c3a6..a0b35bcb19 100644
--- a/python/tvm/topi/gpu/scan.py
+++ b/python/tvm/topi/gpu/scan.py
@@ -17,6 +17,7 @@
# pylint: disable=invalid-name, too-many-locals, too-many-statements
"Scan related operators"
+import operator
from collections.abc import Callable
import tvm
@@ -29,11 +30,13 @@ from ..math import cast, ceil_log2
from ..transform import expand_dims, reshape, squeeze, transpose
from ..utils import ceil_div, get_const_int, prod, swap
+_THRUST_SUM_SCAN = "tvm.contrib.thrust.sum_scan"
+
def _get_thrust_func_name(tvmop):
- tvmop_to_thrust_func_name = {tvm.tirx.generic.add:
"tvm.contrib.thrust.sum_scan"}
- assert tvmop in tvmop_to_thrust_func_name, f"{tvmop} not supported by
thrust"
- return tvmop_to_thrust_func_name[tvmop]
+ if tvmop is not operator.add:
+ raise ValueError(f"{tvmop} not supported by thrust")
+ return _THRUST_SUM_SCAN
def _can_use_scan_thrust(binop):
@@ -43,16 +46,15 @@ def _can_use_scan_thrust(binop):
target = tvm.target.Target.current()
if target is None:
return False
- # pylint: disable=comparison-with-callable
- return binop == tvm.tirx.generic.add and any(
+ return binop is operator.add and any(
[
- can_use_thrust(target, "tvm.contrib.thrust.sum_scan"),
- can_use_rocthrust(target, "tvm.contrib.thrust.sum_scan"),
+ can_use_thrust(target, _THRUST_SUM_SCAN),
+ can_use_rocthrust(target, _THRUST_SUM_SCAN),
]
)
-def exclusive_scan_ir(data, output, reduction=None,
binop=tvm.tirx.generic.add, identity_value=0):
+def exclusive_scan_ir(data, output, reduction=None, binop=operator.add,
identity_value=0):
"""Low level IR to do exclusive sum scan along rows of 2D input.
Parameters
@@ -68,8 +70,8 @@ def exclusive_scan_ir(data, output, reduction=None,
binop=tvm.tirx.generic.add,
binop: function, optional
A binary associative op to use for scan. The function takes two TIR
expressions
- and produce a new TIR expression. By default it uses
tvm.tirx.generic.add to compute
- prefix sum.
+ and produce a new TIR expression. By default it uses ``operator.add``
to compute prefix
+ sum.
identity_value: int or float
A value for the binary operation which provides the identity property.
E.g. if * is
@@ -140,9 +142,7 @@ def exclusive_scan_ir(data, output, reduction=None,
binop=tvm.tirx.generic.add,
T.attr(
bx,
"thread_extent",
- tvm.tirx.generic.cast(
- ceil_div(scan_axis_size, max_threads *
width), "int32"
- ),
+ cast(ceil_div(scan_axis_size, max_threads *
width), "int32"),
),
T.attr(by, "thread_extent", nthread_by),
]
@@ -188,9 +188,7 @@ def exclusive_scan_ir(data, output, reduction=None,
binop=tvm.tirx.generic.add,
T.attr(
bx,
"thread_extent",
- tvm.tirx.generic.cast(
- ceil_div(scan_axis_size, max_threads *
width), "int32"
- ),
+ cast(ceil_div(scan_axis_size, max_threads *
width), "int32"),
),
T.attr(by, "thread_extent", nthread_by),
]
@@ -218,7 +216,7 @@ def exclusive_scan_ir(data, output, reduction=None,
binop=tvm.tirx.generic.add,
return ib.get()
-def get_reduction_from_exclusive_scan(data, ex_scan_output,
binop=tvm.tirx.generic.add):
+def get_reduction_from_exclusive_scan(data, ex_scan_output,
binop=operator.add):
"""Return the sum of the last element of data and the exclusive scan
output.
The is the reduction of data along each row (for 2-D case).
@@ -232,8 +230,8 @@ def get_reduction_from_exclusive_scan(data, ex_scan_output,
binop=tvm.tirx.gener
binop: function, optional
A binary associative op to use for scan. The function takes two TIR
expressions
- and produce a new TIR expression. By default it uses
tvm.tirx.generic.add to compute
- prefix sum.
+ and produce a new TIR expression. By default it uses ``operator.add``
to compute prefix
+ sum.
Returns
-------
@@ -312,7 +310,7 @@ def scan_thrust(
output_dtype,
exclusive=True,
return_reduction=False,
- binop=tvm.tirx.generic.add,
+ binop=operator.add,
workspace=None,
):
"""Do exclusive or inclusive scan on 1D or multidimensional input, using
thrust.
@@ -336,7 +334,7 @@ def scan_thrust(
binop: function, optional
A binary associative op to use for scan. Since we need to lookup the
corresponding
thrust function, arbitrariy callables are not supported. Currently only
- tvm.tirx.generic.add can be passed in.
+ ``operator.add`` can be passed in.
workspace: Optional[tvm.te.Tensor]
A buffer to store intermediate results. The size of the workspace
should be sufficiently
@@ -397,7 +395,7 @@ def exclusive_scan(
axis=-1,
return_reduction=False,
output_dtype=None,
- binop=tvm.tirx.generic.add,
+ binop=operator.add,
identity_value=0,
workspace=None,
):
@@ -422,8 +420,8 @@ def exclusive_scan(
binop: function, optional
A binary associative op to use for scan. The function takes two TIR
expressions
- and produce a new TIR expression. By default it uses
tvm.tirx.generic.add to compute
- prefix sum.
+ and produce a new TIR expression. By default it uses ``operator.add``
to compute prefix
+ sum.
identity_value: int or float
A value for the binary operation which provides the identity property.
E.g. if * is
@@ -534,7 +532,7 @@ def exclusive_scan(
def inclusive_scan(
- data, axis=-1, output_dtype=None, binop=tvm.tirx.generic.add,
identity_value=0, workspace=None
+ data, axis=-1, output_dtype=None, binop=operator.add, identity_value=0,
workspace=None
):
"""Do inclusive scan on 1D or multidimensional input.
@@ -551,8 +549,8 @@ def inclusive_scan(
binop: function, optional
A binary associative op to use for scan. The function takes two TIR
expressions
- and produce a new TIR expression. By default it uses
tvm.tirx.generic.add to compute
- prefix sum.
+ and produce a new TIR expression. By default it uses ``operator.add``
to compute prefix
+ sum.
identity_value: int or float
A value for the binary operation which provides the identity property.
E.g. if * is
@@ -717,7 +715,7 @@ def cumsum(
"""
return scanop(
data=data,
- binop=tvm.tirx.generic.add,
+ binop=operator.add,
identity_value=0,
axis=axis,
dtype=dtype,
@@ -767,7 +765,7 @@ def cumprod(
"""
return scanop(
data=data,
- binop=tvm.tirx.generic.multiply,
+ binop=operator.mul,
identity_value=1,
axis=axis,
dtype=dtype,
diff --git a/python/tvm/topi/gpu/sort.py b/python/tvm/topi/gpu/sort.py
index 95ef358952..2ca4056438 100644
--- a/python/tvm/topi/gpu/sort.py
+++ b/python/tvm/topi/gpu/sort.py
@@ -494,12 +494,12 @@ def _sort_common(
target = tvm.target.Target.current()
if "vulkan" in str(target):
ntx = max_threads
- nbx = tvm.tirx.generic.cast(ceil_div(width, max_threads *
thread_work), "int32")
- nbz = tvm.tirx.generic.cast(ceil_div(size, width), "int32")
+ nbx = cast(ceil_div(width, max_threads * thread_work), "int32")
+ nbz = cast(ceil_div(size, width), "int32")
else:
- ntx = tvm.tirx.generic.cast(tvm.te.min(max_threads, width),
"int32")
- nbx = tvm.tirx.generic.cast(ceil_div(width, max_threads *
thread_work), "int32")
- nbz = tvm.tirx.generic.cast(ceil_div(size, width), "int32")
+ ntx = cast(tvm.te.min(max_threads, width), "int32")
+ nbx = cast(ceil_div(width, max_threads * thread_work), "int32")
+ nbz = cast(ceil_div(size, width), "int32")
tx, bx, by, _, _, _ = _get_threads(ntx, nbx, nthread_by * nbz)
with T.frame_scope(
@@ -635,9 +635,7 @@ def sort_ir(
indices_out,
value_init_func=(
lambda _, tid: (
- tvm.tirx.generic.cast(tid, indices_out_orig.dtype)
- if indices_out is not None
- else None
+ cast(tid, indices_out_orig.dtype) if indices_out
is not None else None
)
),
)
diff --git a/python/tvm/topi/scan.py b/python/tvm/topi/scan.py
index 07dc2cddd5..591eb9b9cd 100644
--- a/python/tvm/topi/scan.py
+++ b/python/tvm/topi/scan.py
@@ -17,6 +17,7 @@
# pylint: disable=invalid-name
"""Scan (cumulative binary) operators"""
+import operator
from collections.abc import Callable
import tvm
@@ -24,7 +25,7 @@ from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder import tirx as T
from ..te import extern
-from ..tirx import decl_buffer, generic
+from ..tirx import decl_buffer
from . import utils
from .math import cast
@@ -186,7 +187,7 @@ def cumsum(
"""
return scanop(
data=data,
- binop=generic.add,
+ binop=operator.add,
identity_value=0,
op_name="cumsum_generic",
axis=axis,
@@ -230,7 +231,7 @@ def cumprod(
"""
return scanop(
data=data,
- binop=generic.multiply,
+ binop=operator.mul,
identity_value=1,
op_name="cumprod_generic",
axis=axis,
