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 093df105bb [REFACTOR][IR] Localize primitive checks and containment
queries (#20518)
093df105bb is described below
commit 093df105bb801cf3065a2781ae689166439c1412
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Oct 2 13:41:23 2026 -0400
[REFACTOR][IR] Localize primitive checks and containment queries (#20518)
Keep operator validation and dtype predicates beside their
implementations, leaving binary type matching as the shared primitive
operation helper. Replace the containment utility with local queries
using the existing TIRx traversal, preserving primitive-expression and
operand boundaries.
---
src/ir/prim/op.cc | 97 ++++++++++++++++++++++++++++---
src/ir/prim/op_utils.h | 72 -----------------------
src/s_tir/schedule/ir_comparator.cc | 13 ++++-
src/tirx/analysis/check_contains.cc | 79 -------------------------
src/tirx/analysis/check_contains.h | 58 ------------------
src/tirx/op/op.cc | 29 ++++-----
src/tirx/transform/common_subexpr_elim.cc | 12 +++-
src/tirx/transform/vectorize_loop.cc | 22 +++++--
8 files changed, 141 insertions(+), 241 deletions(-)
diff --git a/src/ir/prim/op.cc b/src/ir/prim/op.cc
index ce6eb93f15..39be7bb5a8 100644
--- a/src/ir/prim/op.cc
+++ b/src/ir/prim/op.cc
@@ -30,6 +30,39 @@ using prim::is_const_int;
using prim::MakeConst;
using namespace prim::detail;
namespace {
+TVM_FFI_INLINE const PrimTypeNode* GetPrimTypeNode(const PrimExpr& expr) {
+ // Avoid PrimExpr::ty() ObjectRef materialization on binary operator hot
paths.
+ const auto* node = expr.get();
+ TVM_FFI_DCHECK(node != nullptr);
+ TVM_FFI_DCHECK(!node->ExprNode::ty.IsMissing());
+ const auto* prim_ty = node->ExprNode::ty.as<PrimTypeNode>();
+ TVM_FFI_DCHECK(prim_ty != nullptr);
+ return prim_ty;
+}
+
+TVM_FFI_INLINE bool IsFloatType(const PrimType& ty) {
+ return ty.MatchesCode(DLDataTypeCode::kDLFloat);
+}
+
+TVM_FFI_INLINE bool IsBFloat16Type(const PrimType& ty) {
+ return ty.MatchesCode(DLDataTypeCode::kDLBfloat);
+}
+
+TVM_FFI_INLINE bool IsFloat8Type(const PrimType& ty) {
+ return ty.MatchesCode(DLDataTypeCode::kDLFloat8_e3m4,
DLDataTypeCode::kDLFloat8_e4m3,
+ DLDataTypeCode::kDLFloat8_e4m3b11fnuz,
DLDataTypeCode::kDLFloat8_e4m3fn,
+ DLDataTypeCode::kDLFloat8_e4m3fnuz,
DLDataTypeCode::kDLFloat8_e5m2,
+ DLDataTypeCode::kDLFloat8_e5m2fnuz,
DLDataTypeCode::kDLFloat8_e8m0fnu);
+}
+
+TVM_FFI_INLINE bool IsFloat6Type(const PrimType& ty) {
+ return ty.MatchesCode(DLDataTypeCode::kDLFloat6_e2m3fn,
DLDataTypeCode::kDLFloat6_e3m2fn);
+}
+
+TVM_FFI_INLINE bool IsFloat4Type(const PrimType& ty) {
+ return ty.MatchesCode(DLDataTypeCode::kDLFloat4_e2m1fn);
+}
+
// File-local helper: true if `expr` is a call to prim::builtin::vscale().
bool IsVScaleCall(const PrimExpr& expr) {
if (const auto* call = expr.as<CallNode>()) {
@@ -543,21 +576,33 @@ PrimExpr not_equal(PrimExpr a, PrimExpr b, Span span) {
PrimExpr operator&&(PrimExpr a, PrimExpr b) { return logical_and(a, b); }
PrimExpr logical_and(PrimExpr a, PrimExpr b, Span span) {
- type_check_boolean_args(a, b, "&& operator (logical AND)");
+ TVM_FFI_ICHECK(a.ty().MatchesCode(DLDataTypeCode::kDLBool))
+ << "Expected boolean argument as LHS of && operator (logical AND), but
received " << a
+ << " of type " << a.ty();
+ TVM_FFI_ICHECK(b.ty().MatchesCode(DLDataTypeCode::kDLBool))
+ << "Expected boolean argument as RHS of && operator (logical AND), but
received " << b
+ << " of type " << b.ty();
if (auto ret = prim::detail::TryConstFold<prim::And>(a, b)) return
ret.value();
return prim::And(a, b, span);
}
PrimExpr operator||(PrimExpr a, PrimExpr b) { return logical_or(a, b); }
PrimExpr logical_or(PrimExpr a, PrimExpr b, Span span) {
- type_check_boolean_args(a, b, "|| operator (logical OR)");
+ TVM_FFI_ICHECK(a.ty().MatchesCode(DLDataTypeCode::kDLBool))
+ << "Expected boolean argument as LHS of || operator (logical OR), but
received " << a
+ << " of type " << a.ty();
+ TVM_FFI_ICHECK(b.ty().MatchesCode(DLDataTypeCode::kDLBool))
+ << "Expected boolean argument as RHS of || operator (logical OR), but
received " << b
+ << " of type " << b.ty();
if (auto ret = prim::detail::TryConstFold<prim::Or>(a, b)) return
ret.value();
return prim::Or(a, b, span);
}
PrimExpr operator!(PrimExpr a) { return logical_not(a); }
PrimExpr logical_not(PrimExpr a, Span span) {
- type_check_boolean_args(a, "! operator (logical NOT)");
+ TVM_FFI_ICHECK(a.ty().MatchesCode(DLDataTypeCode::kDLBool))
+ << "Expected boolean argument for ! operator (logical NOT), but received
" << a << " of type "
+ << a.ty();
if (auto ret = prim::detail::TryConstFold<prim::Not>(a)) return ret.value();
return prim::Not(a, span);
}
@@ -566,7 +611,12 @@ PrimExpr logical_not(PrimExpr a, Span span) {
PrimExpr operator>>(PrimExpr a, PrimExpr b) { return right_shift(a, b); }
PrimExpr right_shift(PrimExpr a, PrimExpr b, Span span) {
- type_check_integer_args(a, b, ">> operator (right shift)");
+ TVM_FFI_ICHECK(a.ty().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt))
+ << "Expected integer argument as LHS of >> operator (right shift), but
received " << a
+ << " of type " << a.ty();
+ TVM_FFI_ICHECK(b.ty().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt))
+ << "Expected integer argument as RHS of >> operator (right shift), but
received " << b
+ << " of type " << b.ty();
BinaryOpMatchTypes(a, b, span);
TVM_PRIM_INDEX_CONST_PROPAGATION({
@@ -589,7 +639,12 @@ PrimExpr right_shift(PrimExpr a, PrimExpr b, Span span) {
// shift left
PrimExpr operator<<(PrimExpr a, PrimExpr b) { return left_shift(a, b); }
PrimExpr left_shift(PrimExpr a, PrimExpr b, Span span) {
- type_check_integer_args(a, b, "<< operator (left shift)");
+ TVM_FFI_ICHECK(a.ty().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt))
+ << "Expected integer argument as LHS of << operator (left shift), but
received " << a
+ << " of type " << a.ty();
+ TVM_FFI_ICHECK(b.ty().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt))
+ << "Expected integer argument as RHS of << operator (left shift), but
received " << b
+ << " of type " << b.ty();
BinaryOpMatchTypes(a, b, span);
TVM_PRIM_INDEX_CONST_PROPAGATION({
PrimType result_ty = a.ty();
@@ -612,7 +667,14 @@ PrimExpr left_shift(PrimExpr a, PrimExpr b, Span span) {
// bitwise and
PrimExpr operator&(PrimExpr a, PrimExpr b) { return bitwise_and(a, b); }
PrimExpr bitwise_and(PrimExpr a, PrimExpr b, Span span) {
- type_check_int_or_bool_args(a, b, "& operator (bitwise AND)");
+ TVM_FFI_ICHECK(
+ a.ty().MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt,
DLDataTypeCode::kDLBool))
+ << "Expected integer argument as LHS of & operator (bitwise AND), but
received " << a
+ << " of type " << a.ty();
+ TVM_FFI_ICHECK(
+ b.ty().MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt,
DLDataTypeCode::kDLBool))
+ << "Expected integer argument as RHS of & operator (bitwise AND), but
received " << b
+ << " of type " << b.ty();
BinaryOpMatchTypes(a, b, span);
TVM_PRIM_INDEX_CONST_PROPAGATION({
PrimType result_ty = a.ty();
@@ -624,7 +686,14 @@ PrimExpr bitwise_and(PrimExpr a, PrimExpr b, Span span) {
// bitwise_or
PrimExpr operator|(PrimExpr a, PrimExpr b) { return bitwise_or(a, b); }
PrimExpr bitwise_or(PrimExpr a, PrimExpr b, Span span) {
- type_check_int_or_bool_args(a, b, "| operator (bitwise OR)");
+ TVM_FFI_ICHECK(
+ a.ty().MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt,
DLDataTypeCode::kDLBool))
+ << "Expected integer argument as LHS of | operator (bitwise OR), but
received " << a
+ << " of type " << a.ty();
+ TVM_FFI_ICHECK(
+ b.ty().MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt,
DLDataTypeCode::kDLBool))
+ << "Expected integer argument as RHS of | operator (bitwise OR), but
received " << b
+ << " of type " << b.ty();
BinaryOpMatchTypes(a, b, span);
TVM_PRIM_INDEX_CONST_PROPAGATION({
PrimType result_ty = a.ty();
@@ -636,7 +705,14 @@ PrimExpr bitwise_or(PrimExpr a, PrimExpr b, Span span) {
// bitwise_xor
PrimExpr operator^(PrimExpr a, PrimExpr b) { return bitwise_xor(a, b); }
PrimExpr bitwise_xor(PrimExpr a, PrimExpr b, Span span) {
- type_check_int_or_bool_args(a, b, "^ operator (bitwise XOR)");
+ TVM_FFI_ICHECK(
+ a.ty().MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt,
DLDataTypeCode::kDLBool))
+ << "Expected integer argument as LHS of ^ operator (bitwise XOR), but
received " << a
+ << " of type " << a.ty();
+ TVM_FFI_ICHECK(
+ b.ty().MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt,
DLDataTypeCode::kDLBool))
+ << "Expected integer argument as RHS of ^ operator (bitwise XOR), but
received " << b
+ << " of type " << b.ty();
BinaryOpMatchTypes(a, b, span);
TVM_PRIM_INDEX_CONST_PROPAGATION({
PrimType result_ty = a.ty();
@@ -649,7 +725,10 @@ PrimExpr bitwise_xor(PrimExpr a, PrimExpr b, Span span) {
PrimExpr operator~(PrimExpr a) { return bitwise_neg(a); }
PrimExpr bitwise_neg(PrimExpr a, Span span) {
- type_check_int_or_bool_args(a, "~ operator (bitwise NOT)");
+ TVM_FFI_ICHECK(
+ a.ty().MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt,
DLDataTypeCode::kDLBool))
+ << "Expected integer or boolean argument for ~ operator (bitwise NOT),
but received " << a
+ << " of type " << a.ty();
return prim::BitwiseNot(a, span);
}
diff --git a/src/ir/prim/op_utils.h b/src/ir/prim/op_utils.h
index e1cd678eec..ec1c3e34e8 100644
--- a/src/ir/prim/op_utils.h
+++ b/src/ir/prim/op_utils.h
@@ -21,79 +21,7 @@
#define TVM_IR_PRIM_OP_UTILS_H_
#include <tvm/ir/prim/op.h>
namespace tvm::prim::detail {
-TVM_FFI_INLINE const PrimTypeNode* GetPrimTypeNode(const PrimExpr& expr) {
- // Avoid PrimExpr::ty() ObjectRef materialization on binary operator hot
paths.
- const auto* node = expr.get();
- TVM_FFI_DCHECK(node != nullptr);
- TVM_FFI_DCHECK(!node->ExprNode::ty.IsMissing());
- const auto* prim_ty = node->ExprNode::ty.as<PrimTypeNode>();
- TVM_FFI_DCHECK(prim_ty != nullptr);
- return prim_ty;
-}
-
-TVM_FFI_INLINE bool IsFloatType(const PrimType& ty) {
- return ty.MatchesCode(DLDataTypeCode::kDLFloat);
-}
-
-TVM_FFI_INLINE bool IsBFloat16Type(const PrimType& ty) {
- return ty.MatchesCode(DLDataTypeCode::kDLBfloat);
-}
-
-TVM_FFI_INLINE bool IsFloat8Type(const PrimType& ty) {
- return ty.MatchesCode(DLDataTypeCode::kDLFloat8_e3m4,
DLDataTypeCode::kDLFloat8_e4m3,
- DLDataTypeCode::kDLFloat8_e4m3b11fnuz,
DLDataTypeCode::kDLFloat8_e4m3fn,
- DLDataTypeCode::kDLFloat8_e4m3fnuz,
DLDataTypeCode::kDLFloat8_e5m2,
- DLDataTypeCode::kDLFloat8_e5m2fnuz,
DLDataTypeCode::kDLFloat8_e8m0fnu);
-}
-
-TVM_FFI_INLINE bool IsFloat6Type(const PrimType& ty) {
- return ty.MatchesCode(DLDataTypeCode::kDLFloat6_e2m3fn,
DLDataTypeCode::kDLFloat6_e3m2fn);
-}
-
-TVM_FFI_INLINE bool IsFloat4Type(const PrimType& ty) {
- return ty.MatchesCode(DLDataTypeCode::kDLFloat4_e2m1fn);
-}
TVM_DLL void BinaryOpMatchTypes(PrimExpr& lhs, PrimExpr& rhs, Span span);
-inline void type_check_boolean_args(const PrimExpr& arg, const char* op) {
- TVM_FFI_ICHECK(arg.ty().MatchesCode(DLDataTypeCode::kDLBool))
- << "Expected boolean argument for " << op << ", but received " << arg <<
" of type "
- << arg.ty();
-}
-inline void type_check_boolean_args(const PrimExpr& lhs, const PrimExpr& rhs,
const char* op) {
- TVM_FFI_ICHECK(lhs.ty().MatchesCode(DLDataTypeCode::kDLBool))
- << "Expected boolean argument as LHS of " << op << ", but received " <<
lhs << " of type "
- << lhs.ty();
- TVM_FFI_ICHECK(rhs.ty().MatchesCode(DLDataTypeCode::kDLBool))
- << "Expected boolean argument as RHS of " << op << ", but received " <<
rhs << " of type "
- << rhs.ty();
-}
-
-inline void type_check_int_or_bool_args(const PrimExpr& arg, const char* op) {
- TVM_FFI_ICHECK(arg.ty().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt,
- DLDataTypeCode::kDLBool))
- << "Expected integer or boolean argument for " << op << ", but received
" << arg
- << " of type " << arg.ty();
-}
-
-inline void type_check_integer_args(const PrimExpr& lhs, const PrimExpr& rhs,
const char* op) {
- TVM_FFI_ICHECK(lhs.ty().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt))
- << "Expected integer argument as LHS of " << op << ", but received " <<
lhs << " of type "
- << lhs.ty();
- TVM_FFI_ICHECK(rhs.ty().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt))
- << "Expected integer argument as RHS of " << op << ", but received " <<
rhs << " of type "
- << rhs.ty();
-}
-
-inline void type_check_int_or_bool_args(const PrimExpr& lhs, const PrimExpr&
rhs, const char* op) {
- TVM_FFI_ICHECK(lhs.ty().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt,
- DLDataTypeCode::kDLBool))
- << "Expected integer argument as LHS of " << op << ", but received " <<
lhs << " of type "
- << lhs.ty();
- TVM_FFI_ICHECK(rhs.ty().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt,
- DLDataTypeCode::kDLBool))
- << "Expected integer argument as RHS of " << op << ", but received " <<
rhs << " of type "
- << rhs.ty();
-}
} // namespace tvm::prim::detail
#endif // TVM_IR_PRIM_OP_UTILS_H_
diff --git a/src/s_tir/schedule/ir_comparator.cc
b/src/s_tir/schedule/ir_comparator.cc
index fa0f4c074e..d5f78b816a 100644
--- a/src/s_tir/schedule/ir_comparator.cc
+++ b/src/s_tir/schedule/ir_comparator.cc
@@ -22,8 +22,7 @@
#include <tvm/ir/prim/builtin.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/builtin.h>
-
-#include "../../tirx/analysis/check_contains.h"
+#include <tvm/tirx/stmt_functor.h>
namespace tvm {
@@ -38,7 +37,15 @@ bool IsVScaleCall(const PrimExpr& expr) {
// File-local helper: true if `expr` contains a call to
prim::builtin::vscale().
bool ContainsVscaleCall(const PrimExpr& expr) {
- return tirx::CheckContains::ExprContains(expr, IsVScaleCall);
+ struct VScaleFinder : tirx::StmtExprVisitor {
+ ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) final {
+ if (auto expr = value.as<PrimExpr>(); expr && IsVScaleCall(*expr)) {
+ return VisitInterrupt();
+ }
+ return StmtExprVisitor::Visit(value);
+ }
+ };
+ return ffi::make_object<VScaleFinder>()->Visit(expr).has_value();
}
} // namespace
diff --git a/src/tirx/analysis/check_contains.cc
b/src/tirx/analysis/check_contains.cc
deleted file mode 100644
index 9b19aa7c04..0000000000
--- a/src/tirx/analysis/check_contains.cc
+++ /dev/null
@@ -1,79 +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 check_contains.cc
- * \brief Implementation of the analysis that tells if an expression contains
- a node that satisfies a given predicate.
- */
-
-#include "check_contains.h"
-
-#include <tvm/ir/prim/expr.h>
-
-#include <vector>
-
-namespace tvm {
-namespace tirx {
-
-/*!
- * \brief Toplevel (static) function that tells if an expression contains a
subexpression that
- satisfies a given predicate.
- * \param expr The expression to check
- * \param predicate The predicate that must be satisfied
- * \return Whether `expr` contains a subexpression that satisfies `predicate`
- */
-bool CheckContains::ExprContains(const PrimExpr& expr,
- std::function<bool(const PrimExpr&)>
predicate) {
- auto check_contains = ffi::make_object<CheckContains>(predicate);
- check_contains->Visit(expr);
- return check_contains->contains_it_;
-}
-
-/*!
- * \brief Toplevel (static) function that tells if a statement contains a
subexpression that
- satisfies a given predicate.
- * \param stmt The statement to check
- * \param predicate The predicate that must be satisfied
- * \return Whether `stmt` contains a subexpression that satisfies `predicate`
- */
-bool CheckContains::StmtContains(const Stmt& stmt, std::function<bool(const
PrimExpr&)> predicate) {
- auto check_contains = ffi::make_object<CheckContains>(predicate);
- check_contains->Visit(stmt);
- return check_contains->contains_it_;
-}
-
-/*!
- * \brief Protected constructor of CheckContains.
- * \param predicate The predicate that must be satisfied
- */
-CheckContains::CheckContains(std::function<bool(const PrimExpr&)> predicate)
- : predicate_(predicate) {}
-
-ffi::Optional<VisitInterrupt> CheckContains::Visit(ffi::AnyView value) {
- if (auto prim_expr = value.as<PrimExpr>(); prim_expr &&
predicate_(prim_expr.value())) {
- contains_it_ = true;
- return std::nullopt;
- }
- if (value.as<StmtNode>() && contains_it_) return std::nullopt;
- return StmtExprVisitor::Visit(value);
-}
-
-} // namespace tirx
-} // namespace tvm
diff --git a/src/tirx/analysis/check_contains.h
b/src/tirx/analysis/check_contains.h
deleted file mode 100644
index 5b856f7631..0000000000
--- a/src/tirx/analysis/check_contains.h
+++ /dev/null
@@ -1,58 +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 check_contains.h
- * \brief Interface of the analysis that tells if an expression contains
- a node that satisfies a given predicate.
- */
-
-#ifndef TVM_TIR_ANALYSIS_CHECK_CONTAINS_H_
-#define TVM_TIR_ANALYSIS_CHECK_CONTAINS_H_
-
-#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/stmt_functor.h> // For the class StmtExprVisitor
-
-namespace tvm {
-namespace tirx {
-
-/*!
- * \brief Visitor which tells if a given expression or statement contains a
subexpression
- that satisfies a given predicate
- */
-class CheckContains : public StmtExprVisitor {
- public:
- // Toplevel (static) functions
- static bool ExprContains(const PrimExpr& expr, std::function<bool(const
PrimExpr&)> predicate);
- static bool StmtContains(const Stmt& stmt, std::function<bool(const
PrimExpr&)> predicate);
-
- // Constructor
- explicit CheckContains(std::function<bool(const PrimExpr&)> predicate);
-
- ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) override;
-
- private:
- std::function<bool(const PrimExpr&)> predicate_;
- bool contains_it_ = false;
-};
-
-} // namespace tirx
-} // namespace tvm
-
-#endif // TVM_TIR_ANALYSIS_CHECK_CONTAINS_H_
diff --git a/src/tirx/op/op.cc b/src/tirx/op/op.cc
index 68cdd03737..aa662a0aac 100644
--- a/src/tirx/op/op.cc
+++ b/src/tirx/op/op.cc
@@ -38,9 +38,8 @@
#include <tvm/tirx/var.h>
#include <cmath>
-// Shared primitive type matching and dtype predicates.
+// Shared primitive type matching.
#include "../../ir/prim/op_utils.h"
-#include "../analysis/check_contains.h"
namespace tvm::prim {
@@ -147,8 +146,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
PrimExpr logaddexp(PrimExpr a, PrimExpr b, Span span) {
- TVM_FFI_ICHECK(IsFloatType(a.ty())) << a;
- TVM_FFI_ICHECK(IsFloatType(b.ty())) << b;
+ TVM_FFI_ICHECK(a.ty().MatchesCode(DLDataTypeCode::kDLFloat)) << a;
+ TVM_FFI_ICHECK(b.ty().MatchesCode(DLDataTypeCode::kDLFloat)) << b;
BinaryOpMatchTypes(a, b, span);
PrimExpr exp_sum = add(exp(a), exp(b));
PrimExpr log_exp_sum = log(exp_sum);
@@ -160,7 +159,7 @@ PrimExpr infinity(PrimType value_ty, Span span) {
using namespace tirx;
PrimType dtype = value_ty;
TVM_FFI_ICHECK_EQ(dtype.lanes(), 1);
- if (IsFloatType(dtype)) {
+ if (dtype.MatchesCode(DLDataTypeCode::kDLFloat)) {
if (dtype.bits() == 64) {
return FloatImm(value_ty, std::numeric_limits<double>::infinity(), span);
} else if (dtype.bits() == 32 || dtype.bits() == 16) {
@@ -234,7 +233,7 @@ PrimExpr reinterpret(DLDataType t, PrimExpr value, Span
span) {
// pow
PrimExpr pow(PrimExpr x, PrimExpr y, Span span) {
BinaryOpMatchTypes(x, y, span);
- TVM_FFI_ICHECK(IsFloatType(x.ty())) << "power only applies to float";
+ TVM_FFI_ICHECK(x.ty().MatchesCode(DLDataTypeCode::kDLFloat)) << "power only
applies to float";
// If we detect pow(x, 3), suggest using x * x * x
if (y.ty().MatchesCode(DLDataTypeCode::kDLInt)) {
@@ -247,7 +246,7 @@ PrimExpr pow(PrimExpr x, PrimExpr y, Span span) {
"`pow(x, 2) * pow(x, 2) ...`.";
}
}
- } else if (IsFloatType(y.ty())) {
+ } else if (y.ty().MatchesCode(DLDataTypeCode::kDLFloat)) {
const FloatImmNode* fx = y.as<FloatImmNode>();
if (fx) {
if (fx->value >= 3.0) {
@@ -276,7 +275,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
PrimExpr abs(PrimExpr x, Span span) {
if (x.ty().MatchesCode(DLDataTypeCode::kDLInt)) {
return prim::IntegerAbs(x, span);
- } else if (IsFloatType(x.ty()) || IsBFloat16Type(x.ty())) {
+ } else if (x.ty().MatchesCode(DLDataTypeCode::kDLFloat,
DLDataTypeCode::kDLBfloat)) {
const FloatImmNode* fx = x.as<FloatImmNode>();
if (fx) {
return FloatImm(x.ty(), std::fabs(fx->value), fx->span);
@@ -307,7 +306,7 @@ PrimExpr isnan(PrimExpr x, Span span) {
PrimType bool_ty(t);
if (x.ty().MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
return MakeConst(t, false);
- } else if (IsFloatType(x.ty())) {
+ } else if (x.ty().MatchesCode(DLDataTypeCode::kDLFloat)) {
const FloatImmNode* fx = x.as<FloatImmNode>();
if (fx) {
return MakeConst(t, std::isnan(fx->value), fx->span);
@@ -332,7 +331,7 @@ PrimExpr isinf(PrimExpr x, Span span) {
PrimType t = PrimType::Bool(x.ty().lanes());
if (x.ty().MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
return MakeConst(t, false, span);
- } else if (IsFloatType(x.ty())) {
+ } else if (x.ty().MatchesCode(DLDataTypeCode::kDLFloat)) {
PrimExpr infX = infinity(x.ty(), span);
return abs(x, span) == infX && !isnan(x, span);
} else {
@@ -353,7 +352,9 @@ PrimExpr sum(PrimExpr source, ffi::Array<IterVar> rdom,
ffi::Array<PrimExpr> ini
}
PrimExpr all(PrimExpr source, ffi::Array<IterVar> rdom, ffi::Array<PrimExpr>
init, Span span) {
- type_check_boolean_args(source, "tvm::all");
+ TVM_FFI_ICHECK(source.ty().MatchesCode(DLDataTypeCode::kDLBool))
+ << "Expected boolean argument for tvm::all, but received " << source <<
" of type "
+ << source.ty();
PrimVar x("x", source.ty(), span), y("y", source.ty());
PrimExpr result = prim::And(x, y, span);
PrimExpr identity_element = MakeConst(source.ty(), true, span);
@@ -362,7 +363,9 @@ PrimExpr all(PrimExpr source, ffi::Array<IterVar> rdom,
ffi::Array<PrimExpr> ini
}
PrimExpr any(PrimExpr source, ffi::Array<IterVar> rdom, ffi::Array<PrimExpr>
init, Span span) {
- type_check_boolean_args(source, "tvm::any");
+ TVM_FFI_ICHECK(source.ty().MatchesCode(DLDataTypeCode::kDLBool))
+ << "Expected boolean argument for tvm::any, but received " << source <<
" of type "
+ << source.ty();
PrimVar x("x", source.ty(), span), y("y", source.ty(), span);
PrimExpr result = prim::Or(x, y, span);
PrimExpr identity_element = MakeConst(source.ty(), false, span);
@@ -412,7 +415,7 @@ PrimExpr prod(PrimExpr source, ffi::Array<IterVar> rdom,
ffi::Array<PrimExpr> in
// fmod
PrimExpr fmod(PrimExpr x, PrimExpr y, Span span) {
BinaryOpMatchTypes(x, y, span);
- TVM_FFI_ICHECK(IsFloatType(x.ty())) << "fmod only applies to float";
+ TVM_FFI_ICHECK(x.ty().MatchesCode(DLDataTypeCode::kDLFloat)) << "fmod only
applies to float";
static const Op fmod_op = Op::Get("tirx.fmod");
return Call(x.ty(), fmod_op, {x, y}, {}, {}, span).as_or_throw<PrimExpr>();
}
diff --git a/src/tirx/transform/common_subexpr_elim.cc
b/src/tirx/transform/common_subexpr_elim.cc
index 7441ba741c..5002fad3be 100644
--- a/src/tirx/transform/common_subexpr_elim.cc
+++ b/src/tirx/transform/common_subexpr_elim.cc
@@ -87,7 +87,6 @@
#include <vector>
#include "../../support/ordered_map.h"
-#include "../analysis/check_contains.h"
namespace tvm {
namespace tirx {
@@ -302,8 +301,15 @@ class CSEPlanner : public StmtExprVisitor {
// (LT/LE/GT/GE/EQ/NE/And/Or/Not/Cast-to-bool/Select-of-bool).
PrimType expr_ty = expr.ty();
if (expr_ty.MatchesCode(DLDataTypeCode::kDLBool)) return false;
- if (CheckContains::ExprContains(expr, IsForbiddenNode)) return false;
- return true;
+ struct ForbiddenNodeFinder : StmtExprVisitor {
+ ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) final {
+ if (auto expr = value.as<PrimExpr>(); expr && IsForbiddenNode(*expr)) {
+ return VisitInterrupt();
+ }
+ return StmtExprVisitor::Visit(value);
+ }
+ };
+ return !ffi::make_object<ForbiddenNodeFinder>()->Visit(expr).has_value();
}
// ------------------------------------------------------------------
diff --git a/src/tirx/transform/vectorize_loop.cc
b/src/tirx/transform/vectorize_loop.cc
index 0877454572..6c631240bc 100644
--- a/src/tirx/transform/vectorize_loop.cc
+++ b/src/tirx/transform/vectorize_loop.cc
@@ -40,7 +40,6 @@
#include <unordered_map>
#include <vector>
-#include "../../tirx/analysis/check_contains.h"
#include "tvm/ffi/dtype.h"
#include "tvm/tirx/expr.h"
@@ -90,8 +89,15 @@ bool TargetHasVLA(Target target) {
}
bool ContainsCallNode(const Stmt& stmt) {
- return CheckContains::StmtContains(
- stmt, [](const PrimExpr& expr) { return expr.as<CallNode>() != nullptr;
});
+ struct CallFinder : StmtExprVisitor {
+ ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) final {
+ if (auto expr = value.as<PrimExpr>(); expr && expr->as<CallNode>()) {
+ return VisitInterrupt();
+ }
+ return StmtExprVisitor::Visit(value);
+ }
+ };
+ return ffi::make_object<CallFinder>()->Visit(stmt).has_value();
}
PrimType GetTextureElementType(const Expr& texture) {
@@ -1319,7 +1325,15 @@ class LoopVectorizer : public StmtExprMutator {
}
if (!extent_as_int || extent_as_int->value < 1) {
- bool is_scalable_expr = CheckContains::ExprContains(op->extent,
IsVScaleCall);
+ struct VScaleFinder : StmtExprVisitor {
+ ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) final {
+ if (auto expr = value.as<PrimExpr>(); expr && IsVScaleCall(*expr))
{
+ return VisitInterrupt();
+ }
+ return StmtExprVisitor::Visit(value);
+ }
+ };
+ bool is_scalable_expr =
ffi::make_object<VScaleFinder>()->Visit(op->extent).has_value();
TVM_FFI_ICHECK(is_scalable_expr && TargetHasVLA(target_))
<< "Failed to vectorize loop with extent " << op->extent << " for
target " << target_;
}
