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 46c52394a4 [REFACTOR][TE] Represent tensor loads with opaque callees
(#20225)
46c52394a4 is described below
commit 46c52394a47a2841140c412345961940cdd4035d
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Aug 29 13:53:12 2026 -0400
[REFACTOR][TE] Represent tensor loads with opaque callees (#20225)
This change removes the dedicated ProducerLoad/DataProducer path in
favor of ordinary Call nodes whose callee is a TE Tensor.
- Add construction-only OpaqueExpr and OpaqueType markers and reject
them at TE PrimFunc completion.
- Make Tensor an OpaqueExpr and preserve tensor-load traversal through
generic visitor/mutator inheritance fallback.
- Lower Tensor-callee Calls to BufferLoad during PrimFunc creation.
- Remove implicit Tensor-to-PrimExpr conversion and request rank-zero
scalar elements explicitly at affected TOPI packed boundaries.
Validation:
- Clean dependency and incremental builds.
- Pre-commit on all changed files and clean diff/source audits.
- IR, TE, TIRx-base, focused TIRx/Relax visitor, and
all-platform-minimal suites.
- Explicit construction/lowering checks for all affected scalar TOPI
boundaries and non-scalar rank rejection.
---
include/tvm/ir/base_expr.h | 47 +++++++++++++++
include/tvm/relax/expr_functor.h | 2 -
include/tvm/s_tir/stmt.h | 2 +-
include/tvm/te/tensor.h | 28 +++++----
include/tvm/tirx/buffer.h | 43 --------------
include/tvm/tirx/expr.h | 40 -------------
include/tvm/tirx/expr_functor.h | 8 +--
include/tvm/topi/transform.h | 5 +-
python/tvm/ir/__init__.py | 3 +-
python/tvm/ir/expr.py | 5 ++
python/tvm/ir/type.py | 8 +++
python/tvm/relax/expr_functor.py | 2 +-
python/tvm/te/tensor.py | 9 ++-
python/tvm/tirx/__init__.py | 3 +-
python/tvm/tirx/buffer.py | 7 +--
python/tvm/tirx/expr.py | 34 +----------
python/tvm/tirx/expr_functor.py | 39 ++++++-------
python/tvm/tirx/functor.py | 38 ------------
python/tvm/tirx/script/builder/ir.py | 2 -
src/arith/z3_prover.cc | 3 +-
src/ir/expr.cc | 1 +
src/ir/type.cc | 4 ++
src/relax/ir/block_builder.cc | 2 +-
src/relax/ir/expr_functor.cc | 14 ++++-
src/te/operation/compute_op.cc | 4 +-
src/te/operation/create_primfunc.cc | 54 +++++++++++++----
src/te/tensor.cc | 68 +++++++++++++++++-----
src/tirx/analysis/deep_equal.cc | 7 ---
src/tirx/analysis/side_effect.cc | 5 +-
src/tirx/ir/expr.cc | 31 +++-------
src/tirx/ir/expr_functor.cc | 25 ++++----
src/tirx/ir/py_functor.cc | 12 ----
src/tirx/ir/tir_visitor_with_path.cc | 6 +-
src/tirx/ir/tir_visitor_with_path.h | 4 +-
src/tirx/script/printer/buffer.cc | 8 ---
src/tirx/script/printer/expr.cc | 3 +-
src/topi/elemwise.cc | 21 +++++--
src/topi/nn.cc | 23 +++++---
src/topi/transform.cc | 40 +++++++++----
src/topi/utils.cc | 18 ++++--
tests/python/relax/test_expr_functor.py | 14 ++++-
tests/python/te/test_te_create_primfunc.py | 5 +-
tests/python/te/test_te_tensor.py | 9 ++-
tests/python/tirx-base/test_tir_expr_functor.py | 24 ++++----
.../tirx/transform/test_tirx_expr_functor.py | 24 ++++----
45 files changed, 386 insertions(+), 368 deletions(-)
diff --git a/include/tvm/ir/base_expr.h b/include/tvm/ir/base_expr.h
index 443ae42a69..1e39240a7f 100644
--- a/include/tvm/ir/base_expr.h
+++ b/include/tvm/ir/base_expr.h
@@ -85,6 +85,30 @@ class Type : public ffi::ObjectRef {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Type, ffi::ObjectRef,
TypeNode);
};
+/*!
+ * \brief Type marker for opaque construction-time expressions.
+ *
+ * Opaque values may be used while constructing IR, but must be lowered away
+ * before the IR is considered complete.
+ */
+class OpaqueTypeNode final : public TypeNode {
+ public:
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<OpaqueTypeNode>();
+ }
+
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.OpaqueType", OpaqueTypeNode, TypeNode);
+};
+
+/*! \brief Managed reference to OpaqueTypeNode. */
+class OpaqueType final : public Type {
+ public:
+ TVM_DLL OpaqueType();
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(OpaqueType, Type,
OpaqueTypeNode);
+};
+
/*!
* \brief Primitive data types used in the low-level IR.
*
@@ -319,6 +343,29 @@ class Expr : public ffi::ObjectRef {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Expr, ffi::ObjectRef, ExprNode);
};
+/*!
+ * \brief Base node for opaque construction-time expressions.
+ *
+ * Subclasses are passed through by generic expression visitors and mutators.
+ * They must not remain in finished IR.
+ */
+class OpaqueExprNode : public ExprNode {
+ public:
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<OpaqueExprNode>();
+ }
+
+ static constexpr const uint32_t _type_child_slots = 1;
+ TVM_FFI_DECLARE_OBJECT_INFO("ir.OpaqueExpr", OpaqueExprNode, ExprNode);
+};
+
+/*! \brief Managed reference to OpaqueExprNode. */
+class OpaqueExpr : public Expr {
+ public:
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(OpaqueExpr, Expr, OpaqueExprNode);
+};
+
class Call;
/*!
diff --git a/include/tvm/relax/expr_functor.h b/include/tvm/relax/expr_functor.h
index 98fd2fa149..5a65f1bbef 100644
--- a/include/tvm/relax/expr_functor.h
+++ b/include/tvm/relax/expr_functor.h
@@ -156,7 +156,6 @@ class ExprFunctor<R(const Expr& n, Args...)> {
virtual R VisitExpr_(const FunctionNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const CallNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const tirx::BufferLoadNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
- virtual R VisitExpr_(const tirx::ProducerLoadNode* op, Args...)
EXPR_FUNCTOR_DISABLED;
virtual R VisitExpr_(const tirx::LetNode* op, Args...) EXPR_FUNCTOR_DISABLED;
virtual R VisitExpr_(const tirx::ReduceNode* op, Args...)
EXPR_FUNCTOR_DISABLED;
virtual R VisitExpr_(const tirx::AddNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
@@ -211,7 +210,6 @@ class ExprFunctor<R(const Expr& n, Args...)> {
RELAX_EXPR_FUNCTOR_DISPATCH(GlobalVarNode);
RELAX_EXPR_FUNCTOR_DISPATCH(FunctionNode);
RELAX_EXPR_FUNCTOR_DISPATCH(CallNode);
- RELAX_EXPR_FUNCTOR_DISPATCH(tirx::ProducerLoadNode);
RELAX_EXPR_FUNCTOR_DISPATCH(tirx::LetNode);
RELAX_EXPR_FUNCTOR_DISPATCH(tirx::ReduceNode);
RELAX_EXPR_FUNCTOR_DISPATCH(tirx::BufferLoadNode);
diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h
index ae5f2e8124..a05359c99a 100644
--- a/include/tvm/s_tir/stmt.h
+++ b/include/tvm/s_tir/stmt.h
@@ -204,7 +204,7 @@ constexpr const char* warp_execution = "warp_execution";
/*!
* \brief Marks the layout transforms to be used for a tensor.
*
- * Only applies to a DataProducer, as it should be made part of the
+ * Only applies to a tensor-like input, as it should be made part of the
* PrimFunc attributes for TIR.
*/
constexpr const char* layout_transforms = "layout_transforms";
diff --git a/include/tvm/te/tensor.h b/include/tvm/te/tensor.h
index 4ceac38473..5295ff9ed8 100644
--- a/include/tvm/te/tensor.h
+++ b/include/tvm/te/tensor.h
@@ -66,8 +66,8 @@ class Operation : public ffi::ObjectRef {
using ContainerType = OperationNode;
};
-/*! \brief Node to represent a tensor */
-class TensorNode : public DataProducerNode {
+/*! \brief Opaque construction-time node that represents a tensor. */
+class TensorNode : public OpaqueExprNode {
public:
/*! \brief The shape of the tensor */
ffi::Array<PrimExpr> shape;
@@ -80,24 +80,22 @@ class TensorNode : public DataProducerNode {
static void RegisterReflection();
- ffi::Array<PrimExpr> GetShape() const final { return shape; }
+ ffi::Array<PrimExpr> GetShape() const { return shape; }
- PrimType GetDataType() const final { return dtype; }
+ PrimType GetDataType() const { return dtype; }
- TVM_DLL PrimExpr ToPrimExpr() const final;
-
- TVM_DLL ffi::String GetNameHint() const final;
+ TVM_DLL ffi::String GetNameHint() const;
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindConstTreeNode;
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.Tensor", TensorNode, DataProducerNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.Tensor", TensorNode, OpaqueExprNode);
};
/*!
* \brief Tensor structure representing a possible input,
* or intermediate computation result.
*/
-class Tensor : public DataProducer {
+class Tensor : public OpaqueExpr {
private:
/*!
* \brief Helper for indexing operations into tensors
@@ -109,6 +107,7 @@ class Tensor : public DataProducer {
public:
TVM_DLL Tensor(ffi::Array<PrimExpr> shape, PrimType dtype, Operation op, int
value_index);
+
/*!
* \brief check if two tensors equals each other.
* \param other tensor to be checked.
@@ -205,9 +204,18 @@ class Tensor : public DataProducer {
*/
inline Slice operator[](PrimExpr i) const { return Slice(*this, {i}); }
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Tensor, DataProducer, TensorNode);
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Tensor, OpaqueExpr, TensorNode);
};
+/*! \brief Return whether an expression is a Call whose callee is a TE Tensor.
*/
+TVM_DLL bool IsTensorLoad(const Expr& expr);
+
+/*! \brief Recover and validate the Tensor callee of a tensor-load Call. */
+TVM_DLL Tensor GetTensorFromLoad(const Call& call);
+
+/*! \brief Recover and validate the primitive indices of a tensor-load Call. */
+TVM_DLL ffi::Array<PrimExpr> GetTensorLoadIndices(const Call& call);
+
// Implementations of inline functions
inline size_t Tensor::ndim() const { return (*this)->shape.size(); }
diff --git a/include/tvm/tirx/buffer.h b/include/tvm/tirx/buffer.h
index 9413cd6b7f..24cc3b9f1a 100644
--- a/include/tvm/tirx/buffer.h
+++ b/include/tvm/tirx/buffer.h
@@ -352,49 +352,6 @@ TVM_DLL BufferVar decl_buffer(ffi::Array<PrimExpr> shape,
PrimType dtype = PrimT
ffi::String name = "buffer", ffi::String
storage_scope = "",
Span span = Span());
-/*!
- * \brief Base node for data producers.
- *
- * A DataProducer stores necessary information(e.g. a tensor expression) to
produce
- * a multi-dimensional array. The stored information is opaque to the TIR.
- * DataProducer can appear in high-level DSLs that are built on top of the
TIR.
- *
- * A valid TIR PrimFunc should not contain any DataProducer, high level DSLs
should lower
- * all DataProducers to Buffers before TIR transformations.
- *
- * \sa tvm::te::Tensor
- */
-class DataProducerNode : public PrimExprConvertibleNode {
- public:
- /*! \brief destructor. */
- virtual ~DataProducerNode() {}
- /*!
- * \brief Get the shape of the result.
- * \return The shape.
- */
- virtual ffi::Array<PrimExpr> GetShape() const = 0;
- /*!
- * \brief Get the raw element dtype of the result.
- * \return The raw dtype.
- */
- virtual PrimType GetDataType() const = 0;
- /*!
- * \brief Get the name hint of the data producer.
- * \return The data type.
- */
- virtual ffi::String GetNameHint() const = 0;
- TVM_FFI_DECLARE_OBJECT_INFO("tirx.DataProducer", DataProducerNode,
PrimExprConvertibleNode);
-};
-
-/*!
- * \brief Managed reference to DataProducerNode.
- * \sa DataProducerNode
- */
-class DataProducer : public PrimExprConvertible {
- public:
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DataProducer,
PrimExprConvertible, DataProducerNode);
-};
-
/*!
* \brief Creates a TIR buffer for the provided parameters.
* \param shape shape of the buffer
diff --git a/include/tvm/tirx/expr.h b/include/tvm/tirx/expr.h
index 96500d475f..4a4d4b7ce8 100644
--- a/include/tvm/tirx/expr.h
+++ b/include/tvm/tirx/expr.h
@@ -589,44 +589,6 @@ class BufferLoad : public PrimExpr {
TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferLoadNode);
};
-/*!
- * \brief Load value from the result produced by the producer.
- *
- * \note This node only appears in high-level DSLs that are built on top of
the TIR.
- * It should not appear in a valid TIR PrimFunc. A high-level DSL needs
to lower
- * this node before TIR transformations.
- *
- * \sa ProducerLoad, DataProducerNode
- */
-class ProducerLoadNode : public ExprNode {
- public:
- /*! \brief The buffer producer. */
- DataProducer producer;
- /*! \brief The location arguments. */
- ffi::Array<PrimExpr> indices;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ProducerLoadNode>()
- .def_ro("producer", &ProducerLoadNode::producer)
- .def_ro("indices", &ProducerLoadNode::indices);
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.ProducerLoad", ProducerLoadNode,
ExprNode);
-};
-
-/*!
- * \brief Managed reference to ProducerLoadNode.
- * \sa ProducerLoadNode
- */
-class ProducerLoad : public PrimExpr {
- public:
- TVM_DLL explicit ProducerLoad(DataProducer producer, ffi::Array<PrimExpr>
indices,
- Span span = Span());
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ProducerLoad, PrimExpr,
ProducerLoadNode);
- static constexpr bool _type_container_is_exact = true;
- TVM_DEFINE_OBJECT_REF_COW_METHOD(ProducerLoadNode);
-};
-
/*!
* \brief Construct a vector with lanes elements
* where its i-th element equals base + i * stride.
@@ -927,8 +889,6 @@ inline constexpr bool object_ref_contains_v<PrimExpr,
tirx::SelectNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::BufferLoadNode> =
true;
template <>
-inline constexpr bool object_ref_contains_v<PrimExpr, tirx::ProducerLoadNode>
= true;
-template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::RampNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::BroadcastNode> =
true;
diff --git a/include/tvm/tirx/expr_functor.h b/include/tvm/tirx/expr_functor.h
index 170c3a499d..193bdcadc4 100644
--- a/include/tvm/tirx/expr_functor.h
+++ b/include/tvm/tirx/expr_functor.h
@@ -116,7 +116,7 @@ class ExprFunctor<R(const Expr& n, Args...)> {
// Functions that can be overriden by subclass
virtual R VisitExpr_(const VarNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const BufferLoadNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
- virtual R VisitExpr_(const ProducerLoadNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const OpaqueExprNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const TupleNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const TupleGetItemNode* op, Args... args)
EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const LetNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
@@ -160,7 +160,7 @@ class ExprFunctor<R(const Expr& n, Args...)> {
// Set dispatch
IR_EXPR_FUNCTOR_DISPATCH(VarNode);
IR_EXPR_FUNCTOR_DISPATCH(BufferLoadNode);
- IR_EXPR_FUNCTOR_DISPATCH(ProducerLoadNode);
+ IR_EXPR_FUNCTOR_DISPATCH(OpaqueExprNode);
IR_EXPR_FUNCTOR_DISPATCH(TupleNode);
IR_EXPR_FUNCTOR_DISPATCH(TupleGetItemNode);
IR_EXPR_FUNCTOR_DISPATCH(LetNode);
@@ -212,7 +212,7 @@ class TVM_DLL ExprVisitor : public ExprFunctor<void(const
Expr&)> {
// list of functions to override.
void VisitExpr_(const VarNode* op) override;
void VisitExpr_(const BufferLoadNode* op) override;
- void VisitExpr_(const ProducerLoadNode* op) override;
+ void VisitExpr_(const OpaqueExprNode* op) override;
void VisitExpr_(const TupleNode* op) override;
void VisitExpr_(const TupleGetItemNode* op) override;
void VisitExpr_(const LetNode* op) override;
@@ -260,7 +260,7 @@ class TVM_DLL ExprMutator : protected
ExprFunctor<Expr(const Expr&)> {
// list of functions to override.
Expr VisitExpr_(const VarNode* op) override;
Expr VisitExpr_(const BufferLoadNode* op) override;
- Expr VisitExpr_(const ProducerLoadNode* op) override;
+ Expr VisitExpr_(const OpaqueExprNode* op) override;
Expr VisitExpr_(const TupleNode* op) override;
Expr VisitExpr_(const TupleGetItemNode* op) override;
Expr VisitExpr_(const LetNode* op) override;
diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h
index e886c1839a..6badcc1b2d 100644
--- a/include/tvm/topi/transform.h
+++ b/include/tvm/topi/transform.h
@@ -787,9 +787,8 @@ inline Tensor dynamic_strided_slice(const Tensor& x, const
ffi::Array<PrimExpr>&
arith::Analyzer analyzer;
for (size_t i = 0; i < num_slice_axes; ++i) {
- // Check ProducerLoad to keep backward compatibility for Relax.
- if (!begin[i]->IsInstance<ProducerLoadNode>() &&
!end[i]->IsInstance<ProducerLoadNode>() &&
- !strides[i]->IsInstance<ProducerLoadNode>()) {
+ // Dynamic scalar tensor loads cannot be simplified while inferring shape.
+ if (!te::IsTensorLoad(begin[i]) && !te::IsTensorLoad(end[i]) &&
!te::IsTensorLoad(strides[i])) {
out_shape.push_back(
analyzer->Simplify(GetLength(begin[i], end[i], strides[i],
x->shape[i], assume_inbound)));
} else {
diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py
index d907119372..ca83d21f3a 100644
--- a/python/tvm/ir/__init__.py
+++ b/python/tvm/ir/__init__.py
@@ -33,11 +33,12 @@ from .base import (
# Register Type before Expr. Expr's reflected ``ty`` field otherwise creates
# an auto-generated Type wrapper before the concrete Python class is available.
-from .type import FuncType, PointerType, PrimType, TupleType, Type
+from .type import FuncType, OpaqueType, PointerType, PrimType, TupleType, Type
from .expr import (
Call,
Expr,
GlobalVar,
+ OpaqueExpr,
Range,
Tuple,
TupleGetItem,
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 576a02223f..393e825646 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -35,6 +35,11 @@ class Expr(Node):
ty: "tvm.ir.Type"
+@tvm_ffi.register_object("ir.OpaqueExpr")
+class OpaqueExpr(Expr):
+ """Base class for opaque values that must be removed from finished IR."""
+
+
def is_prim_expr(value: object) -> bool:
"""Return whether an expression has a primitive result type."""
return isinstance(value, Expr) and isinstance(value.ty, tvm.ir.PrimType)
diff --git a/python/tvm/ir/type.py b/python/tvm/ir/type.py
index 6b0aebe95f..015232963e 100644
--- a/python/tvm/ir/type.py
+++ b/python/tvm/ir/type.py
@@ -54,6 +54,14 @@ class Type(Node, Scriptable):
return self.is_(other)
+@tvm_ffi.register_object("ir.OpaqueType")
+class OpaqueType(Type):
+ """Type marker for opaque values that must be removed from finished IR."""
+
+ def __init__(self):
+ self.__init_handle_by_constructor__(_ffi_api.OpaqueType)
+
+
@tvm_ffi.register_object("ir.PrimType")
class PrimType(Type):
"""Primitive data type in the low level IR
diff --git a/python/tvm/relax/expr_functor.py b/python/tvm/relax/expr_functor.py
index 688fb6fd6e..84e9ec0884 100644
--- a/python/tvm/relax/expr_functor.py
+++ b/python/tvm/relax/expr_functor.py
@@ -218,7 +218,7 @@ class ExprFunctor:
ret = self.visit_string_imm_(expr)
elif isinstance(expr, DataTypeImm):
ret = self.visit_data_type_imm_(expr)
- elif isinstance(expr, _tirx.Let | _tirx.Reduce | _tirx.ProducerLoad):
+ elif isinstance(expr, _tirx.Let | _tirx.Reduce):
raise TypeError(f"Relax does not support {type(expr).__name__}
expressions")
elif isinstance(expr, Expr):
if is_prim_expr(expr):
diff --git a/python/tvm/te/tensor.py b/python/tvm/te/tensor.py
index ae1fad55a4..c124e7639c 100644
--- a/python/tvm/te/tensor.py
+++ b/python/tvm/te/tensor.py
@@ -19,8 +19,8 @@
# pylint: disable=invalid-name
import tvm_ffi
+from tvm.ir import OpaqueExpr
from tvm.runtime import Object, ObjectConvertible, const
-from tvm.tirx import DataProducer
from tvm.tirx import expr as _expr
from . import _ffi_api, _te_tensor_overload
@@ -249,16 +249,19 @@ class TensorOpBase:
@tvm_ffi.register_object("te.Tensor")
-class Tensor(DataProducer, TensorOpBase):
+class Tensor(OpaqueExpr, TensorOpBase):
"""Tensor object, to construct, see function.Tensor"""
+ def __repr__(self):
+ return f"Tensor(shape={self.shape}, op.name={self.op.name})"
+
def __call__(self, *indices):
ndim = self.ndim
if len(indices) != ndim:
raise ValueError(
f"Need to provide {ndim} index in tensor but {len(indices)}
was provided"
)
- return _expr.ProducerLoad(self, indices)
+ return _ffi_api.TensorLoad(self, indices)
def __getitem__(self, indices):
return TensorSlice(self, indices)
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index a71105044b..ae228360ae 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -30,7 +30,6 @@ from .buffer import (
Buffer,
BufferAccessKind,
BufferType,
- DataProducer,
buffer_data,
buffer_data_pointer_type,
decl_buffer,
@@ -40,7 +39,7 @@ from .expr import convert
from .expr import Var, Reduce, FloatImm, IntImm, StringImm, Cast
from .expr import Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod
from .expr import Min, Max, EQ, NE, LT, LE, GT, GE, And, Or, Not
-from .expr import Select, BufferLoad, ProducerLoad, Ramp, Broadcast, Shuffle
+from .expr import Select, BufferLoad, Ramp, Broadcast, Shuffle
from .expr import CallEffectKind, Let, IterVar, CommReducer
from .stmt import Stmt, Bind, AssertStmt, ForKind, For, While, Return, Break,
Continue
diff --git a/python/tvm/tirx/buffer.py b/python/tvm/tirx/buffer.py
index be4fc25dca..376707a97f 100644
--- a/python/tvm/tirx/buffer.py
+++ b/python/tvm/tirx/buffer.py
@@ -23,7 +23,7 @@ import tvm_ffi
import tvm
from tvm.ir import PointerType, PrimType, Range, Type
-from tvm.runtime import Object, convert
+from tvm.runtime import convert
from . import _buffer_view, _ffi_api
@@ -647,8 +647,3 @@ tvm.ir.Var.data = property(_buffer_data_property)
# particular, ``isinstance(value, Buffer)`` matches every Var. Runtime checks
# must use ``is_buffer_var(value)``.
Buffer = tvm.ir.Var
-
-
-@tvm_ffi.register_object("tirx.DataProducer")
-class DataProducer(Object):
- pass
diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py
index 3e4902c376..6e29d94444 100644
--- a/python/tvm/tirx/expr.py
+++ b/python/tvm/tirx/expr.py
@@ -38,7 +38,7 @@ from tvm.ir.base import Span
from tvm.runtime import DataTypeCode, Object, ObjectConvertible, Scriptable,
const
from . import _ffi_api
-from .buffer import Buffer, DataProducer
+from .buffer import Buffer
def convert(expr) -> Expr:
@@ -590,7 +590,7 @@ class Reduce(ExprWithOp):
The value index.
init : list of Expr
- The initial value for output. This can be an int, float or ProducerLoad
+ The initial value for output. This can be an int, float, or TE
tensor-load Call.
span : Optional[Span]
The location of this expression in the source code.
@@ -1210,36 +1210,6 @@ class BufferLoad(ExprWithOp):
)
-@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.
diff --git a/python/tvm/tirx/expr_functor.py b/python/tvm/tirx/expr_functor.py
index 27dc87c50a..def3b18bda 100644
--- a/python/tvm/tirx/expr_functor.py
+++ b/python/tvm/tirx/expr_functor.py
@@ -24,7 +24,7 @@ from collections.abc import Callable
from typing import TypeVar
import tvm
-from tvm.ir import Expr, Range, Tuple, TupleGetItem
+from tvm.ir import Expr, OpaqueExpr, Range, Tuple, TupleGetItem
from tvm.tirx import IterVar
T = TypeVar("T")
@@ -51,7 +51,6 @@ class ExprFunctor:
self._dispatch_map = {
"tirx.Var": self.visit_var_,
"tirx.BufferLoad": self.visit_buffer_load_,
- "tirx.ProducerLoad": self.visit_producer_load_,
"tirx.Tuple": self.visit_tuple_,
"tirx.TupleGetItem": self.visit_tuple_get_item_,
"tirx.Let": self.visit_let_,
@@ -109,6 +108,9 @@ class ExprFunctor:
if key in self._dispatch_map:
return self._dispatch_map[key](expr)
+ if isinstance(expr, OpaqueExpr):
+ return self.visit_opaque_expr_(expr)
+
return self.visit_expr_default_(expr)
def visit_var_(self, op):
@@ -119,8 +121,8 @@ class ExprFunctor:
"""Default visitor for BufferLoad node."""
return self.visit_expr_default_(op)
- def visit_producer_load_(self, op):
- """Default visitor for ProducerLoad node."""
+ def visit_opaque_expr_(self, op):
+ """Default visitor for an opaque construction-time expression."""
return self.visit_expr_default_(op)
def visit_tuple_(self, op):
@@ -286,13 +288,9 @@ class ExprVisitor(ExprFunctor):
_visit_array(op.indices, _visit_indices)
- def visit_producer_load_(self, op):
- """Visitor implementation for ProducerLoad."""
-
- def _visit_indices(index):
- self.visit_expr(index)
-
- _visit_array(op.indices, _visit_indices)
+ def visit_opaque_expr_(self, op):
+ """Visitor implementation for an opaque construction-time
expression."""
+ pass
def visit_tuple_(self, op):
"""Visitor implementation for Tuple."""
@@ -310,6 +308,9 @@ class ExprVisitor(ExprFunctor):
def visit_call_(self, op):
"""Visitor implementation for Call."""
+ if isinstance(op.op, OpaqueExpr):
+ self.visit_expr(op.op)
+
def _visit_arg(arg):
self.visit_expr(arg)
@@ -473,14 +474,9 @@ class ExprMutator(ExprFunctor):
else:
return tvm.tirx.BufferLoad(op.buffer, indices, op.predicate)
- def visit_producer_load_(self, op):
- """Mutator implementation for ProducerLoad."""
- indices = [self.visit_expr(index) for index in op.indices]
-
- if all(old_index is new_index for old_index, new_index in
zip(op.indices, indices)):
- return op
- else:
- return tvm.tirx.ProducerLoad(op.producer, indices)
+ def visit_opaque_expr_(self, op):
+ """Mutator implementation for an opaque construction-time
expression."""
+ return op
def visit_tuple_(self, op):
"""Mutator implementation for Tuple."""
@@ -510,12 +506,13 @@ class ExprMutator(ExprFunctor):
def visit_call_(self, op):
"""Mutator implementation for Call."""
+ call_op = self.visit_expr(op.op) if isinstance(op.op, OpaqueExpr) else
op.op
args = [self.visit_expr(arg) for arg in op.args]
- if all(old_arg is new_arg for old_arg, new_arg in zip(op.args, args)):
+ if call_op is op.op and all(old_arg is new_arg for old_arg, new_arg in
zip(op.args, args)):
return op
else:
- return tvm.ir.Call(op.op, args, attrs=op.attrs, span=op.span,
ret_ty=op.ty)
+ return tvm.ir.Call(call_op, args, attrs=op.attrs, span=op.span,
ret_ty=op.ty)
def _mutate_binary_op(self, op_cls, op):
"""Helper to mutate binary operators."""
diff --git a/python/tvm/tirx/functor.py b/python/tvm/tirx/functor.py
index 3536f7c474..ed19fffd15 100644
--- a/python/tvm/tirx/functor.py
+++ b/python/tvm/tirx/functor.py
@@ -50,7 +50,6 @@ from .expr import (
Mul,
Not,
Or,
- ProducerLoad,
Ramp,
Reduce,
Select,
@@ -145,7 +144,6 @@ class _PyStmtExprVisitor(tvm_ffi.core.Object):
# Expr
f_visit_var: Callable | None = None,
f_visit_buffer_load: Callable | None = None,
- f_visit_producer_load: Callable | None = None,
f_visit_let: Callable | None = None,
f_visit_call: Callable | None = None,
f_visit_add: Callable | None = None,
@@ -198,7 +196,6 @@ class _PyStmtExprVisitor(tvm_ffi.core.Object):
# Expr
f_visit_var,
f_visit_buffer_load,
- f_visit_producer_load,
f_visit_let,
f_visit_call,
f_visit_add,
@@ -260,7 +257,6 @@ class PyStmtExprVisitor:
# Expr
"visit_var_",
"visit_buffer_load_",
- "visit_producer_load_",
"visit_let_",
"visit_call_",
"visit_add_",
@@ -507,19 +503,6 @@ class PyStmtExprVisitor:
"""
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type:
ignore
- def visit_producer_load_(self, op: ProducerLoad) -> None:
- """Visit ProducerLoad.
-
- Users can customize this function to overwrite
- VisitProducerLoad_(const ProducerLoadNode* op) on the C++ side.
-
- Parameters
- ----------
- op : ProducerLoad
- The ProducerLoad to be visited.
- """
- _ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type:
ignore
-
def visit_let_(self, op: Let) -> None:
"""Visit Let.
@@ -930,7 +913,6 @@ class _PyStmtExprMutator(tvm_ffi.core.Object):
# Expr
f_visit_var: Callable | None = None,
f_visit_buffer_load: Callable | None = None,
- f_visit_producer_load: Callable | None = None,
f_visit_let: Callable | None = None,
f_visit_call: Callable | None = None,
f_visit_add: Callable | None = None,
@@ -983,7 +965,6 @@ class _PyStmtExprMutator(tvm_ffi.core.Object):
# Expr
f_visit_var,
f_visit_buffer_load,
- f_visit_producer_load,
f_visit_let,
f_visit_call,
f_visit_add,
@@ -1045,7 +1026,6 @@ class PyStmtExprMutator:
# Expr
"visit_var_",
"visit_buffer_load_",
- "visit_producer_load_",
"visit_let_",
"visit_call_",
"visit_add_",
@@ -1369,24 +1349,6 @@ class PyStmtExprMutator:
"""
return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)
# type: ignore
- def visit_producer_load_(self, op: ProducerLoad) -> Expr:
- """Visit ProducerLoad.
-
- Users can customize this function to overwrite
- VisitProducerLoad_(const ProducerLoadNode* op) on the C++ side.
-
- Parameters
- ----------
- op : ProducerLoad
- The ProducerLoad to be visited.
-
- Returns
- -------
- result : Expr
- The mutated Expr.
- """
- return _ffi_api.PyStmtExprMutatorDefaultVisitExpr(self._outer(), op)
# type: ignore
-
def visit_let_(self, op: Let) -> Expr:
"""Visit Let.
diff --git a/python/tvm/tirx/script/builder/ir.py
b/python/tvm/tirx/script/builder/ir.py
index e1b87e1f25..d7897647fd 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/tirx/script/builder/ir.py
@@ -76,7 +76,6 @@ from tvm.tirx.expr import (
Mul,
Not,
Or,
- ProducerLoad,
Ramp,
Reduce,
Select,
@@ -3573,7 +3572,6 @@ __all__ = [
"Not",
"Select",
"BufferLoad",
- "ProducerLoad",
"Ramp",
"Broadcast",
"Shuffle",
diff --git a/src/arith/z3_prover.cc b/src/arith/z3_prover.cc
index 49644c78ce..dfff8d7c23 100644
--- a/src/arith/z3_prover.cc
+++ b/src/arith/z3_prover.cc
@@ -736,7 +736,7 @@ class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)> {
return false;
}
return e->IsInstance<CallNode>() || e->IsInstance<BufferLoadNode>() ||
- e->IsInstance<ProducerLoadNode>() || e->IsInstance<ReduceNode>() ||
+ e->IsInstance<ReduceNode>() ||
(e->IsInstance<CastNode>() &&
!IsZ3SupportedExpr(e.as_or_throw<Cast>()->value.get()));
}
@@ -797,7 +797,6 @@ class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)> {
}
z3::expr VisitExpr_(const VarNode* op) override { return Create(op); }
z3::expr VisitExpr_(const BufferLoadNode* op) override { return Create(op); }
- z3::expr VisitExpr_(const ProducerLoadNode* op) override { return
Create(op); }
z3::expr VisitExpr_(const ReduceNode* op) override { return Create(op); }
z3::expr VisitExpr_(const MinNode* op) override {
auto a = VisitInt(op->a);
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index 10714b7953..5f80e20cbd 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -38,6 +38,7 @@ namespace tvm {
TVM_FFI_STATIC_INIT_BLOCK() {
ExprNode::RegisterReflection();
+ OpaqueExprNode::RegisterReflection();
BaseFuncNode::RegisterReflection();
VarNode::RegisterReflection();
GlobalVarNode::RegisterReflection();
diff --git a/src/ir/type.cc b/src/ir/type.cc
index eb3875eb6f..ea4fa4e156 100644
--- a/src/ir/type.cc
+++ b/src/ir/type.cc
@@ -70,6 +70,7 @@ ffi::ObjectPtr<PrimTypeNode> GetCachedPrimTypeNode(DLDataType
dtype) {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
TypeNode::RegisterReflection();
+ OpaqueTypeNode::RegisterReflection();
PrimTypeNode::RegisterReflection();
refl::TypeAttrDef<PrimTypeNode>()
.attr(refl::type_attr::kAnyHash,
reinterpret_cast<void*>(&PrimTypeAnyHash))
@@ -91,6 +92,8 @@ Type Type::Missing() {
bool Type::IsMissing() const { return this->same_as(Type::Missing()); }
+OpaqueType::OpaqueType() : Type(ffi::UnsafeInit{}) { data_ =
ffi::make_object<OpaqueTypeNode>(); }
+
PrimType::PrimType(DLDataType dtype) : Type(ffi::UnsafeInit{}) {
bool is_opaque_handle = dtype.code ==
static_cast<uint8_t>(DLDataTypeCode::kDLOpaqueHandle);
bool is_void = is_opaque_handle && dtype.bits == 0 && dtype.lanes == 0;
@@ -152,6 +155,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef()
.def("ir.TypeMissing", []() { return Type::Missing(); })
.def("ir.TypeIsMissing", [](Type type) { return type.IsMissing(); })
+ .def("ir.OpaqueType", []() { return OpaqueType(); })
.def("ir.PrimType", [](DLDataType dtype) { return PrimType(dtype); });
}
diff --git a/src/relax/ir/block_builder.cc b/src/relax/ir/block_builder.cc
index ec722df110..115e582f0d 100644
--- a/src/relax/ir/block_builder.cc
+++ b/src/relax/ir/block_builder.cc
@@ -581,7 +581,7 @@ class Normalizer : public BlockBuilderImpl, private
ExprFunctor<Expr(const Expr&
Expr VisitExprDefault_(const ffi::Object* op) final {
Expr expr = ffi::GetRef<Expr>(static_cast<const ExprNode*>(op));
- if (expr.as<PrimExpr>()) return expr;
+ if (expr.as<PrimExpr>() || expr.as<OpaqueExpr>()) return expr;
return ExprFunctor::VisitExprDefault_(op);
}
diff --git a/src/relax/ir/expr_functor.cc b/src/relax/ir/expr_functor.cc
index 3746acd758..76c84cdb4d 100644
--- a/src/relax/ir/expr_functor.cc
+++ b/src/relax/ir/expr_functor.cc
@@ -306,7 +306,12 @@ void ExprVisitor::VisitExpr_(const SeqExprNode* op) {
VisitExprDepTypeFieldIfNeeded(this, op->ty);
}
-void ExprVisitor::VisitExprFallback_(const ExprNode* op) {
this->VisitExprDefault_(op); }
+void ExprVisitor::VisitExprFallback_(const ExprNode* op) {
+ if (op->IsInstance<OpaqueExprNode>()) {
+ return;
+ }
+ this->VisitExprDefault_(op);
+}
void ExprVisitor::VisitExpr_(const StringImmNode* op) {
this->VisitSpan(op->span); }
@@ -657,7 +662,12 @@ Expr ExprMutatorBase::VisitExpr_(const TupleGetItemNode*
op) {
}
}
-Expr ExprMutatorBase::VisitExprFallback_(const ExprNode* op) { return
this->VisitExprDefault_(op); }
+Expr ExprMutatorBase::VisitExprFallback_(const ExprNode* op) {
+ if (op->IsInstance<OpaqueExprNode>()) {
+ return ffi::GetRef<Expr>(op);
+ }
+ return this->VisitExprDefault_(op);
+}
Expr ExprMutatorBase::VisitExpr_(const StringImmNode* op) { return
ffi::GetRef<Expr>(op); }
diff --git a/src/te/operation/compute_op.cc b/src/te/operation/compute_op.cc
index 6a461f0d1c..115ea57b87 100644
--- a/src/te/operation/compute_op.cc
+++ b/src/te/operation/compute_op.cc
@@ -164,8 +164,8 @@ ffi::Array<Tensor> ComputeOpNode::InputTensors() const {
std::unordered_set<Tensor> visited;
for (auto& e : body) {
tirx::PostOrderVisit(e, [&ret, &visited](const ffi::ObjectRef& n) {
- if (auto* pload = n.as<tirx::ProducerLoadNode>()) {
- Tensor t = pload->producer.as_or_throw<Tensor>();
+ if (auto call = n.as<Call>(); call.has_value() &&
IsTensorLoad(call.value())) {
+ Tensor t = GetTensorFromLoad(call.value());
if (!visited.count(t)) {
ret.push_back(t);
visited.insert(t);
diff --git a/src/te/operation/create_primfunc.cc
b/src/te/operation/create_primfunc.cc
index daa1b712eb..2332276f10 100644
--- a/src/te/operation/create_primfunc.cc
+++ b/src/te/operation/create_primfunc.cc
@@ -21,6 +21,7 @@
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/unique_name_supply.h>
@@ -45,20 +46,45 @@
namespace tvm {
namespace tirx {
-/*! \brief The helper mutator that transforms ProducerLoad to BufferLoad */
-class ProducerToBufferTransformer : public StmtExprMutator {
+namespace {
+
+void VerifyNoOpaqueArtifacts(const PrimFunc& func) {
+ ffi::String artifact;
+ ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
+ func,
+ [&](const OpaqueExpr& expr) -> ffi::Expected<ffi::WalkResult> {
+ artifact = expr->GetTypeKey();
+ return ffi::WalkResult::Interrupt();
+ },
+ [&](const OpaqueType& type) -> ffi::Expected<ffi::WalkResult> {
+ artifact = type->GetTypeKey();
+ return ffi::WalkResult::Interrupt();
+ });
+ if (!artifact.empty()) {
+ TVM_FFI_THROW(InternalError) << "CreatePrimFunc produced construction-only
opaque artifact "
+ << artifact;
+ }
+}
+
+} // namespace
+
+/*! \brief The helper mutator that transforms Tensor-callee Calls to
BufferLoad. */
+class TensorLoadToBufferTransformer : public StmtExprMutator {
public:
- explicit ProducerToBufferTransformer(
+ explicit TensorLoadToBufferTransformer(
const std::unordered_map<te::Tensor, BufferVar>& tensor2buffers)
: tensor2buffers_(tensor2buffers) {}
- Expr VisitExpr_(const ProducerLoadNode* op) final {
- auto visited_op =
StmtExprMutator::VisitExpr_(op).as_or_throw<ProducerLoad>();
- te::Tensor tensor = visited_op->producer.as_or_throw<te::Tensor>();
+ Expr VisitExpr_(const CallNode* op) final {
+ Call call = StmtExprMutator::VisitExpr_(op).as_or_throw<Call>();
+ if (!te::IsTensorLoad(call)) {
+ return call;
+ }
+ te::Tensor tensor = te::GetTensorFromLoad(call);
auto it = tensor2buffers_.find(tensor);
TVM_FFI_ICHECK(it != tensor2buffers_.end()) << "IndexError: Cannot find
the tensor " << tensor;
const BufferVar& buffer = it->second;
- return BufferLoad(buffer, visited_op->indices);
+ return BufferLoad(buffer, te::GetTensorLoadIndices(call), std::nullopt,
call->span);
}
private:
@@ -110,8 +136,8 @@ struct CreateFuncInfo {
ffi::Array<te::Tensor> arg_list;
/*! \brief The map from each Tensor to its corresponding buffer. */
std::unordered_map<te::Tensor, BufferVar> tensor2buffers;
- /*! \brief The transformer from ProducerLoad to BufferLoad. */
- ProducerToBufferTransformer transformer;
+ /*! \brief The transformer from Tensor-callee Calls to BufferLoad. */
+ TensorLoadToBufferTransformer transformer;
/*! \brief The buffers should be allocated at function root. */
ffi::Array<BufferVar> root_alloc;
/*! \brief The unique name supply to make block name unique. */
@@ -672,7 +698,7 @@ Stmt GenerateStmtFromExternOp(const te::ExternOp&
extern_op, CreateFuncInfo* inf
BufferSubstituter substituter(var_map, input_buffer_map);
Stmt substituted_body = substituter(extern_op->body);
- ProducerToBufferTransformer transformer(info->tensor2buffers);
+ TensorLoadToBufferTransformer transformer(info->tensor2buffers);
Stmt body = transformer(substituted_body);
// Step 4. Generate opaque block as body.
@@ -761,9 +787,10 @@ PrimFunc GenerateAndCompletePrimFunc(const
ffi::Array<te::Tensor>& arg_list,
TVM_FFI_ICHECK(it != info->tensor2buffers.end());
parameters.push_back(it->second.var());
}
+ Stmt body = info->transformer(SeqStmt::Flatten(root_stmts));
PrimFunc func = WithAttrs(
PrimFunc(/*params=*/std::move(parameters),
- /*body=*/SeqStmt::Flatten(root_stmts),
+ /*body=*/std::move(body),
/*ret_type=*/VoidType()),
{{"global_symbol", ffi::String("main")}, {"tirx.noalias", true},
{tvm::attr::kSTir, true}});
const auto fcomplete = tvm::ffi::Function::GetGlobal("script.Complete");
@@ -798,6 +825,7 @@ PrimFunc CreatePrimFunc(const ffi::Array<te::Tensor>&
arg_list,
func =
IndexDataTypeNormalizer(index_dtype_override.value()).Rewrite(std::move(func));
}
auto result = LayoutFreePlaceholdersNormalizer().Process(std::move(func));
+ VerifyNoOpaqueArtifacts(result);
return result;
}
@@ -828,9 +856,10 @@ PrimFunc GenerateAndCompletePrimFunc(const
ffi::Array<ffi::ObjectRef>& arg_tir_v
parameters.push_back(var.value());
}
}
+ Stmt body = info->transformer(SeqStmt::Flatten(root_stmts));
PrimFunc func = WithAttrs(
PrimFunc(/*params=*/std::move(parameters),
- /*body=*/SeqStmt::Flatten(root_stmts),
+ /*body=*/std::move(body),
/*ret_type=*/VoidType()),
{{"global_symbol", ffi::String("main")}, {"tirx.noalias", true},
{tvm::attr::kSTir, true}});
const auto fcomplete = tvm::ffi::Function::GetGlobal("script.Complete");
@@ -870,6 +899,7 @@ PrimFunc CreatePrimFunc(const ffi::Array<ffi::ObjectRef>&
arg_list,
func =
IndexDataTypeNormalizer(index_dtype_override.value()).Rewrite(std::move(func));
}
auto result = LayoutFreePlaceholdersNormalizer().Process(std::move(func));
+ VerifyNoOpaqueArtifacts(result);
return result;
}
diff --git a/src/te/tensor.cc b/src/te/tensor.cc
index 8eeb08b89f..a795bd48ba 100644
--- a/src/te/tensor.cc
+++ b/src/te/tensor.cc
@@ -56,11 +56,9 @@ inline PrimExpr Tensor::IndexTensor(ffi::Array<PrimExpr>
indices,
bool support_negative_indices) const {
ffi::Array<PrimExpr> shape = (*this)->shape;
- if (shape.size() != 0) {
- TVM_FFI_ICHECK_EQ(shape.size(), indices.size())
- << "Tensor dimension mismatch in read "
- << "ndim = " << ndim() << ", indices.size=" << indices.size();
- }
+ TVM_FFI_ICHECK_EQ(shape.size(), indices.size())
+ << "Tensor dimension mismatch in read "
+ << "ndim = " << ndim() << ", indices.size=" << indices.size();
if (support_negative_indices) {
for (size_t i = 0; i < shape.size(); i++) {
@@ -69,7 +67,12 @@ inline PrimExpr Tensor::IndexTensor(ffi::Array<PrimExpr>
indices,
indices.Set(i, new_index);
}
}
- return ProducerLoad((*this), indices);
+ ffi::Array<Expr> args;
+ args.reserve(indices.size());
+ for (const PrimExpr& index : indices) {
+ args.push_back(index);
+ }
+ return PrimExpr(Call((*this)->dtype, *this, args));
}
PrimExpr Tensor::operator()(ffi::Array<PrimVar> indices) const {
@@ -96,19 +99,13 @@ ffi::String TensorNode::GetNameHint() const {
return op->num_outputs() == 1 ? op->name : (op->name + ".v" +
std::to_string(value_index));
}
-PrimExpr TensorNode::ToPrimExpr() const { return ffi::GetRef<Tensor>(this)(); }
-
Tensor Operation::output(size_t i) const {
- auto node = ffi::make_object<TensorNode>();
- node->op = *this;
- node->value_index = i;
- node->dtype = (*this)->output_dtype(i);
- node->shape = (*this)->output_shape(i);
- return Tensor(node);
+ return Tensor((*this)->output_shape(i), (*this)->output_dtype(i), *this,
static_cast<int>(i));
}
Tensor::Tensor(ffi::Array<PrimExpr> shape, PrimType dtype, Operation op, int
value_index) {
auto n = ffi::make_object<TensorNode>();
+ n->ExprNode::ty = OpaqueType();
n->shape = std::move(shape);
n->dtype = dtype;
n->op = op;
@@ -116,6 +113,47 @@ Tensor::Tensor(ffi::Array<PrimExpr> shape, PrimType dtype,
Operation op, int val
data_ = std::move(n);
}
+bool IsTensorLoad(const Expr& expr) {
+ const auto* call = expr.as<CallNode>();
+ return call != nullptr && call->op.as<TensorNode>() != nullptr;
+}
+
+namespace {
+
+ffi::Array<PrimExpr> ValidateTensorLoad(const Call& call, Tensor* tensor_out) {
+ const auto* tensor_node = call->op.as<TensorNode>();
+ TVM_FFI_ICHECK(tensor_node != nullptr) << "Expected a Call whose callee is a
TE Tensor";
+ Tensor tensor = ffi::GetRef<Tensor>(tensor_node);
+ TVM_FFI_ICHECK_EQ(call->args.size(), tensor->shape.size())
+ << "Tensor-load index count must match tensor rank";
+ TVM_FFI_ICHECK(call->ty.as<PrimTypeNode>() != nullptr && call->ty ==
tensor->dtype)
+ << "Tensor-load result type must match the tensor element type";
+
+ ffi::Array<PrimExpr> indices;
+ indices.reserve(call->args.size());
+ for (const Expr& arg : call->args) {
+ auto index = arg.as<PrimExpr>();
+ TVM_FFI_ICHECK(index.has_value()) << "Tensor-load indices must have
primitive type";
+ indices.push_back(index.value());
+ }
+ if (tensor_out != nullptr) {
+ *tensor_out = std::move(tensor);
+ }
+ return indices;
+}
+
+} // namespace
+
+Tensor GetTensorFromLoad(const Call& call) {
+ Tensor tensor;
+ ValidateTensorLoad(call, &tensor);
+ return tensor;
+}
+
+ffi::Array<PrimExpr> GetTensorLoadIndices(const Call& call) {
+ return ValidateTensorLoad(call, nullptr);
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def(
@@ -132,6 +170,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef()
.def_method("te.TensorEqual", &Tensor::operator==)
.def("te.TensorDType", [](Tensor tensor) -> PrimType { return
tensor->dtype; })
+ .def("te.TensorLoad",
+ [](Tensor tensor, ffi::Array<PrimExpr> indices) { return
tensor(indices); })
.def("te.TensorHash",
[](Tensor tensor) -> int64_t {
return static_cast<int64_t>(std::hash<Tensor>()(tensor));
diff --git a/src/tirx/analysis/deep_equal.cc b/src/tirx/analysis/deep_equal.cc
index 48fcb120a2..e915aad1b8 100644
--- a/src/tirx/analysis/deep_equal.cc
+++ b/src/tirx/analysis/deep_equal.cc
@@ -146,13 +146,6 @@ class ExprDeepEqualChecker : private
ExprFunctor<bool(const Expr&, const PrimExp
OptionalDeepEqual(plhs->predicate, prhs->predicate);
}
- bool VisitExpr_(const ProducerLoadNode* plhs, const PrimExpr& rhs) final {
- const auto* prhs = rhs.as<ProducerLoadNode>();
- // run shallow pointer comparison of the producer
- return plhs->ty.as_or_throw<PrimType>() ==
prhs->ty.as_or_throw<PrimType>() &&
- plhs->producer.same_as(prhs->producer) &&
ArrayDeepEqual(plhs->indices, prhs->indices);
- }
-
bool VisitExpr_(const LetNode* plhs, const PrimExpr& rhs) final {
const auto* prhs = rhs.as<LetNode>();
return plhs->ty.as_or_throw<PrimType>() ==
prhs->ty.as_or_throw<PrimType>() &&
diff --git a/src/tirx/analysis/side_effect.cc b/src/tirx/analysis/side_effect.cc
index b5943ba653..eb91637531 100644
--- a/src/tirx/analysis/side_effect.cc
+++ b/src/tirx/analysis/side_effect.cc
@@ -22,6 +22,7 @@
* \brief side effect analysis
*/
#include <tvm/ir/op.h>
+#include <tvm/te/tensor.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/expr_functor.h>
@@ -45,7 +46,9 @@ class ExprSideEffect : public ExprVisitor {
void VisitExpr_(const CallNode* op) final {
static auto op_call_effect =
Op::GetAttrMap<TCallEffectKind>("TCallEffectKind");
- if (auto opt = op->op.as<Op>()) {
+ if (te::IsTensorLoad(ffi::GetRef<Call>(op))) {
+ this->UpdateEffect(CallEffectKind::kReadState);
+ } else if (auto opt = op->op.as<Op>()) {
this->UpdateEffect(static_cast<CallEffectKind>(op_call_effect[opt.value()]));
} else {
this->UpdateEffect(CallEffectKind::kOpaque);
diff --git a/src/tirx/ir/expr.cc b/src/tirx/ir/expr.cc
index a975b31153..860a13f7eb 100644
--- a/src/tirx/ir/expr.cc
+++ b/src/tirx/ir/expr.cc
@@ -22,6 +22,7 @@
*/
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/te/tensor.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>
@@ -95,7 +96,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
NotNode::RegisterReflection();
SelectNode::RegisterReflection();
BufferLoadNode::RegisterReflection();
- ProducerLoadNode::RegisterReflection();
RampNode::RegisterReflection();
BroadcastNode::RegisterReflection();
LetNode::RegisterReflection();
@@ -688,10 +688,13 @@ Reduce::Reduce(CommReducer combiner, ffi::Array<PrimExpr>
source, ffi::Array<Ite
TVM_FFI_ICHECK_EQ(init.size(), source.size()) << "Number of inits should
match number of exprs";
for (size_t i = 0; i < init.size(); i++) {
TVM_FFI_ICHECK(init[i].defined()) << "Init value must be defined";
- TVM_FFI_ICHECK(init[i]->IsInstance<ProducerLoadNode>() ||
init[i]->IsInstance<IntImmNode>() ||
- init[i]->IsInstance<FloatImmNode>())
- << "init can only be a IntImm, FloatImm or ProducerLoad, "
- << "but received " << init[i] << " of type " <<
init[i]->GetTypeKey();
+ if (te::IsTensorLoad(init[i])) {
+ te::GetTensorFromLoad(init[i].as_or_throw<Call>());
+ } else {
+ TVM_FFI_ICHECK(init[i]->IsInstance<IntImmNode>() ||
init[i]->IsInstance<FloatImmNode>())
+ << "init can only be an IntImm, FloatImm or Tensor-load Call, "
+ << "but received " << init[i] << " of type " <<
init[i]->GetTypeKey();
+ }
}
}
n->ExprNode::ty = source[value_index].ty();
@@ -790,23 +793,5 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-// ProducerLoad
-ProducerLoad::ProducerLoad(DataProducer producer, ffi::Array<PrimExpr>
indices, Span span) {
- ffi::ObjectPtr<ProducerLoadNode> node = ffi::make_object<ProducerLoadNode>();
- node->ExprNode::ty = producer->GetDataType();
- node->producer = std::move(producer);
- node->indices = std::move(indices);
- node->span = std::move(span);
- data_ = std::move(node);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.ProducerLoad",
- [](DataProducer producer, ffi::Array<PrimExpr>
indices, Span span) {
- return ProducerLoad(producer, indices, span);
- });
-}
-
} // namespace tirx
} // namespace tvm
diff --git a/src/tirx/ir/expr_functor.cc b/src/tirx/ir/expr_functor.cc
index e465ac970f..9a73caf2c8 100644
--- a/src/tirx/ir/expr_functor.cc
+++ b/src/tirx/ir/expr_functor.cc
@@ -34,9 +34,7 @@ void ExprVisitor::VisitExpr_(const BufferLoadNode* op) {
VisitArray(op->indices, [this](const PrimExpr& e) { this->VisitExpr(e); });
}
-void ExprVisitor::VisitExpr_(const ProducerLoadNode* op) {
- VisitArray(op->indices, [this](const PrimExpr& e) { this->VisitExpr(e); });
-}
+void ExprVisitor::VisitExpr_(const OpaqueExprNode* op) {}
void ExprVisitor::VisitExpr_(const TupleNode* op) {
VisitArray(op->fields, [this](const Expr& e) { this->VisitExpr(e); });
@@ -50,6 +48,9 @@ void ExprVisitor::VisitExpr_(const LetNode* op) {
}
void ExprVisitor::VisitExpr_(const CallNode* op) {
+ if (op->op.as<OpaqueExprNode>()) {
+ this->VisitExpr(op->op);
+ }
VisitArray(op->args, [this](const Expr& e) { this->VisitExpr(e); });
}
@@ -127,15 +128,7 @@ Expr ExprMutator::VisitExpr_(const BufferLoadNode* op) {
}
}
-Expr ExprMutator::VisitExpr_(const ProducerLoadNode* op) {
- auto fmutate = [this](const PrimExpr& e) { return this->VisitPrimExpr(e); };
- ffi::Array<PrimExpr> indices = op->indices.Map(fmutate);
- if (indices.same_as(op->indices)) {
- return ffi::GetRef<PrimExpr>(op);
- } else {
- return ProducerLoad(op->producer, indices);
- }
-}
+Expr ExprMutator::VisitExpr_(const OpaqueExprNode* op) { return
ffi::GetRef<OpaqueExpr>(op); }
Expr ExprMutator::VisitExpr_(const TupleNode* op) {
ffi::Array<Expr> fields =
@@ -160,10 +153,14 @@ Expr ExprMutator::VisitExpr_(const LetNode* op) {
}
Expr ExprMutator::VisitExpr_(const CallNode* op) {
+ Expr call_op = op->op;
+ if (op->op.as<OpaqueExprNode>()) {
+ call_op = this->VisitExpr(op->op);
+ }
ffi::Array<Expr> args =
op->args.Map([this](const Expr& arg) -> Expr { return
this->VisitExpr(arg); });
- if (args.same_as(op->args)) {
+ if (call_op.same_as(op->op) && args.same_as(op->args)) {
return ffi::GetRef<Call>(op);
} else {
Type result_type = op->ExprNode::ty;
@@ -175,7 +172,7 @@ Expr ExprMutator::VisitExpr_(const CallNode* op) {
TVM_FFI_ICHECK(buffer_type);
result_type = buffer_type->DataPointerType();
}
- return Call(result_type, op->op, args, op->attrs, op->ty_args, op->span);
+ return Call(result_type, call_op, args, op->attrs, op->ty_args, op->span);
}
}
diff --git a/src/tirx/ir/py_functor.cc b/src/tirx/ir/py_functor.cc
index 8c20471f7b..ab44d5b0bd 100644
--- a/src/tirx/ir/py_functor.cc
+++ b/src/tirx/ir/py_functor.cc
@@ -104,8 +104,6 @@ class PyStmtExprVisitorNode : public ffi::Object, public
StmtExprVisitor {
ffi::Function f_visit_var{nullptr};
/*! \brief The packed function to the `VisitExpr_(const BufferLoadNode* op)`
function. */
ffi::Function f_visit_buffer_load{nullptr};
- /*! \brief The packed function to the `VisitExpr_(const ProducerLoadNode*
op)` function. */
- ffi::Function f_visit_producer_load{nullptr};
/*! \brief The packed function to the `VisitExpr_(const LetNode* op)`
function. */
ffi::Function f_visit_let{nullptr};
/*! \brief The packed function to the `VisitExpr_(const CallNode* op)`
function. */
@@ -234,7 +232,6 @@ class PyStmtExprVisitorNode : public ffi::Object, public
StmtExprVisitor {
// Expression functions
PY_EXPR_VISITOR_DISPATCH(VarNode, f_visit_var);
PY_EXPR_VISITOR_DISPATCH(BufferLoadNode, f_visit_buffer_load);
- PY_EXPR_VISITOR_DISPATCH(ProducerLoadNode, f_visit_producer_load);
PY_EXPR_VISITOR_DISPATCH(LetNode, f_visit_let);
PY_EXPR_VISITOR_DISPATCH(CallNode, f_visit_call);
PY_EXPR_VISITOR_DISPATCH(AddNode, f_visit_add);
@@ -271,7 +268,6 @@ class PyStmtExprVisitorNode : public ffi::Object, public
StmtExprVisitor {
// Set dispatch
IR_EXPR_VISITOR_DEFAULT_DISPATCH(VarNode);
IR_EXPR_VISITOR_DEFAULT_DISPATCH(BufferLoadNode);
- IR_EXPR_VISITOR_DEFAULT_DISPATCH(ProducerLoadNode);
IR_EXPR_VISITOR_DEFAULT_DISPATCH(TupleNode);
IR_EXPR_VISITOR_DEFAULT_DISPATCH(TupleGetItemNode);
IR_EXPR_VISITOR_DEFAULT_DISPATCH(LetNode);
@@ -353,7 +349,6 @@ class PyStmtExprVisitor : public ffi::ObjectRef {
ffi::Function
f_visit_sblock_realize, //
ffi::Function
f_visit_var, //
ffi::Function
f_visit_buffer_load, //
- ffi::Function
f_visit_producer_load, //
ffi::Function
f_visit_let, //
ffi::Function
f_visit_call, //
ffi::Function
f_visit_add, //
@@ -403,7 +398,6 @@ class PyStmtExprVisitor : public ffi::ObjectRef {
// Set expression functions
n->f_visit_var = std::move(f_visit_var);
n->f_visit_buffer_load = std::move(f_visit_buffer_load);
- n->f_visit_producer_load = std::move(f_visit_producer_load);
n->f_visit_let = std::move(f_visit_let);
n->f_visit_call = std::move(f_visit_call);
n->f_visit_add = std::move(f_visit_add);
@@ -455,8 +449,6 @@ class PyStmtExprMutatorNode : public ffi::Object, public
StmtExprMutator {
ffi::Function f_visit_var{nullptr};
/*! \brief The packed function to the `VisitExpr_(const BufferLoadNode* op)`
function. */
ffi::Function f_visit_buffer_load{nullptr};
- /*! \brief The packed function to the `VisitExpr_(const ProducerLoadNode*
op)` function. */
- ffi::Function f_visit_producer_load{nullptr};
/*! \brief The packed function to the `VisitExpr_(const LetNode* op)`
function. */
ffi::Function f_visit_let{nullptr};
/*! \brief The packed function to the `VisitExpr_(const CallNode* op)`
function. */
@@ -585,7 +577,6 @@ class PyStmtExprMutatorNode : public ffi::Object, public
StmtExprMutator {
// Expression functions
PY_EXPR_MUTATOR_DISPATCH(VarNode, f_visit_var);
PY_EXPR_MUTATOR_DISPATCH(BufferLoadNode, f_visit_buffer_load);
- PY_EXPR_MUTATOR_DISPATCH(ProducerLoadNode, f_visit_producer_load);
PY_EXPR_MUTATOR_DISPATCH(LetNode, f_visit_let);
PY_EXPR_MUTATOR_DISPATCH(CallNode, f_visit_call);
PY_EXPR_MUTATOR_DISPATCH(AddNode, f_visit_add);
@@ -622,7 +613,6 @@ class PyStmtExprMutatorNode : public ffi::Object, public
StmtExprMutator {
// Set dispatch
PY_EXPR_MUTATOR_DEFAULT_DISPATCH(VarNode);
PY_EXPR_MUTATOR_DEFAULT_DISPATCH(BufferLoadNode);
- PY_EXPR_MUTATOR_DEFAULT_DISPATCH(ProducerLoadNode);
PY_EXPR_MUTATOR_DEFAULT_DISPATCH(TupleNode);
PY_EXPR_MUTATOR_DEFAULT_DISPATCH(TupleGetItemNode);
PY_EXPR_MUTATOR_DEFAULT_DISPATCH(LetNode);
@@ -705,7 +695,6 @@ class PyStmtExprMutator : public ffi::ObjectRef {
ffi::Function
f_visit_sblock_realize, //
ffi::Function
f_visit_var, //
ffi::Function
f_visit_buffer_load, //
- ffi::Function
f_visit_producer_load, //
ffi::Function
f_visit_let, //
ffi::Function
f_visit_call, //
ffi::Function
f_visit_add, //
@@ -755,7 +744,6 @@ class PyStmtExprMutator : public ffi::ObjectRef {
// Expression functions
n->f_visit_var = std::move(f_visit_var);
n->f_visit_buffer_load = std::move(f_visit_buffer_load);
- n->f_visit_producer_load = std::move(f_visit_producer_load);
n->f_visit_let = std::move(f_visit_let);
n->f_visit_call = std::move(f_visit_call);
n->f_visit_add = std::move(f_visit_add);
diff --git a/src/tirx/ir/tir_visitor_with_path.cc
b/src/tirx/ir/tir_visitor_with_path.cc
index 882134f600..8954ec7657 100644
--- a/src/tirx/ir/tir_visitor_with_path.cc
+++ b/src/tirx/ir/tir_visitor_with_path.cc
@@ -357,9 +357,7 @@ void TIRVisitorWithPath::VisitExpr_(const BufferLoadNode*
op, AccessPath path) {
Visit(op->indices, path->Attr("indices"));
}
-void TIRVisitorWithPath::VisitExpr_(const ProducerLoadNode* op, AccessPath
path) {
- Visit(op->indices, path->Attr("indices"));
-}
+void TIRVisitorWithPath::VisitExpr_(const OpaqueExprNode* op, AccessPath path)
{}
void TIRVisitorWithPath::VisitExpr_(const TupleNode* op, AccessPath path) {
Visit(op->fields, path->Attr("fields"));
@@ -378,6 +376,8 @@ void TIRVisitorWithPath::VisitExpr_(const LetNode* op,
AccessPath path) {
void TIRVisitorWithPath::VisitExpr_(const CallNode* op, AccessPath path) {
if (auto gvar = op->op.as<GlobalVar>()) {
Visit(gvar.value(), path->Attr("op"));
+ } else if (op->op.as<OpaqueExprNode>()) {
+ Visit(op->op, path->Attr("op"));
}
Visit(op->args, path->Attr("args"));
}
diff --git a/src/tirx/ir/tir_visitor_with_path.h
b/src/tirx/ir/tir_visitor_with_path.h
index d80d44c0cb..05b32bf5cf 100644
--- a/src/tirx/ir/tir_visitor_with_path.h
+++ b/src/tirx/ir/tir_visitor_with_path.h
@@ -66,6 +66,8 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const
Expr&, ffi::reflecti
VisitExpr_(tuple, path);
} else if (auto* tuple_get_item = obj.as<TupleGetItemNode>()) {
VisitExpr_(tuple_get_item, path);
+ } else if (obj.as<OpaqueExprNode>()) {
+ VisitExpr(obj, path);
} else {
TVM_FFI_THROW(TypeError) << "Unsupported non-primitive TIR expression "
<< obj.GetTypeKey();
}
@@ -150,7 +152,7 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const
Expr&, ffi::reflecti
using ExprFunctor::VisitExpr;
void VisitExpr_(const VarNode* op, ffi::reflection::AccessPath path)
override;
void VisitExpr_(const BufferLoadNode* op, ffi::reflection::AccessPath path)
override;
- void VisitExpr_(const ProducerLoadNode* op, ffi::reflection::AccessPath
path) override;
+ void VisitExpr_(const OpaqueExprNode* op, ffi::reflection::AccessPath path)
override;
void VisitExpr_(const TupleNode* op, ffi::reflection::AccessPath path)
override;
void VisitExpr_(const TupleGetItemNode* op, ffi::reflection::AccessPath
path) override;
void VisitExpr_(const LetNode* op, ffi::reflection::AccessPath path)
override;
diff --git a/src/tirx/script/printer/buffer.cc
b/src/tirx/script/printer/buffer.cc
index 950b5acc34..c13a06aa0d 100644
--- a/src/tirx/script/printer/buffer.cc
+++ b/src/tirx/script/printer/buffer.cc
@@ -596,13 +596,6 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
return AssignDoc(lhs, rhs, std::nullopt);
});
-TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
- .set_dispatch<tirx::ProducerLoad>( //
- "", [](tirx::ProducerLoad load, AccessPath p, IRDocsifier d) -> Doc {
- ExprDoc prefix = IdDoc(load->producer->GetNameHint());
- return prefix[BufferIndices(load->indices, p->Attr("indices"), d)];
- });
-
TVM_SCRIPT_REPR(tirx::BufferRegionNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::BufferLoadNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::BufferStoreNode, ReprPrintTIR);
@@ -611,7 +604,6 @@ TVM_SCRIPT_REPR(tirx::IterNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::TileLayoutNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::ComposeLayoutNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::MatchBufferRegionNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::ProducerLoadNode, ReprPrintTIR);
} // namespace printer
} // namespace script
diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc
index b7b4fc6615..b3b2859364 100644
--- a/src/tirx/script/printer/expr.cc
+++ b/src/tirx/script/printer/expr.cc
@@ -408,7 +408,8 @@ Doc PrintTIRCall(Call call, AccessPath call_p, IRDocsifier
d) {
} else if (call->op.as<GlobalVarNode>()) {
prefix = d->AsDoc<ExprDoc>(call->op, call_p->Attr("op"));
} else {
- TVM_FFI_THROW(InternalError) << "call: " << call;
+ TVM_FFI_THROW(TypeError) << "Cannot print a Call whose callee has type "
+ << call->op->GetTypeKey();
}
ffi::Array<ExprDoc> args;
int n_args = call->args.size();
diff --git a/src/topi/elemwise.cc b/src/topi/elemwise.cc
index 4b9d26f276..78f905e6c1 100644
--- a/src/topi/elemwise.cc
+++ b/src/topi/elemwise.cc
@@ -87,8 +87,13 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ffi::Any* rv) { *rv =
negative(args[0].cast<te::Tensor>()); })
.def_packed("topi.clip",
[](ffi::PackedArgs args, ffi::Any* rv) {
- *rv = clip(args[0].cast<te::Tensor>(),
args[1].cast<PrimExpr>(),
- args[2].cast<PrimExpr>());
+ auto a_min_tensor = args[1].try_cast<te::Tensor>();
+ auto a_max_tensor = args[2].try_cast<te::Tensor>();
+ PrimExpr a_min = a_min_tensor ?
a_min_tensor.value()(ffi::Array<PrimExpr>{})
+ : args[1].cast<PrimExpr>();
+ PrimExpr a_max = a_max_tensor ?
a_max_tensor.value()(ffi::Array<PrimExpr>{})
+ : args[2].cast<PrimExpr>();
+ *rv = clip(args[0].cast<te::Tensor>(), a_min, a_max);
})
.def_packed("topi.cast",
[](ffi::PackedArgs args, ffi::Any* rv) {
@@ -106,12 +111,20 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ffi::Any* rv) { *rv =
sign(args[0].cast<te::Tensor>()); })
.def_packed("topi.full",
[](ffi::PackedArgs args, ffi::Any* rv) {
+ auto fill_value_tensor = args[2].try_cast<te::Tensor>();
+ PrimExpr fill_value = fill_value_tensor
+ ?
fill_value_tensor.value()(ffi::Array<PrimExpr>{})
+ : args[2].cast<PrimExpr>();
*rv = full(args[0].cast<ffi::Array<PrimExpr>>(),
args[1].cast<PrimType>(),
- args[2].cast<PrimExpr>());
+ fill_value);
})
.def_packed("topi.full_like",
[](ffi::PackedArgs args, ffi::Any* rv) {
- *rv = full_like(args[0].cast<te::Tensor>(),
args[1].cast<PrimExpr>());
+ auto fill_value_tensor = args[1].try_cast<te::Tensor>();
+ PrimExpr fill_value = fill_value_tensor
+ ?
fill_value_tensor.value()(ffi::Array<PrimExpr>{})
+ : args[1].cast<PrimExpr>();
+ *rv = full_like(args[0].cast<te::Tensor>(), fill_value);
})
.def_packed(
"topi.logical_not",
diff --git a/src/topi/nn.cc b/src/topi/nn.cc
index cd4968a461..32e262ca12 100644
--- a/src/topi/nn.cc
+++ b/src/topi/nn.cc
@@ -62,16 +62,23 @@ TVM_FFI_STATIC_INIT_BLOCK() {
})
.def_packed("topi.nn.pad",
[](ffi::PackedArgs args, ffi::Any* rv) {
+ auto pad_value_tensor = args[3].try_cast<te::Tensor>();
+ PrimExpr pad_value = pad_value_tensor
+ ?
pad_value_tensor.value()(ffi::Array<PrimExpr>{})
+ : args[3].cast<PrimExpr>();
*rv = pad(args[0].cast<te::Tensor>(),
args[1].cast<ffi::Array<PrimExpr>>(),
- args[2].cast<ffi::Array<PrimExpr>>(),
args[3].cast<PrimExpr>());
- })
- .def_packed("topi.nn.space_to_batch_nd",
- [](ffi::PackedArgs args, ffi::Any* rv) {
- *rv = space_to_batch_nd(
- args[0].cast<te::Tensor>(),
args[1].cast<ffi::Array<int64_t>>(),
- args[2].cast<ffi::Array<PrimExpr>>(),
args[3].cast<ffi::Array<PrimExpr>>(),
- args[4].cast<PrimExpr>());
+ args[2].cast<ffi::Array<PrimExpr>>(), pad_value);
})
+ .def_packed(
+ "topi.nn.space_to_batch_nd",
+ [](ffi::PackedArgs args, ffi::Any* rv) {
+ auto pad_value_tensor = args[4].try_cast<te::Tensor>();
+ PrimExpr pad_value = pad_value_tensor ?
pad_value_tensor.value()(ffi::Array<PrimExpr>{})
+ : args[4].cast<PrimExpr>();
+ *rv = space_to_batch_nd(args[0].cast<te::Tensor>(),
args[1].cast<ffi::Array<int64_t>>(),
+ args[2].cast<ffi::Array<PrimExpr>>(),
+ args[3].cast<ffi::Array<PrimExpr>>(),
pad_value);
+ })
.def_packed("topi.nn.batch_to_space_nd",
[](ffi::PackedArgs args, ffi::Any* rv) {
*rv = batch_to_space_nd(
diff --git a/src/topi/transform.cc b/src/topi/transform.cc
index a9d994c2a8..0ecda86956 100644
--- a/src/topi/transform.cc
+++ b/src/topi/transform.cc
@@ -140,8 +140,16 @@ TVM_FFI_STATIC_INIT_BLOCK() {
})
.def_packed("topi.arange",
[](ffi::PackedArgs args, ffi::Any* rv) {
- *rv = arange(args[0].cast<PrimExpr>(),
args[1].cast<PrimExpr>(),
- args[2].cast<PrimExpr>(),
args[3].cast<PrimType>());
+ auto start_tensor = args[0].try_cast<te::Tensor>();
+ auto stop_tensor = args[1].try_cast<te::Tensor>();
+ auto step_tensor = args[2].try_cast<te::Tensor>();
+ PrimExpr start = start_tensor ?
start_tensor.value()(ffi::Array<PrimExpr>{})
+ : args[0].cast<PrimExpr>();
+ PrimExpr stop = stop_tensor ?
stop_tensor.value()(ffi::Array<PrimExpr>{})
+ : args[1].cast<PrimExpr>();
+ PrimExpr step = step_tensor ?
step_tensor.value()(ffi::Array<PrimExpr>{})
+ : args[2].cast<PrimExpr>();
+ *rv = arange(start, stop, step, args[3].cast<PrimType>());
})
.def_packed("topi.meshgrid",
[](ffi::PackedArgs args, ffi::Any* rv) {
@@ -179,9 +187,13 @@ TVM_FFI_STATIC_INIT_BLOCK() {
})
.def_packed("topi.sparse_to_dense",
[](ffi::PackedArgs args, ffi::Any* rv) {
+ auto default_value_tensor = args[3].try_cast<te::Tensor>();
+ PrimExpr default_value =
+ default_value_tensor ?
default_value_tensor.value()(ffi::Array<PrimExpr>{})
+ : args[3].cast<PrimExpr>();
*rv = sparse_to_dense(args[0].cast<te::Tensor>(),
args[1].cast<ffi::Array<PrimExpr>>(),
- args[2].cast<te::Tensor>(),
args[3].cast<PrimExpr>());
+ args[2].cast<te::Tensor>(),
default_value);
})
.def_packed("topi.matmul",
[](ffi::PackedArgs args, ffi::Any* rv) {
@@ -257,14 +269,20 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ffi::Array<PrimExpr> output_shape) {
return relax::dynamic_strided_slice(x, begin, end, strides,
output_shape);
})
- .def_packed("topi.one_hot",
- [](ffi::PackedArgs args, ffi::Any* rv) {
- int depth = args[3].cast<int>();
- int axis = args[4].cast<int>();
- PrimType dtype = args[5].cast<PrimType>();
- *rv = one_hot(args[0].cast<te::Tensor>(),
args[1].cast<PrimExpr>(),
- args[2].cast<PrimExpr>(), depth, axis,
dtype);
- })
+ .def_packed(
+ "topi.one_hot",
+ [](ffi::PackedArgs args, ffi::Any* rv) {
+ int depth = args[3].cast<int>();
+ int axis = args[4].cast<int>();
+ PrimType dtype = args[5].cast<PrimType>();
+ auto on_value_tensor = args[1].try_cast<te::Tensor>();
+ auto off_value_tensor = args[2].try_cast<te::Tensor>();
+ PrimExpr on_value = on_value_tensor ?
on_value_tensor.value()(ffi::Array<PrimExpr>{})
+ : args[1].cast<PrimExpr>();
+ PrimExpr off_value = off_value_tensor ?
off_value_tensor.value()(ffi::Array<PrimExpr>{})
+ : args[2].cast<PrimExpr>();
+ *rv = one_hot(args[0].cast<te::Tensor>(), on_value, off_value,
depth, axis, dtype);
+ })
.def_packed("topi.matrix_set_diag",
[](ffi::PackedArgs args, ffi::Any* rv) {
int k1 = args[2].cast<int>();
diff --git a/src/topi/utils.cc b/src/topi/utils.cc
index 6bc1570bd1..7635d165ef 100644
--- a/src/topi/utils.cc
+++ b/src/topi/utils.cc
@@ -37,14 +37,22 @@ TVM_FFI_STATIC_INIT_BLOCK() {
})
.def_packed("topi.utils.bilinear_sample_nchw",
[](ffi::PackedArgs args, ffi::Any* rv) {
- *rv = detail::bilinear_sample_nchw(
- args[0].cast<te::Tensor>(),
args[1].cast<ffi::Array<PrimExpr>>(),
- args[2].cast<PrimExpr>(), args[3].cast<PrimExpr>());
+ auto y_tensor = args[2].try_cast<te::Tensor>();
+ auto x_tensor = args[3].try_cast<te::Tensor>();
+ PrimExpr y = y_tensor ?
y_tensor.value()(ffi::Array<PrimExpr>{})
+ : args[2].cast<PrimExpr>();
+ PrimExpr x = x_tensor ?
x_tensor.value()(ffi::Array<PrimExpr>{})
+ : args[3].cast<PrimExpr>();
+ *rv =
detail::bilinear_sample_nchw(args[0].cast<te::Tensor>(),
+
args[1].cast<ffi::Array<PrimExpr>>(), y, x);
})
.def_packed("topi.utils.bilinear_sample_nhwc", [](ffi::PackedArgs args,
ffi::Any* rv) {
+ auto y_tensor = args[2].try_cast<te::Tensor>();
+ auto x_tensor = args[3].try_cast<te::Tensor>();
+ PrimExpr y = y_tensor ? y_tensor.value()(ffi::Array<PrimExpr>{}) :
args[2].cast<PrimExpr>();
+ PrimExpr x = x_tensor ? x_tensor.value()(ffi::Array<PrimExpr>{}) :
args[3].cast<PrimExpr>();
*rv = detail::bilinear_sample_nhwc(args[0].cast<te::Tensor>(),
-
args[1].cast<ffi::Array<PrimExpr>>(),
- args[2].cast<PrimExpr>(),
args[3].cast<PrimExpr>());
+
args[1].cast<ffi::Array<PrimExpr>>(), y, x);
});
}
diff --git a/tests/python/relax/test_expr_functor.py
b/tests/python/relax/test_expr_functor.py
index 865848e6e9..532a9f3c79 100644
--- a/tests/python/relax/test_expr_functor.py
+++ b/tests/python/relax/test_expr_functor.py
@@ -19,7 +19,7 @@ import pytest
import tvm
import tvm.testing
-from tvm import relax, tirx
+from tvm import relax, te, tirx
from tvm.ir import Call, Op
from tvm.ir.base import assert_structural_equal
from tvm.relax import PyExprMutator, PyExprVisitor
@@ -445,6 +445,18 @@ def test_call():
"\n".join(["Op", "Var", "Var", "Var", "Var", "ShapeExpr", "Call"]),
)
+ tensor_load = te.placeholder((1,), name="A")(0)
+ BasicVisitor().visit_expr(tensor_load)
+ assert BasicMutator().visit_expr(tensor_load).same_as(tensor_load)
+
+ visitor = ASTPrinter()
+ visitor.visit_expr(tensor_load)
+ assert str(visitor.log) == "\n".join(["Call", "\tExprFallback"])
+
+ mutator = ASTPostPrinterMutator()
+ assert mutator.visit_expr(tensor_load).same_as(tensor_load)
+ assert str(mutator.log) == "\n".join(["ExprFallback", "Call"])
+
def test_if():
if_node = relax.If(x, x, x)
diff --git a/tests/python/te/test_te_create_primfunc.py
b/tests/python/te/test_te_create_primfunc.py
index 9a879e6701..ee1bc60498 100644
--- a/tests/python/te/test_te_create_primfunc.py
+++ b/tests/python/te/test_te_create_primfunc.py
@@ -383,10 +383,7 @@ def test_constant():
M = 11
A = te.placeholder((M,), name="A")
B = te.compute(tuple(), lambda: 2, name="B")
- # Manually craft ProducerLoad because `B[]` is not allowed.
- C = te.compute(
- (M,), lambda x: A[x] + tvm.tirx.expr.ProducerLoad(B, []), name="C",
tag="broadcast"
- )
+ C = te.compute((M,), lambda x: A[x] + B(), name="C", tag="broadcast")
func = te.create_prim_func([C, A])
func = tvm.compile(func)
diff --git a/tests/python/te/test_te_tensor.py
b/tests/python/te/test_te_tensor.py
index 959af34c95..79f6dba0bd 100644
--- a/tests/python/te/test_te_tensor.py
+++ b/tests/python/te/test_te_tensor.py
@@ -33,13 +33,20 @@ def test_tensor():
print(T)
print(T.op.body)
assert tuple(T.shape) == (m, n, l)
+ assert isinstance(A, tvm.ir.OpaqueExpr)
+ assert isinstance(A.ty, tvm.ir.OpaqueType)
assert isinstance(A.op, tvm.te.PlaceholderOp)
assert A == A
assert T.op.output(0) == T
assert T.op.output(0).__hash__() == T.__hash__()
d = {T.op.output(0): 1}
assert d[T] == 1
- assert T[0][0][0].astype("float16").ty == tvm.ir.PrimType("float16")
+ load = T[0][0][0].asobject()
+ assert isinstance(load, tvm.ir.Call)
+ assert load.op.same_as(T)
+ assert list(load.args) == [0, 0, 0]
+ assert load.ty == T.dtype
+ assert load.astype("float16").ty == tvm.ir.PrimType("float16")
def test_rank_zero():
diff --git a/tests/python/tirx-base/test_tir_expr_functor.py
b/tests/python/tirx-base/test_tir_expr_functor.py
index b342782d48..c99550e34c 100644
--- a/tests/python/tirx-base/test_tir_expr_functor.py
+++ b/tests/python/tirx-base/test_tir_expr_functor.py
@@ -18,7 +18,7 @@
import tvm
import tvm.testing
from tvm import tirx as tir
-from tvm.ir import Call, Op, Tuple, TupleGetItem
+from tvm.ir import Call, Op, OpaqueExpr, Tuple, TupleGetItem
from tvm.ir.base import assert_structural_equal
from tvm.tirx.expr import (
EQ,
@@ -44,7 +44,6 @@ from tvm.tirx.expr import (
Mul,
Not,
Or,
- ProducerLoad,
Ramp,
Reduce,
Select,
@@ -104,12 +103,8 @@ class ASTPrinter(ExprVisitor):
self.visit_expr(idx)
self.log.pop_scope()
- def visit_producer_load_(self, op: ProducerLoad) -> None:
- self.log.add("ProducerLoad")
- self.log.push_scope()
- for idx in op.indices:
- self.visit_expr(idx)
- self.log.pop_scope()
+ def visit_opaque_expr_(self, op: OpaqueExpr) -> None:
+ self.log.add("OpaqueExpr")
def visit_tuple_(self, op: Tuple) -> None:
self.log.add("Tuple")
@@ -347,9 +342,9 @@ class ASTPostPrinterMutator(ExprMutator):
self.log.add("BufferLoad")
return result
- def visit_producer_load_(self, op: ProducerLoad) -> tir.Expr:
- result = super().visit_producer_load_(op)
- self.log.add("ProducerLoad")
+ def visit_opaque_expr_(self, op: OpaqueExpr) -> tir.Expr:
+ result = super().visit_opaque_expr_(op)
+ self.log.add("OpaqueExpr")
return result
def visit_tuple_(self, op: Tuple) -> tir.Expr:
@@ -806,6 +801,13 @@ def test_call_visitor_super():
lv.visit_expr(add_node)
assert str(lv.log) == "\n".join(["LeafAdd", "InternalAdd", "InternalVar",
"InternalIntImm"])
+ tensor_load = tvm.te.placeholder((1,), name="A")(0)
+ basic_check(
+ tensor_load,
+ "\n".join(["Call", "\tOpaqueExpr", "\tIntImm"]),
+ "\n".join(["OpaqueExpr", "IntImm", "Call"]),
+ )
+
def test_call_mutator_super():
class InternalMutator(ExprMutator):
diff --git a/tests/python/tirx/transform/test_tirx_expr_functor.py
b/tests/python/tirx/transform/test_tirx_expr_functor.py
index 96bb500111..38845fe6f1 100644
--- a/tests/python/tirx/transform/test_tirx_expr_functor.py
+++ b/tests/python/tirx/transform/test_tirx_expr_functor.py
@@ -18,7 +18,7 @@
import tvm
import tvm.testing
from tvm import tirx as tir
-from tvm.ir import Call, Op
+from tvm.ir import Call, Op, OpaqueExpr
from tvm.ir.base import assert_structural_equal
from tvm.tirx.expr import (
EQ,
@@ -44,7 +44,6 @@ from tvm.tirx.expr import (
Mul,
Not,
Or,
- ProducerLoad,
Ramp,
Reduce,
Select,
@@ -104,12 +103,8 @@ class ASTPrinter(ExprVisitor):
self.visit_expr(idx)
self.log.pop_scope()
- def visit_producer_load_(self, op: ProducerLoad) -> None:
- self.log.add("ProducerLoad")
- self.log.push_scope()
- for idx in op.indices:
- self.visit_expr(idx)
- self.log.pop_scope()
+ def visit_opaque_expr_(self, op: OpaqueExpr) -> None:
+ self.log.add("OpaqueExpr")
def visit_let_(self, op: Let) -> None:
self.log.add("Let")
@@ -334,9 +329,9 @@ class ASTPostPrinterMutator(ExprMutator):
self.log.add("BufferLoad")
return result
- def visit_producer_load_(self, op: ProducerLoad) -> tir.Expr:
- result = super().visit_producer_load_(op)
- self.log.add("ProducerLoad")
+ def visit_opaque_expr_(self, op: OpaqueExpr) -> tir.Expr:
+ result = super().visit_opaque_expr_(op)
+ self.log.add("OpaqueExpr")
return result
def visit_let_(self, op: Let) -> tir.Expr:
@@ -765,6 +760,13 @@ def test_call_visitor_super():
lv.visit_expr(add_node)
assert str(lv.log) == "\n".join(["LeafAdd", "InternalAdd", "InternalVar",
"InternalIntImm"])
+ tensor_load = tvm.te.placeholder((1,), name="A")(0)
+ basic_check(
+ tensor_load,
+ "\n".join(["Call", "\tOpaqueExpr", "\tIntImm"]),
+ "\n".join(["OpaqueExpr", "IntImm", "Call"]),
+ )
+
def test_call_mutator_super():
class InternalMutator(ExprMutator):
