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 e6f512c2ae [CLEANUP][IR][Arith] Move CLZ to shared prim and remove
dialect dependencies (#20367)
e6f512c2ae is described below
commit e6f512c2ae0d4f352546147a201f70ec38f16515
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Sep 16 21:53:32 2026 -0400
[CLEANUP][IR][Arith] Move CLZ to shared prim and remove dialect
dependencies (#20367)
CLZ is a shared integer primitive used by arithmetic simplification.
Move its operator registration and construction helpers into prim, and
use that identity consistently in lowering, dtype rewriting, Python
construction and TOPI.
Remove residual dialect includes and namespace imports from arith,
supply its direct shared and standard-library dependencies, and keep
IterAffineMap constraint-size ordering local. Remove the extra
function-attribute test additions while retaining generic
attribute-update behavior.
---
include/tvm/arith/int_set.h | 1 -
include/tvm/ir/prim/builtin.h | 3 +
include/tvm/ir/prim/op.h | 2 +
include/tvm/s_tir/schedule/instruction.h | 1 -
include/tvm/tirx/op.h | 1 -
python/tvm/ir/prim/__init__.py | 2 +-
python/tvm/ir/prim/op.py | 19 ++-
python/tvm/tirx/op.py | 18 +--
python/tvm/topi/math.py | 2 +-
src/arith/bound_deducer.cc | 2 -
src/arith/canonical_simplify.cc | 2 -
src/arith/const_int_bound.cc | 3 -
src/arith/detect_linear_equation.cc | 7 +-
src/arith/int_set.cc | 2 -
src/arith/iter_affine_map.cc | 69 +++++++++-
src/arith/modular_set.cc | 4 -
src/arith/pattern_match.h | 1 -
src/arith/presburger_set.cc | 1 -
src/arith/rewrite_simplify.cc | 7 +-
src/arith/rewrite_simplify.h | 2 -
src/arith/transitive_comparison_analyzer.cc | 4 +-
src/arith/z3_prover.cc | 8 +-
src/backend/cuda/codegen/intrin_rule_cuda.cc | 2 +-
src/backend/metal/codegen/intrin_rule_metal.cc | 2 +-
src/backend/opencl/codegen/intrin_rule_opencl.cc | 2 +-
src/backend/vulkan/codegen/intrin_rule_spirv.cc | 2 +-
src/ir/prim/builtin.cc | 3 +
src/ir/prim/op.cc | 14 ++
src/target/intrin_rule.h | 3 +-
src/target/llvm/intrin_rule_llvm.cc | 2 +-
src/tirx/ir/data_type_rewriter.cc | 2 +-
src/tirx/op/builtin.cc | 1 +
src/tirx/op/op.cc | 2 -
tests/cpp/function_attrs_test.cc | 154 ----------------------
tests/cpp/pattern_match_test.cc | 1 -
tests/python/arith/test_arith_rewrite_simplify.py | 16 +--
tests/python/ir/test_ir_attrs.py | 37 ------
37 files changed, 131 insertions(+), 273 deletions(-)
diff --git a/include/tvm/arith/int_set.h b/include/tvm/arith/int_set.h
index b2afdc8e9b..cdb5c9af9e 100644
--- a/include/tvm/arith/int_set.h
+++ b/include/tvm/arith/int_set.h
@@ -26,7 +26,6 @@
#include <tvm/ir/expr.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/var.h>
#include <unordered_map>
diff --git a/include/tvm/ir/prim/builtin.h b/include/tvm/ir/prim/builtin.h
index 3c475cda71..87dfece351 100644
--- a/include/tvm/ir/prim/builtin.h
+++ b/include/tvm/ir/prim/builtin.h
@@ -37,6 +37,9 @@ TVM_DLL const Op& ceil();
/*! \brief Base-two logarithm. */
TVM_DLL const Op& log2();
+/*! \brief Count leading zero bits. */
+TVM_DLL const Op& clz();
+
/*! \brief Left shift. */
TVM_DLL const Op& shift_left();
diff --git a/include/tvm/ir/prim/op.h b/include/tvm/ir/prim/op.h
index 7d12045496..a8391753fe 100644
--- a/include/tvm/ir/prim/op.h
+++ b/include/tvm/ir/prim/op.h
@@ -433,6 +433,8 @@ TVM_DLL PrimExpr ceil(PrimExpr x, Span span = Span());
/*! \brief Construct a base-two logarithm with the input's primitive type. */
TVM_DLL PrimExpr log2(PrimExpr x, Span span = Span());
namespace prim {
+/*! \brief Count leading zero bits, preserving the input primitive type. */
+TVM_DLL PrimExpr clz(PrimExpr x, Span span = Span());
/*!
* Query the minimum possible value of dtype.
* \param dtype The primitive type.
diff --git a/include/tvm/s_tir/schedule/instruction.h
b/include/tvm/s_tir/schedule/instruction.h
index 56fdb0ef0c..10a14e1468 100644
--- a/include/tvm/s_tir/schedule/instruction.h
+++ b/include/tvm/s_tir/schedule/instruction.h
@@ -31,7 +31,6 @@ template <typename, typename>
class AttrRegistry;
namespace s_tir {
-using namespace tvm::tirx;
// Forward declaration
class Schedule;
diff --git a/include/tvm/tirx/op.h b/include/tvm/tirx/op.h
index 874719a4e3..c180e89f23 100644
--- a/include/tvm/tirx/op.h
+++ b/include/tvm/tirx/op.h
@@ -339,7 +339,6 @@ TVM_DECLARE_FLOAT_INTRIN_UNARY(atan);
TVM_DECLARE_FLOAT_INTRIN_UNARY(acosh);
TVM_DECLARE_FLOAT_INTRIN_UNARY(asinh);
TVM_DECLARE_FLOAT_INTRIN_UNARY(atanh);
-TVM_DECLARE_INTRIN_UNARY(clz);
#define TVM_DECLARE_INTRIN_BINARY(OpName) \
inline PrimExpr OpName(PrimExpr x, PrimExpr y, Span span = Span()) { \
diff --git a/python/tvm/ir/prim/__init__.py b/python/tvm/ir/prim/__init__.py
index 9a8f388efd..67bd440a0b 100644
--- a/python/tvm/ir/prim/__init__.py
+++ b/python/tvm/ir/prim/__init__.py
@@ -53,7 +53,7 @@ from .expr import (
StringImm,
Sub,
)
-from .op import convert, max_value, min_value
+from .op import clz, convert, max_value, min_value
def expr_deep_equal(lhs: Expr, rhs: Expr) -> bool:
diff --git a/python/tvm/ir/prim/op.py b/python/tvm/ir/prim/op.py
index c56ca21ace..0562e888a4 100644
--- a/python/tvm/ir/prim/op.py
+++ b/python/tvm/ir/prim/op.py
@@ -19,7 +19,7 @@
from typing import Any
from ..base import Span
-from ..expr import Expr
+from ..expr import Call, Expr
from . import _ffi_api
@@ -64,3 +64,20 @@ def max_value(dtype: str, span: Span | None = None) -> Any:
The maximum value of dtype.
"""
return _ffi_api.max_value(dtype, span) # type: ignore
+
+
+def clz(x):
+ """Count leading zero bits of an integer x.
+
+ Parameters
+ ----------
+ x : Expr
+ Input 32 or 64 bit integer.
+ The result is undefined if the input is 0.
+
+ Returns
+ -------
+ y : Expr
+ The result.
+ """
+ return Call("prim.clz", [x], ret_ty="int32")
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index 81a97ee891..cfb875befc 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -27,6 +27,7 @@ import tvm.ir.prim._ffi_api as _prim_ffi_api
from tvm import tirx
from tvm.ir import Call, Expr, ExprWithOp, Op, PointerType, PrimType,
TensorLoad
from tvm.ir.base import Span
+from tvm.ir.prim import clz as clz
from tvm.ir.prim import max_value, min_value
from tvm.ir.type import TensorMapType
from tvm.runtime import const
@@ -1805,23 +1806,6 @@ def rsqrt(x):
return call_intrin(_primexpr_ty(x), "tirx.rsqrt", x)
-def clz(x):
- """Count leading zero bits of an integer x.
-
- Parameters
- ----------
- x : Expr
- Input 32 or 64 bit integer.
- The result is undefined if the input is 0.
-
- Returns
- -------
- y : Expr
- The result.
- """
- return call_intrin("int32", "tirx.clz", x)
-
-
def floor(x: ExprWithOp, span=None):
"""Take floor of float input x.
diff --git a/python/tvm/topi/math.py b/python/tvm/topi/math.py
index 6fd03db7b2..8f680f6a4f 100644
--- a/python/tvm/topi/math.py
+++ b/python/tvm/topi/math.py
@@ -866,7 +866,7 @@ def ceil_log2(x):
if target is not None:
target_name = target.kind.name
if "vulkan" in target_name:
- clz = tvm.tirx.clz(x)
+ clz = tvm.ir.prim.clz(x)
bits = x.ty.dtype.bits
res = tvm.tirx.if_then_else(x & (x - 1) == 0, bits - clz - 1, bits
- clz)
if res.ty != x.ty:
diff --git a/src/arith/bound_deducer.cc b/src/arith/bound_deducer.cc
index 55ddbfe73b..59295ca350 100644
--- a/src/arith/bound_deducer.cc
+++ b/src/arith/bound_deducer.cc
@@ -37,8 +37,6 @@
namespace tvm {
namespace arith {
-using namespace tirx;
-
// Find a target path through structural expression fields, including dynamic
types.
// BoundDeduceInputChecker counts occurrences over the same broader domain and
can
// conservatively decline deduction when a target also appears in type
metadata.
diff --git a/src/arith/canonical_simplify.cc b/src/arith/canonical_simplify.cc
index 93fd0b2c7d..f339d6621e 100644
--- a/src/arith/canonical_simplify.cc
+++ b/src/arith/canonical_simplify.cc
@@ -36,8 +36,6 @@
namespace tvm {
namespace arith {
-using namespace tirx;
-
class SumExpr;
class SplitExpr;
diff --git a/src/arith/const_int_bound.cc b/src/arith/const_int_bound.cc
index 5e1cc86663..0531caa559 100644
--- a/src/arith/const_int_bound.cc
+++ b/src/arith/const_int_bound.cc
@@ -28,7 +28,6 @@
#include <tvm/ir/op.h>
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/builtin.h>
#include <algorithm>
#include <optional>
@@ -40,8 +39,6 @@
namespace tvm {
namespace arith {
-using namespace tirx;
-
TVM_FFI_STATIC_INIT_BLOCK() { ConstIntBoundNode::RegisterReflection(); }
ConstIntBound::ConstIntBound(int64_t min_value, int64_t max_value) {
diff --git a/src/arith/detect_linear_equation.cc
b/src/arith/detect_linear_equation.cc
index daf58918ce..3b890f4b11 100644
--- a/src/arith/detect_linear_equation.cc
+++ b/src/arith/detect_linear_equation.cc
@@ -27,15 +27,14 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr_functor.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/op.h>
+#include <tvm/ir/prim/op.h>
+
+#include <unordered_set>
namespace tvm {
namespace arith {
using namespace tvm::prim;
-using namespace tirx;
-
// Linear equation, the components can be undefined.
struct LinearEqEntry {
PrimExpr base;
diff --git a/src/arith/int_set.cc b/src/arith/int_set.cc
index 2e7655fc05..514da25fab 100644
--- a/src/arith/int_set.cc
+++ b/src/arith/int_set.cc
@@ -395,8 +395,6 @@ IntervalSet ToIntervalSet(IntSet set) {
return IntervalSet::Everything();
}
-using namespace tirx;
-
// Simplified version of int set evaluator that operates on IntervalSet
// We might use better set analysis in the future to replace the intervalset.
class IntervalSetEvaluator : public tvm::ExprFunctor<IntervalSet(const Expr&)>
{
diff --git a/src/arith/iter_affine_map.cc b/src/arith/iter_affine_map.cc
index b4f25a2b13..b16dd1232b 100644
--- a/src/arith/iter_affine_map.cc
+++ b/src/arith/iter_affine_map.cc
@@ -28,9 +28,8 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr_functor.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/op.h>
+#include <unordered_set>
#include <utility>
#include "../support/utils.h"
@@ -42,8 +41,6 @@ namespace tvm {
namespace arith {
using namespace tvm::prim;
-using namespace tirx;
-
TVM_FFI_STATIC_INIT_BLOCK() {
IterMarkNode::RegisterReflection();
IterSplitExprNode::RegisterReflection();
@@ -1295,6 +1292,68 @@ class IterMapRewriter : public tvm::ExprMutator {
}
};
+// Count expression occurrences for constraint ordering, without traversing
type metadata,
+// let bindings, or vector lane descriptors. Shared subexpressions count at
each occurrence.
+class IterConstraintSizeCounter : public tvm::ExprVisitor {
+ public:
+ static size_t Count(const PrimExpr& expr) {
+ auto counter = ffi::make_object<IterConstraintSizeCounter>();
+ counter->Visit(expr);
+ return counter->count_;
+ }
+
+ private:
+ ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) final {
+ if (value.as<ExprNode>()) ++count_;
+ return tvm::ExprVisitor::Visit(value);
+ }
+
+ ffi::Optional<VisitInterrupt> Visit_(const VarNode*) final { return
std::nullopt; }
+
+ ffi::Optional<VisitInterrupt> Visit_(const OpaqueExprNode*) final { return
std::nullopt; }
+
+ ffi::Optional<VisitInterrupt> Visit_(const TupleNode* op) final {
+ return this->Visit(op->fields);
+ }
+
+ ffi::Optional<VisitInterrupt> Visit_(const TupleGetItemNode* op) final {
+ return this->Visit(op->tuple);
+ }
+
+ ffi::Optional<VisitInterrupt> Visit_(const TensorLoadNode* op) final {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->source));
+ return this->Visit(op->indices);
+ }
+
+ ffi::Optional<VisitInterrupt> Visit_(const CallNode* op) final {
+ if (op->op.as<OpaqueExprNode>()) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->op));
+ }
+ return this->Visit(op->args);
+ }
+
+ ffi::Optional<VisitInterrupt> Visit_(const LetNode* op) final {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->value));
+ return this->Visit(op->body);
+ }
+
+ ffi::Optional<VisitInterrupt> Visit_(const RampNode* op) final {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->base));
+ return this->Visit(op->stride);
+ }
+
+ ffi::Optional<VisitInterrupt> Visit_(const BroadcastNode* op) final {
+ return this->Visit(op->value);
+ }
+
+ ffi::Optional<VisitInterrupt> Visit_(const ShuffleNode* op) final {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->indices));
+ return this->Visit(op->vectors);
+ }
+
+ size_t count_{0};
+};
+
/*! \brief An internal struct to represent range extent on iterators(iter <
upper_bound). */
struct IterConstraint {
// The expr of the iter
@@ -1498,7 +1557,7 @@ IterMapResult DetectIterMap(const ffi::Array<PrimExpr>&
indices,
// in the iter var graph has been visited, where the expression of this
iterator will contain the
// expression of its successor, so we sort them by their sizes.
for (IterConstraint& constraint : constraints) {
- constraint.expr_size = CalculateExprComplexity(constraint.iter);
+ constraint.expr_size = IterConstraintSizeCounter::Count(constraint.iter);
}
std::sort(
diff --git a/src/arith/modular_set.cc b/src/arith/modular_set.cc
index 159bd7b4e6..296a239bc4 100644
--- a/src/arith/modular_set.cc
+++ b/src/arith/modular_set.cc
@@ -27,8 +27,6 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr_functor.h>
#include <tvm/ir/prim/builtin.h>
-#include <tvm/tirx/builtin.h>
-#include <tvm/tirx/op.h>
#include <limits>
#include <unordered_map>
@@ -40,8 +38,6 @@ namespace tvm {
namespace arith {
using namespace tvm::prim;
-using namespace tirx;
-
TVM_FFI_STATIC_INIT_BLOCK() { ModularSetNode::RegisterReflection(); }
ModularSet::ModularSet(int64_t coeff, int64_t base) {
diff --git a/src/arith/pattern_match.h b/src/arith/pattern_match.h
index 0b71ff0f6d..674a439758 100644
--- a/src/arith/pattern_match.h
+++ b/src/arith/pattern_match.h
@@ -69,7 +69,6 @@
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/prim/op.h>
-#include <tvm/tirx/builtin.h>
#include <cmath>
#include <tuple>
diff --git a/src/arith/presburger_set.cc b/src/arith/presburger_set.cc
index 35064eac36..23ed696cf0 100644
--- a/src/arith/presburger_set.cc
+++ b/src/arith/presburger_set.cc
@@ -46,7 +46,6 @@ using namespace tvm::prim;
#if defined(TVM_MLIR_VERSION) && TVM_MLIR_VERSION >= 150
TVM_FFI_STATIC_INIT_BLOCK() { PresburgerSetNode::RegisterReflection(); }
-using namespace tirx;
static void Update(const PrimExpr& constraint, PresburgerSetNode* intset) {
auto& space = intset->space;
diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc
index b3be3128a5..086b2b2f22 100644
--- a/src/arith/rewrite_simplify.cc
+++ b/src/arith/rewrite_simplify.cc
@@ -30,9 +30,6 @@
#include <tvm/ir/op.h>
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/builtin.h>
-#include <tvm/tirx/op.h>
#include <algorithm>
#include <tuple>
@@ -53,8 +50,6 @@ TVM_FFI_INLINE bool IsVectorExpr(const ExprNode* expr) {
}
} // namespace
-using namespace tirx;
-
TVM_FFI_STATIC_INIT_BLOCK() {
RewriteSimplifierStatsNode::RegisterReflection(); }
// Note: When using matches_one_of or PMatchesOneOf alongside these
@@ -2425,7 +2420,7 @@ UnchangedOr<Expr> RewriteSimplifier::Impl::Mutate_(const
CallNode* op, InplaceMo
}
static const Op& ceil_op = prim::builtin::ceil();
static const Op& log2_op = prim::builtin::log2();
- static const Op& clz_op = Op::Get("tirx.clz");
+ static const Op& clz_op = prim::builtin::clz();
PrimType ret_ty = op->ty.as_or_throw<PrimType>();
if (op->op.same_as(ceil_op)) {
PrimExpr ceil_arg = op->args[0].as_or_throw<PrimExpr>();
diff --git a/src/arith/rewrite_simplify.h b/src/arith/rewrite_simplify.h
index 23b9edde1d..71586096e7 100644
--- a/src/arith/rewrite_simplify.h
+++ b/src/arith/rewrite_simplify.h
@@ -40,8 +40,6 @@
namespace tvm {
namespace arith {
-using namespace tirx;
-
/* \brief Usage counters for RewriteSimplifier
*
* These are intended for debug and testing purposes, to ensure that
diff --git a/src/arith/transitive_comparison_analyzer.cc
b/src/arith/transitive_comparison_analyzer.cc
index 6b524efcaf..6ade8dc21e 100644
--- a/src/arith/transitive_comparison_analyzer.cc
+++ b/src/arith/transitive_comparison_analyzer.cc
@@ -22,9 +22,9 @@
#include <tvm/arith/analyzer.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/analysis.h>
#include <optional>
+#include <unordered_set>
#include <vector>
#include "constraint_extract.h"
@@ -36,8 +36,6 @@ using namespace tvm::prim;
using prim::is_const_int;
-using namespace tirx;
-
class TransitiveComparisonAnalyzer::Impl {
public:
/* \brief Using previously specified knowns, compare the expressions provided
diff --git a/src/arith/z3_prover.cc b/src/arith/z3_prover.cc
index 4e9556a51a..41a11c6d73 100644
--- a/src/arith/z3_prover.cc
+++ b/src/arith/z3_prover.cc
@@ -34,10 +34,8 @@
#include <tvm/ir/op_attr_types.h>
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/ir/prim/op.h>
#include <tvm/runtime/logging.h>
-#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/builtin.h>
-#include <tvm/tirx/op.h>
#include <algorithm>
#include <climits>
@@ -60,7 +58,6 @@
namespace tvm::arith {
using namespace tvm::prim;
-using namespace tirx;
using namespace ffi;
namespace {
@@ -1024,14 +1021,13 @@ TVM_DLL Z3Prover::~Z3Prover() = default;
#include <tvm/arith/analyzer.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/op.h>
+#include <tvm/ir/prim/op.h>
#include "tvm/ffi/string.h"
#include "tvm/ir/expr.h"
namespace tvm::arith {
-using namespace tirx;
using namespace ffi;
void EnterZ3ContextScope() {}
diff --git a/src/backend/cuda/codegen/intrin_rule_cuda.cc
b/src/backend/cuda/codegen/intrin_rule_cuda.cc
index 4a57f2f469..de2599ebe7 100644
--- a/src/backend/cuda/codegen/intrin_rule_cuda.cc
+++ b/src/backend/cuda/codegen/intrin_rule_cuda.cc
@@ -165,7 +165,7 @@ static PrimExpr DispatchCUDAShuffle(const PrimExpr& e) {
void RegisterCudaIntrinRules() {
// clang-format off
-TVM_REGISTER_OP("tirx.clz")
+TVM_REGISTER_OP("prim.clz")
.set_attr<FLowerIntrinsic>("cuda.FLowerIntrinsic",
DispatchPureExtern<CUDAMath,
/*dtype_from_arg=*/true>);
diff --git a/src/backend/metal/codegen/intrin_rule_metal.cc
b/src/backend/metal/codegen/intrin_rule_metal.cc
index 947ab191b6..93551aced2 100644
--- a/src/backend/metal/codegen/intrin_rule_metal.cc
+++ b/src/backend/metal/codegen/intrin_rule_metal.cc
@@ -58,7 +58,7 @@ static PrimExpr DispatchMetalShuffle(const PrimExpr& e) {
void RegisterMetalIntrinRules() {
// clang-format off
-TVM_REGISTER_OP("tirx.clz")
+TVM_REGISTER_OP("prim.clz")
.set_attr<FLowerIntrinsic>("metal.FLowerIntrinsic",
DispatchPureExtern<Direct>);
TVM_REGISTER_OP("tirx.floor")
diff --git a/src/backend/opencl/codegen/intrin_rule_opencl.cc
b/src/backend/opencl/codegen/intrin_rule_opencl.cc
index 9b3ee23f2e..0d61f6d8ea 100644
--- a/src/backend/opencl/codegen/intrin_rule_opencl.cc
+++ b/src/backend/opencl/codegen/intrin_rule_opencl.cc
@@ -53,7 +53,7 @@ void RegisterOpenCLIntrinRules() {
registered = true;
// clang-format off
-TVM_REGISTER_OP("tirx.clz")
+TVM_REGISTER_OP("prim.clz")
.set_attr<FLowerIntrinsic>("opencl.FLowerIntrinsic",
DispatchPureExtern<Direct>);
TVM_REGISTER_OP("tirx.floor")
diff --git a/src/backend/vulkan/codegen/intrin_rule_spirv.cc
b/src/backend/vulkan/codegen/intrin_rule_spirv.cc
index 13bf738626..bca215bbdb 100644
--- a/src/backend/vulkan/codegen/intrin_rule_spirv.cc
+++ b/src/backend/vulkan/codegen/intrin_rule_spirv.cc
@@ -163,7 +163,7 @@ void RegisterVulkanLegalizeRules() {
registered = true;
// clang-format off
-TVM_REGISTER_OP("tirx.clz")
+TVM_REGISTER_OP("prim.clz")
.set_attr<FLegalize>("vulkan.FLegalize", [](const PrimExpr& e) -> PrimExpr
{
const CallNode* call = e.as<CallNode>();
TVM_FFI_ICHECK(call != nullptr);
diff --git a/src/ir/prim/builtin.cc b/src/ir/prim/builtin.cc
index d8b15a69ad..cc1cc6b6af 100644
--- a/src/ir/prim/builtin.cc
+++ b/src/ir/prim/builtin.cc
@@ -73,6 +73,9 @@ PRIM_DEFINE_BUILTIN_FUNC(log2)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<bool>("TVectorizable", true);
+PRIM_DEFINE_BUILTIN_FUNC(clz).set_num_inputs(1).set_attr<TCallEffectKind>(
+ "TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
+
#undef PRIM_DEFINE_BUILTIN_FUNC
} // namespace builtin
} // namespace prim
diff --git a/src/ir/prim/op.cc b/src/ir/prim/op.cc
index 2b86131b34..8ce00eca71 100644
--- a/src/ir/prim/op.cc
+++ b/src/ir/prim/op.cc
@@ -746,6 +746,20 @@ PrimExpr log2(PrimExpr x, Span span) {
return Call(x_ty, prim::builtin::log2(), {x}, {}, {},
span).as_or_throw<PrimExpr>();
}
+PrimExpr prim::clz(PrimExpr x, Span span) {
+ PrimType x_ty = x.ty();
+ if (x_ty.MatchesElementType(DLDataTypeCode::kDLBfloat, 16)) {
+ PrimType f32_ty = x_ty.IsScalableVector() ?
PrimType::ScalableVector(DLDataTypeCode::kDLFloat,
+ 32,
x_ty.VScaleFactor())
+ : PrimType::Float(32,
x_ty.lanes());
+ PrimExpr x_fp32 = prim::Cast(f32_ty, x, span);
+ PrimExpr result_fp32 =
+ Call(f32_ty, prim::builtin::clz(), {x_fp32}, {}, {},
span).as_or_throw<PrimExpr>();
+ return prim::Cast(x_ty, result_fp32, span);
+ }
+ return Call(x_ty, prim::builtin::clz(), {x}, {}, {},
span).as_or_throw<PrimExpr>();
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef()
.def_packed("node._const",
diff --git a/src/target/intrin_rule.h b/src/target/intrin_rule.h
index f5ddb0eeca..e3c5c042bf 100644
--- a/src/target/intrin_rule.h
+++ b/src/target/intrin_rule.h
@@ -70,7 +70,8 @@ inline PrimExpr DispatchPureExtern(const PrimExpr& e) {
const OpNode* op = call->op.as<OpNode>();
TVM_FFI_ICHECK(op != nullptr);
std::string name = op->name;
- TVM_FFI_ICHECK(name.substr(0, 5) == "tirx." || name == "prim.ceil" || name
== "prim.log2")
+ TVM_FFI_ICHECK(name.substr(0, 5) == "tirx." || name == "prim.ceil" || name
== "prim.log2" ||
+ name == "prim.clz")
<< "Unexpected intrinsic name: " << name;
if (dtype_from_arg) {
TVM_FFI_ICHECK_EQ(call->args.size(), 1U);
diff --git a/src/target/llvm/intrin_rule_llvm.cc
b/src/target/llvm/intrin_rule_llvm.cc
index 3b63a16543..5c46da5704 100644
--- a/src/target/llvm/intrin_rule_llvm.cc
+++ b/src/target/llvm/intrin_rule_llvm.cc
@@ -169,7 +169,7 @@ TVM_REGISTER_OP("tirx.atanh")
return (log(one + x) - log(one - x)) * MakeConst(x_ty, 0.5);
});
-TVM_REGISTER_OP("tirx.clz")
+TVM_REGISTER_OP("prim.clz")
.set_attr<FLegalize>("llvm.FLegalize", [](const PrimExpr& e) -> PrimExpr {
const CallNode* call = e.as<CallNode>();
TVM_FFI_ICHECK(call != nullptr);
diff --git a/src/tirx/ir/data_type_rewriter.cc
b/src/tirx/ir/data_type_rewriter.cc
index 5af0023112..d27daa68b9 100644
--- a/src/tirx/ir/data_type_rewriter.cc
+++ b/src/tirx/ir/data_type_rewriter.cc
@@ -346,7 +346,7 @@ UnchangedOr<Expr> DataTypeLegalizer::Mutate_(const
CallNode* op, InplaceMode inp
return op->args[0].as_or_throw<PrimExpr>() ^
op->args[1].as_or_throw<PrimExpr>();
}
static const Op& pow_op = Op::Get("tirx.pow");
- static const Op& clz_op = Op::Get("tirx.clz");
+ static const Op& clz_op = prim::builtin::clz();
if (op->op.same_as(pow_op)) {
return pow(op->args[0].as_or_throw<PrimExpr>(),
op->args[1].as_or_throw<PrimExpr>());
} else if (op->op.same_as(prim::builtin::if_then_else())) {
diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc
index adaa0b8bd7..391a7508b0 100644
--- a/src/tirx/op/builtin.cc
+++ b/src/tirx/op/builtin.cc
@@ -50,6 +50,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
PRIM_SCRIPT_BUILTIN(vscale);
PRIM_SCRIPT_BUILTIN(ceil);
PRIM_SCRIPT_BUILTIN(log2);
+ PRIM_SCRIPT_BUILTIN(clz);
#undef PRIM_SCRIPT_BUILTIN
}
diff --git a/src/tirx/op/op.cc b/src/tirx/op/op.cc
index 8b1a6a8f1d..828568d16f 100644
--- a/src/tirx/op/op.cc
+++ b/src/tirx/op/op.cc
@@ -520,8 +520,6 @@ TVM_TIR_REGISTER_PURE_UNARY_OP("asinh");
TVM_TIR_REGISTER_PURE_UNARY_OP("atanh");
-TVM_TIR_REGISTER_PURE_UNARY_OP("clz");
-
// binary intrinsics
TVM_TIR_REGISTER_PURE_BINARY_OP("atan2");
diff --git a/tests/cpp/function_attrs_test.cc b/tests/cpp/function_attrs_test.cc
deleted file mode 100644
index bc5edac549..0000000000
--- a/tests/cpp/function_attrs_test.cc
+++ /dev/null
@@ -1,154 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-#include <gtest/gtest.h>
-#include <tvm/ir/function.h>
-#include <tvm/ir/module.h>
-#include <tvm/relax/expr.h>
-#include <tvm/tirx/function.h>
-#include <tvm/tirx/stmt.h>
-
-namespace tvm {
-namespace {
-
-class CustomFuncNode : public BaseFuncNode {
- public:
- ffi::String payload = "preserved";
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.AttrsCustomFunc", CustomFuncNode,
BaseFuncNode);
-};
-
-class MissingCopyFuncNode : public BaseFuncNode {
- public:
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.AttrsMissingCopyFunc",
MissingCopyFuncNode, BaseFuncNode);
-};
-
-class InvalidCopyFuncNode : public BaseFuncNode {
- public:
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.AttrsInvalidCopyFunc",
InvalidCopyFuncNode, BaseFuncNode);
-};
-
-int invalid_copy_mode = 0;
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::ObjectDef<CustomFuncNode>().def_ro("payload",
&CustomFuncNode::payload);
- ffi::reflection::TypeAttrDef<InvalidCopyFuncNode>().def(
- ffi::reflection::type_attr::kShallowCopy, [](BaseFunc func) -> Any {
- if (invalid_copy_mode == 0) return func;
- if (invalid_copy_mode == 1) return
BaseFunc(ffi::make_object<CustomFuncNode>());
- return nullptr;
- });
-}
-
-TEST(FunctionAttrs, GenericSubtypeAndSharedAttributes) {
- BaseFunc original(ffi::make_object<CustomFuncNode>());
- original = WithAttr(std::move(original), "keep", 1);
- DictAttrs attrs = original->attrs;
- auto check_copy = [&](const BaseFunc& result) {
- EXPECT_FALSE(result.same_as(original));
- EXPECT_EQ(result->type_index(), original->type_index());
- EXPECT_EQ(result.as<CustomFuncNode>()->payload, "preserved");
- EXPECT_TRUE(result->ty.same_as(original->ty));
- EXPECT_TRUE(result->span.same_as(original->span));
- EXPECT_EQ(original->attrs->dict.size(), 1);
- EXPECT_TRUE(original->attrs.same_as(attrs));
- EXPECT_EQ(attrs->dict.at("keep").cast<int>(), 1);
- };
- auto added = WithAttr(original, "added", 2);
- check_copy(added);
- EXPECT_EQ(added->attrs->dict.at("added").cast<int>(), 2);
- auto updated = WithAttrs(original, {{"keep", 3}, {"added", 4}});
- check_copy(updated);
- EXPECT_EQ(updated->attrs->dict.at("keep").cast<int>(), 3);
- auto removed = WithoutAttr(original, "keep");
- check_copy(removed);
- EXPECT_TRUE(removed->attrs->dict.empty());
- EXPECT_TRUE(WithAttrs(original, {}).same_as(original));
-}
-
-TEST(FunctionAttrs, UniqueReuseAndSharedDictionary) {
- BaseFunc func(ffi::make_object<CustomFuncNode>());
- const auto* ptr = func.get();
- DictAttrs shared_attrs = func->attrs;
- func = WithAttr(std::move(func), "key", 1);
- EXPECT_EQ(func.get(), ptr);
- EXPECT_TRUE(shared_attrs->dict.empty());
- func = WithAttrs(std::move(func), {{"key", 2}, {"other", 3}});
- EXPECT_EQ(func.get(), ptr);
- func = WithoutAttr(std::move(func), "key");
- EXPECT_EQ(func.get(), ptr);
- EXPECT_FALSE(func->attrs->dict.count("key"));
- EXPECT_EQ(func->attrs->dict.at("other").cast<int>(), 3);
-}
-
-TEST(FunctionAttrs, TypedFunctionsAndModule) {
- tirx::PrimFunc prim({}, tirx::Evaluate(0));
- auto prim_copy = WithAttr(prim, "key", 1);
- EXPECT_FALSE(prim_copy.same_as(prim));
- EXPECT_TRUE(prim_copy->body.same_as(prim->body));
- EXPECT_TRUE(prim_copy->params.same_as(prim->params));
- EXPECT_TRUE(prim_copy->ret_type.same_as(prim->ret_type));
- BaseFunc base = prim;
- auto generic_copy = WithAttrs(base, {{"key", 1}});
- EXPECT_EQ(generic_copy->type_index(), prim->type_index());
- EXPECT_TRUE(generic_copy.as<tirx::PrimFuncNode>()->body.same_as(prim->body));
-
- relax::ExternFunc ext("external_symbol");
- auto ext_copy = WithoutAttr(WithAttr(ext, "key", 1), "key");
- EXPECT_EQ(ext_copy->global_symbol, ext->global_symbol);
- EXPECT_TRUE(ext_copy->attrs->dict.empty());
- IRModule mod = IRModule::FromExpr(prim);
- auto mod_copy = WithAttrs(mod, {{"key", 1}});
- EXPECT_TRUE(mod->attrs->dict.empty());
- EXPECT_EQ(mod_copy->attrs->dict.at("key").cast<int>(), 1);
- EXPECT_TRUE(WithoutAttr(mod_copy, "key")->attrs->dict.empty());
-}
-
-TEST(FunctionAttrs, MissingAndInvalidHooksPreserveInput) {
- BaseFunc missing(ffi::make_object<MissingCopyFuncNode>());
- EXPECT_THROW(WithAttr(missing, "key", 1), ffi::Error);
- EXPECT_THROW(WithAttrs(missing, {{"key", 1}}), ffi::Error);
- EXPECT_THROW(WithoutAttr(missing, "key"), ffi::Error);
- EXPECT_TRUE(WithAttrs(missing, {}).same_as(missing));
- // Unique input needs no copy hook.
- const auto* ptr = missing.get();
- missing = WithAttr(std::move(missing), "key", 1);
- EXPECT_EQ(missing.get(), ptr);
-
- BaseFunc invalid(ffi::make_object<InvalidCopyFuncNode>());
- invalid = WithAttr(std::move(invalid), "key", 1);
- for (invalid_copy_mode = 0; invalid_copy_mode < 3; ++invalid_copy_mode) {
- EXPECT_THROW(WithAttr(invalid, "key", 2), ffi::Error);
- EXPECT_THROW(WithAttrs(invalid, {{"key", 2}}), ffi::Error);
- EXPECT_THROW(WithoutAttr(invalid, "key"), ffi::Error);
- EXPECT_EQ(invalid->attrs->dict.at("key").cast<int>(), 1);
- }
-}
-
-TEST(FunctionAttrs, MovedFromAttributes) {
- BaseFunc func(ffi::make_object<CustomFuncNode>());
- // Simulate a caller moving through the base handle, bypassing DictAttrs'
reset-on-move.
- ffi::ObjectRef moved =
-
std::move(static_cast<ffi::ObjectRef&>(const_cast<BaseFuncNode*>(func.operator->())->attrs));
- EXPECT_EQ(WithAttr(func, "key", 1)->attrs->dict.size(), 1);
- EXPECT_EQ(WithAttrs(func, {{"key", 1}})->attrs->dict.size(), 1);
- EXPECT_TRUE(WithoutAttr(func, "key")->attrs->dict.empty());
- EXPECT_FALSE(func->attrs.defined());
-}
-
-} // namespace
-} // namespace tvm
diff --git a/tests/cpp/pattern_match_test.cc b/tests/cpp/pattern_match_test.cc
index 3a5d22a417..20854b39fc 100644
--- a/tests/cpp/pattern_match_test.cc
+++ b/tests/cpp/pattern_match_test.cc
@@ -24,7 +24,6 @@
TEST(Pattern, Basic) {
using namespace tvm;
- using namespace tvm::tirx;
using namespace tvm::arith;
tvm::PrimVar x("x"), y("y"), z("z");
PrimExpr scalable_lanes = prim::Mul(Call(PrimType::Int(32),
prim::builtin::vscale(), {}), 4);
diff --git a/tests/python/arith/test_arith_rewrite_simplify.py
b/tests/python/arith/test_arith_rewrite_simplify.py
index 5f6eb73391..6dbe2a65a1 100644
--- a/tests/python/arith/test_arith_rewrite_simplify.py
+++ b/tests/python/arith/test_arith_rewrite_simplify.py
@@ -1322,21 +1322,21 @@ class TestIfThenElse(BaseCompare):
class TestCLZ(BaseCompare):
test_case = tvm.testing.parameter(
- TestCase(tvm.tirx.call_intrin("int32", "tirx.clz", 0), T.int32(32)),
- TestCase(tvm.tirx.call_intrin("int32", "tirx.clz", 1), T.int32(31)),
- TestCase(tvm.tirx.call_intrin("int32", "tirx.clz", 2), T.int32(30)),
- TestCase(tvm.tirx.call_intrin("int32", "tirx.clz", 128), T.int32(24)),
+ TestCase(tvm.tirx.call_intrin("int32", "prim.clz", 0), T.int32(32)),
+ TestCase(tvm.tirx.call_intrin("int32", "prim.clz", 1), T.int32(31)),
+ TestCase(tvm.tirx.call_intrin("int32", "prim.clz", 2), T.int32(30)),
+ TestCase(tvm.tirx.call_intrin("int32", "prim.clz", 128), T.int32(24)),
TestCase(
- tvm.tirx.call_intrin("int32", "tirx.clz", tvm.tirx.IntImm("int64",
0)), T.int32(64)
+ tvm.tirx.call_intrin("int32", "prim.clz", tvm.tirx.IntImm("int64",
0)), T.int32(64)
),
TestCase(
- tvm.tirx.call_intrin("int32", "tirx.clz", tvm.tirx.IntImm("int64",
1)), T.int32(63)
+ tvm.tirx.call_intrin("int32", "prim.clz", tvm.tirx.IntImm("int64",
1)), T.int32(63)
),
TestCase(
- tvm.tirx.call_intrin("int32", "tirx.clz", tvm.tirx.IntImm("int64",
2)), T.int32(62)
+ tvm.tirx.call_intrin("int32", "prim.clz", tvm.tirx.IntImm("int64",
2)), T.int32(62)
),
TestCase(
- tvm.tirx.call_intrin("int32", "tirx.clz", tvm.tirx.IntImm("int64",
128)), T.int32(56)
+ tvm.tirx.call_intrin("int32", "prim.clz", tvm.tirx.IntImm("int64",
128)), T.int32(56)
),
)
diff --git a/tests/python/ir/test_ir_attrs.py b/tests/python/ir/test_ir_attrs.py
index 5d8171d152..25480f7265 100644
--- a/tests/python/ir/test_ir_attrs.py
+++ b/tests/python/ir/test_ir_attrs.py
@@ -57,43 +57,6 @@ def test_assert_structural_equal_reports_mismatch():
assert "and rhs at" in message
[email protected]("kind", ["prim", "relax", "extern"])
-def test_function_attribute_copy_preserves_fields(kind):
- span = tvm.ir.Span(tvm.ir.SourceName("attrs"), 1, 2, 3, 4)
- if kind == "prim":
- var = tvm.tirx.Var("x", "int32")
- func = tvm.tirx.PrimFunc([var], tvm.tirx.Evaluate(var), span=span)
- fields = ["params", "body", "ret_type", "ty", "span"]
- elif kind == "relax":
- var = tvm.relax.Var("x", tvm.relax.TensorType([2], "float32"))
- func = tvm.relax.Function([var], var, is_pure=False, span=span)
- fields = ["params", "body", "ret_ty", "ty", "span"]
- else:
- func = tvm.relax.ExternFunc("external_symbol", span=span)
- fields = ["ty", "span"]
-
- func = func.with_attr("keep", 1)
- shared_attrs = func.attrs
- added = func.with_attr("added", 2)
- updated = func.with_attr({"keep": 3, "added": 4})
- removed = func.without_attr("keep")
- assert func.with_attr({}).same_as(func)
- for result in [added, updated, removed]:
- assert type(result) is type(func)
- assert not result.same_as(func)
- for field in fields:
- assert getattr(result, field).same_as(getattr(func, field))
- if kind == "relax":
- assert result.is_pure is False
- if kind == "extern":
- assert result.global_symbol == "external_symbol"
- assert func.attrs.same_as(shared_attrs)
- assert dict(shared_attrs) == {"keep": 1}
- assert dict(added.attrs) == {"keep": 1, "added": 2}
- assert dict(updated.attrs) == {"keep": 3, "added": 4}
- assert not removed.attrs
-
-
if __name__ == "__main__":
test_dict_attrs()
test_attrs_equal()
