Lunderberg commented on code in PR #11646:
URL: https://github.com/apache/tvm/pull/11646#discussion_r894631463
##########
src/arith/rewrite_simplify.cc:
##########
@@ -1640,13 +1640,27 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const
CallNode* op) {
// the operator overload will eagerly constant fold.
return op->args[0] << op->args[1];
}
+ } else if (op->op.same_as(Op::Get("tir.ceil"))) {
+ if (auto as_int = op->args[0].as<IntImmNode>()) {
+ return cast(op->dtype, IntImm(as_int->dtype, as_int->value));
+ } else if (auto as_float = op->args[0].as<FloatImmNode>()) {
+ return cast(op->dtype, FloatImm(as_float->dtype,
std::ceil(as_float->value)));
+ }
+ } else if (op->op.same_as(Op::Get("tir.log2"))) {
+ if (auto as_int = op->args[0].as<IntImmNode>()) {
+ return cast(op->dtype, FloatImm(as_int->dtype,
std::log2(as_int->value)));
+ } else if (auto as_float = op->args[0].as<FloatImmNode>()) {
Review Comment:
Good point. Because this folding would be applied in the `tvm.lower` path
where the target may be unknown, it would be tricky to make it be target-aware.
Since my main goal was for simplifying integer arguments, where
`ceil(log2(n))` is used for loop iteration bounds, we could get similar effects
by restricting this folding to integer arguments to `ceil(log2(n))`, which
would be the same regardless of rounding differences.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
