detect_linear_equation.cc 7.46 KB
Newer Older
1 2 3 4 5 6 7 8
/*
 * 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
9
 *
10
 *   http://www.apache.org/licenses/LICENSE-2.0
11
 *
12 13 14 15 16 17 18 19
 * 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.
 */

20
/*!
21 22
 * \file detect_linear_equation.cc
 * \brief Utility to detect patterns in the expression.
23
 */
24 25 26 27
#include <tvm/tir/expr.h>
#include <tvm/tir/ir_pass.h>
#include <tvm/tir/expr_functor.h>
#include <tvm/tir/stmt_functor.h>
28
#include <tvm/arith/analyzer.h>
29 30 31 32

namespace tvm {
namespace arith {

33
using namespace tir;
34 35 36

// Linear equation, the components can be undefined.
struct LinearEqEntry {
37 38
  PrimExpr base;
  PrimExpr coeff;
39 40
};

41
struct IntervalEntry {
42 43
  PrimExpr min_value;
  PrimExpr max_value;
44 45
};

46
class LinearEqDetector
47
    : public ExprFunctor<LinearEqEntry(const PrimExpr&, const PrimExpr &)> {
48 49 50 51
 public:
  explicit LinearEqDetector(Var var)
      : var_(var) {}

52
  bool Detect(const PrimExpr& e, LinearEqEntry* ret) {
53 54 55
    *ret = VisitExpr(e, e);
    if (fail_) return false;
    if (!ret->base.defined()) {
56
      ret->base = make_zero(var_.dtype());
57
    }
58
    if (!ret->coeff.defined()) {
59
      ret->coeff = make_zero(var_.dtype());
60
    }
61
    return true;
62 63
  }

64
  LinearEqEntry VisitExpr_(const AddNode* op, const PrimExpr& e) final {
65 66 67 68 69 70 71 72
    if (fail_) return LinearEqEntry();
    LinearEqEntry a = VisitExpr(op->a, op->a);
    LinearEqEntry b = VisitExpr(op->b, op->b);
    LinearEqEntry ret;
    ret.base = AddCombine(a.base, b.base);
    ret.coeff = AddCombine(a.coeff, b.coeff);
    return ret;
  }
73

74
  LinearEqEntry VisitExpr_(const SubNode* op, const PrimExpr& e) final {
75 76 77 78 79 80 81 82 83
    if (fail_) return LinearEqEntry();
    LinearEqEntry a = VisitExpr(op->a, op->a);
    LinearEqEntry b = VisitExpr(op->b, op->b);
    LinearEqEntry ret;
    ret.base = SubCombine(a.base, b.base);
    ret.coeff = SubCombine(a.coeff, b.coeff);
    return ret;
  }

84
  LinearEqEntry VisitExpr_(const MulNode* op, const PrimExpr& e) final {
85 86 87 88 89 90 91 92 93 94 95 96 97 98 99
    if (fail_) return LinearEqEntry();
    LinearEqEntry a = VisitExpr(op->a, op->a);
    LinearEqEntry b = VisitExpr(op->b, op->b);
    if (a.coeff.defined()) {
      std::swap(a, b);
    }
    if (a.coeff.defined()) {
      fail_ = true;
      return LinearEqEntry();
    }
    LinearEqEntry ret;
    ret.base = MulCombine(a.base, b.base);
    ret.coeff = MulCombine(a.base, b.coeff);
    return ret;
  }
100
  LinearEqEntry VisitExpr_(const VarNode* op, const PrimExpr& e) final {
101 102
    LinearEqEntry ret;
    if (op == var_.get()) {
103
      ret.coeff = make_const(op->dtype, 1);
104 105 106 107 108
    } else {
      ret.base = e;
    }
    return ret;
  }
109
  LinearEqEntry VisitExprDefault_(const Object* op, const PrimExpr& e) final {
110 111 112 113 114 115 116 117 118 119 120 121 122 123 124
    if (fail_) return LinearEqEntry();
    if (ExprUseVar(e, var_)) {
      fail_ = true;
      return LinearEqEntry();
    } else {
      LinearEqEntry ret;
      ret.base = e;
      return ret;
    }
  }

 private:
  Var var_;
  bool fail_{false};
  // Combine by add
125
  PrimExpr AddCombine(PrimExpr a, PrimExpr b) {
126 127
    if (!a.defined()) return b;
    if (!b.defined()) return a;
128
    return a + b;
129
  }
130
  PrimExpr SubCombine(PrimExpr a, PrimExpr b) {
131
    // Check b first in case they are both undefined
132
    if (!b.defined()) return a;
133
    if (!a.defined()) return -b;
134
    return a - b;
135
  }
136
  PrimExpr MulCombine(PrimExpr a, PrimExpr b) {
137 138
    if (!a.defined()) return a;
    if (!b.defined()) return b;
139
    return a * b;
140 141 142
  }
};

143 144 145 146
Array<PrimExpr> DetectLinearEquation(const PrimExpr& e,
                                          const Array<Var>& vars) {
  PrimExpr base = e;
  Array<PrimExpr> coeff;
147

148 149 150
  for (Var v : vars) {
    LinearEqEntry ret;
    if (!LinearEqDetector(v).Detect(base, &ret)) {
151
      return Array<PrimExpr>();
152
    }
153 154 155
    coeff.push_back(ret.coeff);
    base = std::move(ret.base);
  }
156

157
  std::unordered_set<const VarNode*> vset;
158 159 160 161
  for (size_t i = vars.size(); i > 1; --i) {
    vset.insert(vars[i - 1].get());
    // The previous coeff contains the variable
    if (ExprUseVar(coeff[i - 2], vset)) {
162
      return Array<PrimExpr>();
163 164 165 166
    }
  }
  coeff.push_back(base);
  return coeff;
167 168
}

169 170
// Detect clip condition as min max value
bool DetectClipBound(
171
    const PrimExpr& cond,
172
    std::unordered_map<const VarNode*, IntervalEntry>* bmap) {
173 174
  int flag = 0;
  Var var;
175
  auto fvisit = [&bmap, &flag, &var](const ObjectRef& n) {
176
    if (const VarNode* v = n.as<VarNode>()) {
177 178
      if (bmap->count(v)) {
        if (flag == 0) {
179
          var = Downcast<Var>(n);
180 181 182 183 184 185 186 187 188 189 190 191
          flag = 1;
        } else if (flag == 1) {
          if (!var.same_as(n)) {
            flag = -1;
          }
        }
      }
    }
  };
  PostOrderVisit(cond, fvisit);
  if (flag != 1) return false;
  // canonical form: exp >= 0
192
  PrimExpr canonical;
193
  if (const LTNode* op = cond.as<LTNode>()) {
194 195
    if (!op->a.dtype().is_int()) return false;
    canonical = op->b - op->a - make_const(op->a.dtype(), 1);
196
  } else if (const LENode* op = cond.as<LENode>()) {
197
    if (!op->a.dtype().is_int()) return false;
198
    canonical = op->b - op->a;
199
  } else if (const GTNode* op = cond.as<GTNode>()) {
200 201
    if (!op->a.dtype().is_int()) return false;
    canonical = op->a - op->b - make_const(op->a.dtype(), 1);
202
  } else if (const GENode* op = cond.as<GENode>()) {
203
    if (!op->a.dtype().is_int()) return false;
204 205 206 207 208 209 210 211
    canonical = op->a - op->b;
  } else {
    return false;
  }
  LinearEqEntry ret;
  if (!LinearEqDetector(var).Detect(canonical, &ret)) return false;
  ret.coeff = Simplify(ret.coeff);
  IntervalEntry& p = (*bmap)[var.get()];
212
  if (is_const_int(ret.coeff, 1)) {
213 214
    // var + shift >=0 -> var >= -shift
    if (p.min_value.defined()) {
215
      p.min_value = tir::MaxNode::make(p.min_value, -ret.base);
216 217 218 219 220
    } else {
      p.min_value = -ret.base;
    }
    return true;
  }
221
  if (is_const_int(ret.coeff, -1)) {
222 223
    // -var + shift >=0 -> var <= shift
    if (p.max_value.defined()) {
224
      p.max_value = tir::MinNode::make(p.max_value, ret.base);
225 226 227 228 229 230 231 232 233 234
    } else {
      p.max_value = ret.base;
    }
    return true;
  }
  return false;
}


template<typename OP>
235
void SplitCommExpr(const PrimExpr& e, std::vector<PrimExpr>* ret) {
236 237 238 239 240 241 242 243 244 245
  if (const OP* op = e.as<OP>()) {
    SplitCommExpr<OP>(op->a, ret);
    SplitCommExpr<OP>(op->b, ret);
  } else {
    ret->push_back(e);
  }
}

// Detect the lower and upper bound from the expression.
// e must be connected by and.
246 247
Array<PrimExpr> DetectClipBound(const PrimExpr& e, const Array<Var>& vars) {
  std::vector<PrimExpr> splits;
248
  SplitCommExpr<tir::AndNode>(e, &splits);
249
  std::unordered_map<const VarNode*, IntervalEntry> rmap;
250 251 252
  for (Var v : vars) {
    rmap[v.get()] = IntervalEntry();
  }
253 254
  for (PrimExpr cond : splits) {
    if (!DetectClipBound(cond, &rmap)) return Array<PrimExpr>();
255
  }
256
  Array<PrimExpr> ret;
257 258 259 260 261 262 263 264 265 266 267 268 269 270 271
  for (Var v : vars) {
    IntervalEntry e = rmap[v.get()];
    if (e.min_value.defined()) {
      e.min_value = Simplify(e.min_value);
    }
    if (e.max_value.defined()) {
      e.max_value = Simplify(e.max_value);
    }
    ret.push_back(e.min_value);
    ret.push_back(e.max_value);
  }
  return ret;
}


272 273
}  // namespace arith
}  // namespace tvm