| /* |
| * 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/const_int_bound.cc |
| */ |
| #include <tvm/arith/analyzer.h> |
| #include <tvm/runtime/registry.h> |
| #include <tvm/tir/builtin.h> |
| #include <tvm/tir/expr_functor.h> |
| |
| #include <algorithm> |
| |
| #include "int_operator.h" |
| #include "pattern_match.h" |
| |
| namespace tvm { |
| namespace arith { |
| |
| using namespace tir; |
| |
| TVM_REGISTER_NODE_TYPE(ConstIntBoundNode); |
| |
| ConstIntBound::ConstIntBound(int64_t min_value, int64_t max_value) { |
| auto node = make_object<ConstIntBoundNode>(); |
| node->min_value = min_value; |
| node->max_value = max_value; |
| data_ = std::move(node); |
| } |
| |
| ConstIntBound MakeConstIntBound(int64_t min_value, int64_t max_value) { |
| return ConstIntBound(min_value, max_value); |
| } |
| |
| TVM_REGISTER_GLOBAL("arith.ConstIntBound").set_body_typed(MakeConstIntBound); |
| |
| inline void PrintBoundValue(std::ostream& os, int64_t val) { |
| if (val == ConstIntBound::kPosInf) { |
| os << "pos_inf"; |
| } else if (val == ConstIntBound::kNegInf) { |
| os << "neg_inf"; |
| } else { |
| os << val; |
| } |
| } |
| |
| TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) |
| .set_dispatch<ConstIntBoundNode>([](const ObjectRef& node, ReprPrinter* p) { |
| auto* op = static_cast<const ConstIntBoundNode*>(node.get()); |
| p->stream << "ConstIntBound["; |
| PrintBoundValue(p->stream, op->min_value); |
| p->stream << ','; |
| PrintBoundValue(p->stream, op->max_value); |
| p->stream << ']'; |
| }); |
| |
| // internal entry for const int bound |
| struct ConstIntBoundAnalyzer::Entry { |
| int64_t min_value; |
| int64_t max_value; |
| |
| bool is_const(int64_t value) const { return min_value == max_value && min_value == value; } |
| |
| bool operator==(const Entry& other) const { |
| return min_value == other.min_value && max_value == other.max_value; |
| } |
| }; |
| |
| class ConstIntBoundAnalyzer::Impl |
| : public ExprFunctor<ConstIntBoundAnalyzer::Entry(const PrimExpr&)> { |
| public: |
| /*! \brief additional bound info about expr in bound */ |
| struct BoundInfo { |
| /*! \brief The expr */ |
| PrimExpr expr; |
| /*! \brief The additional bound */ |
| Entry bound; |
| |
| BoundInfo() {} |
| BoundInfo(PrimExpr expr, Entry bound) : expr(expr), bound(bound) {} |
| }; |
| |
| void Bind(const Var& var, const Range& range, bool allow_override) { |
| Entry a = VisitExpr(range->min); |
| Entry b = VisitExpr(range->extent); |
| Entry ret; |
| ret.min_value = a.min_value; |
| ret.max_value = InfAwareAdd(a.max_value, InfAwareAdd(b.max_value, -1)); |
| Update(var, ret, allow_override); |
| } |
| |
| void Update(const Var& var, const Entry& info, bool allow_override) { |
| if (!allow_override) { |
| auto it = var_map_.find(var); |
| if (it != var_map_.end()) { |
| CHECK(it->second == info) << "Trying to update var \'" << var << "\'" |
| << " with a different const bound: " |
| << "original=" |
| << ConstIntBound(it->second.min_value, it->second.max_value) |
| << ", new=" << ConstIntBound(info.min_value, info.max_value); |
| } |
| } |
| var_map_[var] = info; |
| } |
| |
| Entry VisitExpr_(const LetNode* op) final { |
| auto it = var_map_.find(op->var); |
| // if the var has not been binded, update the info. |
| if (it == var_map_.end()) { |
| var_map_[op->var] = this->VisitExpr(op->value); |
| Entry ret = VisitExpr(op->body); |
| var_map_.erase(op->var); |
| return ret; |
| } else { |
| return VisitExpr(op->body); |
| } |
| } |
| |
| void Update(const Var& var, const ConstIntBound& info, bool allow_override) { |
| Update(var, MakeBound(info->min_value, info->max_value), allow_override); |
| } |
| |
| // Override visitor behaviors |
| Entry VisitExprDefault_(const Object* op) final { |
| return Everything(static_cast<const PrimExprNode*>(op)->dtype); |
| } |
| |
| Entry VisitExpr(const PrimExpr& expr) final { |
| Entry res = ExprFunctor::VisitExpr(expr); |
| tir::ExprDeepEqual equal; |
| // a linear search over additional info |
| // assume we won't have a lot of conditions |
| for (const BoundInfo& info : additional_info_) { |
| if (equal(expr, info.expr)) { |
| res = Intersect(res, info.bound); |
| } |
| } |
| if (bound_) { |
| auto val = bound_->find(expr); |
| if (val != bound_->end()) { |
| auto everything = Everything(expr->dtype); |
| CHECK( |
| (val->second->min_value == res.min_value && val->second->max_value == res.max_value) || |
| (val->second->min_value == everything.min_value && |
| val->second->max_value == everything.max_value)) |
| << "Detected bound for " << expr << "conflicts with memorization"; |
| } |
| (*bound_)[expr] = ConstIntBound(res.min_value, res.max_value); |
| } |
| return res; |
| } |
| |
| Entry VisitExpr_(const RampNode* op) final { |
| // op = {base + i * stride | 0 <= i < lanes} |
| // Entry(op) = Union(Entry(base + i * stride) | 0 <= i < lanes) |
| // Note that `base + i * stride` is linear w.r.t. `i` |
| // Entry(op) = Union(Entry(base + i * stride) | i = 0, i = lanes-1) |
| Entry a = VisitExpr(op->base); |
| Entry b = VisitExpr(op->base + (op->lanes - 1) * op->stride); |
| return Union(a, b); |
| } |
| |
| Entry VisitExpr_(const CastNode* op) final { |
| Entry a = VisitExpr(op->value); |
| Entry b = Everything(op->dtype); |
| return Intersect(a, b); |
| } |
| |
| Entry VisitExpr_(const IntImmNode* op) final { return MakeBound(op->value, op->value); } |
| |
| Entry VisitExpr_(const AddNode* op) final { |
| Entry a = VisitExpr(op->a); |
| Entry b = VisitExpr(op->b); |
| Entry ret; |
| ret.min_value = InfAwareAdd(a.min_value, b.min_value); |
| ret.max_value = InfAwareAdd(a.max_value, b.max_value); |
| return ret; |
| } |
| |
| Entry VisitExpr_(const SubNode* op) final { |
| Entry a = VisitExpr(op->a); |
| Entry b = VisitExpr(op->b); |
| Entry ret; |
| ret.min_value = InfAwareAdd(a.min_value, -b.max_value); |
| ret.max_value = InfAwareAdd(a.max_value, -b.min_value); |
| return ret; |
| } |
| |
| Entry VisitExpr_(const MulNode* op) final { |
| Entry a = VisitExpr(op->a); |
| Entry b = VisitExpr(op->b); |
| return BinaryOpBoundary(a, b, InfAwareMul); |
| } |
| |
| Entry VisitExpr_(const DivNode* op) final { |
| Entry a = VisitExpr(op->a); |
| Entry b = VisitExpr(op->b); |
| CHECK(!b.is_const(0)) << "divide by zero"; |
| return HandleDivision(a, b, op->dtype, InfAwareDiv); |
| } |
| |
| Entry VisitExpr_(const ModNode* op) final { |
| Entry a = VisitExpr(op->a); |
| Entry b = VisitExpr(op->b); |
| if (b.min_value > 0) { |
| int64_t b_max_cap = InfAwareAdd(b.max_value, -1); |
| if (a.min_value >= 0) { |
| // 0 <= [a_min, a_max] < b_min |
| if (a.max_value < b.min_value) return a; |
| // other case, we can get close to 0 |
| return MakeBound(0, std::min(a.max_value, b_max_cap)); |
| } else { |
| return MakeBound(std::max(a.min_value, -b_max_cap), |
| std::min(std::max(a.max_value, (int64_t)0), b_max_cap)); |
| } |
| } else { |
| CHECK(!b.is_const(0)) << "mod by zero"; |
| // mod by negative value is rare, |
| // and we just use the simpliest rule. |
| return Everything(op->dtype); |
| } |
| } |
| |
| Entry VisitExpr_(const FloorDivNode* op) final { |
| Entry a = VisitExpr(op->a); |
| Entry b = VisitExpr(op->b); |
| CHECK(!b.is_const(0)) << "floordiv by zero"; |
| return HandleDivision(a, b, op->dtype, InfAwareFloorDiv); |
| } |
| |
| Entry VisitExpr_(const FloorModNode* op) final { |
| Entry a = VisitExpr(op->a); |
| Entry b = VisitExpr(op->b); |
| if (b.min_value > 0) { |
| int64_t b_max_cap = InfAwareAdd(b.max_value, -1); |
| if (a.min_value >= 0) { |
| // 0 <= [a_min, a_max] < b_min |
| if (a.max_value < b.min_value) return a; |
| // other case, we can get close to 0 |
| return MakeBound(0, std::min(a.max_value, b_max_cap)); |
| } else { |
| return MakeBound(0, b_max_cap); |
| } |
| } else { |
| CHECK(!b.is_const(0)) << "floormod by zero"; |
| // mod by negative value is rare, |
| // and we just use the simpliest rule. |
| return Everything(op->dtype); |
| } |
| } |
| |
| Entry VisitExpr_(const MinNode* op) final { |
| Entry a = VisitExpr(op->a); |
| Entry b = VisitExpr(op->b); |
| Entry ret; |
| ret.min_value = std::min(a.min_value, b.min_value); |
| ret.max_value = std::min(a.max_value, b.max_value); |
| return ret; |
| } |
| |
| Entry VisitExpr_(const MaxNode* op) final { |
| Entry a = VisitExpr(op->a); |
| Entry b = VisitExpr(op->b); |
| Entry ret; |
| ret.min_value = std::max(a.min_value, b.min_value); |
| ret.max_value = std::max(a.max_value, b.max_value); |
| return ret; |
| } |
| |
| Entry VisitExpr_(const SelectNode* op) final { |
| Entry a = VisitExpr(op->true_value); |
| Entry b = VisitExpr(op->false_value); |
| return Union(a, b); |
| } |
| |
| Entry VisitExpr_(const CallNode* op) final { |
| // only special handle >> and & which can be |
| // used for index calculation. |
| |
| if (op->op.same_as(tir::builtin::shift_right())) { |
| return VisitRightShift(op); |
| } else if (op->op.same_as(tir::builtin::bitwise_and())) { |
| return VisitBitwiseAnd(op); |
| } else { |
| return Everything(op->dtype); |
| } |
| } |
| |
| Entry VisitExpr_(const VarNode* op) final { |
| Var v = GetRef<Var>(op); |
| auto it = var_map_.find(v); |
| if (it != var_map_.end()) { |
| return it->second; |
| } else { |
| return Everything(op->dtype); |
| } |
| } |
| |
| Entry VisitExpr_(const SizeVarNode* op) final { |
| SizeVar v = GetRef<SizeVar>(op); |
| auto it = var_map_.find(v); |
| if (it != var_map_.end()) { |
| return it->second; |
| } else { |
| return MakeBound(0, kPosInf); |
| } |
| } |
| |
| Entry VisitRightShift(const CallNode* op) { |
| Entry a = VisitExpr(op->args[0]); |
| Entry b = VisitExpr(op->args[1]); |
| return BinaryOpBoundary(a, b, InfAwareRightShift); |
| } |
| |
| Entry VisitBitwiseAnd(const CallNode* op) { |
| Entry a = VisitExpr(op->args[0]); |
| Entry b = VisitExpr(op->args[1]); |
| // handle positive index case. |
| if (a.min_value >= 0 && b.min_value >= 0) { |
| return MakeBound(0, std::min(a.max_value, b.max_value)); |
| } else { |
| if (b.min_value >= 0) { |
| return MakeBound(0, b.max_value); |
| } |
| if (a.min_value >= 0) { |
| return MakeBound(0, a.max_value); |
| } |
| return Everything(op->dtype); |
| } |
| } |
| |
| std::function<void()> EnterConstraint(const PrimExpr& constraint) { |
| std::vector<BoundInfo> info = DetectBoundInfo(constraint); |
| if (info.size() == 0) return nullptr; |
| size_t old_size = additional_info_.size(); |
| additional_info_.insert(additional_info_.end(), info.begin(), info.end()); |
| size_t new_size = old_size + info.size(); |
| auto frecover = [old_size, new_size, this]() { |
| CHECK_EQ(additional_info_.size(), new_size); |
| additional_info_.resize(old_size); |
| }; |
| return frecover; |
| } |
| |
| private: |
| friend class ConstIntBoundAnalyzer; |
| // internal variable map |
| std::unordered_map<Var, Entry, ObjectPtrHash, ObjectPtrEqual> var_map_; |
| // additional bound info |
| std::vector<BoundInfo> additional_info_; |
| // look up table for memorization |
| BoundMapType* bound_{nullptr}; |
| // constants: the limit value means umlimited |
| // NOTE: kNegInf/kPosInf are used to represent infinity. |
| static const constexpr int64_t kNegInf = ConstIntBound::kNegInf; |
| static const constexpr int64_t kPosInf = ConstIntBound::kPosInf; |
| static_assert(-kNegInf == kPosInf, "invariant of inf"); |
| // internal helper functions |
| /*! |
| * \brief Get boundary of binary op who are monotonic wrt to one argument. |
| * \param a The entry of the left operand. |
| * \param b The entry of the right operand. |
| * \param op The operator. |
| * \tparam F the operator function type. |
| * \return The result. |
| */ |
| template <typename F> |
| static Entry BinaryOpBoundary(Entry a, Entry b, const F& op) { |
| Entry ret; |
| // The boundary point must be shihft of the original boundary. |
| int64_t v1 = op(a.min_value, b.min_value); |
| int64_t v2 = op(a.max_value, b.max_value); |
| int64_t v3 = op(a.min_value, b.max_value); |
| int64_t v4 = op(a.max_value, b.min_value); |
| ret.min_value = std::min(std::min(std::min(v1, v2), v3), v4); |
| ret.max_value = std::max(std::max(std::max(v1, v2), v3), v4); |
| return ret; |
| } |
| /*! |
| * \brief Get value boundaries of division (e.g. Div or FloorDiv). |
| * \param a The entry of the left operand. |
| * \param b The entry of the right operand. |
| * \param dt The data type of the division operator. |
| * \param op The division operator. |
| * \tparam F the operator function type. |
| * \return The result. |
| */ |
| template <typename F> |
| static Entry HandleDivision(Entry a, Entry b, DataType dt, const F& op) { |
| // Here we have a / b. |
| // The largest value of the division will be for the smallest (with |
| // respect to the absolute value) value of b. If the range of b starts |
| // at a negative value and ends at a positive one, narrow it down to |
| // be closer to 0, because BinaryOpBoundary only checks end-points of |
| // the domain ranges. |
| |
| // If the range of b contains 0, then some infinity will be involved |
| if (b.min_value <= 0 && 0 <= b.max_value) { |
| Entry b_neg = b.min_value < 0 ? MakeBound(b.min_value, -1) : Everything(dt); |
| Entry b_pos = b.max_value > 0 ? MakeBound(1, b.max_value) : Everything(dt); |
| |
| Entry e_neg = BinaryOpBoundary(a, b_neg, op); |
| Entry e_pos = BinaryOpBoundary(a, b_pos, op); |
| |
| return MakeBound(std::min(e_neg.min_value, e_pos.min_value), |
| std::max(e_neg.max_value, e_pos.max_value)); |
| } |
| // If the range of b does not have 0, use BinaryOpBoundary. |
| return BinaryOpBoundary(a, b, op); |
| } |
| /*! |
| * \brief Compute x + y, aware of inf. |
| * \param x The left operand. |
| * \param y The right operand. |
| * \return the result. |
| */ |
| static int64_t InfAwareAdd(int64_t x, int64_t y) { |
| if (x == kPosInf) { |
| CHECK(y != kNegInf); |
| return kPosInf; |
| } |
| if (x == kNegInf) { |
| CHECK(y != kPosInf); |
| return kNegInf; |
| } |
| if (y == kPosInf || y == kNegInf) return y; |
| if (WillOverflow<AddNode>(x, y, kNegInf, kPosInf)) { |
| if (x > 0) return kPosInf; |
| return kNegInf; |
| } |
| return x + y; |
| } |
| /*! |
| * \brief Compute x * y, aware of inf. |
| * \param x The left operand. |
| * \param y The right operand. |
| * \return the result. |
| */ |
| static int64_t InfAwareMul(int64_t x, int64_t y) { |
| if (!WillOverflow<MulNode>(x, y, kNegInf, kPosInf)) return x * y; |
| if ((x > 0 && y > 0) || (x < 0 && y < 0)) return kPosInf; |
| return kNegInf; |
| } |
| /*! |
| * \brief Compute x / y, aware of inf. |
| * \param x The left operand. |
| * \param y The right operand. |
| * \return the result. |
| */ |
| static int64_t InfAwareDiv(int64_t x, int64_t y) { |
| CHECK_NE(y, 0); |
| if (x == kPosInf || x == kNegInf) { |
| if (y > 0) return x; |
| return -x; |
| } |
| return x / y; |
| } |
| /*! |
| * \brief Compute floodiv(x, y), aware of inf. |
| * \param x The left operand. |
| * \param y The right operand. |
| * \return the result. |
| */ |
| static int64_t InfAwareFloorDiv(int64_t x, int64_t y) { |
| CHECK_NE(y, 0); |
| if (x == kPosInf || x == kNegInf) { |
| if (y > 0) return x; |
| return -x; |
| } |
| return floordiv(x, y); |
| } |
| /*! |
| * \brief Compute x / y, aware of inf. |
| * \param x The left operand. |
| * \param y The right operand. |
| * \return the result. |
| */ |
| static int64_t InfAwareRightShift(int64_t x, int64_t y) { |
| if (x == kPosInf || x == kNegInf) return x; |
| return x >> y; |
| } |
| /*! |
| * \brief Make a new bound entry. |
| */ |
| static Entry MakeBound(int64_t min_value, int64_t max_value) { |
| Entry e; |
| e.min_value = min_value; |
| e.max_value = max_value; |
| return e; |
| } |
| /*! |
| * \brief Create union of two sets. |
| * \param a The left operand. |
| * \param b the right operand. |
| */ |
| static Entry Union(Entry a, Entry b) { |
| Entry ret; |
| ret.min_value = std::min(a.min_value, b.min_value); |
| ret.max_value = std::max(a.max_value, b.max_value); |
| return ret; |
| } |
| /*! |
| * \brief Create intersect of two sets. |
| * \param a The left operand. |
| * \param b the right operand. |
| */ |
| static Entry Intersect(Entry a, Entry b) { |
| Entry ret; |
| ret.min_value = std::max(a.min_value, b.min_value); |
| ret.max_value = std::min(a.max_value, b.max_value); |
| return ret; |
| } |
| /*! |
| * \brief return everything dtype can represent. |
| * \param dtype The data type. |
| * \return Bound that represent everything dtype can represent. |
| */ |
| static Entry Everything(DataType dtype) { |
| if (!dtype.is_int() && !dtype.is_uint()) { |
| return MakeBound(kNegInf, kPosInf); |
| } |
| Entry ret; |
| int64_t vbits = dtype.bits() - static_cast<int>(dtype.is_int()); |
| if (dtype.is_uint()) { |
| ret.min_value = 0; |
| } else { |
| if (vbits >= 63) { |
| ret.min_value = kNegInf; |
| } else { |
| ret.min_value = -(static_cast<int64_t>(1) << vbits); |
| } |
| } |
| if (vbits >= 63) { |
| ret.max_value = kPosInf; |
| } else { |
| ret.max_value = (static_cast<int64_t>(1) << vbits) - 1; |
| } |
| return ret; |
| } |
| |
| /*! |
| * \brief Detect additional constant bound from cond, if any |
| * \param cond The constraint condition. |
| * \return List of detected bounds. |
| */ |
| static std::vector<BoundInfo> DetectBoundInfo(const PrimExpr& cond) { |
| PVar<PrimExpr> x, y; |
| PVar<IntImm> c; |
| // NOTE: canonical form always use <= or < |
| if ((c <= x).Match(cond)) { |
| return {BoundInfo(x.Eval(), MakeBound(c.Eval()->value, kPosInf))}; |
| } |
| if ((c < x).Match(cond)) { |
| return {BoundInfo(x.Eval(), MakeBound(c.Eval()->value + 1, kPosInf))}; |
| } |
| if ((x <= c).Match(cond)) { |
| return {BoundInfo(x.Eval(), MakeBound(kNegInf, c.Eval()->value))}; |
| } |
| if ((x < c).Match(cond)) { |
| return {BoundInfo(x.Eval(), MakeBound(kNegInf, c.Eval()->value - 1))}; |
| } |
| if ((x && y).Match(cond)) { |
| auto ret1 = DetectBoundInfo(x.Eval()); |
| auto ret2 = DetectBoundInfo(y.Eval()); |
| ret1.insert(ret1.end(), ret2.begin(), ret2.end()); |
| return ret1; |
| } |
| return {}; |
| } |
| }; |
| |
| ConstIntBound ConstIntBoundAnalyzer::operator()(const PrimExpr& expr) { |
| Entry ret = impl_->VisitExpr(expr); |
| return ConstIntBound(ret.min_value, ret.max_value); |
| } |
| |
| ConstIntBound ConstIntBoundAnalyzer::operator()(const PrimExpr& expr, BoundMapType* bound) { |
| impl_->bound_ = bound; |
| Entry ret = impl_->VisitExpr(expr); |
| impl_->bound_ = nullptr; |
| return ConstIntBound(ret.min_value, ret.max_value); |
| } |
| |
| void ConstIntBoundAnalyzer::Update(const Var& var, const ConstIntBound& info, bool allow_override) { |
| impl_->Update(var, info, allow_override); |
| } |
| |
| void ConstIntBoundAnalyzer::Bind(const Var& var, const Range& range, bool allow_override) { |
| impl_->Bind(var, range, allow_override); |
| } |
| |
| std::function<void()> ConstIntBoundAnalyzer::EnterConstraint(const PrimExpr& constraint) { |
| return impl_->EnterConstraint(constraint); |
| } |
| |
| ConstIntBoundAnalyzer::ConstIntBoundAnalyzer(Analyzer* parent) : impl_(new Impl()) {} |
| |
| ConstIntBoundAnalyzer::~ConstIntBoundAnalyzer() { delete impl_; } |
| |
| } // namespace arith |
| } // namespace tvm |