analyzer.h

Go to the documentation of this file.

/*
 * 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.
 */

#ifndef TVM_ARITH_ANALYZER_H_
#define TVM_ARITH_ANALYZER_H_

#include <tvm/arith/int_set.h>
#include <tvm/ir/expr.h>
#include <tvm/support/with.h>

#include <limits>
#include <memory>
#include <unordered_map>
#include <vector>

namespace tvm {
namespace arith {
//-------------------------------------------------------
// Base integer analysis API.
//
// We have multiple type of analyzers to do relaxed
// integer set analysis(bound analysis, modulo) and
// equivalence checking and simplification.
//
// Importantly, each analyzer may need result from
// another analyzer.
//-------------------------------------------------------

// Forward declare Analyzer
class Analyzer;

using tir::Var;

enum DivMode {
  kTruncDiv,
  kFloorDiv
};

class ConstIntBoundNode : public Object {
 public:
  int64_t min_value;
  int64_t max_value;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("min_value", &min_value);
    v->Visit("max_value", &max_value);
  }

  bool SEqualReduce(const ConstIntBoundNode* other, SEqualReducer equal) const {
    return equal(min_value, other->min_value) && equal(max_value, other->max_value);
  }

  static const constexpr int64_t kPosInf = std::numeric_limits<int64_t>::max();
  static const constexpr int64_t kNegInf = -kPosInf;

  static constexpr const char* _type_key = "arith.ConstIntBound";
  TVM_DECLARE_FINAL_OBJECT_INFO(ConstIntBoundNode, Object);
};

class ConstIntBound : public ObjectRef {
 public:
  TVM_DLL ConstIntBound(int64_t min_value, int64_t max_value);

  static const constexpr int64_t kPosInf = ConstIntBoundNode::kPosInf;
  static const constexpr int64_t kNegInf = ConstIntBoundNode::kNegInf;
  TVM_DEFINE_OBJECT_REF_METHODS(ConstIntBound, ObjectRef, ConstIntBoundNode);
};

class ConstIntBoundAnalyzer {
 public:
  using BoundMapType = std::unordered_map<PrimExpr, ConstIntBound, ObjectPtrHash, ObjectPtrEqual>;
  TVM_DLL ConstIntBound operator()(const PrimExpr& expr) const;

  TVM_DLL ConstIntBound operator()(const PrimExpr& expr, BoundMapType* bound);

  TVM_DLL void Update(const Var& var, const ConstIntBound& info, bool allow_override = false);
  TVM_DLL void Bind(const Var& var, const Range& range, bool allow_override = false);

 private:
  friend class Analyzer;
  friend class ConstraintContext;
  explicit ConstIntBoundAnalyzer(Analyzer* parent);
  TVM_DLL ~ConstIntBoundAnalyzer();
  std::function<void()> EnterConstraint(const PrimExpr& constraint);
  struct Entry;
  class Impl;
  Impl* impl_;
};

class ModularSetNode : public Object {
 public:
  int64_t coeff;
  int64_t base;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("coeff", &coeff);
    v->Visit("base", &base);
  }

  bool SEqualReduce(const ModularSetNode* other, SEqualReducer equal) const {
    return equal(coeff, other->coeff) && equal(base, other->base);
  }

  static constexpr const char* _type_key = "arith.ModularSet";
  TVM_DECLARE_FINAL_OBJECT_INFO(ModularSetNode, Object);
};

class ModularSet : public ObjectRef {
 public:
  TVM_DLL ModularSet(int64_t coeff, int64_t base);

  TVM_DEFINE_OBJECT_REF_METHODS(ModularSet, ObjectRef, ModularSetNode);
};

class ModularSetAnalyzer {
 public:
  TVM_DLL ModularSet operator()(const PrimExpr& expr);
  TVM_DLL void Update(const Var& var, const ModularSet& info, bool allow_override = false);

 private:
  friend class Analyzer;
  friend class ConstraintContext;
  explicit ModularSetAnalyzer(Analyzer* parent);
  TVM_DLL ~ModularSetAnalyzer();
  std::function<void()> EnterConstraint(const PrimExpr& constraint);
  struct Entry;
  class Impl;
  Impl* impl_;
};

class RewriteSimplifier {
 public:
  TVM_DLL PrimExpr operator()(const PrimExpr& expr);

  TVM_DLL void Update(const Var& var, const PrimExpr& new_expr, bool allow_override = false);

  TVM_DLL std::function<void()> EnterConstraint(const PrimExpr& constraint);

  enum Extension {
    // No extensions enabled
    kNone = 0,

    /* When simplifying an inequality, attempt to use scope-based knowns.
     *
     * Example:
     * if_then_else(i<j && j<k, i<k, false) => if_then_else(i<j && j<k, true, false)
     */
    kTransitivelyProveInequalities = (1 << 0),

    /* When simplifying a boolean expression, convert to an AND of ORs
     * (conjunctive normal form).
     *
     * Example:
     *   (a && b) || c => (a || c) && (b || c)
     */
    kConvertBooleanToAndOfOrs = (1 << 1),

    /* When simplifying a boolean AND or a boolean OR, simplify each
     * branch under the assumption that the other branch does not
     * already dominate the result.  That is, simplify each branch of
     * (A && B) under the assumption that the other branch is true,
     * and simplify each branch of (A || B) under the assumption that
     * the other branch is false.
     *
     * Example:
     *   (n < 10) && (n < 5) => (n < 10)
     *   (n < 10) || (n < 5) => (n < 5)
     */
    kApplyConstraintsToBooleanBranches = (1 << 2),
  };

