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 a5d40bb14a [IR][TE] Share expression effects and register tensor loads
(#20359)
a5d40bb14a is described below
commit a5d40bb14ac6339c25db27c1e8e27f2a66bfbd47
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Sep 16 08:59:34 2026 -0400
[IR][TE] Share expression effects and register tensor loads (#20359)
Move `SideEffect` and call-effect metadata into shared IR so
expression-property queries no longer depend on TIRX or TE. The query
accepts shared `Expr` values, uses registered structural traversal with
dynamic types excluded, and conservatively treats calls with unknown
effects as opaque.
Represent TE loads as calls to the registered `te.tensor_load` operator,
with the tensor followed by its indices. Existing TE helpers validate
and decode this representation for tensor discovery and lowering.
---
include/tvm/ir/expr.h | 18 +++++
include/tvm/{tirx => ir}/op_attr_types.h | 84 ++--------------------
include/tvm/te/tensor.h | 4 +-
include/tvm/tirx/analysis.h | 8 ---
include/tvm/tirx/op_attr_types.h | 71 ------------------
src/arith/simplify_base.cc | 4 +-
src/arith/transitive_comparison_analyzer.cc | 2 +-
src/arith/z3_prover.cc | 2 +-
src/ir/side_effect.cc | 57 +++++++++++++++
src/relax/utils.cc | 8 +--
.../transform/using_assume_to_reduce_branches.cc | 14 ++--
src/te/tensor.cc | 30 +++++---
src/tirx/analysis/side_effect.cc | 78 --------------------
src/tirx/op/op.cc | 2 -
tests/cpp/tir_analysis_side_effect.cc | 14 ++--
tests/python/relax/test_expr_functor.py | 4 +-
tests/python/te/test_te_tensor.py | 5 +-
17 files changed, 132 insertions(+), 273 deletions(-)
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index 62db747e8b..b245795469 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -31,6 +31,7 @@
#include <tvm/ir/attrs.h>
#include <tvm/ir/base_expr.h>
#include <tvm/ir/cow.h>
+#include <tvm/ir/op_attr_types.h>
#include <tvm/ir/source_map.h>
#include <algorithm>
@@ -666,6 +667,22 @@ class Range : public ffi::ObjectRef {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Range, ffi::ObjectRef, RangeNode);
};
+/*!
+ * \brief Analyze the runtime effects of an expression.
+ *
+ * Accumulates effects of evaluated expression children, excluding type and
vector
+ * lane-count metadata. Non-operator callees are conservatively opaque. Missing
+ * operator effect attributes are errors unless an earlier update effect stops
+ * traversal. Returns kPure, kReadState, or kUpdateState.
+ *
+ * Shared IR contains only minimal analyses of expression properties.
Arithmetic
+ * reasoning, constraint solving, and target- or dialect-specific analyses
belong
+ * in their respective modules.
+ * \param expr The expression to inspect.
+ * \return The strongest runtime effect of the expression.
+ */
+TVM_DLL CallEffectKind SideEffect(const Expr& expr);
+
namespace ffi {
template <>
@@ -740,6 +757,7 @@ struct TypeTraits<FloatImm> : public
ObjectRefWithFallbackTraitsBase<FloatImm, d
}
};
} // namespace ffi
+
} // namespace tvm
/* \brief Allow tvm.Var and tvm.GlobalVar as keys in STL tables
diff --git a/include/tvm/tirx/op_attr_types.h b/include/tvm/ir/op_attr_types.h
similarity index 60%
copy from include/tvm/tirx/op_attr_types.h
copy to include/tvm/ir/op_attr_types.h
index 2b8a7428ca..ed04da5326 100644
--- a/include/tvm/tirx/op_attr_types.h
+++ b/include/tvm/ir/op_attr_types.h
@@ -17,87 +17,18 @@
* under the License.
*/
-/*!
- * \file tvm/tirx/op_attr_types.h
- * \brief Attribute types in the Op registry for TIR ops.
- *
- * These attributes can be set via OpRegEntry::set_attr
- *
- * \sa tvm/ir/op.h
+/*! \file tvm/ir/op_attr_types.h
+ * \brief Shared expression operator effects.
*/
-#ifndef TVM_TIR_OP_ATTR_TYPES_H_
-#define TVM_TIR_OP_ATTR_TYPES_H_
+#ifndef TVM_IR_OP_ATTR_TYPES_H_
+#define TVM_IR_OP_ATTR_TYPES_H_
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/string.h>
-#include <tvm/ir/expr.h>
+#include <tvm/ffi/error.h>
+#include <cstdint>
#include <ostream>
namespace tvm {
-namespace tirx {
-/*!
- * \brief Global symbol of the op after lowering.
- */
-using TGlobalSymbol = ffi::String;
-
-/*!
- * \brief Whether the op is overloaded for vector form.
- */
-using TVectorizable = bool;
-
-/*!
- * \brief The intrinsic lowering function for given op.
- */
-using FLowerIntrinsic = ffi::TypedFunction<PrimExpr(PrimExpr)>;
-
-/*!
- * \brief The legalization function for given tirx op.
- */
-using FLegalize = ffi::TypedFunction<PrimExpr(PrimExpr)>;
-
-/*!
- * \brief The operator's name in TVMScript printer
- */
-using TScriptPrinterName = ffi::String;
-
-/*!
- * \brief Specifies that TVMScript printer prints the dtype as the first/last
argument.
- If not specified, dtype will not be printed.
- */
-enum class ScriptDtypePrintLocation : int {
- /*!
- * \brief Do not print dtype as an argument.
- */
- kNone = 0,
- /*!
- * \brief Print dtype as the first argument.
- */
- kFirst = 1,
- /*!
- * \brief FPrint dtype as the last argument.
- */
- kLast = 2,
-};
-
-using TScriptDtypePrintLocation = int64_t;
-
-/*!
- * \brief Broad TIRx op category.
- *
- * Expected values:
- * - "builtin"
- * - "tile_primitive"
- * - "device_intrin"
- */
-using TIRxOpCategory = ffi::String;
-
-/*!
- * \brief Device intrinsic namespace.
- *
- * Expected values include "cuda", "ptx", "nvshmem", "nki", and "metal".
- */
-using TDeviceIntrinsicNamespace = ffi::String;
/*!
* \brief The effect type of the call.
@@ -168,6 +99,5 @@ inline std::ostream& operator<<(std::ostream& os,
CallEffectKind side_effect) {
/*! \brief Use integer to record the kind. */
using TCallEffectKind = int64_t;
-} // namespace tirx
} // namespace tvm
-#endif // TVM_TIR_OP_ATTR_TYPES_H_
+#endif // TVM_IR_OP_ATTR_TYPES_H_
diff --git a/include/tvm/te/tensor.h b/include/tvm/te/tensor.h
index 3f9edda4ea..cbec00fd00 100644
--- a/include/tvm/te/tensor.h
+++ b/include/tvm/te/tensor.h
@@ -207,10 +207,10 @@ class Tensor : public OpaqueExpr {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Tensor, OpaqueExpr, TensorNode);
};
-/*! \brief Return whether an expression is a Call whose callee is a TE Tensor.
*/
+/*! \brief Return whether an expression calls the registered te.tensor_load
operator. */
TVM_DLL bool IsTensorLoad(const Expr& expr);
-/*! \brief Recover and validate the Tensor callee of a tensor-load Call. */
+/*! \brief Recover and validate the first Tensor argument of a tensor-load
Call. */
TVM_DLL Tensor GetTensorFromLoad(const Call& call);
/*! \brief Recover and validate the primitive indices of a tensor-load Call. */
diff --git a/include/tvm/tirx/analysis.h b/include/tvm/tirx/analysis.h
index eb90ca5226..b0e32b65e3 100644
--- a/include/tvm/tirx/analysis.h
+++ b/include/tvm/tirx/analysis.h
@@ -98,14 +98,6 @@ TVM_DLL ffi::Array<Var> UndefinedVars(const PrimExpr& expr);
*/
TVM_DLL ffi::Array<Var> UndefinedVars(const PrimExpr& expr, const
ffi::Array<Var>& defs);
-/*!
- * \brief Analyze the side effect of an expression
- * \param expr The expression to be checked.
- *
- * \return CallEffectKind, can be kPure, kReadState or kUpdateState
- */
-TVM_DLL CallEffectKind SideEffect(const PrimExpr& expr);
-
/*!
* \brief Verifies whether the IR stmt or Expr is in SSA form.
* That is: each Var is defined and assigned once(in Let/For)
diff --git a/include/tvm/tirx/op_attr_types.h b/include/tvm/tirx/op_attr_types.h
index 2b8a7428ca..f333d03887 100644
--- a/include/tvm/tirx/op_attr_types.h
+++ b/include/tvm/tirx/op_attr_types.h
@@ -32,8 +32,6 @@
#include <tvm/ffi/string.h>
#include <tvm/ir/expr.h>
-#include <ostream>
-
namespace tvm {
namespace tirx {
/*!
@@ -99,75 +97,6 @@ using TIRxOpCategory = ffi::String;
*/
using TDeviceIntrinsicNamespace = ffi::String;
-/*!
- * \brief The effect type of the call.
- */
-enum class CallEffectKind : int {
- /*! \brief Function corresponds to an annotation(e.g. likely) and can
translate to identity. */
- kExprAnnotation = 0,
- /*!
- * \brief Pure function that do not interacts
- * with any external state.
- */
- kPure = 1,
- /*!
- * \brief Function's that may read from states(e.g. RAM)
- */
- kReadState = 2,
- /*!
- * \brief Function that may read/write from states(e.g. RAM).
- */
- kUpdateState = 3,
- /*!
- * \brief Opaque function, cannot make any assumption
- */
- kOpaque = kUpdateState,
- /*!
- * \brief Special intrinsic to annotate call arguments info
- * only valid as a direct argument to a call.
- */
- kSpecialCallArg = 4,
- /*!
- * \brief Embed opaque information in the Expr, cannot be codegen.
- */
- kEmbedInfo = 5,
- /*!
- * \brief Function that changes control flow
- */
- kControlJump = 6,
-};
-
-inline std::ostream& operator<<(std::ostream& os, CallEffectKind side_effect) {
- switch (side_effect) {
- case CallEffectKind::kExprAnnotation:
- return os << "kExprAnnotation";
-
- case CallEffectKind::kPure:
- return os << "kPure";
-
- case CallEffectKind::kReadState:
- return os << "kReadState";
-
- case CallEffectKind::kUpdateState:
- return os << "kUpdateState";
-
- case CallEffectKind::kSpecialCallArg:
- return os << "kSpecialCallArg";
-
- case CallEffectKind::kEmbedInfo:
- return os << "kEmbedInfo";
-
- case CallEffectKind::kControlJump:
- return os << "kControlJump";
-
- default:
- TVM_FFI_THROW(InternalError) << "Unknown CallEffectKind: " <<
static_cast<int>(side_effect);
- }
-}
-
-/*! \brief Use integer to record the kind. */
-using TCallEffectKind = int64_t;
-
} // namespace tirx
} // namespace tvm
#endif // TVM_TIR_OP_ATTR_TYPES_H_
diff --git a/src/arith/simplify_base.cc b/src/arith/simplify_base.cc
index 1f18be369b..29d8aaf24b 100644
--- a/src/arith/simplify_base.cc
+++ b/src/arith/simplify_base.cc
@@ -19,8 +19,8 @@
#include "simplify_base.h"
+#include <tvm/ir/expr.h>
#include <tvm/ir/prim/builtin.h>
-#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op.h>
#include "constraint_helpers.h"
@@ -123,7 +123,7 @@ UnchangedOr<Expr> SimplifierBase::Mutate_(const CallNode*
op, InplaceMode inplac
UnchangedOr<PrimExpr> SimplifierBase::Mutate_(const prim::LetNode* op,
InplaceMode inplace_mode) {
PrimExpr value = Mutate(op->value, inplace_mode).ValueOrUnchanged(op->value);
- if (tirx::SideEffect(value) <= tirx::CallEffectKind::kPure) {
+ if (SideEffect(value) <= CallEffectKind::kPure) {
analyzer_->Bind(op->var, value);
}
PrimExpr body = Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body);
diff --git a/src/arith/transitive_comparison_analyzer.cc
b/src/arith/transitive_comparison_analyzer.cc
index 73eb597728..a0c4750633 100644
--- a/src/arith/transitive_comparison_analyzer.cc
+++ b/src/arith/transitive_comparison_analyzer.cc
@@ -558,7 +558,7 @@ std::function<void()>
TransitiveComparisonAnalyzer::EnterConstraint(const PrimEx
void TransitiveComparisonAnalyzer::Impl::AddKnown(const PrimExpr& expr,
std::vector<Comparison>*
vec) {
for (const auto& subexpr : ExtractConstraints(expr, false)) {
- if (tirx::SideEffect(expr) <= tirx::CallEffectKind::kPure) {
+ if (SideEffect(expr) <= CallEffectKind::kPure) {
if (auto cmp = FromExpr(subexpr)) {
vec->push_back(cmp.value());
}
diff --git a/src/arith/z3_prover.cc b/src/arith/z3_prover.cc
index d94c7544f3..7121781a56 100644
--- a/src/arith/z3_prover.cc
+++ b/src/arith/z3_prover.cc
@@ -31,13 +31,13 @@
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/extra/structural_hash.h>
#include <tvm/ir/expr_functor.h>
+#include <tvm/ir/op_attr_types.h>
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/op_attr_types.h>
#include <algorithm>
#include <climits>
diff --git a/src/ir/side_effect.cc b/src/ir/side_effect.cc
new file mode 100644
index 0000000000..468f89b130
--- /dev/null
+++ b/src/ir/side_effect.cc
@@ -0,0 +1,57 @@
+/*
+ * 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.
+ */
+
+/*! \file side_effect.cc
+ * \brief Runtime effect properties of shared expressions.
+ */
+#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/ir/expr.h>
+#include <tvm/ir/op.h>
+
+#include <algorithm>
+
+namespace tvm {
+
+CallEffectKind SideEffect(const Expr& expr) {
+ static auto effects = Op::GetAttrMap<TCallEffectKind>("TCallEffectKind");
+ CallEffectKind kind = CallEffectKind::kPure;
+ ffi::StructuralVisit(
+ expr,
+ [&](const CallNode* node,
+ ffi::StructuralVisitorObj* visitor) ->
ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> {
+ auto effect = static_cast<CallEffectKind>(
+ effects.get(node->op,
static_cast<TCallEffectKind>(CallEffectKind::kOpaque)));
+ kind = std::max(kind, std::min(effect, CallEffectKind::kUpdateState));
+ if (kind == CallEffectKind::kUpdateState) return ffi::VisitInterrupt();
+ return visitor->DefaultVisitExpected(node);
+ },
+ [&](const TensorLoadNode* node,
+ ffi::StructuralVisitorObj* visitor) ->
ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> {
+ kind = std::max(kind, CallEffectKind::kReadState);
+ return visitor->DefaultVisitExpected(node);
+ },
+ [](const TypeNode*,
+ ffi::StructuralVisitorObj*) ->
ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> {
+ // Effects describe evaluated expressions, not expressions embedded in
types.
+ return std::nullopt;
+ });
+ return kind;
+}
+
+} // namespace tvm
diff --git a/src/relax/utils.cc b/src/relax/utils.cc
index f20d048130..740fcca681 100644
--- a/src/relax/utils.cc
+++ b/src/relax/utils.cc
@@ -22,11 +22,11 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/op_attr_types.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/attrs/index.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/utils.h>
-#include <tvm/tirx/op_attr_types.h>
#include <tvm/tirx/stmt_functor.h>
namespace tvm {
@@ -242,11 +242,11 @@ bool IsImpureCall(const Call& call) {
if (purity_map.count(op)) {
return !(purity_map[op]);
}
- static auto effect_map =
Op::GetAttrMap<tirx::TCallEffectKind>("TCallEffectKind");
+ static auto effect_map =
Op::GetAttrMap<TCallEffectKind>("TCallEffectKind");
TVM_FFI_ICHECK(effect_map.count(op))
<< "Cannot find the registered purity or call effect of this op: " <<
op->name;
- auto effect = static_cast<tirx::CallEffectKind>(effect_map[op]);
- return effect > tirx::CallEffectKind::kPure;
+ auto effect = static_cast<CallEffectKind>(effect_map[op]);
+ return effect > CallEffectKind::kPure;
}
// the Type must be FuncType
auto func_ty = GetTypeAs<FuncTypeNode>(call->op);
diff --git a/src/s_tir/transform/using_assume_to_reduce_branches.cc
b/src/s_tir/transform/using_assume_to_reduce_branches.cc
index 27e93c18ed..10412f1c46 100644
--- a/src/s_tir/transform/using_assume_to_reduce_branches.cc
+++ b/src/s_tir/transform/using_assume_to_reduce_branches.cc
@@ -151,10 +151,10 @@ class ParseAssumeAndOvercompute : public
IRMutatorWithAnalyzer {
: self(self), analyzer_context(self->analyzer_, constraint) {
old_num_constraints = self->conditions_.size();
- auto side_effect = tirx::SideEffect(constraint);
- if (side_effect <= tirx::CallEffectKind::kPure) {
+ auto side_effect = SideEffect(constraint);
+ if (side_effect <= CallEffectKind::kPure) {
self->conditions_.push_back(constraint);
- } else if (side_effect <= tirx::CallEffectKind::kReadState) {
+ } else if (side_effect <= CallEffectKind::kReadState) {
assume = constraint;
}
@@ -295,8 +295,8 @@ class ParseAssumeAndOvercompute : public
IRMutatorWithAnalyzer {
std::vector<PrimExpr> buffer_exprs;
for (const auto& expr : arith::ExtractComponents(assumption)) {
- auto side_effect = tirx::SideEffect(expr);
- if (side_effect <= tirx::CallEffectKind::kPure) {
+ auto side_effect = SideEffect(expr);
+ if (side_effect <= CallEffectKind::kPure) {
// Pulling out portions of the assumption that do not depend
// on a buffer value allows the following two forms to be
// treated identically.
@@ -304,7 +304,7 @@ class ParseAssumeAndOvercompute : public
IRMutatorWithAnalyzer {
// Option 1: if i < 3: T.assume(buf[i] == value)
// Option 2: T.assume(i>=3 or buf[i] == value)
additional_predicate = additional_predicate && logical_not(expr);
- } else if (side_effect == tirx::CallEffectKind::kReadState) {
+ } else if (side_effect == CallEffectKind::kReadState) {
buffer_exprs.push_back(expr);
} else {
TVM_FFI_THROW(InternalError)
@@ -357,7 +357,7 @@ class ParseAssumeAndOvercompute : public
IRMutatorWithAnalyzer {
map_buffer_assumption[buf_data.buffer_load->source.as_or_throw<tvm::tirx::BufferVar>()]
=
buf_data;
- auto has_side_effect = tirx::SideEffect(value) >
tirx::CallEffectKind::kPure;
+ auto has_side_effect = SideEffect(value) > CallEffectKind::kPure;
TVM_FFI_ICHECK(!has_side_effect)
<< "BufferVar value in constraint must be pure expression, but was "
<< value;
if (has_side_effect) {
diff --git a/src/te/tensor.cc b/src/te/tensor.cc
index cad803ca72..5efaa4beb8 100644
--- a/src/te/tensor.cc
+++ b/src/te/tensor.cc
@@ -25,6 +25,7 @@
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/op.h>
#include <tvm/te/operation.h>
#include <tvm/te/tensor.h>
@@ -33,19 +34,26 @@ namespace te {
namespace {
+const Op& TensorLoadOp() {
+ static const Op& op = Op::Get("te.tensor_load");
+ return op;
+}
+
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";
+ TVM_FFI_ICHECK(call->op.same_as(TensorLoadOp())) << "Expected a
te.tensor_load Call";
+ TVM_FFI_ICHECK(!call->args.empty()) << "Tensor-load Call requires a Tensor
argument";
+ const auto* tensor_node = call->args[0].as<TensorNode>();
+ TVM_FFI_ICHECK(tensor_node != nullptr) << "Tensor-load first argument must
be a TE Tensor";
Tensor tensor = ffi::GetRef<Tensor>(tensor_node);
- TVM_FFI_ICHECK_EQ(call->args.size(), tensor->shape.size())
+ TVM_FFI_ICHECK_EQ(call->args.size() - 1, 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>();
+ indices.reserve(call->args.size() - 1);
+ for (size_t i = 1; i < call->args.size(); ++i) {
+ auto index = call->args[i].as<PrimExpr>();
TVM_FFI_ICHECK(index.has_value()) << "Tensor-load indices must have
primitive type";
indices.push_back(index.value());
}
@@ -89,6 +97,9 @@ IterVar reduce_axis(Range dom, std::string name) {
PrimVar var(std::string name_hint, PrimType t) { return PrimVar(name_hint, t);
}
+TVM_REGISTER_OP("te.tensor_load")
+ .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kReadState));
+
// Tensor
inline PrimExpr Tensor::IndexTensor(ffi::Array<PrimExpr> indices,
bool support_negative_indices) const {
@@ -106,11 +117,12 @@ inline PrimExpr Tensor::IndexTensor(ffi::Array<PrimExpr>
indices,
}
}
ffi::Array<Expr> args;
- args.reserve(indices.size());
+ args.reserve(indices.size() + 1);
+ args.push_back(*this);
for (const PrimExpr& index : indices) {
args.push_back(index);
}
- return PrimExpr(Call((*this)->dtype, *this, args));
+ return PrimExpr(Call((*this)->dtype, TensorLoadOp(), args));
}
PrimExpr Tensor::operator()(ffi::Array<PrimVar> indices) const {
@@ -163,7 +175,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
bool IsTensorLoad(const Expr& expr) {
const auto* call = expr.as<CallNode>();
- return call != nullptr && call->op.as<TensorNode>() != nullptr;
+ return call != nullptr && call->op.same_as(TensorLoadOp());
}
Tensor GetTensorFromLoad(const Call& call) {
diff --git a/src/tirx/analysis/side_effect.cc b/src/tirx/analysis/side_effect.cc
deleted file mode 100644
index b953d37082..0000000000
--- a/src/tirx/analysis/side_effect.cc
+++ /dev/null
@@ -1,78 +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.
- */
-
-/*!
- * \file side_effect.cc
- * \brief side effect analysis
- */
-#include <tvm/ir/op.h>
-#include <tvm/ir/prim/expr.h>
-#include <tvm/te/tensor.h>
-#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/op_attr_types.h>
-#include <tvm/tirx/stmt_functor.h>
-
-namespace tvm {
-namespace tirx {
-
-class ExprSideEffect : public StmtExprVisitor {
- public:
- ffi::Optional<VisitInterrupt> Visit(ffi::AnyView e) final {
- if (kind_ == CallEffectKind::kUpdateState) return std::nullopt;
- return StmtExprVisitor::Visit(e);
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const TensorLoadNode* op) final {
- this->UpdateEffect(CallEffectKind::kReadState);
- return StmtExprVisitor::Visit_(op);
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const CallNode* op) final {
- static auto op_call_effect =
Op::GetAttrMap<TCallEffectKind>("TCallEffectKind");
-
- 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);
- }
- return StmtExprVisitor::Visit_(op);
- }
-
- void UpdateEffect(CallEffectKind effect_kind) {
- if (effect_kind > CallEffectKind::kUpdateState) {
- effect_kind = CallEffectKind::kUpdateState;
- }
- if (effect_kind > kind_) {
- kind_ = effect_kind;
- }
- }
-
- CallEffectKind kind_{CallEffectKind::kPure};
-};
-
-CallEffectKind SideEffect(const PrimExpr& e) {
- auto visitor = ffi::make_object<ExprSideEffect>();
- visitor->Visit(e);
- return visitor->kind_;
-}
-
-} // namespace tirx
-} // namespace tvm
diff --git a/src/tirx/op/op.cc b/src/tirx/op/op.cc
index 50a163f756..8ccb8eea24 100644
--- a/src/tirx/op/op.cc
+++ b/src/tirx/op/op.cc
@@ -43,11 +43,9 @@
namespace tvm {
namespace builtin = tirx::builtin;
-using tirx::CallEffectKind;
using tirx::is_const_int;
using tirx::IterVar;
using tirx::MakeConst;
-using tirx::TCallEffectKind;
using tirx::TGlobalSymbol;
using tirx::TIRxOpCategory;
using tirx::TScriptPrinterName;
diff --git a/tests/cpp/tir_analysis_side_effect.cc
b/tests/cpp/tir_analysis_side_effect.cc
index 6432c0716d..3af8b846bc 100644
--- a/tests/cpp/tir_analysis_side_effect.cc
+++ b/tests/cpp/tir_analysis_side_effect.cc
@@ -18,20 +18,20 @@
*/
#include <gtest/gtest.h>
+#include <tvm/ir/expr.h>
+#include <tvm/ir/op_attr_types.h>
#include <tvm/ir/prim/builtin.h>
#include <tvm/runtime/logging.h>
#include <tvm/te/operation.h>
-#include <tvm/tirx/analysis.h>
+#include <tvm/tirx/buffer.h>
#include <tvm/tirx/builtin.h>
TEST(SimplePasses, SideEffect) {
using namespace tvm;
auto buf = tirx::decl_buffer({16}, PrimType::Float(32));
auto i = PrimVar("i", PrimType::Int(32));
- TVM_FFI_ICHECK(tirx::SideEffect(tirx::BufferLoad(buf, {i})) ==
tirx::CallEffectKind::kReadState);
- TVM_FFI_ICHECK(tirx::SideEffect(exp(prim::Cast(PrimType::Float(32), i + 1)))
==
- tirx::CallEffectKind::kPure);
- TVM_FFI_ICHECK(tirx::SideEffect(tvm::Call(PrimType::Void(),
tirx::builtin::tvm_storage_sync(), {})
- .as_or_throw<PrimExpr>()) ==
- tirx::CallEffectKind::kUpdateState);
+ TVM_FFI_ICHECK(SideEffect(tirx::BufferLoad(buf, {i})) ==
CallEffectKind::kReadState);
+ TVM_FFI_ICHECK(SideEffect(exp(prim::Cast(PrimType::Float(32), i + 1))) ==
CallEffectKind::kPure);
+ TVM_FFI_ICHECK(SideEffect(tvm::Call(PrimType::Void(),
tirx::builtin::tvm_storage_sync(), {})
+ .as_or_throw<PrimExpr>()) ==
CallEffectKind::kUpdateState);
}
diff --git a/tests/python/relax/test_expr_functor.py
b/tests/python/relax/test_expr_functor.py
index 532a9f3c79..4fd069cd84 100644
--- a/tests/python/relax/test_expr_functor.py
+++ b/tests/python/relax/test_expr_functor.py
@@ -451,11 +451,11 @@ def test_call():
visitor = ASTPrinter()
visitor.visit_expr(tensor_load)
- assert str(visitor.log) == "\n".join(["Call", "\tExprFallback"])
+ assert str(visitor.log) == "\n".join(["Call", "\tOp", "\tExprFallback"])
mutator = ASTPostPrinterMutator()
assert mutator.visit_expr(tensor_load).same_as(tensor_load)
- assert str(mutator.log) == "\n".join(["ExprFallback", "Call"])
+ assert str(mutator.log) == "\n".join(["Op", "ExprFallback", "Call"])
def test_if():
diff --git a/tests/python/te/test_te_tensor.py
b/tests/python/te/test_te_tensor.py
index 79f6dba0bd..1ca7b805c4 100644
--- a/tests/python/te/test_te_tensor.py
+++ b/tests/python/te/test_te_tensor.py
@@ -43,8 +43,9 @@ def test_tensor():
assert d[T] == 1
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.op.same_as(tvm.ir.Op.get("te.tensor_load"))
+ assert load.args[0].same_as(T)
+ assert list(load.args[1:]) == [0, 0, 0]
assert load.ty == T.dtype
assert load.astype("float16").ty == tvm.ir.PrimType("float16")
