lhutton1 commented on code in PR #16523: URL: https://github.com/apache/tvm/pull/16523#discussion_r1478630929
########## src/arith/scalable_expression.cc: ########## @@ -0,0 +1,62 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file tvm/arith/scalable_expression.cc + * \brief Analyze scalable expressions. + */ + +#include "scalable_expression.h" + +#include <tvm/tir/expr.h> +#include <tvm/tir/op.h> + +#include "../tir/transforms/replace_selected_expr.h" Review Comment: nit: I think this include has crept in from the prototype ########## src/arith/rewrite_simplify.cc: ########## @@ -1010,21 +1023,29 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const FloorModNode* op) { } // If all possible indices in ramp are the same. ModularSet bmod = analyzer_->modular_set(b1.Eval()); - int64_t ramp_min = floordiv(bmod->base, c2val); - int64_t ramp_max = floordiv(bmod->base + (lanes.Eval() - 1) * c1val, c2val); - if (ramp_min == ramp_max) { + if (!ScalableLanes(lanes.Eval())) { + int64_t ramp_min = floordiv(bmod->base, c2val); + auto lanes_int = lanes.Eval().as<IntImmNode>()->value; + int64_t ramp_max = floordiv(bmod->base + (lanes_int - 1) * c1val, c2val); + if (ramp_min == ramp_max) { + // If b1 can devide c2 Review Comment: nit: s/devide/divide ########## src/arith/rewrite_simplify.cc: ########## @@ -872,17 +882,20 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const FloorDivNode* op) { return ramp(floordiv(b1, c2), floordiv(c1, c2), lanes).Eval(); } // If all possible indices in ramp are the same. - ModularSet bmod = analyzer_->modular_set(b1.Eval()); - int64_t ramp_min = floordiv(bmod->base, c2val); - int64_t ramp_max = floordiv(bmod->base + (lanes.Eval() - 1) * c1val, c2val); - if (ramp_min == ramp_max) { - // If b1 can devide c2 - if (bmod->coeff % c2val == 0) { - return broadcast(floordiv(b1, c2), lanes).Eval(); - } - // If all indices can be guaranteed to settle inside a coeff range - if (c2val % bmod->coeff == 0 && bmod->base + (lanes.Eval() - 1) * c1val < bmod->coeff) { - return broadcast(floordiv(b1, c2), lanes).Eval(); + if (!ScalableLanes(lanes.Eval())) { + ModularSet bmod = analyzer_->modular_set(b1.Eval()); + int64_t ramp_min = floordiv(bmod->base, c2val); + auto lanes_int = lanes.Eval().as<IntImmNode>()->value; + int64_t ramp_max = floordiv(bmod->base + (lanes_int - 1) * c1val, c2val); + if (ramp_min == ramp_max) { + // If b1 can devide c2 Review Comment: nit: s/devide/divide ########## src/target/spirv/codegen_spirv.cc: ########## @@ -533,7 +535,10 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const RampNode* op) { spirv::Value CodeGenSPIRV::VisitExpr_(const BroadcastNode* op) { std::vector<spirv::Value> values; spirv::Value v = MakeValue(op->value); - for (int i = 0; i < op->lanes; i++) { + // ICHECK(op->lanes->IsInstance<IntImmNode>()) Review Comment: nit: remove ########## src/arith/pattern_match.h: ########## @@ -628,10 +628,11 @@ inline PRampExpr<TBase, TStride, TLanes> ramp(const Pattern<TBase>& base, } template <typename TBase> -inline PRampExpr<TBase, PConstWithTypeLike<TBase>, PConst<int>> ramp(const Pattern<TBase>& base, - int stride, int lanes) { - return PRampExpr<TBase, PConstWithTypeLike<TBase>, PConst<int>>( - base.derived(), PConstWithTypeLike<TBase>(base.derived(), stride), PConst<int>(lanes)); +inline PRampExpr<TBase, PConstWithTypeLike<TBase>, PConstWithTypeLike<TBase>> ramp( + const Pattern<TBase>& base, int stride, int lanes) { + return PRampExpr<TBase, PConstWithTypeLike<TBase>, PConstWithTypeLike<TBase>>( + base.derived(), PConstWithTypeLike<TBase>(base.derived(), stride), Review Comment: I'm not too familiar with the pattern matcher, but just wanted to note something that stood out to me - should `int lanes` have a different type now? ########## src/arith/int_set.cc: ########## @@ -466,14 +466,21 @@ class IntervalSetEvaluator : public ExprFunctor<IntervalSet(const PrimExpr&)> { if (stride.Match(op->stride)) { DataType t = op->base.dtype(); int64_t vstride = stride.Eval()->value; - if (vstride > 0) { - return Combine<Add>(analyzer_, base, - IntervalSet(make_zero(t), make_const(t, vstride * (op->lanes - 1))), - op->dtype); - } else { - return Combine<Add>(analyzer_, base, - IntervalSet(make_const(t, vstride * (op->lanes - 1)), make_zero(t)), - op->dtype); + if (op->lanes->IsInstance<IntImmNode>()) { + int lanes = static_cast<int>(Downcast<IntImm>(op->lanes)->value); + if (vstride > 0) { + return Combine<Add>(analyzer_, base, + IntervalSet(make_zero(t), make_const(t, vstride * (lanes - 1))), + op->dtype); + } else { + return Combine<Add>(analyzer_, base, + IntervalSet(make_const(t, vstride * (lanes - 1)), make_zero(t)), + op->dtype); + } + } else { /* Scalable vector */ + if (vstride > 0) { + return Combine<Add>(analyzer_, base, IntervalSet(make_zero(t), pos_inf()), op->dtype); + } Review Comment: Does `vstride <= 0` need to be handled here? ########## src/arith/rewrite_simplify.h: ########## @@ -221,6 +221,8 @@ class RewriteSimplifier::Impl : public IRMutatorWithAnalyzer { bool CanProveGreaterEqual(const PrimExpr& x, int64_t val) { return analyzer_->CanProveGreaterEqual(x, val); } + // Whether the lanes are scalable + bool ScalableLanes(const PrimExpr& lanes) { return !lanes.as<IntImmNode>(); } Review Comment: I can't see this used anywhere else, it's probably better to move it into the .cc. Additionally, should this have a stricter check to enforce the PrimExpr contains `vscale`? -- 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]
