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);

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

 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_;
};

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;
  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_