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 30813e71bb [REFACTOR][IR] Normalize constant predicate names and
grouping (#20581)
30813e71bb is described below
commit 30813e71bb10a6e6c418fe47903afa4d44684f0f
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 7 07:24:40 2026 -0400
[REFACTOR][IR] Normalize constant predicate names and grouping (#20581)
---
include/tvm/ir/prim/op.h | 167 +++++++++++----------
include/tvm/tirx/op.h | 2 +-
include/tvm/topi/detail/constant_utils.h | 6 +-
src/backend/cuda/codegen/codegen_cuda.cc | 6 +-
src/backend/trn/codegen/codegen_trn.cc | 6 +-
src/backend/vulkan/codegen/codegen_spirv.cc | 2 +-
src/backend/webgpu/codegen/codegen_webgpu.cc | 2 +-
src/ir/expr.cc | 2 +-
src/ir/prim/op.cc | 6 +-
src/relax/analysis/tir_op_pattern_kind.cc | 4 +-
src/relax/distributed/axis_group_graph.cc | 2 +-
.../transform/lower_global_view_to_local_view.cc | 2 +-
src/relax/ir/dataflow_matcher.cc | 2 +-
src/relax/op/nn/attention.cc | 2 +-
src/relax/op/op_common.cc | 2 +-
src/relax/op/tensor/manipulate.cc | 2 +-
src/relax/op/tensor/sampling.cc | 4 +-
src/relax/transform/alter_op_impl.cc | 2 +-
src/relax/transform/fuse_tir.cc | 2 +-
src/relax/transform/split_call_tir_by_pattern.cc | 4 +-
.../transform/split_layout_rewrite_preproc.cc | 2 +-
src/s_tir/analysis/conditional_bounds.cc | 12 +-
src/s_tir/analysis/verify_gpu_code.cc | 2 +-
.../feature_extractor/per_store_feature.cc | 4 +-
.../multi_level_tiling_tensor_core.cc | 2 +-
src/s_tir/schedule/analysis/analysis.cc | 10 +-
src/s_tir/schedule/concrete_schedule.cc | 6 +-
src/s_tir/schedule/ir_comparator.cc | 4 +-
src/s_tir/schedule/primitive/blockize_tensorize.cc | 8 +-
src/s_tir/schedule/primitive/cache_read_write.cc | 2 +-
src/s_tir/schedule/primitive/compute_at.cc | 6 +-
src/s_tir/schedule/primitive/compute_inline.cc | 10 +-
src/s_tir/schedule/primitive/decompose_padding.cc | 2 +-
.../schedule/primitive/layout_transformation.cc | 16 +-
.../schedule/primitive/loop_transformation.cc | 4 +-
src/s_tir/schedule/primitive/pad_einsum.cc | 2 +-
src/s_tir/schedule/primitive/reduction.cc | 4 +-
src/s_tir/schedule/primitive/rolling_buffer.cc | 2 +-
src/s_tir/schedule/transform.cc | 4 +-
src/s_tir/script/ir_builder/frame.cc | 2 +-
src/s_tir/script/printer/stmt.cc | 17 +--
src/s_tir/transform/canonicalize_loop.cc | 2 +-
src/s_tir/transform/inject_software_pipeline.cc | 6 +-
src/s_tir/transform/inject_virtual_thread.cc | 2 +-
src/s_tir/transform/loop_partition.cc | 6 +-
.../transform/lower_cross_thread_reduction.cc | 2 +-
src/s_tir/transform/lower_match_buffer.cc | 4 +-
src/s_tir/transform/lower_opaque_block.cc | 6 +-
src/s_tir/transform/lower_thread_binding.cc | 2 +-
src/s_tir/transform/memhammer_coalesce.cc | 6 +-
src/s_tir/transform/memhammer_rewrite_rule.h | 2 +-
.../remove_weight_layout_rewrite_block.cc | 2 +-
src/sym/analyzer.cc | 6 +-
src/sym/detect_linear_equation.cc | 4 +-
src/sym/int_set.cc | 40 ++---
src/sym/iter_affine_map.cc | 66 ++++----
src/sym/modular_set.cc | 2 +-
src/sym/presburger_set.cc | 4 +-
src/sym/rewrite_simplify.cc | 6 +-
src/sym/simplify_base.cc | 8 +-
src/sym/transitive_comparison_analyzer.cc | 4 +-
src/sym/z3_prover.cc | 2 +-
src/target/llvm/codegen_cpu.cc | 6 +-
src/target/llvm/codegen_llvm.cc | 6 +-
src/target/source/codegen_c.cc | 4 +-
src/te/operation/scan_op.cc | 2 +-
src/tirx/ir/buffer.cc | 10 +-
src/tirx/ir/exec_scope.cc | 2 +-
src/tirx/ir/index_map.cc | 4 +-
src/tirx/ir/ir_mutator_with_analyzer.cc | 12 +-
src/tirx/ir/layout/compose_layout.cc | 2 +-
src/tirx/ir/layout/tile_canonicalize.cc | 4 +-
src/tirx/ir/layout/tile_core.cc | 2 +-
src/tirx/ir/layout/tile_direct_sum_ops.cc | 2 +-
src/tirx/ir/layout/tile_tile_ops.cc | 10 +-
src/tirx/ir/stmt.cc | 2 +-
src/tirx/ir/type.cc | 4 +-
src/tirx/script/ir_builder/ir.cc | 2 +-
src/tirx/script/printer/for_loop.cc | 6 +-
src/tirx/script/printer/stmt.cc | 2 +-
src/tirx/transform/flatten_buffer.cc | 2 +-
src/tirx/transform/lower_intrin.cc | 4 +-
src/tirx/transform/lower_thread_allreduce.h | 8 +-
src/tirx/transform/lower_tirx_opaque.cc | 6 +-
src/tirx/transform/lower_tvm_builtin.cc | 4 +-
src/tirx/transform/remove_no_op.cc | 4 +-
src/tirx/transform/storage_rewrite.cc | 6 +-
src/tirx/transform/tvm_ffi_binder.cc | 32 ++--
src/tirx/transform/vectorize_loop.cc | 4 +-
src/topi/einsum.cc | 2 +-
tests/cpp/pattern_match_test.cc | 6 +-
tests/cpp/sym_simplify_test.cc | 12 +-
92 files changed, 358 insertions(+), 344 deletions(-)
diff --git a/include/tvm/ir/prim/op.h b/include/tvm/ir/prim/op.h
index 9fe4605112..095f60b186 100644
--- a/include/tvm/ir/prim/op.h
+++ b/include/tvm/ir/prim/op.h
@@ -486,82 +486,6 @@ template <typename ValueType,
std::is_same<ValueType,
ffi::BigInt>::value>::type>
inline PrimExpr MakeConst(PrimType dtype, ValueType value, Span span = Span());
-/*!
- * \brief Check whether x is a constant integer expression.
- * \param x The input argument
- * \param value the value to be compared against.
- * \return whether x is constant expression.
- */
-inline bool is_const_int(const PrimExpr& x, int64_t value);
-
-/*!
- * \brief Check whether x is a constant integer 1
- * \param x The input argument.
- * \note This only return true for integer types.
- * \return whether x is constant 1
- */
-inline bool is_one(const PrimExpr& x) { return is_const_int(x, 1); }
-
-/*!
- * \brief Check whether x is a constant integer 0
- * \param x The input argument
- * \return whether x is constant 0
- * \note This only return true for integer types.
- */
-inline bool is_zero(const PrimExpr& x) { return is_const_int(x, 0); }
-
-/*!
- * \brief Check whether x is an integer constant, including wide IntImm values.
- * \return whether x is constant
- */
-inline bool is_const_int(const PrimExpr& x);
-
-/*!
- * \brief Check whether x is an integer/float constant.
- * \note This only return true for integer types.
- * \return whether x is constant
- */
-inline bool is_const_number(const PrimExpr& x);
-
-/*!
- * \brief Check whether x is a constant power of two
- * If x is power of two, write the power to the shift.
- *
- * \param x The input expression.
- * \param shift The output shift if x is power of two.
- * \return whether x is constant power of two
- */
-TVM_DLL bool is_const_power_of_two_integer(const PrimExpr& x, int* shift);
-
-// Implementation details after this
-inline bool is_const_int(const PrimExpr& x) { return x.as<IntImmNode>() !=
nullptr; }
-
-inline bool is_const_number(const PrimExpr& x) {
- if (x.as<IntImmNode>()) {
- return true;
- } else if (x.as<FloatImmNode>()) {
- return true;
- } else if (const auto* op = x.as<prim::BroadcastNode>()) {
- return (op->value->IsInstance<IntImmNode>() ||
op->value->IsInstance<FloatImmNode>());
- }
- return false;
-}
-
-inline bool is_positive_const(const PrimExpr& a) {
- const auto* as_int = a.as<IntImmNode>();
- return as_int && as_int->value > 0;
-}
-
-inline bool is_negative_const(const PrimExpr& a) {
- const auto* as_int = a.as<IntImmNode>();
- return as_int && as_int->value < 0;
-}
-
-inline bool is_const_int(const PrimExpr& x, int64_t value) {
- const auto* as_int = x.as<IntImmNode>();
- return as_int && as_int->value == value;
-}
-
template <typename ValueType>
inline PrimExpr MakeConstScalar(PrimType dtype, ValueType value, Span span =
Span()) {
if constexpr (std::is_enum_v<ValueType>) {
@@ -750,5 +674,96 @@ inline PrimExpr operator%(const PrimExpr& a, const TB& b) {
return a;
}
+// Constant predicates
+namespace prim {
+
+/*!
+ * \brief Check whether x is a constant integer expression.
+ * \param x The input argument
+ * \param value the value to be compared against.
+ * \return whether x is constant expression.
+ */
+inline bool IsConstInt(const PrimExpr& x, int64_t value);
+
+/*!
+ * \brief Check whether x is a constant integer 1
+ * \param x The input argument.
+ * \note This only returns true for integer types.
+ * \return whether x is constant 1
+ */
+inline bool IsOne(const PrimExpr& x);
+
+/*!
+ * \brief Check whether x is a constant integer 0
+ * \param x The input argument
+ * \return whether x is constant 0
+ * \note This only returns true for integer types.
+ */
+inline bool IsZero(const PrimExpr& x);
+
+/*!
+ * \brief Check whether x is an integer constant, including wide IntImm values.
+ * \return whether x is constant
+ */
+inline bool IsConstInt(const PrimExpr& x);
+
+/*!
+ * \brief Check whether x is an integer or floating-point constant.
+ * \note Also recognizes a Broadcast containing an integer or floating-point
constant.
+ * \return whether x is constant
+ */
+inline bool IsConstNumber(const PrimExpr& x);
+
+/*!
+ * \brief Check whether x is a constant power of two
+ * If x is power of two, write the power to the shift.
+ *
+ * \param x The input expression.
+ * \param shift The output shift if x is power of two.
+ * \return whether x is constant power of two
+ */
+TVM_DLL bool IsPowerOfTwoInt(const PrimExpr& x, int* shift);
+
+/*! \brief Check whether a is a positive integer constant. */
+inline bool IsPositiveConst(const PrimExpr& a);
+
+/*! \brief Check whether a is a negative integer constant. */
+inline bool IsNegativeConst(const PrimExpr& a);
+
+// Inline constant predicate implementations.
+inline bool IsConstInt(const PrimExpr& x) { return x.as<IntImmNode>() !=
nullptr; }
+
+inline bool IsConstNumber(const PrimExpr& x) {
+ if (x.as<IntImmNode>()) {
+ return true;
+ } else if (x.as<FloatImmNode>()) {
+ return true;
+ } else if (const auto* op = x.as<prim::BroadcastNode>()) {
+ return (op->value->IsInstance<IntImmNode>() ||
op->value->IsInstance<FloatImmNode>());
+ }
+ return false;
+}
+
+inline bool IsPositiveConst(const PrimExpr& a) {
+ const auto* as_int = a.as<IntImmNode>();
+ return as_int && as_int->value > 0;
+}
+
+inline bool IsNegativeConst(const PrimExpr& a) {
+ const auto* as_int = a.as<IntImmNode>();
+ return as_int && as_int->value < 0;
+}
+
+inline bool IsConstInt(const PrimExpr& x, int64_t value) {
+ const auto* as_int = x.as<IntImmNode>();
+ return as_int && as_int->value == value;
+}
+
+inline bool IsOne(const PrimExpr& x) { return IsConstInt(x, 1); }
+
+inline bool IsZero(const PrimExpr& x) { return IsConstInt(x, 0); }
+
+} // namespace prim
+
} // namespace tvm
#endif // TVM_IR_PRIM_OP_H_
diff --git a/include/tvm/tirx/op.h b/include/tvm/tirx/op.h
index 0940833810..1aaaaa1872 100644
--- a/include/tvm/tirx/op.h
+++ b/include/tvm/tirx/op.h
@@ -384,7 +384,7 @@ inline bool is_no_op(const tirx::Stmt& stmt) {
if (!stmt.defined()) return true;
if (const auto* op = stmt.as<tirx::EvaluateNode>()) {
auto value = op->value.as<PrimExpr>();
- return value && is_const_int(value.value());
+ return value && IsConstInt(value.value());
}
if (const auto* op = stmt.as<tirx::SeqStmtNode>()) {
return op->seq.size() == 0;
diff --git a/include/tvm/topi/detail/constant_utils.h
b/include/tvm/topi/detail/constant_utils.h
index 51f9e910ce..3f850f3479 100644
--- a/include/tvm/topi/detail/constant_utils.h
+++ b/include/tvm/topi/detail/constant_utils.h
@@ -56,11 +56,11 @@ inline bool IsConstInt(PrimExpr expr) { return
expr->IsInstance<tvm::IntImmNode>
* \return true if every element in array is constant int or uint, false
otherwise.
*/
inline bool IsConstIntArray(ffi::Array<PrimExpr> array) {
- bool is_const_int = true;
+ bool IsConstInt = true;
for (auto const& elem : array) {
- is_const_int &= !elem.defined() || elem->IsInstance<tvm::IntImmNode>();
+ IsConstInt &= !elem.defined() || elem->IsInstance<tvm::IntImmNode>();
}
- return is_const_int;
+ return IsConstInt;
}
/*!
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index 01c1511fee..9ff981f207 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -260,7 +260,7 @@ void CodeGenCUDA::InitFuncState(const PrimFunc& func) {
extractor->clusterCtaIdx_y_ext,
extractor->clusterCtaIdx_z_ext};
sym::Analyzer analyzer;
cluster_cta_x_is_linear_rank_ =
- is_one(analyzer->Simplify(launch_dimensions_[4] *
launch_dimensions_[5]));
+ IsOne(analyzer->Simplify(launch_dimensions_[4] * launch_dimensions_[5]));
}
void CodeGenCUDA::DeclareFunction(const GlobalVar& gvar, const PrimFunc& func)
{
@@ -376,7 +376,7 @@ void CodeGenCUDA::Dispatch_(const tirx::ForNode* op) {
// those declarations are printed between the pragma and the for statement,
// nvcc is free to unroll the loop despite disable_unroll.
std::string begin_str = PrintExpr(op->min);
- PrimExpr end = is_zero(op->min) ? op->extent :
sym::Analyzer()->Simplify(op->min + op->extent);
+ PrimExpr end = IsZero(op->min) ? op->extent :
sym::Analyzer()->Simplify(op->min + op->extent);
std::string end_str = PrintExpr(end);
std::string step_str = op->step.has_value() ? PrintExpr(*op->step) : "";
if (op->annotations.count("disable_unroll")) {
@@ -1709,7 +1709,7 @@ void CodeGenCUDA::Dispatch_(const EvaluateNode* op) {
return;
}
}
- if (auto value = op->value.as<PrimExpr>(); value &&
is_const_int(value.value())) return;
+ if (auto value = op->value.as<PrimExpr>(); value &&
IsConstInt(value.value())) return;
CodeGenC::Dispatch_(op);
}
diff --git a/src/backend/trn/codegen/codegen_trn.cc
b/src/backend/trn/codegen/codegen_trn.cc
index e1bc98c4af..4c8ac7e61d 100644
--- a/src/backend/trn/codegen/codegen_trn.cc
+++ b/src/backend/trn/codegen/codegen_trn.cc
@@ -297,7 +297,7 @@ void CodeGenTrainium::Dispatch_(const ForNode* op) {
std::string extent = PrintExpr(op->extent);
PrintIndent();
std::string vid = AllocVarID(op->loop_var.get());
- TVM_FFI_ICHECK(is_zero(op->min));
+ TVM_FFI_ICHECK(IsZero(op->min));
if (ctx_.tensorizing) {
stream << vid << " = nl.arange(" << extent << ")\n";
if (op->annotations.count("nki_dim")) {
@@ -356,7 +356,7 @@ void CodeGenTrainium::Dispatch_(const BufferStoreNode* op) {
}
void CodeGenTrainium::Dispatch_(const EvaluateNode* op) {
- if (auto value = op->value.as<PrimExpr>(); value &&
is_const_int(value.value())) return;
+ if (auto value = op->value.as<PrimExpr>(); value &&
IsConstInt(value.value())) return;
std::string vid = this->PrintExpr(op->value);
if (vid != "") {
this->PrintIndent();
@@ -405,7 +405,7 @@ void CodeGenTrainium::Dispatch_(const CallNode* op,
std::ostream& os) { // NOLI
if (is_op(nki_matmul_op, "tirx.nki.matmul")) {
TVM_FFI_ICHECK_EQ(op->args.size(), 4);
- std::string accum = is_one(op->args[3].as_or_throw<PrimExpr>()) ? " += " :
" = ";
+ std::string accum = IsOne(op->args[3].as_or_throw<PrimExpr>()) ? " += " :
" = ";
os << PrintExpr(op->args[0]) << accum;
ctx_.is_matmul_input = true;
os << "nisa.nc_matmul(" << PrintExpr(op->args[1]) << "," <<
PrintExpr(op->args[2]);
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc
b/src/backend/vulkan/codegen/codegen_spirv.cc
index c9f6d1f5b5..b38b219bad 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.cc
+++ b/src/backend/vulkan/codegen/codegen_spirv.cc
@@ -764,7 +764,7 @@ void CodeGenSPIRV::Dispatch_(const BufferStoreNode* op) {
void CodeGenSPIRV::Dispatch_(const ForNode* op) {
analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent));
spirv::Value init_value = MakeValue(op->min);
- PrimExpr end = is_zero(op->min) ? op->extent : analyzer_->Simplify(op->min +
op->extent);
+ PrimExpr end = IsZero(op->min) ? op->extent : analyzer_->Simplify(op->min +
op->extent);
spirv::Value end_value = MakeValue(end);
// loop step
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc
b/src/backend/webgpu/codegen/codegen_webgpu.cc
index 2ceb0798c9..86eeb8e127 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -834,7 +834,7 @@ void CodeGenWebGPU::DispatchAllocTensor(const BindNode* op,
const CallNode* buff
void CodeGenWebGPU::Dispatch_(const ForNode* op) {
std::string begin_str = PrintExpr(op->min);
- PrimExpr end = is_zero(op->min) ? op->extent :
sym::Analyzer()->Simplify(op->min + op->extent);
+ PrimExpr end = IsZero(op->min) ? op->extent :
sym::Analyzer()->Simplify(op->min + op->extent);
std::string end_str = PrintExpr(end);
std::string step_str = op->step.has_value() ? PrintExpr(*op->step) : "";
std::string vid = AllocVarID(op->loop_var.get());
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index b1b6b7e4f9..2f3587a2cf 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -1128,7 +1128,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// Range
Range::Range(PrimExpr begin, PrimExpr end, Span span)
- : Range(ffi::make_object<RangeNode>(begin, tvm::prim::is_zero(begin) ? end
: (end - begin),
+ : Range(ffi::make_object<RangeNode>(begin, tvm::prim::IsZero(begin) ? end
: (end - begin),
span)) {}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/ir/prim/op.cc b/src/ir/prim/op.cc
index bdbdba9d06..6f5d7e77cd 100644
--- a/src/ir/prim/op.cc
+++ b/src/ir/prim/op.cc
@@ -26,7 +26,7 @@
#include "op_utils.h"
namespace tvm {
-using prim::is_const_int;
+using prim::IsConstInt;
using prim::MakeConst;
using namespace prim::detail;
namespace {
@@ -377,7 +377,7 @@ inline bool ConstPowerHelper(ValueType val, int* shift) {
return true;
}
-bool is_const_power_of_two_integer(const PrimExpr& x, int* shift) {
+bool IsPowerOfTwoInt(const PrimExpr& x, int* shift) {
if (const auto* op = x.as<IntImmNode>()) {
return ConstPowerHelper(op->value, shift);
} else {
@@ -526,7 +526,7 @@ PrimExpr if_then_else(PrimExpr cond, PrimExpr true_value,
PrimExpr false_value,
// likely
PrimExpr likely(PrimExpr cond, Span span) {
- if (is_const_int(cond)) return cond;
+ if (IsConstInt(cond)) return cond;
return Call(cond.ty(), prim::builtin::likely(), {cond}, {}, {},
span).as_or_throw<PrimExpr>();
}
diff --git a/src/relax/analysis/tir_op_pattern_kind.cc
b/src/relax/analysis/tir_op_pattern_kind.cc
index c7147135ca..38ae29db76 100644
--- a/src/relax/analysis/tir_op_pattern_kind.cc
+++ b/src/relax/analysis/tir_op_pattern_kind.cc
@@ -204,8 +204,8 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor {
size_t ndim_store_buf = store->buffer->shape.size();
for (size_t i = 0, j = 0; i < ndim_load_buf; ++i) {
- if
(is_const_int(load->source.as_or_throw<tvm::tirx::TensorVar>()->shape[i], 1) &&
- is_const_int(load->indices[i], 0)) {
+ if
(IsConstInt(load->source.as_or_throw<tvm::tirx::TensorVar>()->shape[i], 1) &&
+ IsConstInt(load->indices[i], 0)) {
// Skip unit load dimensions
// E.g. A[i, j] = B[1, j] is still broadcast
continue;
diff --git a/src/relax/distributed/axis_group_graph.cc
b/src/relax/distributed/axis_group_graph.cc
index 43e685bd39..d58b088c27 100644
--- a/src/relax/distributed/axis_group_graph.cc
+++ b/src/relax/distributed/axis_group_graph.cc
@@ -42,7 +42,7 @@ ffi::Optional<Var> GetShardingVarFromIndex(PrimExpr index,
ffi::Map<Var, Range>
primitive_var_range.Set(var.as_or_throw<PrimVar>(), range);
}
sym::IterSumExpr iter_sum = sym::NormalizeToIterSum(index,
primitive_var_range, analyzer);
- if (!is_zero(iter_sum->base)) {
+ if (!IsZero(iter_sum->base)) {
return std::nullopt;
}
if (iter_sum->args.empty()) {
diff --git a/src/relax/distributed/transform/lower_global_view_to_local_view.cc
b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
index 31aed8a49f..d1d11096ae 100644
--- a/src/relax/distributed/transform/lower_global_view_to_local_view.cc
+++ b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
@@ -231,7 +231,7 @@ class DistributedBufferCompactor : public
s_tir::StmtExprMutator {
int shard = iter_var_shards_[iter_var->var];
if (shard > 1) {
Range dom = iter_var->dom;
- TVM_FFI_ICHECK(is_zero(dom->min));
+ TVM_FFI_ICHECK(IsZero(dom->min));
sym::Analyzer analyzer;
TVM_FFI_ICHECK(analyzer->CanProve(floormod(dom->extent, shard) ==
0));
new_iter_vars.push_back(
diff --git a/src/relax/ir/dataflow_matcher.cc b/src/relax/ir/dataflow_matcher.cc
index 2d162b4fd7..a32489de54 100644
--- a/src/relax/ir/dataflow_matcher.cc
+++ b/src/relax/ir/dataflow_matcher.cc
@@ -485,7 +485,7 @@ static bool ShapeEqual(AnalyzerObj* analyzer, const
ffi::Array<PrimExpr>& lhs,
const ffi::Array<PrimExpr>& rhs) {
if (lhs.size() != rhs.size()) return false;
for (size_t i = 0; i < lhs.size(); ++i)
- if (!tvm::prim::is_one(analyzer->Simplify(lhs[i] == rhs[i]))) return false;
+ if (!tvm::prim::IsOne(analyzer->Simplify(lhs[i] == rhs[i]))) return false;
return true;
}
diff --git a/src/relax/op/nn/attention.cc b/src/relax/op/nn/attention.cc
index 4a6b8796e7..6de038e30b 100644
--- a/src/relax/op/nn/attention.cc
+++ b/src/relax/op/nn/attention.cc
@@ -126,7 +126,7 @@ Type InferTypeAttention(const Call& call, const
BlockBuilder& ctx) {
}
auto diag_equal_or_broadcast = [&](PrimExpr v1, PrimExpr v2, ffi::String
m1, ffi::String m2,
ffi::String dim) {
- if (analyzer->CanProve(v1 != v2) && !tvm::prim::is_one(v2)) {
+ if (analyzer->CanProve(v1 != v2) && !tvm::prim::IsOne(v2)) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "The " << m1 << " " << dim << " and the " << m2 << " " << dim
<< " should be the same or broadcastable. However, the " << dim <<
" of " << m1
diff --git a/src/relax/op/op_common.cc b/src/relax/op/op_common.cc
index 5cbbbc9ebc..5e84cdf892 100644
--- a/src/relax/op/op_common.cc
+++ b/src/relax/op/op_common.cc
@@ -222,7 +222,7 @@ bool CanProveLayoutTransform(const SLayout& input_layout,
const SLayout& desired
ffi::Array<PrimExpr> back_shape = todesired.BackwardShape(desired_shape);
sym::Analyzer analyzer;
for (size_t i = 0; i < shape.size(); ++i) {
- if (tvm::prim::is_const_int(shape[i])) {
+ if (tvm::prim::IsConstInt(shape[i])) {
if (!analyzer->CanProveEqual(shape[i], back_shape[i])) {
can_prove = false;
break;
diff --git a/src/relax/op/tensor/manipulate.cc
b/src/relax/op/tensor/manipulate.cc
index c9f9b7682f..8a9350a9f6 100644
--- a/src/relax/op/tensor/manipulate.cc
+++ b/src/relax/op/tensor/manipulate.cc
@@ -1330,7 +1330,7 @@ InferLayoutOutput InferLayoutSqueeze(
} else {
axis.reserve(ndim);
for (int i = 0; i < ndim; ++i) {
- if (tvm::prim::is_one(shape->values[i])) {
+ if (tvm::prim::IsOne(shape->values[i])) {
axis.push_back(i);
}
}
diff --git a/src/relax/op/tensor/sampling.cc b/src/relax/op/tensor/sampling.cc
index cd5c513cd6..0b09181e1e 100644
--- a/src/relax/op/tensor/sampling.cc
+++ b/src/relax/op/tensor/sampling.cc
@@ -133,8 +133,8 @@ Type InferTypeMultinomialFromUniform(const CallNode*
call_node) {
<< n << "` and the given sample_indices tensor has batch size `"
<< sample_indices_shape->values[0] << "`";
}
- if (!tvm::prim::is_one(uniform_sample_shape->values[1]) ||
- !tvm::prim::is_one(sample_indices_shape->values[1])) {
+ if (!tvm::prim::IsOne(uniform_sample_shape->values[1]) ||
+ !tvm::prim::IsOne(sample_indices_shape->values[1])) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "Multinomial_from_uniform op requires the input uniform_sample and "
"sample_indices to be 2D tensors with the second dimension being 1.
"
diff --git a/src/relax/transform/alter_op_impl.cc
b/src/relax/transform/alter_op_impl.cc
index f98c39d188..2985e7f75e 100644
--- a/src/relax/transform/alter_op_impl.cc
+++ b/src/relax/transform/alter_op_impl.cc
@@ -241,7 +241,7 @@ class AlterOpImplMutator : public ExprMutator {
auto [inverse_index_map, padding_predicate] =
index_map.NonSurjectiveInverse(initial_ranges, analyzer);
- if (tvm::prim::is_zero(padding_predicate)) {
+ if (tvm::prim::IsZero(padding_predicate)) {
return TransformLayout(expr, inverse_index_map);
} else {
auto padded_expr = builder_->Normalize(TransformLayout(expr,
inverse_index_map));
diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc
index 36d4267df6..23d477d7bf 100644
--- a/src/relax/transform/fuse_tir.cc
+++ b/src/relax/transform/fuse_tir.cc
@@ -60,7 +60,7 @@ class SymbolicMatcher : ExprFunctor<void(const Expr& n, const
PrimExpr& other)>
must_prove_ =
analyzer_->Simplify(ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(must_prove_,
f_substitute)
.as_or_throw<PrimExpr>());
- TVM_FFI_ICHECK(!is_zero(must_prove_));
+ TVM_FFI_ICHECK(!IsZero(must_prove_));
}
private:
diff --git a/src/relax/transform/split_call_tir_by_pattern.cc
b/src/relax/transform/split_call_tir_by_pattern.cc
index dd1578ecbe..46915b52da 100644
--- a/src/relax/transform/split_call_tir_by_pattern.cc
+++ b/src/relax/transform/split_call_tir_by_pattern.cc
@@ -269,7 +269,7 @@ class ForMatcher : public TensorizeComparator {
// Build mapping between the loop vars
if (!DefEqual(op->loop_var, rhs->loop_var)) return false;
// Only handle the case where the loop start from 0
- if (!is_zero(op->min) || !is_zero(rhs->min)) return false;
+ if (!IsZero(op->min) || !IsZero(rhs->min)) return false;
if (op->thread_binding.has_value() || rhs->thread_binding.has_value())
return false;
if (op->kind != ForKind::kSerial || op->kind != rhs->kind) return false;
if (!op->annotations.empty() || !rhs->annotations.empty()) return false;
@@ -317,7 +317,7 @@ class ForMatcher : public TensorizeComparator {
if (!rhs->iter_values[i].same_as(loop_stack_rhs_[i]->loop_var)) return
false;
}
// Disallow predicates now
- if (!is_one(op->predicate) || !is_one(rhs->predicate)) return false;
+ if (!IsOne(op->predicate) || !IsOne(rhs->predicate)) return false;
return Dispatch(op->block, rhs->block);
}
diff --git a/src/relax/transform/split_layout_rewrite_preproc.cc
b/src/relax/transform/split_layout_rewrite_preproc.cc
index 2b05b79417..07dad8ee32 100644
--- a/src/relax/transform/split_layout_rewrite_preproc.cc
+++ b/src/relax/transform/split_layout_rewrite_preproc.cc
@@ -175,7 +175,7 @@ class SplitPrimFuncLayoutRewrite : public
s_tir::StmtExprMutator {
.as_or_throw<s_tir::SBlock>();
auto it =
op->annotations.find(s_tir::attr::meta_schedule_layout_rewrite_preproc);
bool is_layout_rewrite_preproc =
- it != op->annotations.end() && is_one((*it).second.cast<PrimExpr>());
+ it != op->annotations.end() && IsOne((*it).second.cast<PrimExpr>());
if (current_subtree_ == 0) {
current_subtree_ = is_layout_rewrite_preproc ? 1 : -1;
diff --git a/src/s_tir/analysis/conditional_bounds.cc
b/src/s_tir/analysis/conditional_bounds.cc
index 810c9ca4b8..3d5fef205e 100644
--- a/src/s_tir/analysis/conditional_bounds.cc
+++ b/src/s_tir/analysis/conditional_bounds.cc
@@ -139,7 +139,7 @@ IntGroupBounds IntGroupBounds::operator+(const Range& r) {
ffi::Array<PrimExpr> lower;
ffi::Array<PrimExpr> upper;
const PrimExpr& coef = this->coef;
- if (tvm::prim::is_one(r->extent)) {
+ if (tvm::prim::IsOne(r->extent)) {
equal.push_back(analyzer->Simplify(r->min * coef));
} else {
lower.push_back(analyzer->Simplify(r->min * coef));
@@ -172,7 +172,7 @@ Range IntGroupBounds::FindBestRange(const ffi::Map<Var,
Range>& vranges_addl) co
uppers.push_back(expr);
}
- if (lowers.size() == 1 && uppers.size() == 1 && tvm::prim::is_one(coef)) {
+ if (lowers.size() == 1 && uppers.size() == 1 && tvm::prim::IsOne(coef)) {
return Range(analyzer->Simplify(lowers[0]), analyzer->Simplify(uppers[0] +
1));
}
@@ -519,11 +519,11 @@ PartialSolvedInequalities SolveLinearInequalities(const
IntConstraints& system_t
ffi::Array<PrimExpr> other_conditions;
for (const PrimExpr& e : current_ineq_set_to_solve) {
PrimExpr e_simp = analyzer->Simplify(e, kSimplifyRewriteCanonicalRewrite);
- if (is_const_int(e_simp, 0)) {
+ if (IsConstInt(e_simp, 0)) {
// contradiction detected
other_conditions = {IntImm::Bool(false)};
break;
- } else if (is_const_int(e_simp, 1)) {
+ } else if (IsConstInt(e_simp, 1)) {
continue;
} else {
other_conditions.push_back(e_simp);
@@ -567,7 +567,7 @@ IntConstraints SolveInequalitiesToRange(const
IntConstraints& inequalities) {
const PrimVar& var = *it;
TVM_FFI_ICHECK(solved_bounds.count(var));
auto bnd = solved_bounds.at(var);
- if (is_one(bnd.coef) && !bnd.equal.empty()) {
+ if (IsOne(bnd.coef) && !bnd.equal.empty()) {
// There is an equation of the form `v == expr`, so this variable can be
completely removed.
// Note that we use the 0-th expression because they are ordered by
complexity,
// so it must be the simplest one.
@@ -622,7 +622,7 @@ ffi::Optional<ffi::Map<Var, Range>>
ConditionalBoundsContext::TrySolveCondition(
// currently only extract simple integral equations which could be solvable.
sym::Analyzer analyzer;
PrimExpr condition = analyzer->Simplify(condition_);
- if (is_const_int(condition)) {
+ if (IsConstInt(condition)) {
return std::nullopt;
}
ffi::Array<PrimExpr> equations;
diff --git a/src/s_tir/analysis/verify_gpu_code.cc
b/src/s_tir/analysis/verify_gpu_code.cc
index 6f0c00735f..28628de8aa 100644
--- a/src/s_tir/analysis/verify_gpu_code.cc
+++ b/src/s_tir/analysis/verify_gpu_code.cc
@@ -232,7 +232,7 @@ class GPUCodeVerifier : public StmtExprVisitor {
for (const auto index : indices) {
if (const auto* ramp = index.as<RampNode>()) {
PrimType ramp_ty = ramp->ty.as_or_throw<PrimType>();
- if (!is_one(ramp->stride) && ramp_ty.IsFixedLengthVector() &&
+ if (!IsOne(ramp->stride) && ramp_ty.IsFixedLengthVector() &&
ElementBytes(ramp_ty) > max_vector_bytes_) {
std::stringstream s;
s << "Number of lanes (" << ramp_ty.lanes() << ") times number of
bytes ("
diff --git a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
index d50c5c6ec6..e790629cb6 100644
--- a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
+++ b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
@@ -308,7 +308,7 @@ Pass SimplifyForFeatureExtraction(bool
normalize_thread_bindings = false) {
updated.CopyOnWrite()->annotations =
std::move(annotations).ValueUnchecked();
return updated;
};
- if (!normalize_thread_bindings_ && is_zero(loop->extent)) {
+ if (!normalize_thread_bindings_ && IsZero(loop->extent)) {
return Evaluate(0);
}
if (normalize_thread_bindings_ && loop->kind == ForKind::kThreadBinding)
{
@@ -349,7 +349,7 @@ Pass SimplifyForFeatureExtraction(bool
normalize_thread_bindings = false) {
thread_bindings_.pop_back();
return result;
}
- if (!normalize_thread_bindings_ && is_zero(loop->min) &&
is_one(loop->extent) &&
+ if (!normalize_thread_bindings_ && IsZero(loop->min) &&
IsOne(loop->extent) &&
loop->kind == ForKind::kSerial && loop->annotations.empty()) {
VarRemapSet(loop->loop_var, MakeConst(loop->loop_var.ty(), 0.0));
return Mutate(loop->body, inplace_mode).ValueOrUnchanged(loop->body);
diff --git
a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
index d7dc2663da..247f2c920e 100644
--- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
@@ -840,7 +840,7 @@ ffi::Optional<LoopRV>
MultiLevelTilingTensorCoreNode::TransformWithTensorIntrin(
std::vector<PrimExpr> sub_index_map_tgt;
const tirx::TensorVar& rhs_buffer =
mapping_info->lhs_buffer_map[lhs_buffer];
for (const Range& range : lhs_region) {
- TVM_FFI_ICHECK(tvm::prim::is_one(range->extent));
+ TVM_FFI_ICHECK(tvm::prim::IsOne(range->extent));
auto var = range->min.as<PrimVar>();
TVM_FFI_ICHECK(var.has_value());
const PrimVar& lhs_representer = lhs_to_index_map_src.at(var.value());
diff --git a/src/s_tir/schedule/analysis/analysis.cc
b/src/s_tir/schedule/analysis/analysis.cc
index af67658153..9fdb3417e0 100644
--- a/src/s_tir/schedule/analysis/analysis.cc
+++ b/src/s_tir/schedule/analysis/analysis.cc
@@ -591,7 +591,7 @@ bool IsAffineBinding(const SBlockRealize& realize, const
ffi::Map<Var, Range>& l
}
for (const sym::IterSumExpr& sum_expr : res->indices) {
const ffi::Array<sym::IterSplitExpr>& args = sum_expr->args;
- if (!args.empty() && !is_one(args[0]->scale)) {
+ if (!args.empty() && !IsOne(args[0]->scale)) {
return false;
}
}
@@ -1053,7 +1053,7 @@ std::pair<ffi::Array<StmtSRef>, std::vector<int>>
CollectComputeLocation(
location_indices.reserve(n_candidate + 2);
bool visited_reduce = false;
for (size_t i = 0; i < n_candidate; ++i) {
- if (is_one(TVM_SREF_TO_FOR(loop_srefs[i])->extent)) {
+ if (IsOne(TVM_SREF_TO_FOR(loop_srefs[i])->extent)) {
continue;
}
@@ -1387,7 +1387,7 @@ bool HasOp(const Stmt& stmt, const ffi::Array<Op>& ops) {
bool HasIfThenElse(const Stmt& stmt) {
auto visit_realize = [](const SBlockRealize& realize) ->
ffi::Expected<ffi::WalkResult> {
- if (!is_one(realize->predicate)) {
+ if (!IsOne(realize->predicate)) {
return ffi::WalkResult::Interrupt(ffi::VisitInterrupt(true));
}
return ffi::WalkResult::Advance();
@@ -1571,7 +1571,7 @@ bool NeedsMultiLevelTiling(const ScheduleState& self,
const StmtSRef& block_sref
std::vector<const VarNode*> spatial_block_vars;
spatial_block_vars.reserve(block->iter_vars.size());
for (const IterVar& block_var : block->iter_vars) {
- bool has_trivial_dom = is_zero(block_var->dom->min) &&
is_one(block_var->dom->extent);
+ bool has_trivial_dom = IsZero(block_var->dom->min) &&
IsOne(block_var->dom->extent);
if (block_var->iter_type == IterVarType::kDataPar && !has_trivial_dom) {
spatial_block_vars.push_back(block_var->var.get());
}
@@ -1905,7 +1905,7 @@ ffi::Optional<TensorizeInfo>
GetTensorizeLoopMapping(const s_tir::ScheduleState&
continue;
}
// padding is allowed only when the block has trivial bindings
- if (allow_padding && !is_zero(residual)) {
+ if (allow_padding && !IsZero(residual)) {
allow_padding = false;
}
diff --git a/src/s_tir/schedule/concrete_schedule.cc
b/src/s_tir/schedule/concrete_schedule.cc
index 36e04d02ca..61f8478642 100644
--- a/src/s_tir/schedule/concrete_schedule.cc
+++ b/src/s_tir/schedule/concrete_schedule.cc
@@ -507,7 +507,7 @@ ffi::Array<LoopRV> ConcreteScheduleNode::Split(const
LoopRV& loop_rv,
infer_index = i;
} else {
PrimExpr factor = this->Get(factor_rvs[i].value());
- if (is_const_int(factor) && !is_positive_const(factor)) {
+ if (IsConstInt(factor) && !IsPositiveConst(factor)) {
throw MakeScheduleError<NonPositiveFactorError>(state_->mod,
factor.as<IntImmNode>()->value,
i);
}
@@ -562,7 +562,7 @@ ffi::Array<LoopRV> ConcreteScheduleNode::LoopPartition(
PrimExpr tot_length = 0;
ffi::Array<StmtSRef> results;
TVM_TIR_SCHEDULE_BEGIN();
- if (!is_const_number(loop->min) || !is_const_number(loop->extent)) {
+ if (!IsConstNumber(loop->min) || !IsConstNumber(loop->extent)) {
throw MakeScheduleError<SymbolicShapeError>(state_->mod,
ffi::GetRef<For>(loop));
}
// infer factor if needed and check validity of factors
@@ -575,7 +575,7 @@ ffi::Array<LoopRV> ConcreteScheduleNode::LoopPartition(
infer_index = i;
} else {
PrimExpr factor = this->Get(factor_rvs[i].value());
- if (is_const_int(factor) && !is_positive_const(factor)) {
+ if (IsConstInt(factor) && !IsPositiveConst(factor)) {
throw MakeScheduleError<NonPositiveFactorError>(state_->mod,
factor.as<IntImmNode>()->value,
i);
}
diff --git a/src/s_tir/schedule/ir_comparator.cc
b/src/s_tir/schedule/ir_comparator.cc
index d447186112..60aa63e1f7 100644
--- a/src/s_tir/schedule/ir_comparator.cc
+++ b/src/s_tir/schedule/ir_comparator.cc
@@ -573,7 +573,7 @@ bool TensorizeComparator::CompareBufferRegion(const
TensorRegion& lhs, const Ten
indices_base.reserve(lhs->region.size());
for (int i = 0; i < offset; i++) {
// High-dim region must be element-wise
- if (!is_one(lhs->region[i]->extent)) {
+ if (!IsOne(lhs->region[i]->extent)) {
if (assert_mode_) {
std::ostringstream os;
os << "CompareBufferRegion returning false because buffer extent
high-dim region must be "
@@ -606,7 +606,7 @@ bool TensorizeComparator::CompareBufferRegion(const
TensorRegion& lhs, const Ten
const std::vector<PrimExpr>& indices_base = it->second;
for (int i = 0; i < offset; i++) {
// High-dim region must be element-wise
- if (!is_one(lhs->region[i]->extent)) {
+ if (!IsOne(lhs->region[i]->extent)) {
if (assert_mode_) {
std::ostringstream os;
os << "CompareBufferRegion returning false because buffer extent
high-dim region must be "
diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc
b/src/s_tir/schedule/primitive/blockize_tensorize.cc
index 2b4799a3e3..c61bae7dcb 100644
--- a/src/s_tir/schedule/primitive/blockize_tensorize.cc
+++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc
@@ -97,7 +97,7 @@ ffi::Array<ffi::Array<sym::IterMark>>
TrivialSubspaceDivision(const ffi::Array<I
const PrimExpr&
predicate,
const
ffi::Array<Var>& outer_iters,
const
ffi::Array<Var>& inner_iters) {
- if (!is_one(predicate)) return {};
+ if (!IsOne(predicate)) return {};
ffi::Array<ffi::Array<sym::IterMark>> res;
std::unordered_set<const VarNode*> outer_loop_vars;
std::unordered_set<const VarNode*> inner_loop_vars;
@@ -256,10 +256,10 @@ ffi::Map<Var, PrimExpr> DeriveBlockBinding(
outer_iter_vars->push_back(outer_iter);
}
PrimExpr sub = [&]() -> PrimExpr {
- if (is_one(inner_mark->extent)) {
+ if (IsOne(inner_mark->extent)) {
// Skip inner var when extent is 1
// substitution
- if (is_one(outer_mark->extent) && !preserve_unit_iters) {
+ if (IsOne(outer_mark->extent) && !preserve_unit_iters) {
// Simplify outer if not preserve_unit_iters
return IntImm(outer_mark->extent.ty(), 0);
} else {
@@ -273,7 +273,7 @@ ffi::Map<Var, PrimExpr> DeriveBlockBinding(
inner_bindings->push_back(NormalizeIterMapToExpr(inner_binding));
inner_iter_vars->push_back(inner_iter);
// substitution
- if (is_one(outer_mark->extent)) {
+ if (IsOne(outer_mark->extent)) {
return inner_iter->var;
} else {
return outer_iter * inner_mark->extent + inner_iter->var;
diff --git a/src/s_tir/schedule/primitive/cache_read_write.cc
b/src/s_tir/schedule/primitive/cache_read_write.cc
index 16059de0ae..2689cca6eb 100644
--- a/src/s_tir/schedule/primitive/cache_read_write.cc
+++ b/src/s_tir/schedule/primitive/cache_read_write.cc
@@ -2015,7 +2015,7 @@ StmtSRef CacheRead(ScheduleState self, const StmtSRef&
block_sref, int read_buff
PrimExpr nested_pred = read_region_opt ? CollectNestedBlockPredicates(
block->body, read_buffer,
BufferIndexType::kRead)
: IntImm::Bool(true);
- if (read_region_opt && !is_one(nested_pred) && block_sref->parent !=
nullptr) {
+ if (read_region_opt && !IsOne(nested_pred) && block_sref->parent !=
nullptr) {
StmtSRef parent_sref = ffi::GetRef<StmtSRef>(block_sref->parent);
return RelaxBufferRegion(self, read_region_opt.value(), block_sref,
parent_sref, scope_sref,
nested_pred);
diff --git a/src/s_tir/schedule/primitive/compute_at.cc
b/src/s_tir/schedule/primitive/compute_at.cc
index 50fe23e9c0..8fac5f6301 100644
--- a/src/s_tir/schedule/primitive/compute_at.cc
+++ b/src/s_tir/schedule/primitive/compute_at.cc
@@ -230,7 +230,7 @@ struct BlockVarDomainInfo {
analyzer->CanProveEqual(bound.max(), intersect.max())) {
dom = bound;
bound = sym::IntSet::Nothing();
- } else if (is_const_int(intersect.min()) && is_const_int(intersect.max()))
{
+ } else if (IsConstInt(intersect.min()) && IsConstInt(intersect.max())) {
// if the bound induce constant iter range, merge bound to loop domain
dom = intersect;
bound = sym::IntSet::Nothing();
@@ -274,7 +274,7 @@ class ScopeReconstructor : public StmtExprMutator {
PrimExpr predicate = IntImm::Bool(true);
for (int i = 0; i < n_iters; ++i) {
Range iter_dom = iter_doms[i].dom.CoverRange(block_->iter_vars[i]->dom);
- if (preserve_unit_loops || !is_one(iter_dom->extent)) {
+ if (preserve_unit_loops || !IsOne(iter_dom->extent)) {
int bits = std::max(iter_dom->min.ty().bits(),
iter_dom->extent.ty().bits());
Var var("ax" + std::to_string(loop_vars.size()), PrimType::Int(bits));
loop_vars.push_back(var);
@@ -515,7 +515,7 @@ void UpdateBlockVarDomainDimwise(
PrimExpr dim_max = max(GetTensorVar(buffer)->shape[i] - 1, 0);
sym::Analyzer analyzer_ref = ffi::GetRef<sym::Analyzer>(analyzer);
- if (provided.CanProveSinglePoint(analyzer_ref) &&
is_const_int(provided.min())) {
+ if (provided.CanProveSinglePoint(analyzer_ref) &&
IsConstInt(provided.min())) {
TVM_FFI_ICHECK(required.CanProveSinglePoint(analyzer_ref) &&
analyzer->CanProveEqual(provided.min(), required.min()));
continue;
diff --git a/src/s_tir/schedule/primitive/compute_inline.cc
b/src/s_tir/schedule/primitive/compute_inline.cc
index d728de4ea4..695239097d 100644
--- a/src/s_tir/schedule/primitive/compute_inline.cc
+++ b/src/s_tir/schedule/primitive/compute_inline.cc
@@ -579,7 +579,7 @@ class ComputeInliner : public BaseInliner {
for (const Var& var : idx_vars_)
prim_idx_vars.push_back(var.as_or_throw<PrimExpr>());
auto inverse_iter_map = sym::InverseAffineIterMap(res->indices,
prim_idx_vars);
for (const auto& iter : producer_block->iter_vars) {
- if (is_const_int(iter->dom->min) &&
analyzer_->CanProveEqual(iter->dom->extent, 1)) {
+ if (IsConstInt(iter->dom->min) &&
analyzer_->CanProveEqual(iter->dom->extent, 1)) {
// fallback mapping for constant iters
inverse_iter_map.Set(iter->var, iter->dom->min);
}
@@ -705,7 +705,7 @@ class ReverseComputeInliner : public BaseInliner {
bool BodyPatternAllowInline(const SBlockRealize& consumer_block_realize) {
const SBlock& consumer_block = consumer_block_realize->block;
- if (!is_one(consumer_block_realize->predicate)) {
+ if (!IsOne(consumer_block_realize->predicate)) {
// Failure: Predicate is the consumer block is not supported
return false;
}
@@ -724,7 +724,7 @@ class ReverseComputeInliner : public BaseInliner {
for (const auto& iter_var : consumer_block->iter_vars) {
consumer_iter_doms.Set(iter_var->var, iter_var->dom);
// Set default mapping for unit iters
- if (is_const_int(iter_var->dom->extent, 1) &&
is_const_int(iter_var->dom->min)) {
+ if (IsConstInt(iter_var->dom->extent, 1) &&
IsConstInt(iter_var->dom->min)) {
VarRemapSet(iter_var->var, iter_var->dom->min);
}
}
@@ -807,7 +807,7 @@ class ReverseComputeInliner : public BaseInliner {
.ValueOrUnchanged(consumer_iter_in_bound_);
// Simplify the predicate using the producer block iter domains
predicate = analyzer_->Simplify(predicate);
- if (is_one(predicate)) {
+ if (IsOne(predicate)) {
return producer_block_realize;
}
if (const auto* if_ = producer_block->body.as<IfThenElseNode>()) {
@@ -1156,7 +1156,7 @@ class ReductionEpilogueFuser : public BaseInliner {
bool ReductionEpilogueFuser::BodyPatternAllowFusion(const SBlockRealize&
epilogue_block_realize) {
// 1. Validate predicate
- if (!is_one(epilogue_block_realize->predicate)) {
+ if (!IsOne(epilogue_block_realize->predicate)) {
// Failure: Predicate in epilogue block is not supported
return false;
}
diff --git a/src/s_tir/schedule/primitive/decompose_padding.cc
b/src/s_tir/schedule/primitive/decompose_padding.cc
index 379ed06b39..cfeea4e9fb 100644
--- a/src/s_tir/schedule/primitive/decompose_padding.cc
+++ b/src/s_tir/schedule/primitive/decompose_padding.cc
@@ -118,7 +118,7 @@ class PaddingInfoAnalyzer {
.as_or_throw<PrimExpr>();
PrimExpr in_bound_value = if_then_else->args[1].as_or_throw<PrimExpr>();
PrimExpr pad_value = if_then_else->args[2].as_or_throw<PrimExpr>();
- if (!is_const_number(pad_value)) {
+ if (!IsConstNumber(pad_value)) {
SetError("Pad value should be constant");
return false;
}
diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc
b/src/s_tir/schedule/primitive/layout_transformation.cc
index dfdb274f2e..c17887ba15 100644
--- a/src/s_tir/schedule/primitive/layout_transformation.cc
+++ b/src/s_tir/schedule/primitive/layout_transformation.cc
@@ -202,7 +202,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
const PrimExpr& buffer_dim = old_buffer_->shape[i];
PrimExpr index =
ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(op->indices[i], f_substitute)
.as_or_throw<PrimExpr>();
- bool is_loop_over_axis = index.same_as(loop->loop_var) &&
is_const_int(loop->min, 0) &&
+ bool is_loop_over_axis = index.same_as(loop->loop_var) &&
IsConstInt(loop->min, 0) &&
prim::ExprDeepEqual()(loop->extent,
buffer_dim) &&
loop->kind == ForKind::kSerial;
if (!is_loop_over_axis) {
@@ -496,7 +496,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
IndexMap inverse, PrimExpr
padding_predicate,
ffi::Optional<IndexMap>
pad_value,
sym::AnalyzerObj* analyzer)
const {
- if (write_info_.size() || is_zero(padding_predicate) ||
!pad_value.has_value()) {
+ if (write_info_.size() || IsZero(padding_predicate) ||
!pad_value.has_value()) {
return std::nullopt;
}
@@ -551,14 +551,14 @@ class TransformLayoutPlanner : public StmtExprVisitor {
PrimExpr
padding_predicate,
ffi::Optional<IndexMap> pad_value,
sym::AnalyzerObj*
analyzer) const {
- if (write_info_.empty() || is_zero(padding_predicate) ||
!pad_value.has_value()) {
+ if (write_info_.empty() || IsZero(padding_predicate) ||
!pad_value.has_value()) {
return std::nullopt;
}
ffi::Map<SBlock, SBlock> new_block_to_old;
auto generate_if_then_else_block = [&](const WriteInfo& info) ->
ffi::Optional<Stmt> {
if (!info.contains_row_major_traversal || !pad_value.has_value() ||
- is_zero(padding_predicate)) {
+ IsZero(padding_predicate)) {
return std::nullopt;
}
@@ -602,7 +602,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
IndexMap inverse, PrimExpr
padding_predicate,
ffi::Optional<IndexMap>
pad_value,
sym::AnalyzerObj* analyzer)
const {
- if (write_info_.empty() || is_zero(padding_predicate) ||
!pad_value.has_value()) {
+ if (write_info_.empty() || IsZero(padding_predicate) ||
!pad_value.has_value()) {
return std::nullopt;
}
@@ -1286,7 +1286,7 @@ void TransformLayout(ScheduleState self, const StmtSRef&
block_sref, int buffer_
}();
}
- bool has_padding = !is_zero(padding_predicate);
+ bool has_padding = !IsZero(padding_predicate);
if (has_padding && !pad_value.has_value()) {
throw MakeScheduleError<TransformationIntroducesPaddingError>(self->mod,
old_buffer, index_map,
padding_predicate);
@@ -1484,7 +1484,7 @@ void TransformBlockLayout(ScheduleState self, const
StmtSRef& block_sref,
block_vars.push_back(iter_var->var);
block_iter_dom.Set(iter_var->var, iter_var->dom);
block_iter_type[iter_var->var.get()] = iter_var->iter_type;
- TVM_FFI_ICHECK(is_zero(iter_var->dom->min));
+ TVM_FFI_ICHECK(IsZero(iter_var->dom->min));
block_iter_range_array.push_back(iter_var->dom->extent);
}
@@ -1503,7 +1503,7 @@ void TransformBlockLayout(ScheduleState self, const
StmtSRef& block_sref,
Var new_block_var{"v" + std::to_string(i),
transformed_block_iters[i].ty()};
new_block_vars.push_back(new_block_var.as_or_throw<PrimExpr>());
IterVarType iter_type;
- if (is_one(new_block_iter_range[i])) {
+ if (IsOne(new_block_iter_range[i])) {
iter_type = kDataPar;
} else {
iter_type = DetectNewBlockIterType(transformed_block_iters[i],
block_iter_type);
diff --git a/src/s_tir/schedule/primitive/loop_transformation.cc
b/src/s_tir/schedule/primitive/loop_transformation.cc
index 3eb0c2c895..40169d1869 100644
--- a/src/s_tir/schedule/primitive/loop_transformation.cc
+++ b/src/s_tir/schedule/primitive/loop_transformation.cc
@@ -1007,12 +1007,12 @@ StmtSRef Fuse(ScheduleState self, const
ffi::Array<StmtSRef>& loop_srefs,
for (int i = static_cast<int>(loops.size()) - 1; i > 0; i--) {
PrimExpr next_lower = analyzer->canonical_simplify(loops[i]->extent *
lower);
substitute_value.Set(
- i, is_one(loops[i]->extent)
+ i, IsOne(loops[i]->extent)
? PrimExpr(0)
: floordiv(floormod(fused_var.as_or_throw<PrimExpr>(),
next_lower), lower));
lower = next_lower;
}
- substitute_value.Set(0, is_one(loops[0]->extent)
+ substitute_value.Set(0, IsOne(loops[0]->extent)
? PrimExpr(0)
: floordiv(fused_var.as_or_throw<PrimExpr>(),
lower));
Stmt new_stmt = loops.back()->body;
diff --git a/src/s_tir/schedule/primitive/pad_einsum.cc
b/src/s_tir/schedule/primitive/pad_einsum.cc
index 392d9ac973..8e8798fbf8 100644
--- a/src/s_tir/schedule/primitive/pad_einsum.cc
+++ b/src/s_tir/schedule/primitive/pad_einsum.cc
@@ -53,7 +53,7 @@ ffi::Optional<ffi::Array<Var>> CheckTrivialBufferAccess(const
TensorRegion& buff
ffi::Array<Var> indices;
indices.reserve(buffer_region->region.size());
for (const Range& range : buffer_region->region) {
- if (!tvm::prim::is_one(range->extent)) {
+ if (!tvm::prim::IsOne(range->extent)) {
return std::nullopt;
}
if (range->min->IsInstance<IntImmNode>()) {
diff --git a/src/s_tir/schedule/primitive/reduction.cc
b/src/s_tir/schedule/primitive/reduction.cc
index b76d254b3e..1bde740159 100644
--- a/src/s_tir/schedule/primitive/reduction.cc
+++ b/src/s_tir/schedule/primitive/reduction.cc
@@ -170,7 +170,7 @@ class LoopHeightError : public ScheduleErrorContextObj {
PrimExpr RewriteInitPredicate(PrimExpr pred,
const std::unordered_set<const VarNode*>&
discarded_loops) {
- if (is_one(pred)) return IntImm::Bool(true);
+ if (IsOne(pred)) return IntImm::Bool(true);
if (const auto* and_node = pred.as<AndNode>()) {
return RewriteInitPredicate(and_node->a, discarded_loops) &&
RewriteInitPredicate(and_node->b, discarded_loops);
@@ -706,7 +706,7 @@ class LoopPropertyError : public ScheduleErrorContextObj {
meet_reduction_loop = true;
}
continue;
- } else if (meet_reduction_loop && !is_one(loop->extent)) {
+ } else if (meet_reduction_loop && !IsOne(loop->extent)) {
throw MakeScheduleError<LoopPropertyError>(self->mod, loop,
kUnboundLoopUnderReductionLoop);
}
}
diff --git a/src/s_tir/schedule/primitive/rolling_buffer.cc
b/src/s_tir/schedule/primitive/rolling_buffer.cc
index edace22ceb..7f2f721cf9 100644
--- a/src/s_tir/schedule/primitive/rolling_buffer.cc
+++ b/src/s_tir/schedule/primitive/rolling_buffer.cc
@@ -214,7 +214,7 @@ class RollingBufferInfoCollector {
// If the bound is just a Var, that implies the stride is 1
iter_var = p_var.Eval();
stride = 1;
- } else if (is_const_int(bound->min)) {
+ } else if (IsConstInt(bound->min)) {
// If the bound is an int, we can't roll over it
iter_var = std::nullopt;
} else {
diff --git a/src/s_tir/schedule/transform.cc b/src/s_tir/schedule/transform.cc
index 11eb90e767..22c935166f 100644
--- a/src/s_tir/schedule/transform.cc
+++ b/src/s_tir/schedule/transform.cc
@@ -527,7 +527,7 @@ ffi::Optional<ffi::ObjectRef> NormalizePrimFunc(Schedule
sch) {
if (binds[i].get() != loop->loop_var.get()) {
return std::nullopt;
}
- if (!is_zero(loop->min)) {
+ if (!IsZero(loop->min)) {
return std::nullopt;
}
}
@@ -544,7 +544,7 @@ ffi::Optional<ffi::ObjectRef> NormalizePrimFunc(Schedule
sch) {
for (const IterVar& iter : sch->Get(block)->iter_vars) {
PrimVar var = iter->var.CopyWithSuffix("");
index_map_inputs.push_back(var);
- if (!is_one(iter->dom->extent)) {
+ if (!IsOne(iter->dom->extent)) {
index_map_outputs.push_back(var);
if (iter->iter_type == IterVarType::kDataPar) {
has_spatial_iter = true;
diff --git a/src/s_tir/script/ir_builder/frame.cc
b/src/s_tir/script/ir_builder/frame.cc
index 89bc00ae35..b1a99e6944 100644
--- a/src/s_tir/script/ir_builder/frame.cc
+++ b/src/s_tir/script/ir_builder/frame.cc
@@ -40,7 +40,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
} else if (auto flag = value.as<int64_t>()) {
enabled = *flag != 0;
} else if (auto flag = value.as<tvm::IntImm>()) {
- enabled = !tvm::prim::is_zero(*flag);
+ enabled = !tvm::prim::IsZero(*flag);
}
TVM_FFI_CHECK(enabled, ValueError) << "Ts.prim_func cannot disable the
s_tir attribute";
});
diff --git a/src/s_tir/script/printer/stmt.cc b/src/s_tir/script/printer/stmt.cc
index 9db4b7e8e6..43110e662a 100644
--- a/src/s_tir/script/printer/stmt.cc
+++ b/src/s_tir/script/printer/stmt.cc
@@ -53,14 +53,13 @@ ffi::Array<StmtDoc> SBlockBody(DocTranslatorObj* d, const
s_tir::SBlockNode* blo
default:
TVM_FFI_THROW(TypeError) << "printer unsupported SBlock iter var
kind";
}
- ExprDoc domain =
- prim::is_zero(iter->dom->min) && iter->dom->min.ty() ==
iter->dom->extent.ty()
- ? d->Translate(iter->dom->extent).value()
- : NamespaceDoc("ir")
- ->Attr("Range")
- ->Attr("from_min_extent")
- ->Call({d->Translate(iter->dom->min).value(),
- d->Translate(iter->dom->extent).value()});
+ ExprDoc domain = prim::IsZero(iter->dom->min) && iter->dom->min.ty() ==
iter->dom->extent.ty()
+ ? d->Translate(iter->dom->extent).value()
+ : NamespaceDoc("ir")
+ ->Attr("Range")
+ ->Attr("from_min_extent")
+ ->Call({d->Translate(iter->dom->min).value(),
+
d->Translate(iter->dom->extent).value()});
ffi::Array<ExprDoc> args = {domain};
if (realize)
args.push_back(d->Translate(realize->iter_values[i]).value());
IdDoc lhs = VarDoc(d, iter->var);
@@ -76,7 +75,7 @@ ffi::Array<StmtDoc> SBlockBody(DocTranslatorObj* d, const
s_tir::SBlockNode* blo
std::nullopt));
}
for (size_t i = 0; i < axes.size(); ++i) d->Emit(axes[i],
block->iter_vars[i]);
- if (realize && !tvm::prim::is_one(realize->predicate)) {
+ if (realize && !tvm::prim::IsOne(realize->predicate)) {
d->Emit(ExprStmtDoc(NamespaceDoc("s_tir")->Attr("where")->Call(
{d->Translate(realize->predicate).value()})),
realize->predicate);
diff --git a/src/s_tir/transform/canonicalize_loop.cc
b/src/s_tir/transform/canonicalize_loop.cc
index f7f7bdb847..203c7a20ed 100644
--- a/src/s_tir/transform/canonicalize_loop.cc
+++ b/src/s_tir/transform/canonicalize_loop.cc
@@ -47,7 +47,7 @@ class LoopCanonicalizer : public StmtExprMutator {
private:
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
- if (is_zero(op->min) && op->HasTrivialStep()) {
+ if (IsZero(op->min) && op->HasTrivialStep()) {
return StmtExprMutator::Mutate_(op, inplace_mode);
}
const auto* loop_var = op->loop_var.get();
diff --git a/src/s_tir/transform/inject_software_pipeline.cc
b/src/s_tir/transform/inject_software_pipeline.cc
index e23d368138..fa88a2b74f 100644
--- a/src/s_tir/transform/inject_software_pipeline.cc
+++ b/src/s_tir/transform/inject_software_pipeline.cc
@@ -74,7 +74,7 @@ ffi::Optional<Var> GetBufferDataVar(const ffi::Any& data) {
*/
SBlock MakeSBlock(const Stmt& body, const ffi::Map<Var, TensorVar>&
buffer_data_to_buffer) {
if (const SBlockRealizeNode* block_realize = body.as<SBlockRealizeNode>()) {
- if (is_one(block_realize->predicate)) {
+ if (IsOne(block_realize->predicate)) {
// no need to create a new block
return block_realize->block;
}
@@ -411,7 +411,7 @@ class PipelineRewriter : public StmtExprMutator {
PrimExpr epigogue_start = pipeline_loop_->min + pipeline_loop_->extent;
ffi::Optional<PrimExpr> extra_epilogue_lower_bound = std::nullopt;
if (max_stage_ > 1 &&
!analyzer_->CanProveGreaterEqual(pipeline_loop_->extent, max_stage_)) {
- if (is_const_int(epigogue_start)) {
+ if (IsConstInt(epigogue_start)) {
epigogue_start = max(epigogue_start, pipeline_loop_->min + max_stage_);
} else {
// for dynamic case, introduce extra lowerbound as loop predicate
@@ -1186,7 +1186,7 @@ class PipelineInjector : public StmtExprMutator {
};
for (size_t i = 0; i < pipeline_body_seq->seq.size(); i++) {
const auto* nested_block_realize =
pipeline_body_seq->seq[i].as<SBlockRealizeNode>();
- if (nested_block_realize && is_one(nested_block_realize->predicate) &&
+ if (nested_block_realize && IsOne(nested_block_realize->predicate) &&
nested_block_realize->block->body->IsInstance<SeqStmtNode>()) {
const SBlock& nested_pipeline_block = nested_block_realize->block;
TVM_FFI_ICHECK(
diff --git a/src/s_tir/transform/inject_virtual_thread.cc
b/src/s_tir/transform/inject_virtual_thread.cc
index c5b50328a0..1dc7740c7c 100644
--- a/src/s_tir/transform/inject_virtual_thread.cc
+++ b/src/s_tir/transform/inject_virtual_thread.cc
@@ -480,7 +480,7 @@ class VTInjector : public s_tir::IRMutatorWithAnalyzer {
}
// For
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
- TVM_FFI_ICHECK(is_zero(op->min));
+ TVM_FFI_ICHECK(IsZero(op->min));
auto extent_result = this->Mutate(op->extent, inplace_mode);
bool extent_unchanged = extent_result.UnchangedOrSameAs(op->extent);
PrimExpr extent = std::move(extent_result).ValueOrUnchanged(op->extent);
diff --git a/src/s_tir/transform/loop_partition.cc
b/src/s_tir/transform/loop_partition.cc
index 9c093ca872..e7d48f0529 100644
--- a/src/s_tir/transform/loop_partition.cc
+++ b/src/s_tir/transform/loop_partition.cc
@@ -116,7 +116,7 @@ static bool IsVirtualThread(const ForNode* op) {
static bool HasPartitionHint(const ForNode* op) {
auto it = op->annotations.find(attr::loop_partition_hint);
- return it != op->annotations.end() &&
is_one((*it).second.as_or_throw<PrimExpr>());
+ return it != op->annotations.end() &&
IsOne((*it).second.as_or_throw<PrimExpr>());
}
// Resolve expression-valued hints before partitioning clones or renames loops.
@@ -170,7 +170,7 @@ class CandidateSelector final : public StmtExprVisitor {
return StmtExprVisitor::Visit_(op);
}
// partition const loop when sets partition_const_loop_
- if (!is_const_int(op->min) || !is_const_int(op->extent) ||
partition_const_loop_) {
+ if (!IsConstInt(op->min) || !IsConstInt(op->extent) ||
partition_const_loop_) {
record_.insert({var, false});
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
if (record_.at(var) && !no_split_) {
@@ -190,7 +190,7 @@ class CandidateSelector final : public StmtExprVisitor {
runtime::ThreadScope scope =
runtime::ThreadScope::Create(op->args[0].as_or_throw<StringImm>()->value);
PrimExpr extent = op->args[1].as_or_throw<PrimExpr>();
- if ((scope.rank == 0) && (!is_const_int(extent) ||
partition_const_loop_)) {
+ if ((scope.rank == 0) && (!IsConstInt(extent) || partition_const_loop_))
{
record_.insert({var.get(), false});
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
if (record_.at(var.get()) && !no_split_) {
diff --git a/src/s_tir/transform/lower_cross_thread_reduction.cc
b/src/s_tir/transform/lower_cross_thread_reduction.cc
index fde0ac24f7..03c52c16ef 100644
--- a/src/s_tir/transform/lower_cross_thread_reduction.cc
+++ b/src/s_tir/transform/lower_cross_thread_reduction.cc
@@ -883,7 +883,7 @@ class CrossThreadReductionTransformer : public
StmtExprMutator {
// - the block-realize has a non-constant-true predicate.
bool need_in_thread_reduction =
n_bound_reduction_loops < static_cast<int>(reduction_loops.size()) ||
- !is_one(realize->predicate);
+ !IsOne(realize->predicate);
// Step 3. Create intermediate buffers, storing them in `ct_buffers` and
// `it_buffers`. Let the scope block allocate these new buffers.
ffi::Array<TensorVar>& new_buffers =
block2new_buffers_[block_stack_.back()];
diff --git a/src/s_tir/transform/lower_match_buffer.cc
b/src/s_tir/transform/lower_match_buffer.cc
index fbc26e95be..c6a121a7aa 100644
--- a/src/s_tir/transform/lower_match_buffer.cc
+++ b/src/s_tir/transform/lower_match_buffer.cc
@@ -197,8 +197,8 @@ class MatchBufferLower : public StmtExprMutator {
<< " required alignment=" << buffer->data_alignment
<< ", provided alignment=" << source_buffer->data_alignment;
}
- if (is_zero(buffer->elem_offset)) {
- TVM_FFI_ICHECK(is_zero(source_buffer->elem_offset))
+ if (IsZero(buffer->elem_offset)) {
+ TVM_FFI_ICHECK(IsZero(source_buffer->elem_offset))
<< "Trying to bind a TensorVar with offset into one without offset "
<< " required elem_offset=" << buffer->elem_offset
<< ", provided elem_offset=" << source_buffer->elem_offset;
diff --git a/src/s_tir/transform/lower_opaque_block.cc
b/src/s_tir/transform/lower_opaque_block.cc
index d6f276861f..9cb88a22b6 100644
--- a/src/s_tir/transform/lower_opaque_block.cc
+++ b/src/s_tir/transform/lower_opaque_block.cc
@@ -74,7 +74,7 @@ class OpaqueBlockLower : public StmtExprMutator {
PrimExpr predicate = this->Mutate(op->predicate,
inplace_mode).ValueOrUnchanged(op->predicate);
// Step 2. Transform the `predicate` to if-then-else
Stmt body = new_block->body;
- if (!is_one(predicate)) {
+ if (!IsOne(predicate)) {
body = IfThenElse(predicate, std::move(body));
}
// Step 3. Handle allocations in reverse order
@@ -111,7 +111,7 @@ class OpaqueBlockLower : public StmtExprMutator {
.as_or_throw<UnchangedOr<ffi::Optional<PrimExpr>>>()
.ValueOrUnchanged(op->step);
auto previous_remap = VarRemapGet(op->loop_var);
- if (op->kind != ForKind::kThreadBinding && is_one(extent) &&
op->annotations.empty()) {
+ if (op->kind != ForKind::kThreadBinding && IsOne(extent) &&
op->annotations.empty()) {
// handling unit loop
VarRemapSet(op->loop_var, prim::cast(op->loop_var.ty(), min));
}
@@ -137,7 +137,7 @@ class OpaqueBlockLower : public StmtExprMutator {
VarRemapSet(op->loop_var, previous_remap);
// Step 2. Keep thread-binding loops until LowerThreadBinding.
- if (op->kind != ForKind::kThreadBinding && is_one(extent) &&
op->annotations.empty() &&
+ if (op->kind != ForKind::kThreadBinding && IsOne(extent) &&
op->annotations.empty() &&
!op->annotations.count(s_tir::attr::irregular_loop_mark)) {
return body;
}
diff --git a/src/s_tir/transform/lower_thread_binding.cc
b/src/s_tir/transform/lower_thread_binding.cc
index d6d3dd87a9..4270e3f807 100644
--- a/src/s_tir/transform/lower_thread_binding.cc
+++ b/src/s_tir/transform/lower_thread_binding.cc
@@ -56,7 +56,7 @@ class ThreadBindingLowerer : public StmtExprMutator {
}
PrimExpr min = Mutate(op->min, inplace_mode).ValueOrUnchanged(op->min);
PrimExpr extent = Mutate(op->extent,
inplace_mode).ValueOrUnchanged(op->extent);
- TVM_FFI_ICHECK(is_zero(min)) << "Thread binding loops must start at zero";
+ TVM_FFI_ICHECK(IsZero(min)) << "Thread binding loops must start at zero";
TVM_FFI_ICHECK(op->thread_binding.has_value());
TVM_FFI_ICHECK(!op->annotations.count("loop_partition_hint") ||
op->annotations.at("loop_partition_hint") == nullptr)
diff --git a/src/s_tir/transform/memhammer_coalesce.cc
b/src/s_tir/transform/memhammer_coalesce.cc
index 5843bfc27c..865383b8b1 100644
--- a/src/s_tir/transform/memhammer_coalesce.cc
+++ b/src/s_tir/transform/memhammer_coalesce.cc
@@ -178,7 +178,7 @@ ffi::Array<PrimExpr> GetMapping(const Stmt& stmt, const
ConstraintSet& constrain
sym::Analyzer analyzer;
for (int i = 0; i < static_cast<int>(write_region->region.size()); i++) {
PrimExpr pattern = analyzer->Simplify(write_index[i] -
write_region->region[i]->min);
- if (!is_zero(pattern)) {
+ if (!IsZero(pattern)) {
result.push_back(pattern);
}
}
@@ -212,7 +212,7 @@ Stmt InverseMapping::Rewrite(const Stmt& stmt, const
ConstraintSet& constraints,
ffi::Map<Var, PrimExpr> substitute_map;
// Step 3.1 construct target buffer indices
for (int i = 0, j = 0; i < static_cast<int>(write_region->region.size());
i++) {
- if (is_one(write_region->region[i]->extent)) {
+ if (IsOne(write_region->region[i]->extent)) {
write_index.push_back(write_region->region[i]->min);
} else {
PrimVar var =
loop_vars[j].as_or_throw<PrimVar>().CopyWithSuffix("_inverse");
@@ -228,7 +228,7 @@ Stmt InverseMapping::Rewrite(const Stmt& stmt, const
ConstraintSet& constraints,
};
// Step 3.2 construct source buffer indices
for (int i = 0, j = 0; i < static_cast<int>(read_region->region.size());
i++) {
- if (is_one(read_region->region[i]->extent)) {
+ if (IsOne(read_region->region[i]->extent)) {
read_index.push_back(read_region->region[i]->min);
} else {
PrimExpr inverse = inverse_mapping[loop_vars[j++].as_or_throw<Var>()];
diff --git a/src/s_tir/transform/memhammer_rewrite_rule.h
b/src/s_tir/transform/memhammer_rewrite_rule.h
index bbf87732bd..92712861ae 100644
--- a/src/s_tir/transform/memhammer_rewrite_rule.h
+++ b/src/s_tir/transform/memhammer_rewrite_rule.h
@@ -171,7 +171,7 @@ class CreateLocalStage : public RewriteRule {
TensorVar tgt_buffer =
constraints.write_region->source.as_or_throw<tvm::tirx::TensorVar>();
return IsCopyBetweenScope(src_buffer, tgt_buffer,
runtime::StorageRank::kGlobal,
runtime::StorageRank::kShared) &&
- is_one(constraints.add_local_stage);
+ IsOne(constraints.add_local_stage);
}
};
diff --git a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc
b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc
index 78dd491931..aa25b8e055 100644
--- a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc
+++ b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc
@@ -64,7 +64,7 @@ class RemoveLayoutRewriteBlock : public StmtExprMutator {
.as_or_throw<SBlock>();
auto it =
block->annotations.find(s_tir::attr::meta_schedule_layout_rewrite_preproc);
- if (it == block->annotations.end() ||
!is_one((*it).second.cast<PrimExpr>())) {
+ if (it == block->annotations.end() ||
!IsOne((*it).second.cast<PrimExpr>())) {
// The block is not a weight layout block
// Remove allocates if needed
ffi::Array<TensorVar> alloc_buffers;
diff --git a/src/sym/analyzer.cc b/src/sym/analyzer.cc
index d99043e0ea..8fad2e0b32 100644
--- a/src/sym/analyzer.cc
+++ b/src/sym/analyzer.cc
@@ -58,7 +58,7 @@ void AnalyzerObj::Bind(const Var& var, const PrimExpr& expr,
bool allow_override
void AnalyzerObj::Bind(const Var& var, const Range& range, bool
allow_override) {
TVM_FFI_ICHECK(range.defined());
- if (prim::is_one(range->extent)) {
+ if (prim::IsOne(range->extent)) {
this->Bind(var, range->min, allow_override);
} else {
this->const_int_bound.Bind(var, range, allow_override);
@@ -182,7 +182,7 @@ bool
AnalyzerObj::CanProveLessEqualThanSymbolicShapeValue(const PrimExpr& lhs,
const PrimExpr&
shape) {
if (this->CanProve(lhs <= shape, ProofStrength::kSymbolicBound)) return true;
// no need to do further attempt if shape is already a constant.
- if (prim::is_const_int(shape)) return false;
+ if (prim::IsConstInt(shape)) return false;
// collect constant scale and ignore symbolic part
// so 32 * n => cscale = 32
ffi::BigInt cscale = 1;
@@ -258,7 +258,7 @@ PrimExpr AnalyzerObj::Simplify(const PrimExpr& expr, int
steps) {
res = this->canonical_simplify(res);
for (int i = 0; i < steps; ++i) {
- if (prim::is_const_int(res)) {
+ if (prim::IsConstInt(res)) {
return res;
}
if (i % 2 == 0) {
diff --git a/src/sym/detect_linear_equation.cc
b/src/sym/detect_linear_equation.cc
index 190c19c5e6..bf7292fc35 100644
--- a/src/sym/detect_linear_equation.cc
+++ b/src/sym/detect_linear_equation.cc
@@ -232,13 +232,13 @@ bool DetectClipBound(const PrimExpr& cond,
ffi::Optional<PrimExpr> min_value;
ffi::Optional<PrimExpr> max_value;
- if (is_const_int(ret.coeff.value(), 1)) {
+ if (IsConstInt(ret.coeff.value(), 1)) {
// var + shift >=0 -> var >= -shift
min_value = -ret.base.value();
if (is_eq) {
max_value = min_value;
}
- } else if (is_const_int(ret.coeff.value(), -1)) {
+ } else if (IsConstInt(ret.coeff.value(), -1)) {
// -var + shift >=0 -> var <= shift
max_value = ret.base;
if (is_eq) {
diff --git a/src/sym/int_set.cc b/src/sym/int_set.cc
index b9050d3152..8af288fcb4 100644
--- a/src/sym/int_set.cc
+++ b/src/sym/int_set.cc
@@ -46,8 +46,8 @@ namespace tvm {
namespace sym {
using namespace tvm::prim;
-using prim::is_one;
-using prim::is_zero;
+using prim::IsOne;
+using prim::IsZero;
using prim::MakeConst;
TVM_FFI_STATIC_INIT_BLOCK() { IntervalSetNode::RegisterReflection(); }
@@ -182,8 +182,8 @@ inline IntervalSet Combine<prim::Mul>(AnalyzerObj*
analyzer, IntervalSet a, Inte
std::swap(a, b);
}
if (b->IsSinglePoint()) {
- if (is_zero(b->min_value)) return b;
- if (is_one(b->min_value)) return a;
+ if (IsZero(b->min_value)) return b;
+ if (IsOne(b->min_value)) return a;
if (analyzer->CanProveGreaterEqual(b->min_value, 0)) {
PrimExpr min_value = a->HasLowerBound() ? a->min_value * b->min_value :
neg_inf();
PrimExpr max_value = a->HasUpperBound() ? a->max_value * b->min_value :
pos_inf();
@@ -212,10 +212,10 @@ inline IntervalSet Combine<prim::Div>(AnalyzerObj*
analyzer, IntervalSet a, Inte
if (a->IsEmpty()) return a;
if (b->IsEmpty()) return b;
if (b->IsSinglePoint()) {
- if (is_zero(b->min_value)) {
+ if (IsZero(b->min_value)) {
TVM_FFI_THROW(InternalError) << "Divide by zero in CombineInterval Div";
}
- if (is_one(b->min_value)) return a;
+ if (IsOne(b->min_value)) return a;
// no relaxation is needed in here due to set is inclusive
if (analyzer->CanProveGreaterEqual(b->min_value, 0)) {
PrimExpr min_value = a->HasLowerBound() ? a->min_value / b->min_value :
neg_inf();
@@ -247,7 +247,7 @@ inline IntervalSet Combine<prim::Mod>(AnalyzerObj*
analyzer, IntervalSet a, Inte
if (b->IsSinglePoint()) {
const PrimExpr& divisor = b->min_value;
- if (is_zero(divisor)) {
+ if (IsZero(divisor)) {
TVM_FFI_THROW(InternalError) << "Modular by zero in CombineInterval Mod";
}
// We need to add more bound constraints throughout the code.
@@ -274,10 +274,10 @@ inline IntervalSet Combine<prim::FloorDiv>(AnalyzerObj*
analyzer, IntervalSet a,
if (a->IsEmpty()) return a;
if (b->IsEmpty()) return b;
if (b->IsSinglePoint()) {
- if (is_zero(b->min_value)) {
+ if (IsZero(b->min_value)) {
TVM_FFI_THROW(InternalError) << "Divide by zero in CombineInterval Div";
}
- if (is_one(b->min_value)) return a;
+ if (IsOne(b->min_value)) return a;
// no relaxation is needed in here due to set is inclusive
if (analyzer->CanProveGreaterEqual(b->min_value, 0)) {
PrimExpr min_value = a->HasLowerBound() ? floordiv(a->min_value,
b->min_value) : neg_inf();
@@ -309,7 +309,7 @@ inline IntervalSet Combine<prim::FloorMod>(AnalyzerObj*
analyzer, IntervalSet a,
if (b->IsSinglePoint()) {
const PrimExpr& divisor = b->min_value;
- if (is_zero(divisor)) {
+ if (IsZero(divisor)) {
TVM_FFI_THROW(InternalError) << "Modular by zero in CombineInterval Mod";
}
if (analyzer->CanProveGreaterEqual(divisor, 0)) {
@@ -831,20 +831,20 @@ bool IntSet::CanProveSinglePoint(const Analyzer& ana)
const {
bool IntSet::CanProvePositive() const {
Analyzer analyzer;
const IntervalSetNode* s_int = (*this).as<IntervalSetNode>();
- return (s_int && is_positive_const(analyzer->Simplify(s_int->min_value)));
+ return (s_int && IsPositiveConst(analyzer->Simplify(s_int->min_value)));
}
bool IntSet::CanProveNegative() const {
Analyzer analyzer;
const IntervalSetNode* s_int = (*this).as<IntervalSetNode>();
- return (s_int && is_negative_const(analyzer->Simplify(s_int->max_value)));
+ return (s_int && IsNegativeConst(analyzer->Simplify(s_int->max_value)));
}
bool IntSet::CanProveNonPositive() const {
Analyzer analyzer;
if (const auto* s_int = (*this).as<IntervalSetNode>()) {
auto max = analyzer->Simplify(s_int->max_value);
- return is_zero(max) || is_negative_const(max);
+ return IsZero(max) || IsNegativeConst(max);
}
return false;
}
@@ -853,7 +853,7 @@ bool IntSet::CanProveNonNegative() const {
Analyzer analyzer;
if (const IntervalSetNode* s_int = (*this).as<IntervalSetNode>()) {
auto min = analyzer->Simplify(s_int->min_value);
- return is_zero(min) || is_positive_const(min);
+ return IsZero(min) || IsPositiveConst(min);
}
return false;
}
@@ -877,7 +877,7 @@ SignType IntSet::GetSignType() const {
return kPositive;
} else if (CanProveNegative()) {
return kNegative;
- } else if (IsSinglePoint() && is_zero(PointValue())) {
+ } else if (IsSinglePoint() && IsZero(PointValue())) {
return kZero;
} else {
return kUnknown;
@@ -904,11 +904,11 @@ IntSet IntSet::Interval(PrimExpr min, PrimExpr max) {
// Range related code
inline bool ProveEqual(AnalyzerObj* analyzer, PrimExpr lhs, PrimExpr rhs) {
- return is_zero(analyzer->Simplify(lhs - rhs));
+ return IsZero(analyzer->Simplify(lhs - rhs));
}
IntSet IntSet::FromMinExtent(PrimExpr min, PrimExpr extent) {
- if (is_one(extent)) {
+ if (IsOne(extent)) {
return IntSet::SinglePoint(min);
}
return IntervalSet(min, extent + min - 1);
@@ -916,7 +916,7 @@ IntSet IntSet::FromMinExtent(PrimExpr min, PrimExpr extent)
{
IntSet IntSet::FromRange(Range r) {
// must make sure it can be matched back by MatchRange.
- if (is_one(r->extent)) {
+ if (IsOne(r->extent)) {
return IntSet::SinglePoint(r->min);
}
return IntervalSet(r->min, r->extent + r->min - 1);
@@ -1175,7 +1175,7 @@ ffi::Optional<ffi::Array<IntSet>>
EstimateRegionStrictBound(const ffi::Array<Ran
ffi::Array<PrimExpr> affine_indices;
affine_indices.reserve(ndim);
for (const Range& range : region) {
- if (!is_const_number(range->extent)) {
+ if (!IsConstNumber(range->extent)) {
// dynamic extent is not supported yet.
return std::nullopt;
}
@@ -1239,7 +1239,7 @@ ffi::Array<IntSet> EstimateRegionUpperBound(const
ffi::Array<Range>& region,
// dynamic extent is not supported yet.
PrimExpr extent = range->extent;
- if (!is_const_number(extent)) {
+ if (!IsConstNumber(extent)) {
IntSet relaxed = EvalSet(extent, AsIntSet(var_dom));
TVM_FFI_ICHECK(relaxed.HasUpperBound());
extent = relaxed.max();
diff --git a/src/sym/iter_affine_map.cc b/src/sym/iter_affine_map.cc
index d2f152e299..182f0f6a25 100644
--- a/src/sym/iter_affine_map.cc
+++ b/src/sym/iter_affine_map.cc
@@ -176,9 +176,9 @@ class IterMapRewriter : public tvm::ExprMutator {
for (auto kv : input_iters) {
const PrimVar& var = kv.first;
const Range& vrng = kv.second;
- if (simplify_trivial_iterators && is_one(vrng->extent)) {
+ if (simplify_trivial_iterators && IsOne(vrng->extent)) {
var_map_.insert_or_assign(var, IterSumExpr({}, vrng->min));
- } else if (is_zero(vrng->min)) {
+ } else if (IsZero(vrng->min)) {
IterMark mark(var.as_or_throw<PrimExpr>(), vrng->extent);
var_map_.insert_or_assign(var, IterSplitExpr(mark));
input_marks_.push_back(mark);
@@ -264,7 +264,7 @@ class IterMapRewriter : public tvm::ExprMutator {
if (check_level == IterMapLevel::Bijective) {
// all input marks must be visited
for (const IterMark& mark : input_marks_) {
- if (collector.visited_.count(mark) == 0 && !is_one(mark->extent)) {
+ if (collector.visited_.count(mark) == 0 && !IsOne(mark->extent)) {
return false;
}
}
@@ -682,7 +682,7 @@ class IterMapRewriter : public tvm::ExprMutator {
ffi::Optional<PrimExpr>
predicate_induced_max) {
// normalize to zero base
PrimExpr base = expr->base;
- if (!is_zero(base)) {
+ if (!IsZero(base)) {
expr.CopyOnWrite()->base = 0;
if (predicate_induced_min.has_value())
predicate_induced_min = predicate_induced_min.value() - base;
@@ -692,7 +692,7 @@ class IterMapRewriter : public tvm::ExprMutator {
ffi::Optional<IterSumExpr> opt = TryFuseIters(expr, check_level_, false);
TVM_FFI_ICHECK(!opt.has_value() || opt.value()->args.size() == 1);
// scale should be 1
- if (opt.has_value() && is_one(opt.value()->args[0]->scale)) {
+ if (opt.has_value() && IsOne(opt.value()->args[0]->scale)) {
const IterSplitExpr split = opt.value()->args[0];
IterSumExpr structured_form =
split->source->source.as_or_throw<IterSumExpr>();
// get the flattened form
@@ -730,7 +730,7 @@ class IterMapRewriter : public tvm::ExprMutator {
}
// When iter_min_delta is present, we need to normalize the structured
form to have minimum of
// 0, and add the delta to the mark_offset
- if (!is_zero(iter_min_delta)) {
+ if (!IsZero(iter_min_delta)) {
// structured form's offset should be updated
flattened_map_.erase(structured_form);
structured_form.CopyOnWrite()->base -= iter_min_delta;
@@ -880,7 +880,7 @@ class IterMapRewriter : public tvm::ExprMutator {
} else if (op->value == min_const_scale) {
// for ties, we want to look into 1 extent trivial iters
// prioritize trivial iterators
- if (is_one(expr->args[i]->extent) &&
!is_one(expr->args[base_index]->extent)) {
+ if (IsOne(expr->args[i]->extent) &&
!IsOne(expr->args[base_index]->extent)) {
base_index = static_cast<int>(i);
}
}
@@ -915,7 +915,7 @@ class IterMapRewriter : public tvm::ExprMutator {
*/
int FindFirstPossibleUnitExtentIndex(const IterSumExpr& expr) {
for (size_t i = 0; i < expr->args.size(); ++i) {
- if (is_one(expr->args[i]->extent)) return static_cast<int>(i);
+ if (IsOne(expr->args[i]->extent)) return static_cast<int>(i);
}
return static_cast<int>(expr->args.size());
}
@@ -944,7 +944,7 @@ class IterMapRewriter : public tvm::ExprMutator {
const PrimExpr& cur_scale = expr->args[j]->scale;
// for bijective mapping, the matched scale must equal to expected scale
if (analyzer_->CanProveEqual(cur_scale, expected_scale)) {
- if (is_one(expr->args[j]->extent)) return j;
+ if (IsOne(expr->args[j]->extent)) return j;
// if extent is not one and there is a possible extent=1 split
// further out, we need to extent the search
// extent=1 gets higher priority since they don't change the scale
@@ -1137,7 +1137,7 @@ class IterMapRewriter : public tvm::ExprMutator {
if (matched_pos == -1) {
// if exact scale is not possible, try to find an iter with scale
// that is smaller but closest to the scale.
- if (check_level != IterMapLevel::Bijective && is_const_int(base_scale,
1)) {
+ if (check_level != IterMapLevel::Bijective && IsConstInt(base_scale,
1)) {
matched_pos =
FindIterSmallerClosestToScale(expr, visited, expected_scale,
&opt_matched_scale);
}
@@ -1435,9 +1435,9 @@ bool MatchBoundConstraints(PrimExpr pred,
ffi::Map<PrimVar, Range>* input_iters,
bool bound_at_left;
if (lhs_uses_itervar || rhs_uses_itervar) {
// At least it uses one input iter
- if (is_const_int(lhs_expr) || !lhs_uses_itervar) {
+ if (IsConstInt(lhs_expr) || !lhs_uses_itervar) {
bound_at_left = true;
- } else if (is_const_int(rhs_expr) || !rhs_uses_itervar) {
+ } else if (IsConstInt(rhs_expr) || !rhs_uses_itervar) {
bound_at_left = false;
} else {
bound_at_left = false; // accumulate bound to rhs
@@ -1538,7 +1538,7 @@ IterMapResult DetectIterMap(const ffi::Array<PrimExpr>&
indices,
}
ffi::Map<PrimVar, Range> constrained_input_iters = input_iters;
std::vector<IterConstraint> constraints;
- if (!is_one(predicate) &&
+ if (!IsOne(predicate) &&
!MatchBoundConstraints(predicate, &constrained_input_iters,
&constraints)) {
result->errors.push_back("Could not parse predicate as constraints on the
input iterators.");
return result;
@@ -1858,7 +1858,7 @@ std::optional<std::pair<IterSplitExpr, PrimExpr>>
IterMapRewriter::PadDividendTo
ApproxLeastCommonMultiple(info.padding_factor, divisor *
split->lower_factor, analyzer_);
// If the split itself require no padding, return directly.
- if (is_zero(left_pad) && is_zero(right_pad)) {
+ if (IsZero(left_pad) && IsZero(right_pad)) {
return std::make_pair(split, PrimExpr(0));
}
@@ -1883,7 +1883,7 @@ std::optional<std::pair<IterSplitExpr, PrimExpr>>
IterMapRewriter::PadDividendTo
return std::make_pair(split, left_pad);
}
auto& info = it->second;
- if (is_zero(info.left_pad) && CanProveDivisible(mark->extent,
info.padding_factor)) {
+ if (IsZero(info.left_pad) && CanProveDivisible(mark->extent,
info.padding_factor)) {
// the iter mark requires no padding
return std::make_pair(split, left_pad);
}
@@ -1951,8 +1951,8 @@ ffi::Optional<PrimExpr>
IterMapRewriter::SplitFloorDivConst(IterSplitExpr lhs, P
PrimExpr rhs) {
// (lhs + base) // rhs
- if (is_one(rhs)) {
- if (is_zero(base)) {
+ if (IsOne(rhs)) {
+ if (IsZero(base)) {
// floordiv(x, 1) = x
return lhs;
@@ -1962,8 +1962,8 @@ ffi::Optional<PrimExpr>
IterMapRewriter::SplitFloorDivConst(IterSplitExpr lhs, P
}
}
- if (!is_one(lhs->scale)) {
- if (CanProveDivisible(lhs->scale, rhs) && is_zero(base)) {
+ if (!IsOne(lhs->scale)) {
+ if (CanProveDivisible(lhs->scale, rhs) && IsZero(base)) {
// floordiv(x*c1*c2, c2) = x*c1, c1=scale/rhs
lhs.CopyOnWrite()->scale = floordiv(lhs->scale, rhs);
return lhs;
@@ -1972,7 +1972,7 @@ ffi::Optional<PrimExpr>
IterMapRewriter::SplitFloorDivConst(IterSplitExpr lhs, P
// floordiv(x*c1*c2 + y*c2, c2) = x*c1 + y, c1=scale/rhs
lhs.CopyOnWrite()->scale = floordiv(lhs->scale, rhs);
return IterSumExpr({lhs}, floordiv(base, rhs));
- } else if (CanProveDivisible(rhs, lhs->scale) && is_zero(base)) {
+ } else if (CanProveDivisible(rhs, lhs->scale) && IsZero(base)) {
// floordiv(x*c1, c1*c2) = floordiv(x, c2), c2=rhs/scale
rhs = floordiv(rhs, lhs->scale);
lhs.CopyOnWrite()->scale = IntImm(rhs.ty(), 1);
@@ -2011,7 +2011,7 @@ ffi::Optional<PrimExpr>
IterMapRewriter::SplitFloorDivConst(IterSplitExpr lhs, P
/* lower_factor = */ padded->lower_factor * rhs,
/* extent = */
analyzer_->Simplify(floordiv(padded->extent, rhs)),
/* scale = */ padded->scale);
- } else if (is_one(padded->lower_factor) &&
+ } else if (IsOne(padded->lower_factor) &&
analyzer_->CanProveEqual(padded->extent, padded->source->extent))
{
// floordiv(floormod(floordiv(iter, lower_factor), ext), c)
// = floordiv(iter, c)
@@ -2028,7 +2028,7 @@ ffi::Optional<PrimExpr>
IterMapRewriter::SplitFloorDivConst(IterSplitExpr lhs, P
}
auto new_base = analyzer_->Simplify(floordiv(base - left_pad, rhs), 6);
- if (is_zero(new_base)) {
+ if (IsZero(new_base)) {
return new_split;
} else {
@@ -2084,16 +2084,16 @@ ffi::Optional<PrimExpr>
IterMapRewriter::SplitFloorModConst(IterSplitExpr lhs, P
PrimExpr rhs) {
// (lhs + base) % rhs
- if (is_one(rhs)) {
+ if (IsOne(rhs)) {
// floormod(x, 1) = 0
return IntImm(lhs.ty(), 0);
}
- if (!is_one(lhs->scale)) {
+ if (!IsOne(lhs->scale)) {
if (CanProveDivisible(lhs->scale, rhs) && CanProveDivisible(base, rhs)) {
// floormod(x*c1*c2, c1) = 0
return IntImm(lhs.ty(), 0);
- } else if (CanProveDivisible(rhs, lhs->scale) && is_zero(base)) {
+ } else if (CanProveDivisible(rhs, lhs->scale) && IsZero(base)) {
// floormod(x*c1, c1*c2) = (floormod(x, c2)) * c1, where c2 = rhs/scale
rhs = floordiv(rhs, lhs->scale);
} else if (CanProveDivisible(rhs, lhs->scale) && CanProveDivisible(base,
lhs->scale)) {
@@ -2120,7 +2120,7 @@ ffi::Optional<PrimExpr>
IterMapRewriter::SplitFloorModConst(IterSplitExpr lhs, P
// Right padding only contains values excluded by padding_predicate_, so it
// cannot make the inner floormod wrap over the original iterator domain.
// Keep left-padded marks conservative because padding shifts their values.
- if (is_zero(padding_it->second.left_pad)) {
+ if (IsZero(padding_it->second.left_pad)) {
source_upper_bound = origin_it->second->extent;
}
}
@@ -2240,14 +2240,14 @@ class IterMapToExprNormalizer : public tvm::ExprMutator
{
if (!expr->source.unique()) source_mode = InplaceMode::kDisallow;
source = Mutate(expr->source->source,
source_mode).ValueOrUnchanged(expr->source->source);
}
- if (analyzer_->CanProve(expr->extent == expr->source->extent) &&
is_one(expr->lower_factor)) {
+ if (analyzer_->CanProve(expr->extent == expr->source->extent) &&
IsOne(expr->lower_factor)) {
return source * expr->scale;
} else if (analyzer_->CanProve(expr->source->extent == expr->lower_factor
* expr->extent) ||
analyzer_->CanProve(expr->source->extent == expr->extent *
expr->lower_factor)) {
// Simplify if `expr` is always 0. The 2nd condition guarantess that we
do not aggressively
// simplify trivial iters like `vi \in [0, 1)`, which can be useful for
subsequent analysis
// like tensorization.
- if (is_one(expr->extent) && !is_one(expr->source->extent)) {
+ if (IsOne(expr->extent) && !IsOne(expr->source->extent)) {
return IntImm(expr->extent.ty(), 0);
}
return floordiv(source, expr->lower_factor) * expr->scale;
@@ -2299,7 +2299,7 @@ ffi::Array<PrimExpr> IterMapSimplify(const
ffi::Array<PrimExpr>& indices,
/*simplify_trivial_iterators=*/simplify_trivial_iterators);
ffi::Array<IterSumExpr> rewrite = res->indices;
- if (rewrite.empty() && !is_one(input_pred) && check_level !=
IterMapLevel::Bijective) {
+ if (rewrite.empty() && !IsOne(input_pred) && check_level !=
IterMapLevel::Bijective) {
// The input predicate may cause detect iter map to fail
// but we can still detect the iter map without the input predicate
// in which case an unpadded iter map is valid and can be used for
@@ -2309,7 +2309,7 @@ ffi::Array<PrimExpr> IterMapSimplify(const
ffi::Array<PrimExpr>& indices,
// A padded fallback is not equivalent over the original iterator domain
unless its
// padding predicate is also preserved. IterMapSimplify only returns
expressions, so it
// cannot carry that predicate to callers.
- if (!fallback->indices.empty() &&
is_zero(fallback->padding_predicate.value())) {
+ if (!fallback->indices.empty() &&
IsZero(fallback->padding_predicate.value())) {
rewrite = fallback->indices;
}
}
@@ -2435,12 +2435,12 @@ class SubspaceDivider {
IterSumExpr({}, expr->base), IntImm(dtype, 1));
} else if (expr->args.size() == 1) {
// arg + base, if arg=Y*E(X)+X, then arg+base = Y*E(X)+(X+base)
- if (!is_one(expr->args[0]->scale)) {
+ if (!IsOne(expr->args[0]->scale)) {
unresolved_count_++;
return DivisionResult::Failure();
}
DivisionResult res = DivideIterSplitExpr(expr->args[0]);
- if (!is_zero(expr->base)) res = AddBase(res, expr->base);
+ if (!IsZero(expr->base)) res = AddBase(res, expr->base);
return res;
}
// arg1 + arg2 + ... + argn + base
@@ -2452,7 +2452,7 @@ class SubspaceDivider {
// we check in inverse order so we can visit from inner to outer
for (auto it = expr->args.rbegin(); it != expr->args.rend(); ++it) {
const IterSplitExpr& arg = *it;
- if (is_one(arg->scale)) scale_is_one = true;
+ if (IsOne(arg->scale)) scale_is_one = true;
DivisionResult arg_division = DivideIterSplitExpr(arg);
IterSplitExpr new_arg{ffi::UnsafeInit{}};
if (arg_division.IsInner()) {
diff --git a/src/sym/modular_set.cc b/src/sym/modular_set.cc
index c2aafe1820..3df172c03f 100644
--- a/src/sym/modular_set.cc
+++ b/src/sym/modular_set.cc
@@ -293,7 +293,7 @@ class ModularSetAnalyzer::Impl : public
tvm::ExprFunctor<ModularSetAnalyzer::Ent
Entry b = Dispatch(op->b);
if (b.is_const() && b.base >= 0 && b.base <
std::numeric_limits<int64_t>::max()) {
int shift;
- if (is_const_power_of_two_integer(IntImm::Int64(b.base + 1), &shift)) {
+ if (IsPowerOfTwoInt(IntImm::Int64(b.base + 1), &shift)) {
return ModByConst(op->a, static_cast<int64_t>(1) << shift, true);
}
}
diff --git a/src/sym/presburger_set.cc b/src/sym/presburger_set.cc
index ce0d2a82b4..4ec7d884e1 100644
--- a/src/sym/presburger_set.cc
+++ b/src/sym/presburger_set.cc
@@ -151,7 +151,7 @@ PrimExpr PresburgerSetNode::GenerateConstraint() const {
#else
auto coeff = disjunct.atEq(i, j);
#endif
- if (coeff >= 0 || is_zero(linear_eq)) {
+ if (coeff >= 0 || IsZero(linear_eq)) {
linear_eq = linear_eq + IntImm::Int64(coeff) * vars[j];
} else {
linear_eq = linear_eq - IntImm::Int64(-coeff) * vars[j];
@@ -175,7 +175,7 @@ PrimExpr PresburgerSetNode::GenerateConstraint() const {
#else
auto coeff = disjunct.atIneq(i, j);
#endif
- if (coeff >= 0 || is_zero(linear_eq)) {
+ if (coeff >= 0 || IsZero(linear_eq)) {
linear_eq = linear_eq + IntImm::Int64(coeff) * vars[j];
} else {
linear_eq = linear_eq - IntImm::Int64(-coeff) * vars[j];
diff --git a/src/sym/rewrite_simplify.cc b/src/sym/rewrite_simplify.cc
index 4b77e46fa6..d16a216189 100644
--- a/src/sym/rewrite_simplify.cc
+++ b/src/sym/rewrite_simplify.cc
@@ -2448,7 +2448,7 @@ UnchangedOr<Expr> RewriteSimplifier::Impl::Mutate_(const
CallNode* op, InplaceMo
if (op == nullptr) return ret;
if (op->op.same_as(prim::builtin::likely()) &&
- prim::is_const_int(op->args[0].as_or_throw<PrimExpr>())) {
+ prim::IsConstInt(op->args[0].as_or_throw<PrimExpr>())) {
return op->args[0].as_or_throw<PrimExpr>();
}
static const Op& ceil_op = prim::builtin::ceil();
@@ -2508,7 +2508,7 @@ UnchangedOr<Expr> RewriteSimplifier::Impl::Mutate_(const
CallNode* op, InplaceMo
PrimExpr inner_then_expr = inner_call->args[1].as_or_throw<PrimExpr>();
PrimExpr inner_else_expr = inner_call->args[2].as_or_throw<PrimExpr>();
// Only check constant cases to avoid recursion
- if (prim::is_const_number(inner_else_expr) &&
prim::is_const_number(else_expr) &&
+ if (prim::IsConstNumber(inner_else_expr) &&
prim::IsConstNumber(else_expr) &&
analyzer_->CanProve(inner_else_expr == else_expr)) {
return Call(ret_ty, op->op, {cond && inner_cond, inner_then_expr,
else_expr}, op->attrs,
op->ty_args, op->span)
@@ -2554,7 +2554,7 @@ UnchangedOr<PrimExpr>
RewriteSimplifier::Impl::Mutate_(const prim::CastNode* op,
bool RewriteSimplifier::Impl::CanInlineLet(const prim::LetNode* op) {
// Only inline trivial bindings to avoid deep expression explosion
// when we need let to construct complicated expressions.
- if (prim::is_const_number(op->value)) return true;
+ if (prim::IsConstNumber(op->value)) return true;
if (op->value.as<PrimVar>()) return true;
return false;
}
diff --git a/src/sym/simplify_base.cc b/src/sym/simplify_base.cc
index 6ad018f92d..5892294386 100644
--- a/src/sym/simplify_base.cc
+++ b/src/sym/simplify_base.cc
@@ -100,8 +100,8 @@ UnchangedOr<Expr> SimplifierBase::Mutate_(const CallNode*
op, InplaceMode inplac
.ValueOrUnchanged((*args)[2])
.as_or_throw<Expr>();
}();
- if (prim::is_zero(cond)) return false_value;
- if (prim::is_one(cond)) return true_value;
+ if (prim::IsZero(cond)) return false_value;
+ if (prim::IsOne(cond)) return true_value;
if (cond.same_as(op->args[0]) && true_value.same_as(op->args[1]) &&
false_value.same_as(op->args[2])) {
// Reuse the original node identity; there is no replacement to process.
@@ -151,8 +151,8 @@ UnchangedOr<PrimExpr> SimplifierBase::Mutate_(const
prim::SelectNode* op,
})
.ValueOrUnchanged(op->false_value);
}();
- if (prim::is_zero(cond)) return false_value;
- if (prim::is_one(cond)) return true_value;
+ if (prim::IsZero(cond)) return false_value;
+ if (prim::IsOne(cond)) return true_value;
if (cond.same_as(op->condition) && true_value.same_as(op->true_value) &&
false_value.same_as(op->false_value)) {
// Reuse the original node identity; there is no replacement to process.
diff --git a/src/sym/transitive_comparison_analyzer.cc
b/src/sym/transitive_comparison_analyzer.cc
index d9bfd24e7c..570a8659a2 100644
--- a/src/sym/transitive_comparison_analyzer.cc
+++ b/src/sym/transitive_comparison_analyzer.cc
@@ -34,7 +34,7 @@ namespace tvm {
namespace sym {
using namespace tvm::prim;
-using prim::is_const_int;
+using prim::IsConstInt;
class TransitiveComparisonAnalyzer::Impl {
public:
@@ -597,7 +597,7 @@ void TransitiveComparisonAnalyzer::Impl::Bind(const Var&
var, const Range& range
prev_bindings_.Set(var, range);
PrimExpr prim_var = var.as_or_throw<PrimExpr>();
- if (is_const_int(range->extent, 1)) {
+ if (IsConstInt(range->extent, 1)) {
AddKnown(prim_var == range->min, &knowns_);
} else {
AddKnown(prim_var >= range->min, &knowns_);
diff --git a/src/sym/z3_prover.cc b/src/sym/z3_prover.cc
index 288de36201..6fffc537a2 100644
--- a/src/sym/z3_prover.cc
+++ b/src/sym/z3_prover.cc
@@ -426,7 +426,7 @@ class Z3Prover::Impl : tvm::ExprFunctor<z3::expr(const
Expr&)> {
// free var
//
// NOTE: min + extent builds a fresh AddNode that is not folded, so we
must
- // test is_const_int on min and extent individually and add the two
constants
+ // test IsConstInt on min and extent individually and add the two
constants
// in C++. Otherwise this fast path is never taken and we always emit
the more expensive
// symbolic constraint below.
if (auto min_imm = min.as<IntImm>(), extent_imm = extent.as<IntImm>();
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index 7ec79f6672..ce8bb6a329 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -1150,7 +1150,7 @@ void CodeGenCPU::Dispatch_(const ForNode* op) {
if (op->kind == ForKind::kSerial || op->kind == ForKind::kUnrolled) {
CodeGenLLVM::Dispatch_(op);
} else if (op->kind == ForKind::kParallel) {
- TVM_FFI_ICHECK(is_zero(op->min))
+ TVM_FFI_ICHECK(IsZero(op->min))
<< "Parallel launch require canonical loop with zero start index";
TVM_FFI_ICHECK(op->HasTrivialStep())
<< "Parallel launch require canonical loop with trivial loop step";
@@ -1169,7 +1169,7 @@ void CodeGenCPU::Dispatch_(const ForNode* op) {
TVM_FFI_ICHECK(!parallel_env_.in_parallel_loop)
<< "Nested parallel loop is not supported by threadpool, try fuse
them instead";
parallel_env_.in_parallel_loop = true;
- PrimExpr end = is_zero(op->min) ? op->extent :
analyzer_->Simplify(op->min + op->extent);
+ PrimExpr end = IsZero(op->min) ? op->extent :
analyzer_->Simplify(op->min + op->extent);
bool stride_pattern = false;
if (auto it = op->annotations.find("parallel_stride_pattern"); it !=
op->annotations.end()) {
ffi::Any annotation = (*it).second;
@@ -1177,7 +1177,7 @@ void CodeGenCPU::Dispatch_(const ForNode* op) {
stride_pattern = value.value();
} else if (auto value = annotation.try_cast<IntImm>();
value && value.value().ty() == PrimType::Bool()) {
- stride_pattern = is_one(value.value());
+ stride_pattern = IsOne(value.value());
} else {
TVM_FFI_THROW(ValueError) << "parallel_stride_pattern must be a
constant boolean";
}
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index dfd8d97dd7..c4e896ea63 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -1387,7 +1387,7 @@ llvm::Value* CodeGenLLVM::CreateIntrinsic(const CallNode*
op) {
PrimType(load->ty.as_or_throw<PrimType>()->dtype));
return buffer_ptr.addr;
} else if (op->op.same_as(tirx::builtin::reinterpret()) &&
args[0].as<PrimExpr>() &&
- is_zero(args[0].as<PrimExpr>().value())) {
+ IsZero(args[0].as<PrimExpr>().value())) {
llvm::Type* target = GetLLVMType(ret_type);
TVM_FFI_ICHECK(target->isPointerTy())
<< "A zero reinterpret shortcut requires pointer result type, but got
" << ret_type;
@@ -1761,7 +1761,7 @@ void CodeGenLLVM::BufferAccessHelper(
// If the buffer index is a contiguous ramp node, we only need to
// access the first element, then cast to the value type.
if (const prim::RampNode* ramp_index = last_index.as<prim::RampNode>()) {
- if (is_one(ramp_index->stride)) {
+ if (IsOne(ramp_index->stride)) {
last_index = ramp_index->base;
last_index_lanes =
GetLanesOrVScaleFactor(PrimType(last_index.ty()->dtype));
}
@@ -2104,7 +2104,7 @@ void CodeGenLLVM::Dispatch_(const ForNode* op) {
TVM_FFI_ICHECK(op->kind == ForKind::kSerial);
}
PrimExpr step = op->step.value_or(IntImm(op->extent.ty(), 1));
- PrimExpr end = is_zero(op->min) ? op->extent : analyzer_->Simplify(op->min +
op->extent);
+ PrimExpr end = IsZero(op->min) ? op->extent : analyzer_->Simplify(op->min +
op->extent);
llvm::Value* begin_value = MakeValue(op->min);
llvm::Value* end_value = MakeValue(end);
CreateSerialFor(begin_value, end_value, MakeValue(step), op->loop_var,
op->body);
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index be5bdd5ae0..11ec216935 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -1397,7 +1397,7 @@ void CodeGenC::Dispatch_(const AssertStmtNode* op) {
void CodeGenC::Dispatch_(const ForNode* op) {
std::string begin_str = PrintExpr(op->min);
- PrimExpr end = is_zero(op->min) ? op->extent :
sym::Analyzer()->Simplify(op->min + op->extent);
+ PrimExpr end = IsZero(op->min) ? op->extent :
sym::Analyzer()->Simplify(op->min + op->extent);
std::string end_str = PrintExpr(end);
std::string step_str = op->step.has_value() ? PrintExpr(*op->step) : "";
PrintIndent();
@@ -1480,7 +1480,7 @@ void CodeGenC::Dispatch_(const SeqStmtNode* op) {
}
void CodeGenC::Dispatch_(const EvaluateNode* op) {
- if (auto value = op->value.as<PrimExpr>(); value &&
is_const_int(value.value())) return;
+ if (auto value = op->value.as<PrimExpr>(); value &&
IsConstInt(value.value())) return;
const CallNode* call = op->value.as<CallNode>();
if (call) {
if (call->op.same_as(tirx::builtin::assume_aligned())) {
diff --git a/src/te/operation/scan_op.cc b/src/te/operation/scan_op.cc
index bd2496d794..84ea312f6b 100644
--- a/src/te/operation/scan_op.cc
+++ b/src/te/operation/scan_op.cc
@@ -57,7 +57,7 @@ ScanOp::ScanOp(std::string name, std::string tag,
TVM_FFI_ICHECK_EQ(init.size(), state_placeholder.size());
sym::Analyzer analyzer;
auto prove_equal = [&](PrimExpr lhs, PrimExpr rhs) {
- return is_zero(analyzer->Simplify(lhs - rhs));
+ return IsZero(analyzer->Simplify(lhs - rhs));
};
for (size_t i = 0; i < init.size(); ++i) {
diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc
index 2cce2f3f4f..374c1ff29f 100644
--- a/src/tirx/ir/buffer.cc
+++ b/src/tirx/ir/buffer.cc
@@ -71,7 +71,7 @@ ffi::ObjectRef RealizeBufferSubscript(
ffi::Optional<PrimExpr>>>()) {
all_points = false;
ffi::Optional<PrimExpr> step = descriptor.value().get<2>();
- TVM_FFI_CHECK(!step.has_value() || is_one(step.value()), ValueError)
+ TVM_FFI_CHECK(!step.has_value() || IsOne(step.value()), ValueError)
<< "Buffer slices with a non-unit step are not supported";
}
}
@@ -123,7 +123,7 @@ ffi::ObjectRef RealizeBufferRegionSubscript(Expr value,
SubscriptSlice slice, Sp
ffi::Optional<PrimExpr>>>()) {
all_points = false;
ffi::Optional<PrimExpr> step = descriptor.value().get<2>();
- TVM_FFI_CHECK(!step.has_value() || is_one(step.value()), ValueError)
+ TVM_FFI_CHECK(!step.has_value() || IsOne(step.value()), ValueError)
<< "TensorRegion slices with a non-unit step are not supported";
}
}
@@ -448,7 +448,7 @@ ffi::Array<PrimExpr>
TensorTypeNode::ElemOffset(ffi::Array<PrimExpr> input_indic
}
}
- if (elem_offset.defined() && !is_zero(elem_offset) && !inner) {
+ if (elem_offset.defined() && !IsZero(elem_offset) && !inner) {
output_index = output_index + elem_offset;
}
@@ -596,11 +596,11 @@ TensorVar TensorVar::MakeSlice(ffi::Array<PrimExpr>
begins, ffi::Array<PrimExpr>
// check if stride is needed.
for (size_t i = 0; i < extents.size(); ++i) {
if (!can_relax) {
- if (!is_zero(begins[i]) || !is_zero(ana->Simplify(extents[i] -
n->shape[i]))) {
+ if (!IsZero(begins[i]) || !IsZero(ana->Simplify(extents[i] -
n->shape[i]))) {
need_stride = true;
}
}
- if (!is_one(extents[i])) can_relax = false;
+ if (!IsOne(extents[i])) can_relax = false;
}
// make stride.
if (need_stride) {
diff --git a/src/tirx/ir/exec_scope.cc b/src/tirx/ir/exec_scope.cc
index 76cf2ad936..4205b1879a 100644
--- a/src/tirx/ir/exec_scope.cc
+++ b/src/tirx/ir/exec_scope.cc
@@ -315,7 +315,7 @@ static ffi::Optional<ScopeIdDef> Compliment(const
ScopeIdDef& lhs, const ScopeId
if (lhs->scope == ScopeBinding::kClusterCtaPair || rhs->scope ==
ScopeBinding::kClusterCtaPair) {
return std::nullopt;
}
- if (is_zero(rhs.fused_extent())) return std::nullopt;
+ if (IsZero(rhs.fused_extent())) return std::nullopt;
sym::Analyzer ana;
auto try_compliment = [&](PrimExpr lhs_ext, PrimExpr rhs_ext,
ScopeBinding scope) -> ffi::Optional<ScopeIdDef> {
diff --git a/src/tirx/ir/index_map.cc b/src/tirx/ir/index_map.cc
index 0b04541b0f..e4eb390466 100644
--- a/src/tirx/ir/index_map.cc
+++ b/src/tirx/ir/index_map.cc
@@ -126,7 +126,7 @@ std::pair<IndexMap, PrimExpr> IndexMapInverseImpl(const
IndexMap& self,
for (int i = 0, n = self->initial_indices.size(); i < n; ++i) {
PrimVar index = self->initial_indices[i];
PrimExpr expr{ffi::UnsafeInit{}};
- if (is_one(initial_ranges[i]->extent) && !inverse_exprs_map.count(index)) {
+ if (IsOne(initial_ranges[i]->extent) && !inverse_exprs_map.count(index)) {
expr = initial_ranges[i]->min;
} else {
expr = inverse_exprs_map.at(index);
@@ -299,7 +299,7 @@ ffi::Array<PrimExpr> IndexMapNode::MapShape(const
ffi::Array<PrimExpr>& shape,
ffi::Array<PrimExpr> output;
for (auto& range : mapped) {
- TVM_FFI_ICHECK(is_zero(range->min));
+ TVM_FFI_ICHECK(IsZero(range->min));
output.push_back(range->extent);
}
diff --git a/src/tirx/ir/ir_mutator_with_analyzer.cc
b/src/tirx/ir/ir_mutator_with_analyzer.cc
index 4b72595516..20312625d4 100644
--- a/src/tirx/ir/ir_mutator_with_analyzer.cc
+++ b/src/tirx/ir/ir_mutator_with_analyzer.cc
@@ -186,8 +186,8 @@ UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const
IfThenElseNode* op,
.ValueOrUnchanged(op->else_case.value());
});
}
- if (is_one(real_condition)) return then_case;
- if (is_zero(real_condition)) {
+ if (IsOne(real_condition)) return then_case;
+ if (IsZero(real_condition)) {
return else_case.value_or(Evaluate(0));
}
@@ -270,10 +270,10 @@ UnchangedOr<Expr> IRMutatorWithAnalyzer::Mutate_(const
CallNode* op, InplaceMode
});
});
}
- if (is_zero(cond)) {
+ if (IsZero(cond)) {
return false_value;
}
- if (is_one(cond)) {
+ if (IsOne(cond)) {
return true_value;
}
if (cond.same_as(op->args[0]) && true_value.same_as(op->args[1]) &&
@@ -323,10 +323,10 @@ UnchangedOr<PrimExpr>
IRMutatorWithAnalyzer::Mutate_(const prim::SelectNode* op,
false_value = Mutate(op->false_value,
inplace_mode).ValueOrUnchanged(op->false_value);
});
}
- if (is_zero(cond)) {
+ if (IsZero(cond)) {
return false_value;
}
- if (is_one(cond)) {
+ if (IsOne(cond)) {
return true_value;
}
// normal path
diff --git a/src/tirx/ir/layout/compose_layout.cc
b/src/tirx/ir/layout/compose_layout.cc
index 8a19e3fdc5..11d1eda8c3 100644
--- a/src/tirx/ir/layout/compose_layout.cc
+++ b/src/tirx/ir/layout/compose_layout.cc
@@ -47,7 +47,7 @@ PrimExpr ApplyFullSwizzle(const ComposeLayoutNode* layout,
const PrimExpr& m) {
}
void AddExpr(std::optional<PrimExpr>* sum, const PrimExpr& term, const
sym::Analyzer& analyzer) {
- if (is_zero(term)) return;
+ if (IsZero(term)) return;
if (sum->has_value()) {
*sum = analyzer->Simplify(sum->value() + term);
} else {
diff --git a/src/tirx/ir/layout/tile_canonicalize.cc
b/src/tirx/ir/layout/tile_canonicalize.cc
index ed6724cf21..3cdf4d44e6 100644
--- a/src/tirx/ir/layout/tile_canonicalize.cc
+++ b/src/tirx/ir/layout/tile_canonicalize.cc
@@ -33,7 +33,7 @@ TileLayout RemoveUnitIters(TileLayout layout) {
auto new_layout = layout.CopyOnWrite();
std::vector<Iter> new_shard;
std::copy_if(layout->shard.begin(), layout->shard.end(),
std::back_inserter(new_shard),
- [](const Iter& iter) { return !is_one(iter->extent); });
+ [](const Iter& iter) { return !IsOne(iter->extent); });
// if new_shard is empty, add a unit iter (using axis from original shard)
if (new_shard.empty() && !layout->shard.empty()) {
new_shard.push_back(Iter(1, 1, layout->shard[0]->axis));
@@ -46,7 +46,7 @@ TileLayout RemoveZeroOffsets(TileLayout layout) {
auto new_layout = layout.CopyOnWrite();
ffi::Map<Axis, PrimExpr> new_offset;
for (const auto& [axis, off] : layout->offset) {
- if (!is_zero(off)) {
+ if (!IsZero(off)) {
new_offset.Set(axis, off);
}
}
diff --git a/src/tirx/ir/layout/tile_core.cc b/src/tirx/ir/layout/tile_core.cc
index 9d22fdab8d..b3522ff119 100644
--- a/src/tirx/ir/layout/tile_core.cc
+++ b/src/tirx/ir/layout/tile_core.cc
@@ -347,7 +347,7 @@ ffi::Array<PrimExpr> TileLayoutNode::GetShardShape() const {
bool TileLayoutNode::IsTrivial() const {
if (shard.size() > 1) return false;
if (shard.size() == 1) {
- if (!shard[0]->axis->IsMemoryAxis() || !is_one(shard[0]->stride)) return
false;
+ if (!shard[0]->axis->IsMemoryAxis() || !IsOne(shard[0]->stride)) return
false;
}
return replica.size() == 0 && offset.size() == 0;
}
diff --git a/src/tirx/ir/layout/tile_direct_sum_ops.cc
b/src/tirx/ir/layout/tile_direct_sum_ops.cc
index a7804c03d2..dc2414de0c 100644
--- a/src/tirx/ir/layout/tile_direct_sum_ops.cc
+++ b/src/tirx/ir/layout/tile_direct_sum_ops.cc
@@ -73,7 +73,7 @@ Layout TileLayoutNode::DirectSum(const TileLayout& left_in,
const Array<PrimExpr
static bool IterEqualRelaxUnit(const Iter& a, const Iter& b, sym::AnalyzerObj*
analyzer) {
if (!(*analyzer).CanProveEqual(a->extent, b->extent)) return false;
- if (!is_one(a->extent)) {
+ if (!IsOne(a->extent)) {
if (!(*analyzer).CanProveEqual(a->stride, b->stride)) return false;
if (!a->axis.same_as(b->axis)) return false;
}
diff --git a/src/tirx/ir/layout/tile_tile_ops.cc
b/src/tirx/ir/layout/tile_tile_ops.cc
index 69c9f54afa..2822bf7780 100644
--- a/src/tirx/ir/layout/tile_tile_ops.cc
+++ b/src/tirx/ir/layout/tile_tile_ops.cc
@@ -59,7 +59,7 @@ std::pair<TileLayout, std::vector<int64_t>> Group(TileLayout
layout,
seps.push_back(new_shard.size());
}
extent_i = analyzer->Simplify(extent_i);
- if (!is_one(extent_i)) {
+ if (!IsOne(extent_i)) {
TVM_FFI_ICHECK(shape_idx < shape.size())
<< "layout " << layout << " can not be grouped by shape " << shape;
new_shard.push_back(Iter(extent_i, analyzer->Simplify(stride_i),
layout->shard[i]->axis));
@@ -288,7 +288,7 @@ std::optional<std::pair<TileLayout, std::vector<int64_t>>>
TryGroup(
seps.push_back(new_shard.size());
}
extent_i = analyzer->Simplify(extent_i);
- if (!is_one(extent_i)) {
+ if (!IsOne(extent_i)) {
if (shape_idx >= shape.size()) return std::nullopt;
new_shard.push_back(Iter(extent_i, analyzer->Simplify(stride_i),
layout->shard[i]->axis));
}
@@ -494,7 +494,7 @@ ffi::Optional<TileLayout> TileLayoutNode::IsTileInner(
auto inner_span_map = BuildSpanMap(layout);
auto rescale_by_inner_span = [&](const Iter& iter) -> ffi::Optional<Iter> {
auto it = inner_span_map.find(iter->axis.name());
- if (it != inner_span_map.end() && !is_one(iter->extent)) {
+ if (it != inner_span_map.end() && !IsOne(iter->extent)) {
if (!analyzer->CanProveEqual(floormod(iter->stride, (*it).second), 0)) {
return std::nullopt;
}
@@ -529,7 +529,7 @@ ffi::Optional<TileLayout> TileLayoutNode::IsTileInner(
Iter inner_iter = grouped_layout->shard[inner_seps[i] + j];
Iter tiled_iter = grouped_tiled->shard[tiled_seps_even[i + 1] -
inner_count + j];
if (!analyzer->CanProveEqual(inner_iter->extent, tiled_iter->extent) ||
- (!is_one(inner_iter->extent) &&
+ (!IsOne(inner_iter->extent) &&
!(analyzer->CanProveEqual(inner_iter->stride, tiled_iter->stride) &&
inner_iter->axis.same_as(tiled_iter->axis)))) {
return std::nullopt;
@@ -620,7 +620,7 @@ ffi::Optional<Layout> TileLayoutNode::IsTileOuter(const
Layout& tile_layout,
Iter outer_iter = grouped_layout->shard[outer_seps[i] + j];
Iter tiled_iter = grouped_tiled->shard[tiled_seps_even[i] + j];
if (!analyzer->CanProveEqual(outer_iter->extent, tiled_iter->extent) ||
- (!is_one(outer_iter->extent) &&
!outer_iter->axis.same_as(tiled_iter->axis))) {
+ (!IsOne(outer_iter->extent) &&
!outer_iter->axis.same_as(tiled_iter->axis))) {
return std::nullopt;
}
}
diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc
index 4ab20240b1..dfa9335f0d 100644
--- a/src/tirx/ir/stmt.cc
+++ b/src/tirx/ir/stmt.cc
@@ -924,7 +924,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-bool ForNode::HasTrivialStep() const { return !step.has_value() ||
is_one(*step); }
+bool ForNode::HasTrivialStep() const { return !step.has_value() ||
IsOne(*step); }
std::ostream& operator<<(std::ostream& out, ForKind type) { // NOLINT(*)
switch (type) {
diff --git a/src/tirx/ir/type.cc b/src/tirx/ir/type.cc
index 3af4f18d85..d95db688e1 100644
--- a/src/tirx/ir/type.cc
+++ b/src/tirx/ir/type.cc
@@ -36,8 +36,8 @@ namespace tvm::tirx {
bool TensorTypeNode::IsScalar(bool alloc_or_decl) const {
// TODO(@bohan): logical scope is not considered
- return shape.size() == 1 && tvm::prim::is_one(shape[0]) && strides.empty() &&
- (!alloc_or_decl || tvm::prim::is_zero(elem_offset)) && data_alignment
== 64 &&
+ return shape.size() == 1 && tvm::prim::IsOne(shape[0]) && strides.empty() &&
+ (!alloc_or_decl || tvm::prim::IsZero(elem_offset)) && data_alignment
== 64 &&
offset_factor == 1 && allocated_addr.empty() && layout.has_value() &&
ffi::StructuralEqual()(layout.value(),
TileLayoutNode::DefaultLayout({1}));
}
diff --git a/src/tirx/script/ir_builder/ir.cc b/src/tirx/script/ir_builder/ir.cc
index 8d3c8c73fa..23423066b4 100644
--- a/src/tirx/script/ir_builder/ir.cc
+++ b/src/tirx/script/ir_builder/ir.cc
@@ -334,7 +334,7 @@ ForFrame ThreadBinding(PrimExpr start, PrimExpr stop,
ffi::String thread,
Stmt body, Span span) ->
For {
TVM_FFI_ICHECK_EQ(vars.size(), 1);
TVM_FFI_ICHECK_EQ(doms.size(), 1);
- TVM_FFI_ICHECK(steps.size() == 1 && (!steps[0].has_value() ||
is_one(*steps[0])));
+ TVM_FFI_ICHECK(steps.size() == 1 && (!steps[0].has_value() ||
IsOne(*steps[0])));
IterVar iter_var(Range(nullptr), tvm::PrimVar("iter", dtype),
IterVarType::kThreadIndex,
thread);
return For(vars[0].as_or_throw<tvm::PrimVar>(), doms[0]->min,
doms[0]->extent,
diff --git a/src/tirx/script/printer/for_loop.cc
b/src/tirx/script/printer/for_loop.cc
index 0a8517bfd3..8ddc76d4bc 100644
--- a/src/tirx/script/printer/for_loop.cc
+++ b/src/tirx/script/printer/for_loop.cc
@@ -63,7 +63,7 @@ ffi::Optional<ExprDoc> ForDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
ExprDoc min = d->Translate(loop->min).value();
ExprDoc extent = d->Translate(loop->extent).value();
ExprDoc end = OperationDoc(OperationDocNode::Kind::kAdd, {min, extent});
- if (prim::is_zero(loop->min)) {
+ if (prim::IsZero(loop->min)) {
end = extent;
} else if (loop->min.as<IntImmNode>() && loop->extent.as<IntImmNode>()) {
// Python integer addition can widen the endpoint before the builder sees
@@ -120,7 +120,7 @@ ffi::Optional<ExprDoc> ForDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
{
IdDoc var = VarDoc(d, loop->loop_var);
ffi::Array<ExprDoc> bounds = {min, end};
- if (loop->kind == tirx::ForKind::kThreadBinding &&
prim::is_zero(loop->min) &&
+ if (loop->kind == tirx::ForKind::kThreadBinding && prim::IsZero(loop->min)
&&
loop->min.ty() == loop->extent.ty()) {
bounds = {end};
}
@@ -133,7 +133,7 @@ ffi::Optional<ExprDoc> ForDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
bounds.push_back(values.back());
keys.pop_back();
values.pop_back();
- } else if (prim::is_zero(loop->min) && loop->min.ty() ==
PrimType::Int(32)) {
+ } else if (prim::IsZero(loop->min) && loop->min.ty() ==
PrimType::Int(32)) {
bounds.erase(bounds.begin());
}
}
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index 08695fb90a..c0e8494e68 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/src/tirx/script/printer/stmt.cc
@@ -355,7 +355,7 @@ ffi::Optional<ExprDoc>
SeqStmtDocTranslate(DocTranslatorObj* d, ffi::AnyView inp
continue;
auto scalar = docs.back().as<AssignDoc>();
if (!IsScalarBuffer(d, alloc->var) || !scalar.has_value() ||
- !std::all_of(store->indices.begin(), store->indices.end(),
tvm::prim::is_zero))
+ !std::all_of(store->indices.begin(), store->indices.end(),
tvm::prim::IsZero))
continue;
bool reads_allocation = false;
ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(
diff --git a/src/tirx/transform/flatten_buffer.cc
b/src/tirx/transform/flatten_buffer.cc
index 068da6966f..51d7591ab0 100644
--- a/src/tirx/transform/flatten_buffer.cc
+++ b/src/tirx/transform/flatten_buffer.cc
@@ -158,7 +158,7 @@ class BufferFlattener : public IRMutatorWithAnalyzer {
type->shape.Set(i, analyzer_->canonical_simplify(type->shape[i]));
}
type->layout = std::nullopt;
- if (type->elem_offset.defined() && !is_zero(type->elem_offset)) {
+ if (type->elem_offset.defined() && !IsZero(type->elem_offset)) {
type->elem_offset =
IntImm(type->elem_offset.ty().as_or_throw<PrimType>(), 0);
}
// Body-local buffers keep their identity when flattening changes nothing.
diff --git a/src/tirx/transform/lower_intrin.cc
b/src/tirx/transform/lower_intrin.cc
index a0f08480fc..7644268324 100644
--- a/src/tirx/transform/lower_intrin.cc
+++ b/src/tirx/transform/lower_intrin.cc
@@ -234,7 +234,7 @@ class IntrinInjecter : public IRMutatorWithAnalyzer {
PrimType dtype = op->ty.as_or_throw<PrimType>();
TVM_FFI_ICHECK(dtype.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt));
- if (support_bitwise_op_ && is_const_power_of_two_integer(op->b, &shift)) {
+ if (support_bitwise_op_ && IsPowerOfTwoInt(op->b, &shift)) {
// lower to right shift if possible.
return op->a >> IntImm(dtype, shift);
}
@@ -299,7 +299,7 @@ class IntrinInjecter : public IRMutatorWithAnalyzer {
PrimType dtype = op->ty.as_or_throw<PrimType>();
TVM_FFI_ICHECK(dtype.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt));
- if (support_bitwise_op_ && is_const_power_of_two_integer(op->b, &shift)) {
+ if (support_bitwise_op_ && IsPowerOfTwoInt(op->b, &shift)) {
// lower to masking if possible.
ffi::BigInt mask = (ffi::BigInt(1) << shift) - 1;
return op->a & IntImm(dtype, mask);
diff --git a/src/tirx/transform/lower_thread_allreduce.h
b/src/tirx/transform/lower_thread_allreduce.h
index 9b8f5b7740..ca4af28e47 100644
--- a/src/tirx/transform/lower_thread_allreduce.h
+++ b/src/tirx/transform/lower_thread_allreduce.h
@@ -164,7 +164,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
GetAllocationKey(op->source.as_or_throw<tvm::tirx::TensorVar>().get());
if (auto it = load_remap_.find(allocation); it != load_remap_.end()) {
for (const auto& index : op->indices) {
- TVM_FFI_ICHECK(is_zero(index));
+ TVM_FFI_ICHECK(IsZero(index));
}
return it->second;
}
@@ -190,7 +190,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
const auto* replacement = it->second.template as<TensorLoadNode>();
TVM_FFI_ICHECK(replacement);
for (const auto& index : store->indices) {
- TVM_FFI_ICHECK(is_zero(index));
+ TVM_FFI_ICHECK(IsZero(index));
}
auto* writer = store.CopyOnWrite();
writer->buffer = replacement->source.template
as_or_throw<tvm::tirx::TensorVar>();
@@ -239,7 +239,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
PrimExpr cond = call->args[3].as_or_throw<PrimExpr>();
for (size_t idx = 0; idx < size; ++idx) {
values.push_back(inputs[idx].as_or_throw<PrimExpr>());
- if (!is_one(cond)) {
+ if (!IsOne(cond)) {
values[idx] = Select(cond, values[idx],
inits[idx].as_or_throw<PrimExpr>());
}
dtypes.push_back(values[idx].ty());
@@ -788,7 +788,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
}
// The local buffer index.
PrimExpr BufIndex(PrimExpr reduce_index, PrimExpr group_index, int
reduce_extent) {
- if (!is_zero(group_index)) {
+ if (!IsZero(group_index)) {
return analyzer_->Simplify(group_index * reduce_extent + reduce_index);
} else {
return reduce_index;
diff --git a/src/tirx/transform/lower_tirx_opaque.cc
b/src/tirx/transform/lower_tirx_opaque.cc
index 21aa83bb6d..5b352b699f 100644
--- a/src/tirx/transform/lower_tirx_opaque.cc
+++ b/src/tirx/transform/lower_tirx_opaque.cc
@@ -72,10 +72,10 @@ class TIRxOpaqueLower : public StmtExprMutator {
ffi::Any previous_remap = VarRemapGet(op->loop_var);
PrimVar launch_var(ffi::UnsafeInit{});
if (op->kind == ForKind::kThreadBinding) {
- TVM_FFI_ICHECK(is_zero(min)) << "Thread binding must have zero minimum";
+ TVM_FFI_ICHECK(IsZero(min)) << "Thread binding must have zero minimum";
launch_var = PrimVar(op->loop_var->name, extent.ty());
VarRemapSet(op->loop_var, prim::cast(op->loop_var.ty(), launch_var));
- } else if (is_one(extent) && op->annotations.empty()) {
+ } else if (IsOne(extent) && op->annotations.empty()) {
VarRemapSet(op->loop_var, prim::cast(op->loop_var.ty(), min));
}
@@ -108,7 +108,7 @@ class TIRxOpaqueLower : public StmtExprMutator {
<< "Run LoopPartition before opaque lowering of a thread-binding
loop with "
"loop_partition_hint";
TVM_FFI_ICHECK(op->thread_binding.has_value());
- } else if (is_one(extent) && op->annotations.empty()) {
+ } else if (IsOne(extent) && op->annotations.empty()) {
// Case 2. Unit loop elimination
return body;
} else {
diff --git a/src/tirx/transform/lower_tvm_builtin.cc
b/src/tirx/transform/lower_tvm_builtin.cc
index eb5e21ee6e..92ed136afc 100644
--- a/src/tirx/transform/lower_tvm_builtin.cc
+++ b/src/tirx/transform/lower_tvm_builtin.cc
@@ -517,7 +517,7 @@ class BuiltinLower : public StmtExprMutator {
prep_seq.emplace_back(
TVMStructSet(scope.stack_array, idx, builtin::kDLTensorShape,
op->args[1]));
Expr strides = op->args[2];
- if (auto prim_strides = strides.as<PrimExpr>(); prim_strides &&
is_zero(prim_strides.value())) {
+ if (auto prim_strides = strides.as<PrimExpr>(); prim_strides &&
IsZero(prim_strides.value())) {
strides = ConstHandle(0);
}
prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx,
builtin::kDLTensorStrides, strides));
@@ -534,7 +534,7 @@ class BuiltinLower : public StmtExprMutator {
int data_bytes = GetVectorBytes(dtype);
PrimExpr elem_offset = op->args[5].as_or_throw<PrimExpr>();
PrimExpr byte_offset = elem_offset;
- if (!is_zero(elem_offset)) {
+ if (!IsZero(elem_offset)) {
byte_offset = elem_offset * IntImm(elem_offset.ty(), data_bytes);
} else {
byte_offset = elem_offset;
diff --git a/src/tirx/transform/remove_no_op.cc
b/src/tirx/transform/remove_no_op.cc
index 8235e58c26..902c41292a 100644
--- a/src/tirx/transform/remove_no_op.cc
+++ b/src/tirx/transform/remove_no_op.cc
@@ -129,7 +129,7 @@ class NoOpRemover : public IRMutatorWithAnalyzer {
Stmt stmt = Parent::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
var_range_map_.erase(op->loop_var.get());
op = stmt.as<ForNode>();
- if (is_zero(op->extent)) {
+ if (IsZero(op->extent)) {
return Evaluate(0);
}
return is_no_op(op->body) ? MakeEvaluate({op->min, op->extent}) : stmt;
@@ -172,7 +172,7 @@ class NoOpRemover : public IRMutatorWithAnalyzer {
PrimExpr stores_existing_value =
store->value - MakeTensorLoad(store->buffer, store->indices) == 0;
stores_existing_value = analyzer_->Simplify(stores_existing_value);
- if (is_one(stores_existing_value)) {
+ if (IsOne(stores_existing_value)) {
return only_side_effects();
}
diff --git a/src/tirx/transform/storage_rewrite.cc
b/src/tirx/transform/storage_rewrite.cc
index 24ac6e62e6..42cf133f4f 100644
--- a/src/tirx/transform/storage_rewrite.cc
+++ b/src/tirx/transform/storage_rewrite.cc
@@ -1655,7 +1655,7 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
// vectorized.
if (indices.size()) {
const prim::RampNode* ramp_index = indices[indices.size() -
1].as<prim::RampNode>();
- if (ramp_index && is_one(ramp_index->stride)) {
+ if (ramp_index && IsOne(ramp_index->stride)) {
if (ramp_index->lanes->IsInstance<IntImmNode>()) {
int lanes =
ramp_index->lanes.as_or_throw<IntImm>()->value.as<int>().value();
sym::ModularSet me = analyzer_->modular_set(ramp_index->base);
@@ -1837,7 +1837,7 @@ class VectorTypeRewriter : public StmtExprMutator {
return {node, shuffle_index};
}
- if (ramp_index && is_one(ramp_index->stride) &&
ramp_index->lanes->IsInstance<IntImmNode>()) {
+ if (ramp_index && IsOne(ramp_index->stride) &&
ramp_index->lanes->IsInstance<IntImmNode>()) {
int lanes =
ramp_index->lanes.as_or_throw<IntImm>()->value.as<int>().value();
PrimExpr new_index = ramp_index->base / MakeConst(ramp_index->base.ty(),
lanes);
if (lanes != info.factor()) {
@@ -1882,7 +1882,7 @@ class VectorTypeRewriter : public StmtExprMutator {
return {node, shuffle_index};
}
- if (ramp_index && is_one(ramp_index->stride) &&
ramp_index->lanes->IsInstance<IntImmNode>()) {
+ if (ramp_index && IsOne(ramp_index->stride) &&
ramp_index->lanes->IsInstance<IntImmNode>()) {
int lanes =
ramp_index->lanes.as_or_throw<IntImm>()->value.as<int>().value();
PrimExpr new_index = ramp_index->base / MakeConst(ramp_index->base.ty(),
lanes);
if (lanes != info.factor()) {
diff --git a/src/tirx/transform/tvm_ffi_binder.cc
b/src/tirx/transform/tvm_ffi_binder.cc
index 08ca86da02..1a49058efb 100644
--- a/src/tirx/transform/tvm_ffi_binder.cc
+++ b/src/tirx/transform/tvm_ffi_binder.cc
@@ -190,11 +190,11 @@ bool TVMFFIABIBuilder::BindScalar(const PrimExpr& arg,
const PrimExpr& value,
// Duplicate bind: create rich assertion with both paths
PrimExpr prev_value = it->second.value.as_or_throw<PrimExpr>();
PrimExpr scond = analyzer_->Simplify(prev_value == value);
- if (is_zero(scond)) {
+ if (IsZero(scond)) {
TVM_FFI_THROW(InternalError) << "Bind have an unmet assertion: " <<
prev_value
<< " == " << value << " at " <<
RenderAccessPath(path);
}
- if (!is_one(scond)) {
+ if (!IsOne(scond)) {
ffi::String current_path_str = RenderAccessPath(path);
ffi::String first_path_str =
RenderAccessPath(it->second.first_def_path);
int param_index = GetParamIndex(path);
@@ -220,11 +220,11 @@ bool TVMFFIABIBuilder::BindScalar(const PrimExpr& arg,
const PrimExpr& value,
// Non-Var expression (e.g. batch_size + 1): defer assertion to Finalize()
// so display-var substitution can render human-readable names.
PrimExpr scond = analyzer_->Simplify(arg == value);
- if (is_zero(scond)) {
+ if (IsZero(scond)) {
TVM_FFI_THROW(InternalError) << "Bind have an unmet assertion: " << arg
<< " == " << value
<< " at " << RenderAccessPath(path);
}
- if (!is_one(scond)) {
+ if (!IsOne(scond)) {
pending_const_asserts_.push_back({scond, path, arg});
}
}
@@ -249,11 +249,11 @@ bool TVMFFIABIBuilder::BindPointer(const Var& arg, const
Expr& value,
PrimExpr prev_value = pointer_as_uint(it->second.value);
PrimExpr current_value = pointer_as_uint(value);
PrimExpr condition = analyzer_->Simplify(prev_value == current_value);
- if (is_zero(condition)) {
+ if (IsZero(condition)) {
TVM_FFI_THROW(InternalError) << "Bind has an unmet pointer assertion at "
<< RenderAccessPath(path);
}
- if (!is_one(condition)) {
+ if (!IsOne(condition)) {
ffi::String current_path = RenderAccessPath(path);
ffi::String first_path = RenderAccessPath(it->second.first_def_path);
ffi::Array<StringImm> parts{StringImm("Mismatched "),
StringImm(current_path)};
@@ -409,8 +409,8 @@ void TVMFFIABIBuilder::BindBuffer(const TensorVar& arg,
const TensorVar& value,
}
if (value->elem_offset.defined()) {
- if (is_zero(arg->elem_offset)) {
- TVM_FFI_ICHECK(is_zero(value->elem_offset))
+ if (IsZero(arg->elem_offset)) {
+ TVM_FFI_ICHECK(IsZero(value->elem_offset))
<< "Trying to bind a TensorVar with offset into one without offset "
<< " required elem_offset=" << arg->elem_offset
<< ", provided elem_offset=" << value->elem_offset;
@@ -423,11 +423,11 @@ void TVMFFIABIBuilder::BindBuffer(const TensorVar& arg,
const TensorVar& value,
PrimExpr factor = IntImm(offset.ty(), arg->offset_factor);
PrimExpr zero = IntImm(offset.ty(), 0);
PrimExpr acond = analyzer_->Simplify(truncmod(offset, factor) == zero);
- if (is_zero(acond)) {
+ if (IsZero(acond)) {
TVM_FFI_THROW(InternalError)
<< "Bind have an unmet assertion at " <<
RenderAccessPath(offset_path);
}
- if (!is_one(acond)) {
+ if (!IsOne(acond)) {
int param_index = GetParamIndex(base_path);
int data_bytes =
((((arg->dtype->dtype).bits *
static_cast<int16_t>((arg->dtype->dtype).lanes)) + 7) /
@@ -448,7 +448,7 @@ void TVMFFIABIBuilder::BindBuffer(const TensorVar& arg,
const TensorVar& value,
TVM_FFI_ICHECK(fuzzy_match) << "TensorVar size mismatch at " <<
RenderAccessPath(base_path);
size_t diff = value->shape.size() - arg->shape.size();
for (size_t i = 0; i < diff; ++i) {
- TVM_FFI_ICHECK(is_one(analyzer_->Simplify(value->shape[i])))
+ TVM_FFI_ICHECK(IsOne(analyzer_->Simplify(value->shape[i])))
<< "TensorVar shape mismatch at " << RenderAccessPath(base_path) <<
": " << arg->shape
<< " vs " << value->shape;
}
@@ -633,8 +633,8 @@ void TVMFFIABIBuilder::DecodeAllParams() {
{buffer.value(),
IntImm::Int32(buffer.value()->data_alignment)}));
// Empty tensors may have arbitrary data pointers. Their checks must not
// be strengthened into an unconditional alignment promise.
- if (is_const_int(size)) {
- if (!is_zero(size)) decl_buffers_.push_back(alignment);
+ if (IsConstInt(size)) {
+ if (!IsZero(size)) decl_buffers_.push_back(alignment);
} else {
decl_buffers_.push_back(IfThenElse(size != 0, alignment));
}
@@ -798,11 +798,11 @@ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const
TensorVar& buffer, const PrimEx
PrimExpr factor = IntImm(offset.ty(), buffer->offset_factor);
PrimExpr zero = IntImm(offset.ty(), 0);
PrimExpr acond = analyzer_->Simplify(truncmod(offset, factor) == zero);
- if (is_zero(acond)) {
+ if (IsZero(acond)) {
TVM_FFI_THROW(InternalError)
<< "Bind have an unmet assertion at " <<
RenderAccessPath(byte_offset_path);
}
- if (!is_one(acond)) {
+ if (!IsOne(acond)) {
EmitAssert(acond, "ValueError", //
"Misaligned Tensor data on argument #",
std::to_string(param_index),
when_calling_imm_, sig_imm_, "`,\n expected data
alignment=",
@@ -819,7 +819,7 @@ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const TensorVar&
buffer, const PrimEx
// Use custom assertion for device_type to show human-readable device name
if (const auto* const_dt = device_type_.as<IntImmNode>()) {
PrimExpr cond = analyzer_->Simplify(IntImm::Int32(const_dt->value) ==
actual_device_type);
- if (!is_one(cond)) {
+ if (!IsOne(cond)) {
std::string device_name =
runtime::DLDeviceType2Str(const_dt->value.as<int>().value());
EmitAssert(cond, "ValueError", //
"Mismatched ", buf_name, ".device_type on argument #",
diff --git a/src/tirx/transform/vectorize_loop.cc
b/src/tirx/transform/vectorize_loop.cc
index 39f2673552..62c57bed4d 100644
--- a/src/tirx/transform/vectorize_loop.cc
+++ b/src/tirx/transform/vectorize_loop.cc
@@ -1063,7 +1063,7 @@ class Vectorizer : public StmtExprMutator {
if (op->kind == ForKind::kVectorized) {
LOG(WARNING) << "Detect vectorize inside vectorized loop, ignoring...";
}
- TVM_FFI_ICHECK(is_zero(op->min));
+ TVM_FFI_ICHECK(IsZero(op->min));
TVM_FFI_ICHECK(!op->extent.ty().IsScalableVector() &&
!op->extent.ty().IsFixedLengthVector());
auto extent_update = this->Mutate(op->extent, inplace_mode);
bool extent_unchanged = extent_update.UnchangedOrSameAs(op->extent);
@@ -1315,7 +1315,7 @@ class LoopVectorizer : public StmtExprMutator {
if (op->kind == ForKind::kVectorized) {
auto* extent_as_int = op->extent.as<IntImmNode>();
- TVM_FFI_ICHECK(is_zero(op->min));
+ TVM_FFI_ICHECK(IsZero(op->min));
// General calls still have vectorization paths that query a compile-time
// lane count, so keep them on the existing fixed-width path for now.
if (extent_as_int && extent_as_int->value > 1 && TargetHasRVV(target_) &&
diff --git a/src/topi/einsum.cc b/src/topi/einsum.cc
index c912e7a2ee..308bce6a0b 100644
--- a/src/topi/einsum.cc
+++ b/src/topi/einsum.cc
@@ -127,7 +127,7 @@ PrimExpr GetIndexForBroadcastedDim(const PrimVar& index,
const PrimExpr& extent,
const PrimExpr& broadcasted_extent) {
// Check if current dimension is being broadcasted to `broadcasted_extent`
(symbolic shape is
// handled)
- if (is_one(extent) && !is_one(broadcasted_extent)) {
+ if (IsOne(extent) && !IsOne(broadcasted_extent)) {
return IntImm(index.ty(), 0);
}
return index;
diff --git a/tests/cpp/pattern_match_test.cc b/tests/cpp/pattern_match_test.cc
index c1b45328a4..15582f10f4 100644
--- a/tests/cpp/pattern_match_test.cc
+++ b/tests/cpp/pattern_match_test.cc
@@ -72,7 +72,7 @@ TEST(Pattern, Basic) {
// bit intrinsics
{
TVM_FFI_ICHECK((px >> pz).Match(x >> 1));
- TVM_FFI_ICHECK(prim::is_const_int(pz.Eval(), 1));
+ TVM_FFI_ICHECK(prim::IsConstInt(pz.Eval(), 1));
}
TVM_FFI_ICHECK(!(px >> pz).Match(x << 1));
TVM_FFI_ICHECK((px << pz).Match(x << 1));
@@ -83,7 +83,7 @@ TEST(Pattern, Basic) {
// select
{
TVM_FFI_ICHECK(select(px > pz, py, py + pz).Match(prim::Select(x > 1, y, y
+ 1)));
- TVM_FFI_ICHECK(prim::is_const_int(pz.Eval(), 1));
+ TVM_FFI_ICHECK(prim::IsConstInt(pz.Eval(), 1));
}
TVM_FFI_ICHECK(!select(px > pz, py, py + pz).Match(prim::Select(x > 2, y, y
+ 1)));
TVM_FFI_ICHECK(!select(px > pz, py, py).Match(prim::Select(x > 2, y, y +
1)));
@@ -94,7 +94,7 @@ TEST(Pattern, Basic) {
// if_then_else
{
TVM_FFI_ICHECK(if_then_else(px > pz, py, py + pz).Match(if_then_else(x >
1, y, y + 1)));
- TVM_FFI_ICHECK(prim::is_const_int(pz.Eval(), 1));
+ TVM_FFI_ICHECK(prim::IsConstInt(pz.Eval(), 1));
}
// cast pattern
{
diff --git a/tests/cpp/sym_simplify_test.cc b/tests/cpp/sym_simplify_test.cc
index fce5703bc9..0e5ea4717c 100644
--- a/tests/cpp/sym_simplify_test.cc
+++ b/tests/cpp/sym_simplify_test.cc
@@ -29,11 +29,11 @@ TEST(Simplify, MinMax) {
auto x = tvm::te::var("x");
auto e1 = (tvm::max(x, 1) - tvm::max(x, 1));
auto e1s = ana->canonical_simplify(e1);
- TVM_FFI_ICHECK(tvm::prim::is_zero(e1s));
+ TVM_FFI_ICHECK(tvm::prim::IsZero(e1s));
auto e2 = (x * tvm::min(x, 1)) - (x * tvm::min(x, 1));
auto e2s = ana->canonical_simplify(e2);
- TVM_FFI_ICHECK(tvm::prim::is_zero(e2s));
+ TVM_FFI_ICHECK(tvm::prim::IsZero(e2s));
}
TEST(Simplify, Mul) {
@@ -41,7 +41,7 @@ TEST(Simplify, Mul) {
auto x = tvm::te::var("x");
auto e = (x * x) - (x * x);
auto es = ana->canonical_simplify(e);
- TVM_FFI_ICHECK(tvm::prim::is_zero(es));
+ TVM_FFI_ICHECK(tvm::prim::IsZero(es));
}
TEST(Simplify, Mod) {
@@ -53,7 +53,7 @@ TEST(Simplify, Mod) {
// and therefore, the constant folding will be attempted in CanonicalSimplify
auto mod = ana->canonical_simplify(tvm::prim::Mod(x, y));
auto es = ana->canonical_simplify(mod - x);
- TVM_FFI_ICHECK(tvm::prim::is_zero(es));
+ TVM_FFI_ICHECK(tvm::prim::IsZero(es));
}
TEST(AnalyzerObjectRef, CopySharesMutableState) {
@@ -107,12 +107,12 @@ TEST(Simplify, AssumeConstraintKeepsTensorLoadStable) {
{
auto exit_scope = analyzer->rewrite_simplify.EnterConstraint(constraint);
- EXPECT_FALSE(tvm::prim::is_one(analyzer->rewrite_simplify(constraint)));
+ EXPECT_FALSE(tvm::prim::IsOne(analyzer->rewrite_simplify(constraint)));
exit_scope();
}
{
auto exit_scope = analyzer->rewrite_simplify.EnterConstraint(constraint,
true);
- EXPECT_TRUE(tvm::prim::is_one(analyzer->rewrite_simplify(constraint)));
+ EXPECT_TRUE(tvm::prim::IsOne(analyzer->rewrite_simplify(constraint)));
exit_scope();
}
