MLIR  17.0.0git
AffineStructures.cpp
Go to the documentation of this file.
1 //===- AffineStructures.cpp - MLIR Affine Structures Class-----------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // Structures for affine/polyhedral analysis of affine dialect ops.
10 //
11 //===----------------------------------------------------------------------===//
12 
21 #include "mlir/IR/IntegerSet.h"
22 #include "mlir/Support/LLVM.h"
24 #include "llvm/ADT/STLExtras.h"
25 #include "llvm/ADT/SmallPtrSet.h"
26 #include "llvm/ADT/SmallVector.h"
27 #include "llvm/Support/Debug.h"
28 #include "llvm/Support/raw_ostream.h"
29 #include <optional>
30 
31 #define DEBUG_TYPE "affine-structures"
32 
33 using namespace mlir;
34 using namespace presburger;
35 
36 namespace {
37 
38 // See comments for SimpleAffineExprFlattener.
39 // An AffineExprFlattener extends a SimpleAffineExprFlattener by recording
40 // constraint information associated with mod's, floordiv's, and ceildiv's
41 // in FlatAffineValueConstraints 'localVarCst'.
42 struct AffineExprFlattener : public SimpleAffineExprFlattener {
43 public:
44  // Constraints connecting newly introduced local variables (for mod's and
45  // div's) to existing (dimensional and symbolic) ones. These are always
46  // inequalities.
47  IntegerPolyhedron localVarCst;
48 
49  AffineExprFlattener(unsigned nDims, unsigned nSymbols)
50  : SimpleAffineExprFlattener(nDims, nSymbols),
51  localVarCst(PresburgerSpace::getSetSpace(nDims, nSymbols)) {}
52 
53 private:
54  // Add a local variable (needed to flatten a mod, floordiv, ceildiv expr).
55  // The local variable added is always a floordiv of a pure add/mul affine
56  // function of other variables, coefficients of which are specified in
57  // `dividend' and with respect to the positive constant `divisor'. localExpr
58  // is the simplified tree expression (AffineExpr) corresponding to the
59  // quantifier.
60  void addLocalFloorDivId(ArrayRef<int64_t> dividend, int64_t divisor,
61  AffineExpr localExpr) override {
62  SimpleAffineExprFlattener::addLocalFloorDivId(dividend, divisor, localExpr);
63  // Update localVarCst.
64  localVarCst.addLocalFloorDiv(dividend, divisor);
65  }
66 };
67 
68 } // namespace
69 
70 // Flattens the expressions in map. Returns failure if 'expr' was unable to be
71 // flattened (i.e., semi-affine expressions not handled yet).
72 static LogicalResult
74  unsigned numSymbols,
75  std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
76  FlatAffineValueConstraints *localVarCst) {
77  if (exprs.empty()) {
78  localVarCst->reset(numDims, numSymbols);
79  return success();
80  }
81 
82  AffineExprFlattener flattener(numDims, numSymbols);
83  // Use the same flattener to simplify each expression successively. This way
84  // local variables / expressions are shared.
85  for (auto expr : exprs) {
86  if (!expr.isPureAffine())
87  return failure();
88 
89  flattener.walkPostOrder(expr);
90  }
91 
92  assert(flattener.operandExprStack.size() == exprs.size());
93  flattenedExprs->clear();
94  flattenedExprs->assign(flattener.operandExprStack.begin(),
95  flattener.operandExprStack.end());
96 
97  if (localVarCst)
98  localVarCst->clearAndCopyFrom(flattener.localVarCst);
99 
100  return success();
101 }
102 
103 // Flattens 'expr' into 'flattenedExpr'. Returns failure if 'expr' was unable to
104 // be flattened (semi-affine expressions not handled yet).
107  unsigned numSymbols,
108  SmallVectorImpl<int64_t> *flattenedExpr,
109  FlatAffineValueConstraints *localVarCst) {
110  std::vector<SmallVector<int64_t, 8>> flattenedExprs;
111  LogicalResult ret = ::getFlattenedAffineExprs({expr}, numDims, numSymbols,
112  &flattenedExprs, localVarCst);
113  *flattenedExpr = flattenedExprs[0];
114  return ret;
115 }
116 
117 /// Flattens the expressions in map. Returns failure if 'expr' was unable to be
118 /// flattened (i.e., semi-affine expressions not handled yet).
120  AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
121  FlatAffineValueConstraints *localVarCst) {
122  if (map.getNumResults() == 0) {
123  localVarCst->reset(map.getNumDims(), map.getNumSymbols());
124  return success();
125  }
127  map.getNumSymbols(), flattenedExprs,
128  localVarCst);
129 }
130 
132  IntegerSet set, std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
133  FlatAffineValueConstraints *localVarCst) {
134  if (set.getNumConstraints() == 0) {
135  localVarCst->reset(set.getNumDims(), set.getNumSymbols());
136  return success();
137  }
139  set.getNumSymbols(), flattenedExprs,
140  localVarCst);
141 }
142 
143 //===----------------------------------------------------------------------===//
144 // FlatAffineConstraints / FlatAffineValueConstraints.
145 //===----------------------------------------------------------------------===//
146 
147 std::unique_ptr<FlatAffineValueConstraints>
149  return std::make_unique<FlatAffineValueConstraints>(*this);
150 }
151 
152 // Construct from an IntegerSet.
154  ValueRange operands)
155  : IntegerPolyhedron(set.getNumInequalities(), set.getNumEqualities(),
156  set.getNumDims() + set.getNumSymbols() + 1,
157  PresburgerSpace::getSetSpace(set.getNumDims(),
158  set.getNumSymbols(),
159  /*numLocals=*/0)) {
160  // Populate values.
161  if (operands.empty()) {
162  values.resize(getNumDimAndSymbolVars(), std::nullopt);
163  } else {
164  assert(set.getNumInputs() == operands.size() && "operand count mismatch");
165  values.assign(operands.begin(), operands.end());
166  }
167 
168  // Flatten expressions and add them to the constraint system.
169  std::vector<SmallVector<int64_t, 8>> flatExprs;
170  FlatAffineValueConstraints localVarCst;
171  if (failed(getFlattenedAffineExprs(set, &flatExprs, &localVarCst))) {
172  assert(false && "flattening unimplemented for semi-affine integer sets");
173  return;
174  }
175  assert(flatExprs.size() == set.getNumConstraints());
176  insertVar(VarKind::Local, getNumVarKind(VarKind::Local),
177  /*num=*/localVarCst.getNumLocalVars());
178 
179  for (unsigned i = 0, e = flatExprs.size(); i < e; ++i) {
180  const auto &flatExpr = flatExprs[i];
181  assert(flatExpr.size() == getNumCols());
182  if (set.getEqFlags()[i]) {
183  addEquality(flatExpr);
184  } else {
185  addInequality(flatExpr);
186  }
187  }
188  // Add the other constraints involving local vars from flattening.
189  append(localVarCst);
190 }
191 
192 // Construct a hyperrectangular constraint set from ValueRanges that represent
193 // induction variables, lower and upper bounds. `ivs`, `lbs` and `ubs` are
194 // expected to match one to one. The order of variables and constraints is:
195 //
196 // ivs | lbs | ubs | eq/ineq
197 // ----+-----+-----+---------
198 // 1 -1 0 >= 0
199 // ----+-----+-----+---------
200 // -1 0 1 >= 0
201 //
202 // All dimensions as set as VarKind::SetDim.
205  ValueRange ubs) {
207  unsigned nIvs = ivs.size();
208  assert(nIvs == lbs.size() && "expected as many lower bounds as ivs");
209  assert(nIvs == ubs.size() && "expected as many upper bounds as ivs");
210 
211  if (nIvs == 0)
212  return res;
213 
214  res.appendDimVar(ivs);
215  unsigned lbsStart = res.appendDimVar(lbs);
216  unsigned ubsStart = res.appendDimVar(ubs);
217 
218  MLIRContext *ctx = ivs.front().getContext();
219  for (int ivIdx = 0, e = nIvs; ivIdx < e; ++ivIdx) {
220  // iv - lb >= 0
221  AffineMap lb = AffineMap::get(/*dimCount=*/3 * nIvs, /*symbolCount=*/0,
222  getAffineDimExpr(lbsStart + ivIdx, ctx));
223  if (failed(res.addBound(BoundType::LB, ivIdx, lb)))
224  llvm_unreachable("Unexpected FlatAffineValueConstraints creation error");
225  // -iv + ub >= 0
226  AffineMap ub = AffineMap::get(/*dimCount=*/3 * nIvs, /*symbolCount=*/0,
227  getAffineDimExpr(ubsStart + ivIdx, ctx));
228  if (failed(res.addBound(BoundType::UB, ivIdx, ub)))
229  llvm_unreachable("Unexpected FlatAffineValueConstraints creation error");
230  }
231  return res;
232 }
233 
234 void FlatAffineValueConstraints::reset(unsigned numReservedInequalities,
235  unsigned numReservedEqualities,
236  unsigned newNumReservedCols,
237  unsigned newNumDims,
238  unsigned newNumSymbols,
239  unsigned newNumLocals) {
240  assert(newNumReservedCols >= newNumDims + newNumSymbols + newNumLocals + 1 &&
241  "minimum 1 column");
242  *this = FlatAffineValueConstraints(numReservedInequalities,
243  numReservedEqualities, newNumReservedCols,
244  newNumDims, newNumSymbols, newNumLocals);
245 }
246 
247 void FlatAffineValueConstraints::reset(unsigned newNumDims,
248  unsigned newNumSymbols,
249  unsigned newNumLocals) {
250  reset(/*numReservedInequalities=*/0, /*numReservedEqualities=*/0,
251  /*numReservedCols=*/newNumDims + newNumSymbols + newNumLocals + 1,
252  newNumDims, newNumSymbols, newNumLocals);
253 }
254 
256  unsigned numReservedInequalities, unsigned numReservedEqualities,
257  unsigned newNumReservedCols, unsigned newNumDims, unsigned newNumSymbols,
258  unsigned newNumLocals, ArrayRef<Value> valArgs) {
259  assert(newNumReservedCols >= newNumDims + newNumSymbols + newNumLocals + 1 &&
260  "minimum 1 column");
262  if (!valArgs.empty())
263  newVals.assign(valArgs.begin(), valArgs.end());
264 
266  numReservedInequalities, numReservedEqualities, newNumReservedCols,
267  newNumDims, newNumSymbols, newNumLocals, newVals);
268 }
269 
270 void FlatAffineValueConstraints::reset(unsigned newNumDims,
271  unsigned newNumSymbols,
272  unsigned newNumLocals,
273  ArrayRef<Value> valArgs) {
274  reset(0, 0, newNumDims + newNumSymbols + newNumLocals + 1, newNumDims,
275  newNumSymbols, newNumLocals, valArgs);
276 }
277 
279  unsigned pos = getNumDimVars();
280  return insertVar(VarKind::SetDim, pos, vals);
281 }
282 
284  unsigned pos = getNumSymbolVars();
285  return insertVar(VarKind::Symbol, pos, vals);
286 }
287 
289  ValueRange vals) {
290  return insertVar(VarKind::SetDim, pos, vals);
291 }
292 
294  ValueRange vals) {
295  return insertVar(VarKind::Symbol, pos, vals);
296 }
297 
298 unsigned FlatAffineValueConstraints::insertVar(VarKind kind, unsigned pos,
299  unsigned num) {
300  unsigned absolutePos = IntegerPolyhedron::insertVar(kind, pos, num);
301 
302  if (kind != VarKind::Local) {
303  values.insert(values.begin() + absolutePos, num, std::nullopt);
304  assert(values.size() == getNumDimAndSymbolVars());
305  }
306 
307  return absolutePos;
308 }
309 
310 unsigned FlatAffineValueConstraints::insertVar(VarKind kind, unsigned pos,
311  ValueRange vals) {
312  assert(!vals.empty() && "expected ValueRange with Values.");
313  assert(kind != VarKind::Local &&
314  "values cannot be attached to local variables.");
315  unsigned num = vals.size();
316  unsigned absolutePos = IntegerPolyhedron::insertVar(kind, pos, num);
317 
318  // If a Value is provided, insert it; otherwise use None.
319  for (unsigned i = 0; i < num; ++i)
320  values.insert(values.begin() + absolutePos + i,
321  vals[i] ? std::optional<Value>(vals[i]) : std::nullopt);
322 
323  assert(values.size() == getNumDimAndSymbolVars());
324  return absolutePos;
325 }
326 
328  return llvm::any_of(
329  values, [](const std::optional<Value> &var) { return var.has_value(); });
330 }
331 
332 /// Checks if two constraint systems are in the same space, i.e., if they are
333 /// associated with the same set of variables, appearing in the same order.
335  const FlatAffineValueConstraints &b) {
336  return a.getNumDimVars() == b.getNumDimVars() &&
337  a.getNumSymbolVars() == b.getNumSymbolVars() &&
338  a.getNumVars() == b.getNumVars() &&
339  a.getMaybeValues().equals(b.getMaybeValues());
340 }
341 
342 /// Calls areVarsAligned to check if two constraint systems have the same set
343 /// of variables in the same order.
345  const FlatAffineValueConstraints &other) {
346  return areVarsAligned(*this, other);
347 }
348 
349 /// Checks if the SSA values associated with `cst`'s variables in range
350 /// [start, end) are unique.
351 static bool LLVM_ATTRIBUTE_UNUSED areVarsUnique(
352  const FlatAffineValueConstraints &cst, unsigned start, unsigned end) {
353 
354  assert(start <= cst.getNumDimAndSymbolVars() &&
355  "Start position out of bounds");
356  assert(end <= cst.getNumDimAndSymbolVars() && "End position out of bounds");
357 
358  if (start >= end)
359  return true;
360 
361  SmallPtrSet<Value, 8> uniqueVars;
362  ArrayRef<std::optional<Value>> maybeValues =
363  cst.getMaybeValues().slice(start, end - start);
364  for (std::optional<Value> val : maybeValues) {
365  if (val && !uniqueVars.insert(*val).second)
366  return false;
367  }
368  return true;
369 }
370 
371 /// Checks if the SSA values associated with `cst`'s variables are unique.
372 static bool LLVM_ATTRIBUTE_UNUSED
374  return areVarsUnique(cst, 0, cst.getNumDimAndSymbolVars());
375 }
376 
377 /// Checks if the SSA values associated with `cst`'s variables of kind `kind`
378 /// are unique.
379 static bool LLVM_ATTRIBUTE_UNUSED
381 
382  if (kind == VarKind::SetDim)
383  return areVarsUnique(cst, 0, cst.getNumDimVars());
384  if (kind == VarKind::Symbol)
385  return areVarsUnique(cst, cst.getNumDimVars(),
386  cst.getNumDimAndSymbolVars());
387  llvm_unreachable("Unexpected VarKind");
388 }
389 
390 /// Merge and align the variables of A and B starting at 'offset', so that
391 /// both constraint systems get the union of the contained variables that is
392 /// dimension-wise and symbol-wise unique; both constraint systems are updated
393 /// so that they have the union of all variables, with A's original
394 /// variables appearing first followed by any of B's variables that didn't
395 /// appear in A. Local variables in B that have the same division
396 /// representation as local variables in A are merged into one.
397 // E.g.: Input: A has ((%i, %j) [%M, %N]) and B has (%k, %j) [%P, %N, %M])
398 // Output: both A, B have (%i, %j, %k) [%M, %N, %P]
399 static void mergeAndAlignVars(unsigned offset, FlatAffineValueConstraints *a,
401  assert(offset <= a->getNumDimVars() && offset <= b->getNumDimVars());
402  // A merge/align isn't meaningful if a cst's vars aren't distinct.
403  assert(areVarsUnique(*a) && "A's values aren't unique");
404  assert(areVarsUnique(*b) && "B's values aren't unique");
405 
406  assert(llvm::all_of(
407  llvm::drop_begin(a->getMaybeValues(), offset),
408  [](const std::optional<Value> &var) { return var.has_value(); }));
409 
410  assert(llvm::all_of(
411  llvm::drop_begin(b->getMaybeValues(), offset),
412  [](const std::optional<Value> &var) { return var.has_value(); }));
413 
414  SmallVector<Value, 4> aDimValues;
415  a->getValues(offset, a->getNumDimVars(), &aDimValues);
416 
417  {
418  // Merge dims from A into B.
419  unsigned d = offset;
420  for (auto aDimValue : aDimValues) {
421  unsigned loc;
422  if (b->findVar(aDimValue, &loc)) {
423  assert(loc >= offset && "A's dim appears in B's aligned range");
424  assert(loc < b->getNumDimVars() &&
425  "A's dim appears in B's non-dim position");
426  b->swapVar(d, loc);
427  } else {
428  b->insertDimVar(d, aDimValue);
429  }
430  d++;
431  }
432  // Dimensions that are in B, but not in A, are added at the end.
433  for (unsigned t = a->getNumDimVars(), e = b->getNumDimVars(); t < e; t++) {
434  a->appendDimVar(b->getValue(t));
435  }
436  assert(a->getNumDimVars() == b->getNumDimVars() &&
437  "expected same number of dims");
438  }
439 
440  // Merge and align symbols of A and B
441  a->mergeSymbolVars(*b);
442  // Merge and align locals of A and B
443  a->mergeLocalVars(*b);
444 
445  assert(areVarsAligned(*a, *b) && "IDs expected to be aligned");
446 }
447 
448 // Call 'mergeAndAlignVars' to align constraint systems of 'this' and 'other'.
450  unsigned offset, FlatAffineValueConstraints *other) {
451  mergeAndAlignVars(offset, this, other);
452 }
453 
456  return composeMatchingMap(
457  computeAlignedMap(vMap->getAffineMap(), vMap->getOperands()));
458 }
459 
460 // Similar to `composeMap` except that no Values need be associated with the
461 // constraint system nor are they looked at -- the dimensions and symbols of
462 // `other` are expected to correspond 1:1 to `this` system.
464  assert(other.getNumDims() == getNumDimVars() && "dim mismatch");
465  assert(other.getNumSymbols() == getNumSymbolVars() && "symbol mismatch");
466 
467  std::vector<SmallVector<int64_t, 8>> flatExprs;
468  if (failed(flattenAlignedMapAndMergeLocals(other, &flatExprs)))
469  return failure();
470  assert(flatExprs.size() == other.getNumResults());
471 
472  // Add dimensions corresponding to the map's results.
473  insertDimVar(/*pos=*/0, /*num=*/other.getNumResults());
474 
475  // We add one equality for each result connecting the result dim of the map to
476  // the other variables.
477  // E.g.: if the expression is 16*i0 + i1, and this is the r^th
478  // iteration/result of the value map, we are adding the equality:
479  // d_r - 16*i0 - i1 = 0. Similarly, when flattening (i0 + 1, i0 + 8*i2), we
480  // add two equalities: d_0 - i0 - 1 == 0, d1 - i0 - 8*i2 == 0.
481  for (unsigned r = 0, e = flatExprs.size(); r < e; r++) {
482  const auto &flatExpr = flatExprs[r];
483  assert(flatExpr.size() >= other.getNumInputs() + 1);
484 
485  SmallVector<int64_t, 8> eqToAdd(getNumCols(), 0);
486  // Set the coefficient for this result to one.
487  eqToAdd[r] = 1;
488 
489  // Dims and symbols.
490  for (unsigned i = 0, f = other.getNumInputs(); i < f; i++) {
491  // Negate `eq[r]` since the newly added dimension will be set to this one.
492  eqToAdd[e + i] = -flatExpr[i];
493  }
494  // Local columns of `eq` are at the beginning.
495  unsigned j = getNumDimVars() + getNumSymbolVars();
496  unsigned end = flatExpr.size() - 1;
497  for (unsigned i = other.getNumInputs(); i < end; i++, j++) {
498  eqToAdd[j] = -flatExpr[i];
499  }
500 
501  // Constant term.
502  eqToAdd[getNumCols() - 1] = -flatExpr[flatExpr.size() - 1];
503 
504  // Add the equality connecting the result of the map to this constraint set.
505  addEquality(eqToAdd);
506  }
507 
508  return success();
509 }
510 
511 // Turn a symbol into a dimension.
513  unsigned pos;
514  if (cst->findVar(value, &pos) && pos >= cst->getNumDimVars() &&
515  pos < cst->getNumDimAndSymbolVars()) {
516  cst->swapVar(pos, cst->getNumDimVars());
517  cst->setDimSymbolSeparation(cst->getNumSymbolVars() - 1);
518  }
519 }
520 
521 /// Merge and align symbols of `this` and `other` such that both get union of
522 /// of symbols that are unique. Symbols in `this` and `other` should be
523 /// unique. Symbols with Value as `None` are considered to be inequal to all
524 /// other symbols.
527 
528  assert(areVarsUnique(*this, VarKind::Symbol) && "Symbol vars are not unique");
529  assert(areVarsUnique(other, VarKind::Symbol) && "Symbol vars are not unique");
530 
531  SmallVector<Value, 4> aSymValues;
533 
534  // Merge symbols: merge symbols into `other` first from `this`.
535  unsigned s = other.getNumDimVars();
536  for (Value aSymValue : aSymValues) {
537  unsigned loc;
538  // If the var is a symbol in `other`, then align it, otherwise assume that
539  // it is a new symbol
540  if (other.findVar(aSymValue, &loc) && loc >= other.getNumDimVars() &&
541  loc < other.getNumDimAndSymbolVars())
542  other.swapVar(s, loc);
543  else
544  other.insertSymbolVar(s - other.getNumDimVars(), aSymValue);
545  s++;
546  }
547 
548  // Symbols that are in other, but not in this, are added at the end.
549  for (unsigned t = other.getNumDimVars() + getNumSymbolVars(),
550  e = other.getNumDimAndSymbolVars();
551  t < e; t++)
553 
554  assert(getNumSymbolVars() == other.getNumSymbolVars() &&
555  "expected same number of symbols");
556  assert(areVarsUnique(*this, VarKind::Symbol) && "Symbol vars are not unique");
557  assert(areVarsUnique(other, VarKind::Symbol) && "Symbol vars are not unique");
558 }
559 
560 // Changes all symbol variables which are loop IVs to dim variables.
562  // Gather all symbols which are loop IVs.
563  SmallVector<Value, 4> loopIVs;
564  for (unsigned i = getNumDimVars(), e = getNumDimAndSymbolVars(); i < e; i++) {
566  loopIVs.push_back(getValue(i));
567  }
568  // Turn each symbol in 'loopIVs' into a dim variable.
569  for (auto iv : loopIVs) {
570  turnSymbolIntoDim(this, iv);
571  }
572 }
573 
575  if (containsVar(val))
576  return;
577 
578  // Caller is expected to fully compose map/operands if necessary.
579  assert((isTopLevelValue(val) || isAffineForInductionVar(val)) &&
580  "non-terminal symbol / loop IV expected");
581  // Outer loop IVs could be used in forOp's bounds.
582  if (auto loop = getForInductionVarOwner(val)) {
583  appendDimVar(val);
584  if (failed(this->addAffineForOpDomain(loop)))
585  LLVM_DEBUG(
586  loop.emitWarning("failed to add domain info to constraint system"));
587  return;
588  }
589  // Add top level symbol.
590  appendSymbolVar(val);
591  // Check if the symbol is a constant.
592  if (auto constOp = val.getDefiningOp<arith::ConstantIndexOp>())
593  addBound(BoundType::EQ, val, constOp.value());
594 }
595 
598  unsigned pos;
599  // Pre-condition for this method.
600  if (!findVar(forOp.getInductionVar(), &pos)) {
601  assert(false && "Value not found");
602  return failure();
603  }
604 
605  int64_t step = forOp.getStep();
606  if (step != 1) {
607  if (!forOp.hasConstantLowerBound())
608  LLVM_DEBUG(forOp.emitWarning("domain conservatively approximated"));
609  else {
610  // Add constraints for the stride.
611  // (iv - lb) % step = 0 can be written as:
612  // (iv - lb) - step * q = 0 where q = (iv - lb) / step.
613  // Add local variable 'q' and add the above equality.
614  // The first constraint is q = (iv - lb) floordiv step
615  SmallVector<int64_t, 8> dividend(getNumCols(), 0);
616  int64_t lb = forOp.getConstantLowerBound();
617  dividend[pos] = 1;
618  dividend.back() -= lb;
619  addLocalFloorDiv(dividend, step);
620  // Second constraint: (iv - lb) - step * q = 0.
622  eq[pos] = 1;
623  eq.back() -= lb;
624  // For the local var just added above.
625  eq[getNumCols() - 2] = -step;
626  addEquality(eq);
627  }
628  }
629 
630  if (forOp.hasConstantLowerBound()) {
631  addBound(BoundType::LB, pos, forOp.getConstantLowerBound());
632  } else {
633  // Non-constant lower bound case.
634  if (failed(addBound(BoundType::LB, pos, forOp.getLowerBoundMap(),
635  forOp.getLowerBoundOperands())))
636  return failure();
637  }
638 
639  if (forOp.hasConstantUpperBound()) {
640  addBound(BoundType::UB, pos, forOp.getConstantUpperBound() - 1);
641  return success();
642  }
643  // Non-constant upper bound case.
644  return addBound(BoundType::UB, pos, forOp.getUpperBoundMap(),
645  forOp.getUpperBoundOperands());
646 }
647 
649  AffineParallelOp parallelOp) {
650  size_t ivPos = 0;
651  for (auto iv : parallelOp.getIVs()) {
652  unsigned pos;
653  if (!findVar(iv, &pos)) {
654  assert(false && "variable expected for the IV value");
655  return failure();
656  }
657 
658  AffineMap lowerBound = parallelOp.getLowerBoundMap(ivPos);
659  if (lowerBound.isConstant())
660  addBound(BoundType::LB, pos, lowerBound.getSingleConstantResult());
661  else if (failed(addBound(BoundType::LB, pos, lowerBound,
662  parallelOp.getLowerBoundsOperands())))
663  return failure();
664 
665  auto upperBound = parallelOp.getUpperBoundMap(ivPos);
666  if (upperBound.isConstant())
667  addBound(BoundType::UB, pos, upperBound.getSingleConstantResult());
668  else if (failed(addBound(BoundType::UB, pos, upperBound,
669  parallelOp.getUpperBoundsOperands())))
670  return failure();
671  }
672  return success();
673 }
674 
677  ArrayRef<AffineMap> ubMaps,
678  ArrayRef<Value> operands) {
679  assert(lbMaps.size() == ubMaps.size());
680  assert(lbMaps.size() <= getNumDimVars());
681 
682  for (unsigned i = 0, e = lbMaps.size(); i < e; ++i) {
683  AffineMap lbMap = lbMaps[i];
684  AffineMap ubMap = ubMaps[i];
685  assert(!lbMap || lbMap.getNumInputs() == operands.size());
686  assert(!ubMap || ubMap.getNumInputs() == operands.size());
687 
688  // Check if this slice is just an equality along this dimension. If so,
689  // retrieve the existing loop it equates to and add it to the system.
690  if (lbMap && ubMap && lbMap.getNumResults() == 1 &&
691  ubMap.getNumResults() == 1 &&
692  lbMap.getResult(0) + 1 == ubMap.getResult(0) &&
693  // The condition above will be true for maps describing a single
694  // iteration (e.g., lbMap.getResult(0) = 0, ubMap.getResult(0) = 1).
695  // Make sure we skip those cases by checking that the lb result is not
696  // just a constant.
697  !lbMap.getResult(0).isa<AffineConstantExpr>()) {
698  // Limited support: we expect the lb result to be just a loop dimension.
699  // Not supported otherwise for now.
700  AffineDimExpr result = lbMap.getResult(0).dyn_cast<AffineDimExpr>();
701  if (!result)
702  return failure();
703 
704  AffineForOp loop =
705  getForInductionVarOwner(operands[result.getPosition()]);
706  if (!loop)
707  return failure();
708 
709  if (failed(addAffineForOpDomain(loop)))
710  return failure();
711  continue;
712  }
713 
714  // This slice refers to a loop that doesn't exist in the IR yet. Add its
715  // bounds to the system assuming its dimension variable position is the
716  // same as the position of the loop in the loop nest.
717  if (lbMap && failed(addBound(BoundType::LB, i, lbMap, operands)))
718  return failure();
719  if (ubMap && failed(addBound(BoundType::UB, i, ubMap, operands)))
720  return failure();
721  }
722  return success();
723 }
724 
726  IntegerSet set = ifOp.getIntegerSet();
727  // Canonicalize set and operands to ensure unique values for
728  // FlatAffineValueConstraints below and for early simplification.
729  SmallVector<Value> operands(ifOp.getOperands());
730  canonicalizeSetAndOperands(&set, &operands);
731 
732  // Create the base constraints from the integer set attached to ifOp.
733  FlatAffineValueConstraints cst(set, operands);
734 
735  // Merge the constraints from ifOp to the current domain. We need first merge
736  // and align the IDs from both constraints, and then append the constraints
737  // from the ifOp into the current one.
739  append(cst);
740 }
741 
743  return IntegerPolyhedron::hasConsistentState() &&
744  values.size() == getNumDimAndSymbolVars();
745 }
746 
748  unsigned varLimit) {
749  IntegerPolyhedron::removeVarRange(kind, varStart, varLimit);
750  unsigned offset = getVarKindOffset(kind);
751 
752  if (kind != VarKind::Local) {
753  values.erase(values.begin() + varStart + offset,
754  values.begin() + varLimit + offset);
755  }
756 }
757 
758 // Determine whether the variable at 'pos' (say var_r) can be expressed as
759 // modulo of another known variable (say var_n) w.r.t a constant. For example,
760 // if the following constraints hold true:
761 // ```
762 // 0 <= var_r <= divisor - 1
763 // var_n - (divisor * q_expr) = var_r
764 // ```
765 // where `var_n` is a known variable (called dividend), and `q_expr` is an
766 // `AffineExpr` (called the quotient expression), `var_r` can be written as:
767 //
768 // `var_r = var_n mod divisor`.
769 //
770 // Additionally, in a special case of the above constaints where `q_expr` is an
771 // variable itself that is not yet known (say `var_q`), it can be written as a
772 // floordiv in the following way:
773 //
774 // `var_q = var_n floordiv divisor`.
775 //
776 // Returns true if the above mod or floordiv are detected, updating 'memo' with
777 // these new expressions. Returns false otherwise.
778 static bool detectAsMod(const FlatAffineValueConstraints &cst, unsigned pos,
779  int64_t lbConst, int64_t ubConst,
781  MLIRContext *context) {
782  assert(pos < cst.getNumVars() && "invalid position");
783 
784  // Check if a divisor satisfying the condition `0 <= var_r <= divisor - 1` can
785  // be determined.
786  if (lbConst != 0 || ubConst < 1)
787  return false;
788  int64_t divisor = ubConst + 1;
789 
790  // Check for the aforementioned conditions in each equality.
791  for (unsigned curEquality = 0, numEqualities = cst.getNumEqualities();
792  curEquality < numEqualities; curEquality++) {
793  int64_t coefficientAtPos = cst.atEq64(curEquality, pos);
794  // If current equality does not involve `var_r`, continue to the next
795  // equality.
796  if (coefficientAtPos == 0)
797  continue;
798 
799  // Constant term should be 0 in this equality.
800  if (cst.atEq64(curEquality, cst.getNumCols() - 1) != 0)
801  continue;
802 
803  // Traverse through the equality and construct the dividend expression
804  // `dividendExpr`, to contain all the variables which are known and are
805  // not divisible by `(coefficientAtPos * divisor)`. Hope here is that the
806  // `dividendExpr` gets simplified into a single variable `var_n` discussed
807  // above.
808  auto dividendExpr = getAffineConstantExpr(0, context);
809 
810  // Track the terms that go into quotient expression, later used to detect
811  // additional floordiv.
812  unsigned quotientCount = 0;
813  int quotientPosition = -1;
814  int quotientSign = 1;
815 
816  // Consider each term in the current equality.
817  unsigned curVar, e;
818  for (curVar = 0, e = cst.getNumDimAndSymbolVars(); curVar < e; ++curVar) {
819  // Ignore var_r.
820  if (curVar == pos)
821  continue;
822  int64_t coefficientOfCurVar = cst.atEq64(curEquality, curVar);
823  // Ignore vars that do not contribute to the current equality.
824  if (coefficientOfCurVar == 0)
825  continue;
826  // Check if the current var goes into the quotient expression.
827  if (coefficientOfCurVar % (divisor * coefficientAtPos) == 0) {
828  quotientCount++;
829  quotientPosition = curVar;
830  quotientSign = (coefficientOfCurVar * coefficientAtPos) > 0 ? 1 : -1;
831  continue;
832  }
833  // Variables that are part of dividendExpr should be known.
834  if (!memo[curVar])
835  break;
836  // Append the current variable to the dividend expression.
837  dividendExpr = dividendExpr + memo[curVar] * coefficientOfCurVar;
838  }
839 
840  // Can't construct expression as it depends on a yet uncomputed var.
841  if (curVar < e)
842  continue;
843 
844  // Express `var_r` in terms of the other vars collected so far.
845  if (coefficientAtPos > 0)
846  dividendExpr = (-dividendExpr).floorDiv(coefficientAtPos);
847  else
848  dividendExpr = dividendExpr.floorDiv(-coefficientAtPos);
849 
850  // Simplify the expression.
851  dividendExpr = simplifyAffineExpr(dividendExpr, cst.getNumDimVars(),
852  cst.getNumSymbolVars());
853  // Only if the final dividend expression is just a single var (which we call
854  // `var_n`), we can proceed.
855  // TODO: Handle AffineSymbolExpr as well. There is no reason to restrict it
856  // to dims themselves.
857  auto dimExpr = dividendExpr.dyn_cast<AffineDimExpr>();
858  if (!dimExpr)
859  continue;
860 
861  // Express `var_r` as `var_n % divisor` and store the expression in `memo`.
862  if (quotientCount >= 1) {
863  auto ub = cst.getConstantBound64(
864  FlatAffineValueConstraints::BoundType::UB, dimExpr.getPosition());
865  // If `var_n` has an upperbound that is less than the divisor, mod can be
866  // eliminated altogether.
867  if (ub && *ub < divisor)
868  memo[pos] = dimExpr;
869  else
870  memo[pos] = dimExpr % divisor;
871  // If a unique quotient `var_q` was seen, it can be expressed as
872  // `var_n floordiv divisor`.
873  if (quotientCount == 1 && !memo[quotientPosition])
874  memo[quotientPosition] = dimExpr.floorDiv(divisor) * quotientSign;
875 
876  return true;
877  }
878  }
879  return false;
880 }
881 
882 /// Check if the pos^th variable can be expressed as a floordiv of an affine
883 /// function of other variables (where the divisor is a positive constant)
884 /// given the initial set of expressions in `exprs`. If it can be, the
885 /// corresponding position in `exprs` is set as the detected affine expr. For
886 /// eg: 4q <= i + j <= 4q + 3 <=> q = (i + j) floordiv 4. An equality can
887 /// also yield a floordiv: eg. 4q = i + j <=> q = (i + j) floordiv 4. 32q + 28
888 /// <= i <= 32q + 31 => q = i floordiv 32.
890  unsigned pos, MLIRContext *context,
892  assert(pos < cst.getNumVars() && "invalid position");
893 
894  // Get upper-lower bound pair for this variable.
895  SmallVector<bool, 8> foundRepr(cst.getNumVars(), false);
896  for (unsigned i = 0, e = cst.getNumVars(); i < e; ++i)
897  if (exprs[i])
898  foundRepr[i] = true;
899 
900  SmallVector<int64_t, 8> dividend(cst.getNumCols());
901  unsigned divisor;
902  auto ulPair = computeSingleVarRepr(cst, foundRepr, pos, dividend, divisor);
903 
904  // No upper-lower bound pair found for this var.
905  if (ulPair.kind == ReprKind::None || ulPair.kind == ReprKind::Equality)
906  return false;
907 
908  // Construct the dividend expression.
909  auto dividendExpr = getAffineConstantExpr(dividend.back(), context);
910  for (unsigned c = 0, f = cst.getNumVars(); c < f; c++)
911  if (dividend[c] != 0)
912  dividendExpr = dividendExpr + dividend[c] * exprs[c];
913 
914  // Successfully detected the floordiv.
915  exprs[pos] = dividendExpr.floorDiv(divisor);
916  return true;
917 }
918 
919 std::pair<AffineMap, AffineMap>
921  unsigned pos, unsigned offset, unsigned num, unsigned symStartPos,
922  ArrayRef<AffineExpr> localExprs, MLIRContext *context) const {
923  assert(pos + offset < getNumDimVars() && "invalid dim start pos");
924  assert(symStartPos >= (pos + offset) && "invalid sym start pos");
925  assert(getNumLocalVars() == localExprs.size() &&
926  "incorrect local exprs count");
927 
928  SmallVector<unsigned, 4> lbIndices, ubIndices, eqIndices;
929  getLowerAndUpperBoundIndices(pos + offset, &lbIndices, &ubIndices, &eqIndices,
930  offset, num);
931 
932  /// Add to 'b' from 'a' in set [0, offset) U [offset + num, symbStartPos).
933  auto addCoeffs = [&](ArrayRef<int64_t> a, SmallVectorImpl<int64_t> &b) {
934  b.clear();
935  for (unsigned i = 0, e = a.size(); i < e; ++i) {
936  if (i < offset || i >= offset + num)
937  b.push_back(a[i]);
938  }
939  };
940 
943  unsigned dimCount = symStartPos - num;
944  unsigned symCount = getNumDimAndSymbolVars() - symStartPos;
945  lbExprs.reserve(lbIndices.size() + eqIndices.size());
946  // Lower bound expressions.
947  for (auto idx : lbIndices) {
948  auto ineq = getInequality64(idx);
949  // Extract the lower bound (in terms of other coeff's + const), i.e., if
950  // i - j + 1 >= 0 is the constraint, 'pos' is for i the lower bound is j
951  // - 1.
952  addCoeffs(ineq, lb);
953  std::transform(lb.begin(), lb.end(), lb.begin(), std::negate<int64_t>());
954  auto expr =
955  getAffineExprFromFlatForm(lb, dimCount, symCount, localExprs, context);
956  // expr ceildiv divisor is (expr + divisor - 1) floordiv divisor
957  int64_t divisor = std::abs(ineq[pos + offset]);
958  expr = (expr + divisor - 1).floorDiv(divisor);
959  lbExprs.push_back(expr);
960  }
961 
963  ubExprs.reserve(ubIndices.size() + eqIndices.size());
964  // Upper bound expressions.
965  for (auto idx : ubIndices) {
966  auto ineq = getInequality64(idx);
967  // Extract the upper bound (in terms of other coeff's + const).
968  addCoeffs(ineq, ub);
969  auto expr =
970  getAffineExprFromFlatForm(ub, dimCount, symCount, localExprs, context);
971  expr = expr.floorDiv(std::abs(ineq[pos + offset]));
972  // Upper bound is exclusive.
973  ubExprs.push_back(expr + 1);
974  }
975 
976  // Equalities. It's both a lower and a upper bound.
978  for (auto idx : eqIndices) {
979  auto eq = getEquality64(idx);
980  addCoeffs(eq, b);
981  if (eq[pos + offset] > 0)
982  std::transform(b.begin(), b.end(), b.begin(), std::negate<int64_t>());
983 
984  // Extract the upper bound (in terms of other coeff's + const).
985  auto expr =
986  getAffineExprFromFlatForm(b, dimCount, symCount, localExprs, context);
987  expr = expr.floorDiv(std::abs(eq[pos + offset]));
988  // Upper bound is exclusive.
989  ubExprs.push_back(expr + 1);
990  // Lower bound.
991  expr =
992  getAffineExprFromFlatForm(b, dimCount, symCount, localExprs, context);
993  expr = expr.ceilDiv(std::abs(eq[pos + offset]));
994  lbExprs.push_back(expr);
995  }
996 
997  auto lbMap = AffineMap::get(dimCount, symCount, lbExprs, context);
998  auto ubMap = AffineMap::get(dimCount, symCount, ubExprs, context);
999 
1000  return {lbMap, ubMap};
1001 }
1002 
1003 /// Computes the lower and upper bounds of the first 'num' dimensional
1004 /// variables (starting at 'offset') as affine maps of the remaining
1005 /// variables (dimensional and symbolic variables). Local variables are
1006 /// themselves explicitly computed as affine functions of other variables in
1007 /// this process if needed.
1009  unsigned offset, unsigned num, MLIRContext *context,
1011  bool getClosedUB) {
1012  assert(num < getNumDimVars() && "invalid range");
1013 
1014  // Basic simplification.
1016 
1017  LLVM_DEBUG(llvm::dbgs() << "getSliceBounds for first " << num
1018  << " variables\n");
1019  LLVM_DEBUG(dump());
1020 
1021  // Record computed/detected variables.
1023  // Initialize dimensional and symbolic variables.
1024  for (unsigned i = 0, e = getNumDimVars(); i < e; i++) {
1025  if (i < offset)
1026  memo[i] = getAffineDimExpr(i, context);
1027  else if (i >= offset + num)
1028  memo[i] = getAffineDimExpr(i - num, context);
1029  }
1030  for (unsigned i = getNumDimVars(), e = getNumDimAndSymbolVars(); i < e; i++)
1031  memo[i] = getAffineSymbolExpr(i - getNumDimVars(), context);
1032 
1033  bool changed;
1034  do {
1035  changed = false;
1036  // Identify yet unknown variables as constants or mod's / floordiv's of
1037  // other variables if possible.
1038  for (unsigned pos = 0; pos < getNumVars(); pos++) {
1039  if (memo[pos])
1040  continue;
1041 
1042  auto lbConst = getConstantBound64(BoundType::LB, pos);
1043  auto ubConst = getConstantBound64(BoundType::UB, pos);
1044  if (lbConst.has_value() && ubConst.has_value()) {
1045  // Detect equality to a constant.
1046  if (*lbConst == *ubConst) {
1047  memo[pos] = getAffineConstantExpr(*lbConst, context);
1048  changed = true;
1049  continue;
1050  }
1051 
1052  // Detect an variable as modulo of another variable w.r.t a
1053  // constant.
1054  if (detectAsMod(*this, pos, *lbConst, *ubConst, memo, context)) {
1055  changed = true;
1056  continue;
1057  }
1058  }
1059 
1060  // Detect an variable as a floordiv of an affine function of other
1061  // variables (divisor is a positive constant).
1062  if (detectAsFloorDiv(*this, pos, context, memo)) {
1063  changed = true;
1064  continue;
1065  }
1066 
1067  // Detect an variable as an expression of other variables.
1068  unsigned idx;
1069  if (!findConstraintWithNonZeroAt(pos, /*isEq=*/true, &idx)) {
1070  continue;
1071  }
1072 
1073  // Build AffineExpr solving for variable 'pos' in terms of all others.
1074  auto expr = getAffineConstantExpr(0, context);
1075  unsigned j, e;
1076  for (j = 0, e = getNumVars(); j < e; ++j) {
1077  if (j == pos)
1078  continue;
1079  int64_t c = atEq64(idx, j);
1080  if (c == 0)
1081  continue;
1082  // If any of the involved IDs hasn't been found yet, we can't proceed.
1083  if (!memo[j])
1084  break;
1085  expr = expr + memo[j] * c;
1086  }
1087  if (j < e)
1088  // Can't construct expression as it depends on a yet uncomputed
1089  // variable.
1090  continue;
1091 
1092  // Add constant term to AffineExpr.
1093  expr = expr + atEq64(idx, getNumVars());
1094  int64_t vPos = atEq64(idx, pos);
1095  assert(vPos != 0 && "expected non-zero here");
1096  if (vPos > 0)
1097  expr = (-expr).floorDiv(vPos);
1098  else
1099  // vPos < 0.
1100  expr = expr.floorDiv(-vPos);
1101  // Successfully constructed expression.
1102  memo[pos] = expr;
1103  changed = true;
1104  }
1105  // This loop is guaranteed to reach a fixed point - since once an
1106  // variable's explicit form is computed (in memo[pos]), it's not updated
1107  // again.
1108  } while (changed);
1109 
1110  int64_t ubAdjustment = getClosedUB ? 0 : 1;
1111 
1112  // Set the lower and upper bound maps for all the variables that were
1113  // computed as affine expressions of the rest as the "detected expr" and
1114  // "detected expr + 1" respectively; set the undetected ones to null.
1115  std::optional<FlatAffineValueConstraints> tmpClone;
1116  for (unsigned pos = 0; pos < num; pos++) {
1117  unsigned numMapDims = getNumDimVars() - num;
1118  unsigned numMapSymbols = getNumSymbolVars();
1119  AffineExpr expr = memo[pos + offset];
1120  if (expr)
1121  expr = simplifyAffineExpr(expr, numMapDims, numMapSymbols);
1122 
1123  AffineMap &lbMap = (*lbMaps)[pos];
1124  AffineMap &ubMap = (*ubMaps)[pos];
1125 
1126  if (expr) {
1127  lbMap = AffineMap::get(numMapDims, numMapSymbols, expr);
1128  ubMap = AffineMap::get(numMapDims, numMapSymbols, expr + ubAdjustment);
1129  } else {
1130  // TODO: Whenever there are local variables in the dependence
1131  // constraints, we'll conservatively over-approximate, since we don't
1132  // always explicitly compute them above (in the while loop).
1133  if (getNumLocalVars() == 0) {
1134  // Work on a copy so that we don't update this constraint system.
1135  if (!tmpClone) {
1136  tmpClone.emplace(FlatAffineValueConstraints(*this));
1137  // Removing redundant inequalities is necessary so that we don't get
1138  // redundant loop bounds.
1139  tmpClone->removeRedundantInequalities();
1140  }
1141  std::tie(lbMap, ubMap) = tmpClone->getLowerAndUpperBound(
1142  pos, offset, num, getNumDimVars(), /*localExprs=*/{}, context);
1143  }
1144 
1145  // If the above fails, we'll just use the constant lower bound and the
1146  // constant upper bound (if they exist) as the slice bounds.
1147  // TODO: being conservative for the moment in cases that
1148  // lead to multiple bounds - until getConstDifference in LoopFusion.cpp is
1149  // fixed (b/126426796).
1150  if (!lbMap || lbMap.getNumResults() > 1) {
1151  LLVM_DEBUG(llvm::dbgs()
1152  << "WARNING: Potentially over-approximating slice lb\n");
1153  auto lbConst = getConstantBound64(BoundType::LB, pos + offset);
1154  if (lbConst.has_value()) {
1155  lbMap = AffineMap::get(numMapDims, numMapSymbols,
1156  getAffineConstantExpr(*lbConst, context));
1157  }
1158  }
1159  if (!ubMap || ubMap.getNumResults() > 1) {
1160  LLVM_DEBUG(llvm::dbgs()
1161  << "WARNING: Potentially over-approximating slice ub\n");
1162  auto ubConst = getConstantBound64(BoundType::UB, pos + offset);
1163  if (ubConst.has_value()) {
1164  ubMap = AffineMap::get(
1165  numMapDims, numMapSymbols,
1166  getAffineConstantExpr(*ubConst + ubAdjustment, context));
1167  }
1168  }
1169  }
1170  LLVM_DEBUG(llvm::dbgs()
1171  << "lb map for pos = " << Twine(pos + offset) << ", expr: ");
1172  LLVM_DEBUG(lbMap.dump(););
1173  LLVM_DEBUG(llvm::dbgs()
1174  << "ub map for pos = " << Twine(pos + offset) << ", expr: ");
1175  LLVM_DEBUG(ubMap.dump(););
1176  }
1177 }
1178 
1180  AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs) {
1181  FlatAffineValueConstraints localCst;
1182  if (failed(getFlattenedAffineExprs(map, flattenedExprs, &localCst))) {
1183  LLVM_DEBUG(llvm::dbgs()
1184  << "composition unimplemented for semi-affine maps\n");
1185  return failure();
1186  }
1187 
1188  // Add localCst information.
1189  if (localCst.getNumLocalVars() > 0) {
1190  unsigned numLocalVars = getNumLocalVars();
1191  // Insert local dims of localCst at the beginning.
1192  insertLocalVar(/*pos=*/0, /*num=*/localCst.getNumLocalVars());
1193  // Insert local dims of `this` at the end of localCst.
1194  localCst.appendLocalVar(/*num=*/numLocalVars);
1195  // Dimensions of localCst and this constraint set match. Append localCst to
1196  // this constraint set.
1197  append(localCst);
1198  }
1199 
1200  return success();
1201 }
1202 
1204  AffineMap boundMap,
1205  bool isClosedBound) {
1206  assert(boundMap.getNumDims() == getNumDimVars() && "dim mismatch");
1207  assert(boundMap.getNumSymbols() == getNumSymbolVars() && "symbol mismatch");
1208  assert(pos < getNumDimAndSymbolVars() && "invalid position");
1209  assert((type != BoundType::EQ || isClosedBound) &&
1210  "EQ bound must be closed.");
1211 
1212  // Equality follows the logic of lower bound except that we add an equality
1213  // instead of an inequality.
1214  assert((type != BoundType::EQ || boundMap.getNumResults() == 1) &&
1215  "single result expected");
1216  bool lower = type == BoundType::LB || type == BoundType::EQ;
1217 
1218  std::vector<SmallVector<int64_t, 8>> flatExprs;
1219  if (failed(flattenAlignedMapAndMergeLocals(boundMap, &flatExprs)))
1220  return failure();
1221  assert(flatExprs.size() == boundMap.getNumResults());
1222 
1223  // Add one (in)equality for each result.
1224  for (const auto &flatExpr : flatExprs) {
1225  SmallVector<int64_t> ineq(getNumCols(), 0);
1226  // Dims and symbols.
1227  for (unsigned j = 0, e = boundMap.getNumInputs(); j < e; j++) {
1228  ineq[j] = lower ? -flatExpr[j] : flatExpr[j];
1229  }
1230  // Invalid bound: pos appears in `boundMap`.
1231  // TODO: This should be an assertion. Fix `addDomainFromSliceMaps` and/or
1232  // its callers to prevent invalid bounds from being added.
1233  if (ineq[pos] != 0)
1234  continue;
1235  ineq[pos] = lower ? 1 : -1;
1236  // Local columns of `ineq` are at the beginning.
1237  unsigned j = getNumDimVars() + getNumSymbolVars();
1238  unsigned end = flatExpr.size() - 1;
1239  for (unsigned i = boundMap.getNumInputs(); i < end; i++, j++) {
1240  ineq[j] = lower ? -flatExpr[i] : flatExpr[i];
1241  }
1242  // Make the bound closed in if flatExpr is open. The inequality is always
1243  // created in the upper bound form, so the adjustment is -1.
1244  int64_t boundAdjustment = (isClosedBound || type == BoundType::EQ) ? 0 : -1;
1245  // Constant term.
1246  ineq[getNumCols() - 1] = (lower ? -flatExpr[flatExpr.size() - 1]
1247  : flatExpr[flatExpr.size() - 1]) +
1248  boundAdjustment;
1249  type == BoundType::EQ ? addEquality(ineq) : addInequality(ineq);
1250  }
1251 
1252  return success();
1253 }
1254 
1256  AffineMap boundMap) {
1257  return addBound(type, pos, boundMap, /*isClosedBound=*/type != BoundType::UB);
1258 }
1259 
1260 AffineMap
1262  ValueRange operands) const {
1263  assert(map.getNumInputs() == operands.size() && "number of inputs mismatch");
1264 
1265  SmallVector<Value> dims, syms;
1266 #ifndef NDEBUG
1267  SmallVector<Value> newSyms;
1268  SmallVector<Value> *newSymsPtr = &newSyms;
1269 #else
1270  SmallVector<Value> *newSymsPtr = nullptr;
1271 #endif // NDEBUG
1272 
1273  dims.reserve(getNumDimVars());
1274  syms.reserve(getNumSymbolVars());
1275  for (unsigned i = getVarKindOffset(VarKind::SetDim),
1276  e = getVarKindEnd(VarKind::SetDim);
1277  i < e; ++i)
1278  dims.push_back(values[i] ? *values[i] : Value());
1279  for (unsigned i = getVarKindOffset(VarKind::Symbol),
1280  e = getVarKindEnd(VarKind::Symbol);
1281  i < e; ++i)
1282  syms.push_back(values[i] ? *values[i] : Value());
1283 
1284  AffineMap alignedMap =
1285  alignAffineMapWithValues(map, operands, dims, syms, newSymsPtr);
1286  // All symbols are already part of this FlatAffineConstraints.
1287  assert(syms.size() == newSymsPtr->size() && "unexpected new/missing symbols");
1288  assert(std::equal(syms.begin(), syms.end(), newSymsPtr->begin()) &&
1289  "unexpected new/missing symbols");
1290  return alignedMap;
1291 }
1292 
1294  AffineMap boundMap,
1295  ValueRange boundOperands) {
1296  // Fully compose map and operands; canonicalize and simplify so that we
1297  // transitively get to terminal symbols or loop IVs.
1298  auto map = boundMap;
1299  SmallVector<Value, 4> operands(boundOperands.begin(), boundOperands.end());
1300  fullyComposeAffineMapAndOperands(&map, &operands);
1301  map = simplifyAffineMap(map);
1302  canonicalizeMapAndOperands(&map, &operands);
1303  for (auto operand : operands)
1305  return addBound(type, pos, computeAlignedMap(map, operands));
1306 }
1307 
1308 // Adds slice lower bounds represented by lower bounds in 'lbMaps' and upper
1309 // bounds in 'ubMaps' to each value in `values' that appears in the constraint
1310 // system. Note that both lower/upper bounds share the same operand list
1311 // 'operands'.
1312 // This function assumes 'values.size' == 'lbMaps.size' == 'ubMaps.size', and
1313 // skips any null AffineMaps in 'lbMaps' or 'ubMaps'.
1314 // Note that both lower/upper bounds use operands from 'operands'.
1315 // Returns failure for unimplemented cases such as semi-affine expressions or
1316 // expressions with mod/floordiv.
1318  ArrayRef<Value> values, ArrayRef<AffineMap> lbMaps,
1319  ArrayRef<AffineMap> ubMaps, ArrayRef<Value> operands) {
1320  assert(values.size() == lbMaps.size());
1321  assert(lbMaps.size() == ubMaps.size());
1322 
1323  for (unsigned i = 0, e = lbMaps.size(); i < e; ++i) {
1324  unsigned pos;
1325  if (!findVar(values[i], &pos))
1326  continue;
1327 
1328  AffineMap lbMap = lbMaps[i];
1329  AffineMap ubMap = ubMaps[i];
1330  assert(!lbMap || lbMap.getNumInputs() == operands.size());
1331  assert(!ubMap || ubMap.getNumInputs() == operands.size());
1332 
1333  // Check if this slice is just an equality along this dimension.
1334  if (lbMap && ubMap && lbMap.getNumResults() == 1 &&
1335  ubMap.getNumResults() == 1 &&
1336  lbMap.getResult(0) + 1 == ubMap.getResult(0)) {
1337  if (failed(addBound(BoundType::EQ, pos, lbMap, operands)))
1338  return failure();
1339  continue;
1340  }
1341 
1342  // If lower or upper bound maps are null or provide no results, it implies
1343  // that the source loop was not at all sliced, and the entire loop will be a
1344  // part of the slice.
1345  if (lbMap && lbMap.getNumResults() != 0 && ubMap &&
1346  ubMap.getNumResults() != 0) {
1347  if (failed(addBound(BoundType::LB, pos, lbMap, operands)))
1348  return failure();
1349  if (failed(addBound(BoundType::UB, pos, ubMap, operands)))
1350  return failure();
1351  } else {
1352  auto loop = getForInductionVarOwner(values[i]);
1353  if (failed(this->addAffineForOpDomain(loop)))
1354  return failure();
1355  }
1356  }
1357  return success();
1358 }
1359 
1360 bool FlatAffineValueConstraints::findVar(Value val, unsigned *pos) const {
1361  unsigned i = 0;
1362  for (const auto &mayBeVar : values) {
1363  if (mayBeVar && *mayBeVar == val) {
1364  *pos = i;
1365  return true;
1366  }
1367  i++;
1368  }
1369  return false;
1370 }
1371 
1373  return llvm::any_of(values, [&](const std::optional<Value> &mayBeVar) {
1374  return mayBeVar && *mayBeVar == val;
1375  });
1376 }
1377 
1378 void FlatAffineValueConstraints::swapVar(unsigned posA, unsigned posB) {
1379  IntegerPolyhedron::swapVar(posA, posB);
1380 
1381  if (getVarKindAt(posA) == VarKind::Local &&
1382  getVarKindAt(posB) == VarKind::Local)
1383  return;
1384 
1385  // Treat value of a local variable as std::nullopt.
1386  if (getVarKindAt(posA) == VarKind::Local)
1387  values[posB] = std::nullopt;
1388  else if (getVarKindAt(posB) == VarKind::Local)
1389  values[posA] = std::nullopt;
1390  else
1391  std::swap(values[posA], values[posB]);
1392 }
1393 
1395  int64_t value) {
1396  unsigned pos;
1397  if (!findVar(val, &pos))
1398  // This is a pre-condition for this method.
1399  assert(0 && "var not found");
1400  addBound(type, pos, value);
1401 }
1402 
1403 void FlatAffineValueConstraints::printSpace(raw_ostream &os) const {
1405  os << "(";
1406  for (unsigned i = 0, e = getNumDimAndSymbolVars(); i < e; i++) {
1407  if (hasValue(i))
1408  os << "Value ";
1409  else
1410  os << "None ";
1411  }
1412  for (unsigned i = getVarKindOffset(VarKind::Local),
1413  e = getVarKindEnd(VarKind::Local);
1414  i < e; ++i)
1415  os << "Local ";
1416  os << " const)\n";
1417 }
1418 
1420  const IntegerRelation &other) {
1421 
1422  if (auto *otherValueSet =
1423  dyn_cast<const FlatAffineValueConstraints>(&other)) {
1424  *this = *otherValueSet;
1425  } else {
1426  *static_cast<IntegerRelation *>(this) = other;
1427  values.clear();
1428  values.resize(getNumDimAndSymbolVars(), std::nullopt);
1429  }
1430 }
1431 
1433  unsigned pos, bool darkShadow, bool *isResultIntegerExact) {
1435  if (getVarKindAt(pos) != VarKind::Local)
1436  newVals.erase(newVals.begin() + pos);
1437  // Note: Base implementation discards all associated Values.
1438  IntegerPolyhedron::fourierMotzkinEliminate(pos, darkShadow,
1439  isResultIntegerExact);
1440  values = newVals;
1441  assert(values.size() == getNumDimAndSymbolVars());
1442 }
1443 
1445  unsigned pos;
1446  bool ret = findVar(val, &pos);
1447  assert(ret);
1448  (void)ret;
1450 }
1451 
1453  const FlatAffineValueConstraints &otherCst) {
1454  assert(otherCst.getNumDimVars() == getNumDimVars() && "dims mismatch");
1455  assert(otherCst.getMaybeValues()
1456  .slice(0, getNumDimVars())
1457  .equals(getMaybeValues().slice(0, getNumDimVars())) &&
1458  "dim values mismatch");
1459  assert(otherCst.getNumLocalVars() == 0 && "local vars not supported here");
1460  assert(getNumLocalVars() == 0 && "local vars not supported yet here");
1461 
1462  // Align `other` to this.
1463  if (!areVarsAligned(*this, otherCst)) {
1464  FlatAffineValueConstraints otherCopy(otherCst);
1465  mergeAndAlignVars(/*offset=*/getNumDimVars(), this, &otherCopy);
1466  return IntegerPolyhedron::unionBoundingBox(otherCopy);
1467  }
1468 
1469  return IntegerPolyhedron::unionBoundingBox(otherCst);
1470 }
1471 
1472 /// Compute an explicit representation for local vars. For all systems coming
1473 /// from MLIR integer sets, maps, or expressions where local vars were
1474 /// introduced to model floordivs and mods, this always succeeds.
1477  MLIRContext *context) {
1478  unsigned numDims = cst.getNumDimVars();
1479  unsigned numSyms = cst.getNumSymbolVars();
1480 
1481  // Initialize dimensional and symbolic variables.
1482  for (unsigned i = 0; i < numDims; i++)
1483  memo[i] = getAffineDimExpr(i, context);
1484  for (unsigned i = numDims, e = numDims + numSyms; i < e; i++)
1485  memo[i] = getAffineSymbolExpr(i - numDims, context);
1486 
1487  bool changed;
1488  do {
1489  // Each time `changed` is true at the end of this iteration, one or more
1490  // local vars would have been detected as floordivs and set in memo; so the
1491  // number of null entries in memo[...] strictly reduces; so this converges.
1492  changed = false;
1493  for (unsigned i = 0, e = cst.getNumLocalVars(); i < e; ++i)
1494  if (!memo[numDims + numSyms + i] &&
1495  detectAsFloorDiv(cst, /*pos=*/numDims + numSyms + i, context, memo))
1496  changed = true;
1497  } while (changed);
1498 
1499  ArrayRef<AffineExpr> localExprs =
1500  ArrayRef<AffineExpr>(memo).take_back(cst.getNumLocalVars());
1501  return success(
1502  llvm::all_of(localExprs, [](AffineExpr expr) { return expr; }));
1503 }
1504 
1506  unsigned pos, unsigned ineqPos, AffineValueMap &vmap,
1507  MLIRContext *context) const {
1508  unsigned numDims = getNumDimVars();
1509  unsigned numSyms = getNumSymbolVars();
1510 
1511  assert(pos < numDims && "invalid position");
1512  assert(ineqPos < getNumInequalities() && "invalid inequality position");
1513 
1514  // Get expressions for local vars.
1516  if (failed(computeLocalVars(*this, memo, context)))
1517  assert(false &&
1518  "one or more local exprs do not have an explicit representation");
1519  auto localExprs = ArrayRef<AffineExpr>(memo).take_back(getNumLocalVars());
1520 
1521  // Compute the AffineExpr lower/upper bound for this inequality.
1522  SmallVector<int64_t, 8> inequality = getInequality64(ineqPos);
1524  bound.reserve(getNumCols() - 1);
1525  // Everything other than the coefficient at `pos`.
1526  bound.append(inequality.begin(), inequality.begin() + pos);
1527  bound.append(inequality.begin() + pos + 1, inequality.end());
1528 
1529  if (inequality[pos] > 0)
1530  // Lower bound.
1531  std::transform(bound.begin(), bound.end(), bound.begin(),
1532  std::negate<int64_t>());
1533  else
1534  // Upper bound (which is exclusive).
1535  bound.back() += 1;
1536 
1537  // Convert to AffineExpr (tree) form.
1538  auto boundExpr = getAffineExprFromFlatForm(bound, numDims - 1, numSyms,
1539  localExprs, context);
1540 
1541  // Get the values to bind to this affine expr (all dims and symbols).
1542  SmallVector<Value, 4> operands;
1543  getValues(0, pos, &operands);
1544  SmallVector<Value, 4> trailingOperands;
1545  getValues(pos + 1, getNumDimAndSymbolVars(), &trailingOperands);
1546  operands.append(trailingOperands.begin(), trailingOperands.end());
1547  vmap.reset(AffineMap::get(numDims - 1, numSyms, boundExpr), operands);
1548 }
1549 
1550 IntegerSet
1552  if (getNumConstraints() == 0)
1553  // Return universal set (always true): 0 == 0.
1555  getAffineConstantExpr(/*constant=*/0, context),
1556  /*eqFlags=*/true);
1557 
1558  // Construct local references.
1560 
1561  if (failed(computeLocalVars(*this, memo, context))) {
1562  // Check if the local variables without an explicit representation have
1563  // zero coefficients everywhere.
1564  SmallVector<unsigned> noLocalRepVars;
1565  unsigned numDimsSymbols = getNumDimAndSymbolVars();
1566  for (unsigned i = numDimsSymbols, e = getNumVars(); i < e; ++i) {
1567  if (!memo[i] && !isColZero(/*pos=*/i))
1568  noLocalRepVars.push_back(i - numDimsSymbols);
1569  }
1570  if (!noLocalRepVars.empty()) {
1571  LLVM_DEBUG({
1572  llvm::dbgs() << "local variables at position(s) ";
1573  llvm::interleaveComma(noLocalRepVars, llvm::dbgs());
1574  llvm::dbgs() << " do not have an explicit representation in:\n";
1575  this->dump();
1576  });
1577  return IntegerSet();
1578  }
1579  }
1580 
1581  ArrayRef<AffineExpr> localExprs =
1582  ArrayRef<AffineExpr>(memo).take_back(getNumLocalVars());
1583 
1584  // Construct the IntegerSet from the equalities/inequalities.
1585  unsigned numDims = getNumDimVars();
1586  unsigned numSyms = getNumSymbolVars();
1587 
1589  std::fill(eqFlags.begin(), eqFlags.begin() + getNumEqualities(), true);
1590  std::fill(eqFlags.begin() + getNumEqualities(), eqFlags.end(), false);
1591 
1593  exprs.reserve(getNumConstraints());
1594 
1595  for (unsigned i = 0, e = getNumEqualities(); i < e; ++i)
1596  exprs.push_back(getAffineExprFromFlatForm(getEquality64(i), numDims,
1597  numSyms, localExprs, context));
1598  for (unsigned i = 0, e = getNumInequalities(); i < e; ++i)
1599  exprs.push_back(getAffineExprFromFlatForm(getInequality64(i), numDims,
1600  numSyms, localExprs, context));
1601  return IntegerSet::get(numDims, numSyms, exprs, eqFlags);
1602 }
1603 
1605  ValueRange dims, ValueRange syms,
1606  SmallVector<Value> *newSyms) {
1607  assert(operands.size() == map.getNumInputs() &&
1608  "expected same number of operands and map inputs");
1609  MLIRContext *ctx = map.getContext();
1610  Builder builder(ctx);
1611  SmallVector<AffineExpr> dimReplacements(map.getNumDims(), {});
1612  unsigned numSymbols = syms.size();
1613  SmallVector<AffineExpr> symReplacements(map.getNumSymbols(), {});
1614  if (newSyms) {
1615  newSyms->clear();
1616  newSyms->append(syms.begin(), syms.end());
1617  }
1618 
1619  for (const auto &operand : llvm::enumerate(operands)) {
1620  // Compute replacement dim/sym of operand.
1621  AffineExpr replacement;
1622  auto dimIt = std::find(dims.begin(), dims.end(), operand.value());
1623  auto symIt = std::find(syms.begin(), syms.end(), operand.value());
1624  if (dimIt != dims.end()) {
1625  replacement =
1626  builder.getAffineDimExpr(std::distance(dims.begin(), dimIt));
1627  } else if (symIt != syms.end()) {
1628  replacement =
1629  builder.getAffineSymbolExpr(std::distance(syms.begin(), symIt));
1630  } else {
1631  // This operand is neither a dimension nor a symbol. Add it as a new
1632  // symbol.
1633  replacement = builder.getAffineSymbolExpr(numSymbols++);
1634  if (newSyms)
1635  newSyms->push_back(operand.value());
1636  }
1637  // Add to corresponding replacements vector.
1638  if (operand.index() < map.getNumDims()) {
1639  dimReplacements[operand.index()] = replacement;
1640  } else {
1641  symReplacements[operand.index() - map.getNumDims()] = replacement;
1642  }
1643  }
1644 
1645  return map.replaceDimsAndSymbols(dimReplacements, symReplacements,
1646  dims.size(), numSymbols);
1647 }
1648 
1650  FlatAffineValueConstraints domain = *this;
1651  // Convert all range variables to local variables.
1652  domain.convertToLocal(VarKind::SetDim, getNumDomainDims(),
1654  return domain;
1655 }
1656 
1658  FlatAffineValueConstraints range = *this;
1659  // Convert all domain variables to local variables.
1660  range.convertToLocal(VarKind::SetDim, 0, getNumDomainDims());
1661  return range;
1662 }
1663 
1665  assert(getNumDomainDims() == other.getNumRangeDims() &&
1666  "Domain of this and range of other do not match");
1667  assert(std::equal(values.begin(), values.begin() + getNumDomainDims(),
1668  other.values.begin() + other.getNumDomainDims()) &&
1669  "Domain of this and range of other do not match");
1670 
1671  FlatAffineRelation rel = other;
1672 
1673  // Convert `rel` from
1674  // [otherDomain] -> [otherRange]
1675  // to
1676  // [otherDomain] -> [otherRange thisRange]
1677  // and `this` from
1678  // [thisDomain] -> [thisRange]
1679  // to
1680  // [otherDomain thisDomain] -> [thisRange].
1681  unsigned removeDims = rel.getNumRangeDims();
1684 
1685  // Merge symbol and local variables.
1686  mergeSymbolVars(rel);
1687  mergeLocalVars(rel);
1688 
1689  // Convert `rel` from [otherDomain] -> [otherRange thisRange] to
1690  // [otherDomain] -> [thisRange] by converting first otherRange range vars
1691  // to local vars.
1692  rel.convertToLocal(VarKind::SetDim, rel.getNumDomainDims(),
1693  rel.getNumDomainDims() + removeDims);
1694  // Convert `this` from [otherDomain thisDomain] -> [thisRange] to
1695  // [otherDomain] -> [thisRange] by converting last thisDomain domain vars
1696  // to local vars.
1697  convertToLocal(VarKind::SetDim, getNumDomainDims() - removeDims,
1698  getNumDomainDims());
1699 
1700  auto thisMaybeValues = getMaybeValues(VarKind::SetDim);
1701  auto relMaybeValues = rel.getMaybeValues(VarKind::SetDim);
1702 
1703  // Add and match domain of `rel` to domain of `this`.
1704  for (unsigned i = 0, e = rel.getNumDomainDims(); i < e; ++i)
1705  if (relMaybeValues[i].has_value())
1706  setValue(i, *relMaybeValues[i]);
1707  // Add and match range of `this` to range of `rel`.
1708  for (unsigned i = 0, e = getNumRangeDims(); i < e; ++i) {
1709  unsigned rangeIdx = rel.getNumDomainDims() + i;
1710  if (thisMaybeValues[rangeIdx].has_value())
1711  rel.setValue(rangeIdx, *thisMaybeValues[rangeIdx]);
1712  }
1713 
1714  // Append `this` to `rel` and simplify constraints.
1715  rel.append(*this);
1717 
1718  *this = rel;
1719 }
1720 
1722  unsigned oldDomain = getNumDomainDims();
1723  unsigned oldRange = getNumRangeDims();
1724  // Add new range vars.
1725  appendRangeVar(oldDomain);
1726  // Swap new vars with domain.
1727  for (unsigned i = 0; i < oldDomain; ++i)
1728  swapVar(i, oldDomain + oldRange + i);
1729  // Remove the swapped domain.
1730  removeVarRange(0, oldDomain);
1731  // Set domain and range as inverse.
1732  numDomainDims = oldRange;
1733  numRangeDims = oldDomain;
1734 }
1735 
1736 void FlatAffineRelation::insertDomainVar(unsigned pos, unsigned num) {
1737  assert(pos <= getNumDomainDims() &&
1738  "Var cannot be inserted at invalid position");
1739  insertDimVar(pos, num);
1740  numDomainDims += num;
1741 }
1742 
1743 void FlatAffineRelation::insertRangeVar(unsigned pos, unsigned num) {
1744  assert(pos <= getNumRangeDims() &&
1745  "Var cannot be inserted at invalid position");
1746  insertDimVar(getNumDomainDims() + pos, num);
1747  numRangeDims += num;
1748 }
1749 
1752  numDomainDims += num;
1753 }
1754 
1756  insertDimVar(getNumDimVars(), num);
1757  numRangeDims += num;
1758 }
1759 
1760 void FlatAffineRelation::removeVarRange(VarKind kind, unsigned varStart,
1761  unsigned varLimit) {
1762  assert(varLimit <= getNumVarKind(kind));
1763  if (varStart >= varLimit)
1764  return;
1765 
1766  FlatAffineValueConstraints::removeVarRange(kind, varStart, varLimit);
1767 
1768  // If kind is not SetDim, domain and range don't need to be updated.
1769  if (kind != VarKind::SetDim)
1770  return;
1771 
1772  // Compute number of domain and range variables to remove. This is done by
1773  // intersecting the range of domain/range vars with range of vars to remove.
1774  unsigned intersectDomainLHS = std::min(varLimit, getNumDomainDims());
1775  unsigned intersectDomainRHS = varStart;
1776  unsigned intersectRangeLHS = std::min(varLimit, getNumDimVars());
1777  unsigned intersectRangeRHS = std::max(varStart, getNumDomainDims());
1778 
1779  if (intersectDomainLHS > intersectDomainRHS)
1780  numDomainDims -= intersectDomainLHS - intersectDomainRHS;
1781  if (intersectRangeLHS > intersectRangeRHS)
1782  numRangeDims -= intersectRangeLHS - intersectRangeRHS;
1783 }
1784 
1786  FlatAffineRelation &rel) {
1787  // Get flattened affine expressions.
1788  std::vector<SmallVector<int64_t, 8>> flatExprs;
1789  FlatAffineValueConstraints localVarCst;
1790  if (failed(getFlattenedAffineExprs(map, &flatExprs, &localVarCst)))
1791  return failure();
1792 
1793  unsigned oldDimNum = localVarCst.getNumDimVars();
1794  unsigned oldCols = localVarCst.getNumCols();
1795  unsigned numRangeVars = map.getNumResults();
1796  unsigned numDomainVars = map.getNumDims();
1797 
1798  // Add range as the new expressions.
1799  localVarCst.appendDimVar(numRangeVars);
1800 
1801  // Add equalities between source and range.
1802  SmallVector<int64_t, 8> eq(localVarCst.getNumCols());
1803  for (unsigned i = 0, e = map.getNumResults(); i < e; ++i) {
1804  // Zero fill.
1805  std::fill(eq.begin(), eq.end(), 0);
1806  // Fill equality.
1807  for (unsigned j = 0, f = oldDimNum; j < f; ++j)
1808  eq[j] = flatExprs[i][j];
1809  for (unsigned j = oldDimNum, f = oldCols; j < f; ++j)
1810  eq[j + numRangeVars] = flatExprs[i][j];
1811  // Set this dimension to -1 to equate lhs and rhs and add equality.
1812  eq[numDomainVars + i] = -1;
1813  localVarCst.addEquality(eq);
1814  }
1815 
1816  // Create relation and return success.
1817  rel = FlatAffineRelation(numDomainVars, numRangeVars, localVarCst);
1818  return success();
1819 }
1820 
1822  FlatAffineRelation &rel) {
1823 
1824  AffineMap affineMap = map.getAffineMap();
1825  if (failed(getRelationFromMap(affineMap, rel)))
1826  return failure();
1827 
1828  // Set symbol values for domain dimensions and symbols.
1829  for (unsigned i = 0, e = rel.getNumDomainDims(); i < e; ++i)
1830  rel.setValue(i, map.getOperand(i));
1831  for (unsigned i = rel.getNumDimVars(), e = rel.getNumDimAndSymbolVars();
1832  i < e; ++i)
1833  rel.setValue(i, map.getOperand(i - rel.getNumRangeDims()));
1834 
1835  return success();
1836 }
1837 
1840  MultiAffineFunction &multiAff) {
1842  std::vector<SmallVector<int64_t, 8>> flattenedExprs;
1843  LogicalResult result = getFlattenedAffineExprs(map, &flattenedExprs, &cst);
1844 
1845  if (result.failed())
1846  return failure();
1847 
1848  DivisionRepr divs = cst.getLocalReprs();
1849  assert(divs.hasAllReprs() &&
1850  "AffineMap cannot produce divs without local representation");
1851 
1852  // TODO: We shouldn't have to do this conversion.
1853  Matrix mat(map.getNumResults(), map.getNumInputs() + divs.getNumDivs() + 1);
1854  for (unsigned i = 0, e = flattenedExprs.size(); i < e; ++i)
1855  for (unsigned j = 0, f = flattenedExprs[i].size(); j < f; ++j)
1856  mat(i, j) = flattenedExprs[i][j];
1857 
1858  multiAff = MultiAffineFunction(
1859  PresburgerSpace::getRelationSpace(map.getNumDims(), map.getNumResults(),
1860  map.getNumSymbols(), divs.getNumDivs()),
1861  mat, divs);
1862 
1863  return success();
1864 }
static bool areVarsAligned(const FlatAffineValueConstraints &a, const FlatAffineValueConstraints &b)
Checks if two constraint systems are in the same space, i.e., if they are associated with the same se...
static void turnSymbolIntoDim(FlatAffineValueConstraints *cst, Value value)
static bool detectAsMod(const FlatAffineValueConstraints &cst, unsigned pos, int64_t lbConst, int64_t ubConst, SmallVectorImpl< AffineExpr > &memo, MLIRContext *context)
static bool LLVM_ATTRIBUTE_UNUSED areVarsUnique(const FlatAffineValueConstraints &cst, unsigned start, unsigned end)
Checks if the SSA values associated with cst's variables in range [start, end) are unique.
static LogicalResult getFlattenedAffineExprs(ArrayRef< AffineExpr > exprs, unsigned numDims, unsigned numSymbols, std::vector< SmallVector< int64_t, 8 >> *flattenedExprs, FlatAffineValueConstraints *localVarCst)
static bool detectAsFloorDiv(const FlatAffineValueConstraints &cst, unsigned pos, MLIRContext *context, SmallVectorImpl< AffineExpr > &exprs)
Check if the pos^th variable can be expressed as a floordiv of an affine function of other variables ...
static void mergeAndAlignVars(unsigned offset, FlatAffineValueConstraints *a, FlatAffineValueConstraints *b)
Merge and align the variables of A and B starting at 'offset', so that both constraint systems get th...
static LogicalResult computeLocalVars(const FlatAffineValueConstraints &cst, SmallVectorImpl< AffineExpr > &memo, MLIRContext *context)
Compute an explicit representation for local vars.
@ None
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
static Value min(ImplicitLocOpBuilder &builder, Value value, Value bound)
An integer constant appearing in affine expression.
Definition: AffineExpr.h:232
A dimensional identifier appearing in an affine expression.
Definition: AffineExpr.h:216
unsigned getPosition() const
Definition: AffineExpr.cpp:325
Base type for affine expression.
Definition: AffineExpr.h:68
constexpr bool isa() const
Definition: AffineExpr.h:270
U dyn_cast() const
Definition: AffineExpr.h:281
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:43
int64_t getSingleConstantResult() const
Returns the constant result of this map.
Definition: AffineMap.cpp:306
MLIRContext * getContext() const
Definition: AffineMap.cpp:266
bool isConstant() const
Returns true if this affine map has only constant results.
Definition: AffineMap.cpp:300
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
unsigned getNumSymbols() const
Definition: AffineMap.cpp:323
unsigned getNumDims() const
Definition: AffineMap.cpp:319
ArrayRef< AffineExpr > getResults() const
Definition: AffineMap.cpp:332
unsigned getNumResults() const
Definition: AffineMap.cpp:327
AffineMap replaceDimsAndSymbols(ArrayRef< AffineExpr > dimReplacements, ArrayRef< AffineExpr > symReplacements, unsigned numResultDims, unsigned numResultSyms) const
This method substitutes any uses of dimensions and symbols (e.g.
Definition: AffineMap.cpp:421
unsigned getNumInputs() const
Definition: AffineMap.cpp:328
AffineExpr getResult(unsigned idx) const
Definition: AffineMap.cpp:336
void dump() const
An AffineValueMap is an affine map plus its ML value operands and results for analysis purposes.
Value getOperand(unsigned i) const
ArrayRef< Value > getOperands() const
AffineMap getAffineMap() const
void reset(AffineMap map, ValueRange operands, ValueRange results={})
This class is a general helper class for creating context-global objects like types,...
Definition: Builders.h:50
AffineExpr getAffineSymbolExpr(unsigned position)
Definition: Builders.cpp:331
AffineExpr getAffineDimExpr(unsigned position)
Definition: Builders.cpp:327
A FlatAffineRelation represents a set of ordered pairs (domain -> range) where "domain" and "range" a...
unsigned getNumRangeDims() const
void appendDomainVar(unsigned num=1)
Append num variables of the specified kind after the last variable of that kind.
void compose(const FlatAffineRelation &other)
Given affine relation other: (domainOther -> rangeOther), this operation takes the composition of oth...
unsigned getNumDomainDims() const
Returns the number of variables corresponding to domain/range of relation.
void inverse()
Swap domain and range of the relation.
FlatAffineValueConstraints getDomainSet() const
Returns a set corresponding to the domain/range of the affine relation.
void removeVarRange(VarKind kind, unsigned varStart, unsigned varLimit) override
Removes variables in the column range [varStart, varLimit), and copies any remaining valid data into ...
void appendRangeVar(unsigned num=1)
FlatAffineValueConstraints getRangeSet() const
void insertRangeVar(unsigned pos, unsigned num=1)
void insertDomainVar(unsigned pos, unsigned num=1)
Insert num variables of the specified kind after the pos variable of that kind.
FlatAffineValueConstraints represents an extension of IntegerPolyhedron where each non-local variable...
void clearAndCopyFrom(const IntegerRelation &other) override
Replaces the contents of this FlatAffineValueConstraints with other.
unsigned appendLocalVar(unsigned num=1)
void swapVar(unsigned posA, unsigned posB) override
Swap the posA^th variable with the posB^th variable.
void addInductionVarOrTerminalSymbol(Value val)
Add the specified values as a dim or symbol var depending on its nature, if it already doesn't exist ...
unsigned insertSymbolVar(unsigned pos, unsigned num=1)
void addAffineIfOpDomain(AffineIfOp ifOp)
Adds constraints imposed by the affine.if operation.
LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap, bool isClosedBound)
Adds a bound for the variable at the specified position with constraints being drawn from the specifi...
void fourierMotzkinEliminate(unsigned pos, bool darkShadow=false, bool *isResultIntegerExact=nullptr) override
Eliminates the variable at the specified position using Fourier-Motzkin variable elimination,...
bool hasConsistentState() const override
Returns false if the fields corresponding to various variable counts, or equality/inequality buffer s...
bool hasValues() const
Returns true if at least one variable has an associated Value.
unsigned insertVar(presburger::VarKind kind, unsigned pos, unsigned num=1) override
Insert num variables of the specified kind at position pos.
Value getValue(unsigned pos) const
Returns the Value associated with the pos^th variable.
void convertLoopIVSymbolsToDims()
Changes all symbol variables which are loop IVs to dim variables.
LogicalResult addDomainFromSliceMaps(ArrayRef< AffineMap > lbMaps, ArrayRef< AffineMap > ubMaps, ArrayRef< Value > operands)
Adds constraints (lower and upper bounds) for each loop in the loop nest described by the bound maps ...
unsigned insertLocalVar(unsigned pos, unsigned num=1)
ArrayRef< std::optional< Value > > getMaybeValues() const
std::unique_ptr< FlatAffineValueConstraints > clone() const
Clones this object.
LogicalResult addAffineParallelOpDomain(AffineParallelOp parallelOp)
Add constraints (lower and upper bounds) for the specified 'affine.parallel' operation's Value using ...
bool hasValue(unsigned pos) const
Returns true if the pos^th variable has an associated Value.
void setValue(unsigned pos, Value val)
Sets the Value associated with the pos^th variable.
AffineMap computeAlignedMap(AffineMap map, ValueRange operands) const
Align map with this constraint system based on operands.
IntegerSet getAsIntegerSet(MLIRContext *context) const
Returns the constraint system as an integer set.
std::pair< AffineMap, AffineMap > getLowerAndUpperBound(unsigned pos, unsigned offset, unsigned num, unsigned symStartPos, ArrayRef< AffineExpr > localExprs, MLIRContext *context) const
Gets the lower and upper bound of the offset + posth variable treating [0, offset) U [offset + num,...
void projectOut(Value val)
Projects out the variable that is associate with Value.
LogicalResult addAffineForOpDomain(AffineForOp forOp)
Adds constraints (lower and upper bounds) for the specified 'affine.for' operation's Value using IR i...
void mergeAndAlignVarsWithOther(unsigned offset, FlatAffineValueConstraints *other)
Merge and align the variables of this and other starting at offset, so that both constraint systems g...
SmallVector< std::optional< Value >, 8 > values
Values corresponding to the (column) non-local variables of this constraint system appearing in the o...
unsigned appendDimVar(ValueRange vals)
Append variables of the specified kind after the last variable of that kind.
void removeVarRange(presburger::VarKind kind, unsigned varStart, unsigned varLimit) override
Removes variables in the column range [varStart, varLimit), and copies any remaining valid data into ...
void mergeSymbolVars(FlatAffineValueConstraints &other)
Merge and align symbols of this and other such that both get union of of symbols that are unique.
void getValues(unsigned start, unsigned end, SmallVectorImpl< Value > *values) const
Returns the Values associated with variables in range [start, end).
bool findVar(Value val, unsigned *pos) const
Looks up the position of the variable with the specified Value.
void reset(unsigned numReservedInequalities, unsigned numReservedEqualities, unsigned numReservedCols, unsigned numDims, unsigned numSymbols, unsigned numLocals=0)
Clears any existing data and reserves memory for the specified constraints.
void printSpace(raw_ostream &os) const override
Prints the number of constraints, dimensions, symbols and locals in the FlatAffineConstraints.
bool areVarsAlignedWithOther(const FlatAffineValueConstraints &other)
Returns true if this constraint system and other are in the same space, i.e., if they are associated ...
void getIneqAsAffineValueMap(unsigned pos, unsigned ineqPos, AffineValueMap &vmap, MLIRContext *context) const
Returns the bound for the variable at pos from the inequality at ineqPos as a 1-d affine value map (a...
void getSliceBounds(unsigned offset, unsigned num, MLIRContext *context, SmallVectorImpl< AffineMap > *lbMaps, SmallVectorImpl< AffineMap > *ubMaps, bool getClosedUB=false)
Computes the lower and upper bounds of the first num dimensional variables (starting at offset) as an...
unsigned insertDimVar(unsigned pos, unsigned num=1)
Insert variables of the specified kind at position pos.
static FlatAffineValueConstraints getHyperrectangular(ValueRange ivs, ValueRange lbs, ValueRange ubs)
LogicalResult addSliceBounds(ArrayRef< Value > values, ArrayRef< AffineMap > lbMaps, ArrayRef< AffineMap > ubMaps, ArrayRef< Value > operands)
Adds slice lower bounds represented by lower bounds in lbMaps and upper bounds in ubMaps to each vari...
LogicalResult composeMap(const AffineValueMap *vMap)
Composes the affine value map with this FlatAffineValueConstrains, adding the results of the map as d...
FlatAffineValueConstraints(unsigned numReservedInequalities, unsigned numReservedEqualities, unsigned numReservedCols, unsigned numDims, unsigned numSymbols, unsigned numLocals, ArrayRef< std::optional< Value >> valArgs={})
Constructs a constraint system reserving memory for the specified number of constraints and variables...
LogicalResult composeMatchingMap(AffineMap other)
Composes an affine map whose dimensions and symbols match one to one with the dimensions and symbols ...
bool containsVar(Value val) const
Returns true if an variable with the specified Value exists, false otherwise.
unsigned appendSymbolVar(ValueRange vals)
LogicalResult flattenAlignedMapAndMergeLocals(AffineMap map, std::vector< SmallVector< int64_t, 8 >> *flattenedExprs)
Given an affine map that is aligned with this constraint system:
LogicalResult unionBoundingBox(const FlatAffineValueConstraints &other)
Updates the constraints to be the smallest bounding (enclosing) box that contains the points of this ...
An integer set representing a conjunction of one or more affine equalities and inequalities.
Definition: IntegerSet.h:44
unsigned getNumDims() const
Definition: IntegerSet.cpp:15
static IntegerSet get(unsigned dimCount, unsigned symbolCount, ArrayRef< AffineExpr > constraints, ArrayRef< bool > eqFlags)
unsigned getNumInputs() const
Definition: IntegerSet.cpp:17
unsigned getNumConstraints() const
Definition: IntegerSet.cpp:21
ArrayRef< AffineExpr > getConstraints() const
Definition: IntegerSet.cpp:41
ArrayRef< bool > getEqFlags() const
Returns the equality bits, which specify whether each of the constraints is an equality or inequality...
Definition: IntegerSet.cpp:51
unsigned getNumSymbols() const
Definition: IntegerSet.cpp:16
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:56
virtual void addLocalFloorDivId(ArrayRef< int64_t > dividend, int64_t divisor, AffineExpr localExpr)
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:350
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:93
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Definition: Value.cpp:20
Specialization of arith.constant op that returns an integer of index type.
Definition: Arith.h:89
Class storing division representation of local variables of a constraint system.
Definition: Utils.h:118
bool hasAllReprs() const
Definition: Utils.h:134
unsigned getNumDivs() const
Definition: Utils.h:126
An IntegerPolyhedron represents the set of points from a PresburgerSpace that satisfy a list of affin...
An IntegerRelation represents the set of points from a PresburgerSpace that satisfy a list of affine ...
int64_t atEq64(unsigned i, unsigned j) const
The same, but casts to int64_t.
unsigned getVarKindEnd(VarKind kind) const
Return the index at Which the specified kind of vars ends.
void addInequality(ArrayRef< MPInt > inEq)
Adds an inequality (>= 0) from the coefficients specified in inEq.
void addEquality(ArrayRef< MPInt > eq)
Adds an equality from the coefficients specified in eq.
void normalizeConstraintsByGCD()
Normalized each constraints by the GCD of its coefficients.
SmallVector< int64_t, 8 > getEquality64(unsigned idx) const
The same, but casts to int64_t.
std::optional< int64_t > getConstantBound64(BoundType type, unsigned pos) const
The same, but casts to int64_t.
unsigned getNumVarKind(VarKind kind) const
Get the number of vars of the specified kind.
VarKind getVarKindAt(unsigned pos) const
Return the VarKind of the var at the specified position.
void convertToLocal(VarKind kind, unsigned varStart, unsigned varLimit)
void addLocalFloorDiv(ArrayRef< MPInt > dividend, const MPInt &divisor)
Adds a new local variable as the floordiv of an affine function of other variables,...
void append(const IntegerRelation &other)
Appends constraints from other into this.
void setDimSymbolSeparation(unsigned newSymbolCount)
Changes the partition between dimensions and symbols.
unsigned getNumCols() const
Returns the number of columns in the constraint system.
void getLowerAndUpperBoundIndices(unsigned pos, SmallVectorImpl< unsigned > *lbIndices, SmallVectorImpl< unsigned > *ubIndices, SmallVectorImpl< unsigned > *eqIndices=nullptr, unsigned offset=0, unsigned num=0) const
Gather positions of all lower and upper bounds of the variable at pos, and optionally any equalities ...
BoundType
The type of bound: equal, lower bound or upper bound.
DivisionRepr getLocalReprs(std::vector< MaybeLocalRepr > *repr=nullptr) const
Returns a DivisonRepr representing the division representation of local variables in the constraint s...
void removeRedundantLocalVars()
Removes local variables using equalities.
unsigned mergeLocalVars(IntegerRelation &other)
Adds additional local vars to the sets such that they both have the union of the local vars in each s...
SmallVector< int64_t, 8 > getInequality64(unsigned idx) const
bool findConstraintWithNonZeroAt(unsigned colIdx, bool isEq, unsigned *rowIdx) const
Searches for a constraint with a non-zero coefficient at colIdx in equality (isEq=true) or inequality...
bool isColZero(unsigned pos) const
Returns true if the pos^th column is all zero for both inequalities and equalities.
unsigned getVarKindOffset(VarKind kind) const
Return the index at which the specified kind of vars starts.
This is a class to represent a resizable matrix.
Definition: Matrix.h:35
This class represents a multi-affine function with the domain as Z^d, where d is the number of domain...
Definition: PWMAFunction.h:41
PresburgerSpace is the space of all possible values of a tuple of integer valued variables/variables.
static PresburgerSpace getSetSpace(unsigned numDims=0, unsigned numSymbols=0, unsigned numLocals=0)
static void printSpace(std::ostream &os, int count)
Definition: RunnerUtils.h:93
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:223
VarKind
Kind of variable.
LLVM_ATTRIBUTE_ALWAYS_INLINE MPInt abs(const MPInt &x)
Definition: MPInt.h:370
MaybeLocalRepr computeSingleVarRepr(const IntegerRelation &cst, ArrayRef< bool > foundRepr, unsigned pos, MutableArrayRef< MPInt > dividend, MPInt &divisor)
Returns the MaybeLocalRepr struct which contains the indices of the constraints that can be expressed...
Definition: Utils.cpp:223
Include the generated interface declarations.
AffineMap simplifyAffineMap(AffineMap map)
Simplifies an affine map by simplifying its underlying AffineExpr results.
Definition: AffineMap.cpp:663
void canonicalizeSetAndOperands(IntegerSet *set, SmallVectorImpl< Value > *operands)
Canonicalizes an integer set the same way canonicalizeMapAndOperands does for affine maps.
Definition: AffineOps.cpp:1245
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
int64_t floorDiv(int64_t lhs, int64_t rhs)
Returns the result of MLIR's floordiv operation on constants.
Definition: MathExtras.h:33
AffineMap alignAffineMapWithValues(AffineMap map, ValueRange operands, ValueRange dims, ValueRange syms, SmallVector< Value > *newSyms=nullptr)
Re-indexes the dimensions and symbols of an affine map with given operands values to align with dims ...
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
LogicalResult getRelationFromMap(AffineMap &map, FlatAffineRelation &rel)
Builds a relation from the given AffineMap/AffineValueMap map, containing all pairs of the form opera...
AffineForOp getForInductionVarOwner(Value val)
Returns the loop parent of an induction variable.
Definition: AffineOps.cpp:2310
AffineExpr getAffineExprFromFlatForm(ArrayRef< int64_t > flatExprs, unsigned numDims, unsigned numSymbols, ArrayRef< AffineExpr > localExprs, MLIRContext *context)
Constructs an affine expression from a flat ArrayRef.
Definition: AffineExpr.cpp:904
LogicalResult getFlattenedAffineExprs(AffineMap map, std::vector< SmallVector< int64_t, 8 >> *flattenedExprs, FlatAffineValueConstraints *cst=nullptr)
Flattens the result expressions of the map to their corresponding flattened forms and set in 'flatten...
bool isTopLevelValue(Value value)
TODO: These should be renamed if they are on the mlir namespace.
Definition: AffineOps.cpp:232
AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context)
Definition: AffineExpr.cpp:527
void fullyComposeAffineMapAndOperands(AffineMap *map, SmallVectorImpl< Value > *operands)
Given an affine map map and its input operands, this method composes into map, maps of AffineApplyOps...
Definition: AffineOps.cpp:875
LogicalResult getFlattenedAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols, SmallVectorImpl< int64_t > *flattenedExpr, FlatAffineValueConstraints *cst=nullptr)
Flattens 'expr' into 'flattenedExpr', which contains the coefficients of the dimensions,...
AffineExpr simplifyAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols)
Simplify an affine expression by flattening and some amount of simple analysis.
bool isAffineForInductionVar(Value val)
Returns true if the provided value is the induction variable of an AffineForOp.
Definition: AffineOps.cpp:2304
void canonicalizeMapAndOperands(AffineMap *map, SmallVectorImpl< Value > *operands)
Modifies both map and operands in-place so as to:
Definition: AffineOps.cpp:1240
AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context)
These free functions allow clients of the API to not use classes in detail.
Definition: AffineExpr.cpp:502
LogicalResult getMultiAffineFunctionFromMap(AffineMap map, presburger::MultiAffineFunction &multiAff)
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value.
Definition: LogicalResult.h:72
AffineExpr getAffineSymbolExpr(unsigned position, MLIRContext *context)
Definition: AffineExpr.cpp:512
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
bool failed() const
Returns true if the provided LogicalResult corresponds to a failure value.
Definition: LogicalResult.h:44
Eliminates variable at the specified position using Fourier-Motzkin variable elimination.