text: /*
   * 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.
   */
  
  /*!
   *  Copyright (c) 2019 by Contributors
   * \file tvm/arithmetic/const_int_bound.cc
   */
  #include <tvm/arithmetic.h>
  #include <tvm/ir_functor_ext.h>
  #include <algorithm>
  #include "int_operator.h"
  #include "pattern_match.h"
  
  namespace tvm {
  namespace arith {
  
  using namespace ir;
  
  TVM_REGISTER_NODE_TYPE(ConstIntBoundNode);
  
  ConstIntBound::ConstIntBound(
      int64_t min_value, int64_t max_value) {
    auto node = make_node<ConstIntBoundNode>();
    node->min_value = min_value;
    node->max_value = max_value;
    data_ = std::move(node);
  }
  
  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(IRPrinter, vtable)
  .set_dispatch<ConstIntBoundNode>([](const ObjectRef& node, IRPrinter* 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 Expr&)> {
   public:
    /*! \brief additional bound info about expr \in bound */
    struct BoundInfo {
      /*! \brief The expr */
      Expr expr;
      /*! \brief The additional bound */
      Entry bound;
  
      BoundInfo() {}
      BoundInfo(Expr expr, Entry bound)
          : expr(expr), bound(bound) {
      }
    };
  
    void Bind(const Var& var, const Range& range) {
      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, false);
    }
  
    void Update(const Var& var,
                const Entry& info,
                bool override) {
      if (!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;
    }
  
    void Update(const Var& var,
                const ConstIntBound& info,
                bool override) {
      Update(var, MakeBound(info->min_value, info->max_value), override);
    }
  
    // Override visitor behaviors
    Entry VisitExprDefault_(const Node* op) final {
      return Everything(
          static_cast<const ExprNode*>(op)->type);
    }
  
    Entry VisitExpr(const Expr& expr) final {
      Entry res = ExprFunctor::VisitExpr(expr);
      // a linear search over additional info
      // assume we won't have a lot of conditions
      for (const BoundInfo& info : additional_info_) {
        if (ir::Equal(expr, info.expr)) {
          res = Intersect(res, info.bound);
        }
      }
      return res;
    }
  
    Entry VisitExpr_(const Cast* op) final {
      Entry a = VisitExpr(op->value);
      Entry b = Everything(op->type);
      return Intersect(a, b);
    }
  
    Entry VisitExpr_(const IntImm* op) final {
      return MakeBound(op->value, op->value);
    }
  
    Entry VisitExpr_(const UIntImm* op) final {
      if (op->value <= static_cast<uint64_t>(kPosInf)) {
        return MakeBound(op->value, op->value);
      } else {
        return Everything(op->type);
      }
    }
  
    Entry VisitExpr_(const Add* 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 Sub* 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 Mul* op) final {
      Entry a = VisitExpr(op->a);
      Entry b = VisitExpr(op->b);
      return BinaryOpBoundry(a, b, InfAwareMul);
    }
  
    Entry VisitExpr_(const Div* op) final {
      Entry a = VisitExpr(op->a);
      Entry b = VisitExpr(op->b);
      CHECK(!b.is_const(0)) << "divide by zero";
      // assume no division by 0
      if (b.min_value == 0) b.min_value = 1;
      if (b.max_value == 0) b.max_value = -1;
      return BinaryOpBoundry(a, b, InfAwareDiv);
    }
  
    Entry VisitExpr_(const Mod* 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->type);
      }
    }
  
    Entry VisitExpr_(const FloorDiv* op) final {
      Entry a = VisitExpr(op->a);
      Entry b = VisitExpr(op->b);
      CHECK(!b.is_const(0)) << "floordiv by zero";
      // assume no division by 0
      if (b.min_value == 0) b.min_value = 1;
      if (b.max_value == 0) b.max_value = -1;
      return BinaryOpBoundry(a, b, InfAwareFloorDiv);
    }
  
    Entry VisitExpr_(const FloorMod* 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->type);
      }
    }
  
    Entry VisitExpr_(const Min* 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 Max* 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 Select* op) final {
      Entry a = VisitExpr(op->true_value);
      Entry b = VisitExpr(op->false_value);
      return Union(a, b);
    }
  
    Entry VisitExpr_(const Call* op) final {
      // only special handle >> and & which can be
      // used for index calculation.
      if (op->is_intrinsic(Call::shift_right)) {
        return VisitRightShift(op);
      } else if (op->is_intrinsic(Call::bitwise_and)) {
        return VisitBitwiseAnd(op);
      } else {
        return Everything(op->type);
      }
    }
  
    Entry VisitExpr_(const Variable* 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->type);
      }
    }
  
    Entry VisitRightShift(const Call* op) {
      Entry a = VisitExpr(op->args[0]);
      Entry b = VisitExpr(op->args[1]);
      return BinaryOpBoundry(a, b, InfAwareRightShift);
    }
  
    Entry VisitBitwiseAnd(const Call* 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->type);
      }
    }
  
    std::function<void()> EnterConstraint(const Expr& 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:
    // internal variable map
    std::unordered_map<Var, Entry, ExprHash, ExprEqual> var_map_;
    // additional bound info
    std::vector<BoundInfo> additional_info_;
    // 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 param a The entry of the left operand.
     * \param param a The entry of the right operand.
     * \param op The operator.
     * \tparam F the operator function type.
     * \return The result.
     */
    template<typename F>
    static Entry BinaryOpBoundry(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 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<Add>(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<Mul>(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(Type 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 Expr& cond) {
      PVar<Expr> x, y;
      PVar<Integer> 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 Expr& expr) {
    Entry ret = impl_->VisitExpr(expr);
    return ConstIntBound(ret.min_value, ret.max_value);
  }
  
  void ConstIntBoundAnalyzer::Update(const Var& var,
                                     const ConstIntBound& info,
                                     bool override) {
    impl_->Update(var, info, override);
  }
  
  void ConstIntBoundAnalyzer::Bind(const Var& var, const Range& range) {
    impl_->Bind(var, range);
  }
  
  std::function<void()> ConstIntBoundAnalyzer::EnterConstraint(const Expr& constraint) {
    return impl_->EnterConstraint(constraint);
  }
  
  ConstIntBoundAnalyzer::ConstIntBoundAnalyzer(Analyzer* parent)
      : impl_(new Impl()) {
  }
  
  ConstIntBoundAnalyzer::~ConstIntBoundAnalyzer() {
    delete impl_;
  }
  
  }  // namespace arith
  }  // namespace tvm
  