arithmetic.h 9.71 KB
Newer Older
1 2
/*!
 *  Copyright (c) 2016 by Contributors
3
 * \file arithmetic.h
4
 * \brief Algebra and set operations and simplifications.
5
 */
6 7
#ifndef TVM_ARITHMETIC_H_
#define TVM_ARITHMETIC_H_
8

9
#include <vector>
10 11 12
#include <unordered_map>
#include <memory>
#include "./expr.h"
13 14

namespace tvm {
15 16 17

class Tensor;

18
/*! \brief namespace of arithmetic */
19
namespace arith {
20 21 22
/*!
 * \brief Sign of an expression or set.
 */
23 24 25 26 27 28 29
enum SignType {
  kPositive,
  kNegative,
  kZero,
  kUnknown
};

30
// internal node container of int set.
31
struct IntSetNode;
32

33
/*!
34
 * \brief Integer set class, represent a set of integers in one dimension.
35
 */
36
class IntSet : public NodeRef {
37
 public:
38 39
  /*! \brief constructor */
  IntSet() {}
40
  // constructor from not container.
41
  explicit IntSet(std::shared_ptr<Node> n) : NodeRef(n) {}
42
  /*!
43 44 45 46 47
   * \brief access the internal node container
   * \return the pointer to the internal node container
   */
  inline const IntSetNode* operator->() const;
  /*!
48 49 50 51 52 53 54 55
   * \brief Find a range that covers the region.
   * \param max_range The range to be covered.
   * \return The covering range.
   */
  Range cover_range(Range max_range) const;
  /*!
   * \brief find an interval that covers the set.
   * \return The covering interval set.
56
   */
57
  IntSet cover_interval() const;
58 59 60 61 62 63
  /*! \return Lower bound of the set */
  Expr min() const;
  /*! \return upper bound of the set */
  Expr max() const;
  /*! \return Whether the set represent nothing  */
  bool is_nothing() const;
64 65 66 67
  /*! \return Whether the set represent everything  */
  bool is_everything() const;
  /*! \return Whether the set is a single point */
  bool is_single_point() const;
68 69
  /*! \return Whether the set is proved to be bigger than 0 */
  bool can_prove_positive() const;
70 71 72 73
  /*! \return Whether the set is proved to be smaller than 0 */
  bool can_prove_negative() const;
  /*! \return The sign of the elements in the integer set */
  SignType sign_type() const;
74 75 76 77 78 79 80 81 82 83 84 85
  /*!
   * \brief The single point value, call only if is_single_point is true
   * \return The point value.
   */
  Expr point_value() const;
  /*!
   * \brief Try to match IntSet with range r.
   *
   * \note It is guanrateed that IntSet::range(r).match_range(r) == true
   * \return true if we can prove they are the same.
   */
  bool match_range(const Range& r) const;
86 87 88
  /*! \return The set contains nothing */
  static IntSet nothing();
  /*! \return The set contains everything */
89
  static IntSet everything();
90
  /*!
91 92 93
   * \brief construct a point set.
   * \param point The point in the set.
   * \return construct a single point set
94
   */
95
  static IntSet single_point(Expr point);
96
  /*!
97 98 99 100 101 102
   * \brief construct a integer set from vector expression.
   * \param vec The vector expression, can also be single point.
   * \return The result set containing the indices in the vector.
   */
  static IntSet vector(Expr vec);
  /*!
103 104 105
   * \brief Construct a set representing a range.
   * \param r The range
   * \return constructed set.
106
   */
107
  static IntSet range(Range r);
108 109 110 111 112 113 114
  /*!
   * \brief Construct a set representing a interval.
   * \param min The minimum value of the interval.
   * \param max The maximum value of the interval.
   * \return constructed set.
   */
  static IntSet interval(Expr min, Expr max);
115 116 117
};

/*!
118 119 120 121 122 123 124 125 126 127 128 129 130
 * \brief Range of a linear integer function.
 *  Use to do specify the possible index values.
 *
 *  set = { base + coeff * x | x in Z }
 *
 *  When coeff != 0, it can also be written as
 *  set = { n | n % coeff == base }
 *
 *  This is useful to decide if the index is dividable by certain value.
 *  For example, if index = 0 + 4 x, then we know it can be divided by 4.
 */
struct ModularEntry {
  /*! \brief The base */
131
  int base{0};
132
  /*! \brief linear co-efficient */
133
  int coeff{1};
134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152