  TVM_DLL void SetEnabledExtensions(Extension flags);

  TVM_DLL Extension GetEnabledExtensions() const;

 private:
  friend class Analyzer;
  friend class ConstraintContext;
  friend class CanonicalSimplifier;
  explicit RewriteSimplifier(Analyzer* parent);
  TVM_DLL ~RewriteSimplifier();
  class Impl;
  Impl* impl_;
};

class CanonicalSimplifier {
 public:
  TVM_DLL PrimExpr operator()(const PrimExpr& expr);

  TVM_DLL void Update(const Var& var, const PrimExpr& new_expr, bool allow_override = false);

 private:
  friend class Analyzer;
  friend class ConstraintContext;
  explicit CanonicalSimplifier(Analyzer* parent);
  TVM_DLL ~CanonicalSimplifier();
  class Impl;
  Impl* impl_;
};

enum class CompareResult : int {
  kInconsistent = 0,
  kEQ = 1,
  kLT = 2,
  kLE = 3,
  kGT = 4,
  kGE = 5,
  kNE = 6,
  kUnknown = 7
};

inline constexpr CompareResult operator&(CompareResult lhs, CompareResult rhs) {
  return CompareResult(static_cast<int>(lhs) & static_cast<int>(rhs));
}
inline constexpr CompareResult operator|(CompareResult lhs, CompareResult rhs) {
  return CompareResult(static_cast<int>(lhs) | static_cast<int>(rhs));
}

class TransitiveComparisonAnalyzer {
 public:
  /* \brief Using previously specified knowns, compare the expressions provided
   *
   * \param lhs The left-hand side of the comparison
   *
   * \param rhs The right-hand side of the comparison
   *
   * \param propagate_inequalities If true, attempt to find a sequence
   * of transitive inequalities that allow the lhs and rhs to be
   * compared.  If false, only use the known comparison that have been
   * directly provided.  Using `propagate_inequalities = false` is
   * roughly equivalent to comparing against all known inequality
   * expressions using `ExprDeepEqual`, but also allows for constant
   * offsets on either side of the inequality.
   *
   * \return The most specific result that can be proven about the
   * comparison.  If nothing can be proven, returns kUnknown.
   */
  TVM_DLL CompareResult TryCompare(const PrimExpr& lhs, const PrimExpr& rhs,
                                   bool propagate_inequalities = true);

  TVM_DLL void Bind(const Var& var, const PrimExpr& expr, bool allow_override = false);

  TVM_DLL void Bind(const Var& var, const Range& range, bool allow_override = false);

  TVM_DLL std::function<void()> EnterConstraint(const PrimExpr& constraint);

 private:
  friend class Analyzer;
  friend class ConstraintContext;
  TransitiveComparisonAnalyzer();
  TVM_DLL ~TransitiveComparisonAnalyzer();
  class Impl;
  std::unique_ptr<Impl> impl_;
};

class ConstraintContext {
 private:
  // declare friend to enable with.
  friend class With<ConstraintContext>;
  ConstraintContext(Analyzer* analyzer, PrimExpr constraint)
      : analyzer_(analyzer), constraint_(constraint) {}
  // enter the scope.
  void EnterWithScope();
  // exit the scope.
  void ExitWithScope();
  Analyzer* analyzer_;
  PrimExpr constraint_;
  std::vector<std::function<void()>> recovery_functions_;
};

class IntSetAnalyzer {
 public:
  TVM_DLL IntSet operator()(const PrimExpr& expr, const Map<Var, IntSet>& dom_map);

  TVM_DLL IntSet operator()(const PrimExpr& expr);

  TVM_DLL void Update(const Var& var, const IntSet& new_interval_set, bool allow_override = false);

  TVM_DLL void Bind(const Var& var, const Range& new_range, bool allow_override = false);

  std::function<void()> EnterConstraint(const PrimExpr& constraint);

 private:
  friend class Analyzer;
  explicit IntSetAnalyzer(Analyzer* parent);
  TVM_DLL ~IntSetAnalyzer();
  class Impl;
  Impl* impl_;
};

class TVM_DLL Analyzer {
 public:
  /*
   * Disable copy constructor.
   */
  Analyzer(const Analyzer&) = delete;
  Analyzer& operator=(const Analyzer&) = delete;
  ConstIntBoundAnalyzer const_int_bound;
  ModularSetAnalyzer modular_set;
  RewriteSimplifier rewrite_simplify;
  CanonicalSimplifier canonical_simplify;
  IntSetAnalyzer int_set;
  TransitiveComparisonAnalyzer transitive_comparisons;
  Analyzer();
  void Bind(const Var& var, const PrimExpr& expr, bool allow_override = false);
  void Bind(const Var& var, const Range& range, bool allow_override = false);
  void Bind(const Map<Var, Range>& variables, bool allow_override = false);
  bool CanProveGreaterEqual(const PrimExpr& expr, int64_t lower_bound);
  bool CanProveLess(const PrimExpr& expr, int64_t upper_bound);
  bool CanProveEqual(const PrimExpr& lhs, const PrimExpr& rhs);
  bool CanProve(const PrimExpr& cond);
  PrimExpr Simplify(const PrimExpr& expr, int steps = 2);
};

}  // namespace arith
}  // namespace tvm
#endif  // TVM_ARITH_ANALYZER_H_