  /*! \return entry represent everything */
  static ModularEntry everything() {
    // always safe to set 0 + x, so it can be everything.
    ModularEntry e;
    e.base = 0; e.coeff = 1;
    return e;
  }
  /*!
   * \brief Add two modular entries together to get a new modular entry.
   * \param a The left operand.
   * \param b The right operand.
   * \return The combined modular entry.
   */
  static ModularEntry Add(const ModularEntry& a,
                          const ModularEntry& b);
};

/*!
153 154 155
 * \brief Base class of all IntSet containers.
 */
struct IntSetNode : public Node {
156 157
  static constexpr const char* _type_key = "IntSet";
  TVM_DECLARE_BASE_NODE_INFO(IntSetNode, Node);
158 159
};

160 161 162 163 164 165 166 167 168

/*!
 * \brief Detect if e can be rewritten as e = base + var * coeff
 *  Where coeff and base are invariant of var.
 *
 * \return [base, coeff] if it is possible, empty array if it is not.
 */
Array<Expr> DetectLinearEquation(Expr e, Var var);

169 170 171 172 173 174 175 176
/*!
 * \brief Find an symbolic integer set that contains all possible values of
 *  e given the domain of each iteration variables.
 *
 * \param e The expression to be evaluated.
 * \param dom_map The domain of each variable.
 * \return An integer set that can cover all the possible values of e.
 */
177 178
IntSet EvalSet(Expr e,
               const Map<IterVar, IntSet>& dom_map);
179 180 181 182 183 184 185
/*!
 * \brief Same as EvalSet, but takes unordered_map
 *
 * \param e The expression to be evaluated.
 * \param dom_map The domain of each variable.
 * \return An integer set that can cover all the possible values of e.
 */
186 187
IntSet EvalSet(Expr e,
               const std::unordered_map<const Variable*, IntSet>& dom_map);
188 189 190 191 192 193 194 195 196 197 198

/*!
 * \brief Find an symbolic integer set that contains is union over
 *  all the possible conditional values in dom_map.
 *
 * \param r The initial range.
 * \param dom_map The domain of each variable.
 * \return An integer set that can cover all the possible values.
 */
IntSet EvalSet(Range r,
               const Map<IterVar, IntSet>& dom_map);
199 200 201 202 203 204 205 206 207 208 209

/*!
 * \brief Find an symbolic integer set that contains is union over
 *  all the possible conditional values in dom_map.
 *
 * \param s The initial set.
 * \param dom_map The domain of each variable.
 * \return An integer set that can cover all the possible values.
 */
IntSet EvalSet(IntSet s,
               const std::unordered_map<const Variable*, IntSet>& dom_map);
210 211 212 213 214 215 216
/*!
 * \brief Same as EvalSet, but takes unordered_map
 *
 * \param r The range to be evaluated.
 * \param dom_map The domain of each variable.
 * \return An integer set that can cover all the possible values of e.
 */
217 218 219
IntSet EvalSet(Range r,
               const std::unordered_map<const Variable*, IntSet>& dom_map);

220 221
/*! \brief Map from Expr to IntSet */
using ExprIntSetMap = std::unordered_map<Expr, IntSet, ExprHash, ExprEqual>;
222 223 224 225 226 227 228 229
/*!
 * \brief Find the integer set of every sub-expression, given the
 *  domain of each iteration variables.
 *
 * \param e The expression to be evaluated.
 * \param dom_map The domain of each variable.
 * \return the map from the expression to its possible value.
 */
230 231
ExprIntSetMap EvalSetForEachSubExpr(
    Expr e,
232 233
    const std::unordered_map<const Variable*, IntSet>& dom_map);

234 235 236 237 238 239 240
/*!
 * \brief Create an union set of all sets
 * \param sets The sets to be unioned
 * \return the set after union
 */
IntSet Union(const Array<IntSet>& sets);

241 242 243 244 245 246 247
/*!
 * \brief Create an union set of all sets
 * \param sets The sets to be intersected
 * \return the set after intersected
 */
IntSet Intersect(const Array<IntSet>& sets);

248 249 250 251 252 253 254
/*!
 * \brief Deduce the bound of the target variable in a expression,
 *  give the domain of each variables. Return undefined IntSet to
 *  represent failure.
 *
 * \param v The target variable to be deduced.
 * \param cond The conditional expression.
255
 * \param hint_map The domain of variable, used to help deduce.
256 257
 * \param relax_map The domain of each variable, used to relax the domain,
 *        The deduce bound mush implies e for all value in relax_map
258 259
 * \return An integer set that can cover all the possible values.
 */
260 261 262
IntSet DeduceBound(Expr v, Expr cond,
                   const Map<Var, IntSet>& hint_map,
                   const Map<Var, IntSet>& relax_map);
263 264 265 266 267 268 269 270 271 272 273 274 275
/*!
 * \brief Same as DeduceBound with  unordered_map signature.
 *
 * \param v The target variable to be deduced.
 * \param cond The conditional expression.
 * \param hint_map The domain of variable, used to help deduce.
 * \param relax_map The domain of each variable, used to relax the domain,
 *        The deduce bound mush implies e for all value in relax_map
 * \return An integer set that can cover all the possible values.
 */
IntSet DeduceBound(Expr v, Expr cond,
                   const std::unordered_map<const Variable*, IntSet>& hint_map,
                   const std::unordered_map<const Variable*, IntSet>& relax_map);
276

277
/*!
278 279 280 281 282 283 284 285 286 287
 * \brief Infer a regular domain that covers all the calls or provides within the given statement.
 * \param body The given statement.
 * \param tensor The name of the calls or provides.
 * \param consider_calls If calls (read) are considered.
 * \param consider_provides If provides (write) are considered.
 * \return The domain that covers all the calls or provides within the given statement.
 */
Domain DomainTouched(Stmt body, const Tensor &tensor, bool consider_calls, bool consider_provides);

/*!
288 289 290 291 292 293 294 295
 * \brief Evaluate the expression with modular analysis
 * \param e The expression to be evaluated.
 * \param mod_map Map of modular statistics of known variables.
 * \return The ModularEntry covering all possible value of e.
 */
ModularEntry EvalModular(
    const Expr& e,
    const std::unordered_map<const Variable*, ModularEntry>& mod_map);
296

297 298 299 300 301 302 303 304 305 306 307 308
/*!
 * \brief Same as EvalModular, used by front-end.
 * \param e The expression to be evaluated.
 * \param mod_map Map of modular statistics of known variables.
 * \return A ModularSet covering all possible value of e.
 */
IntSet EvalModular(const Expr& e,
                   const Map<Var, IntSet>& mod_map);
// implementation
inline const IntSetNode* IntSet::operator->() const {
  return static_cast<const IntSetNode*>(node_.get());
}
309
}  // namespace arith
310
}  // namespace tvm
311
#endif  // TVM_ARITHMETIC_H_