AffineStructures.h

Go to the documentation of this file.

//===- AffineStructures.h - MLIR Affine Structures Class --------*- C++ -*-===//
//
// Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// Structures for affine/polyhedral analysis of ML functions.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_ANALYSIS_AFFINE_STRUCTURES_H
#define MLIR_ANALYSIS_AFFINE_STRUCTURES_H

#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/Support/LogicalResult.h"

namespace mlir {

class AffineApplyOp;
class AffineBound;
class AffineCondition;
class AffineMap;
class AffineForOp;
class IntegerSet;
class MLIRContext;
class Value;
class HyperRectangularSet;
class MemRefType;

struct MutableAffineMap {
public:
  MutableAffineMap() {}
  MutableAffineMap(AffineMap map);

  ArrayRef<AffineExpr> getResults() const { return results; }
  AffineExpr getResult(unsigned idx) const { return results[idx]; }
  void setResult(unsigned idx, AffineExpr result) { results[idx] = result; }
  unsigned getNumResults() const { return results.size(); }
  unsigned getNumDims() const { return numDims; }
  void setNumDims(unsigned d) { numDims = d; }
  unsigned getNumSymbols() const { return numSymbols; }
  void setNumSymbols(unsigned d) { numSymbols = d; }
  MLIRContext *getContext() const { return context; }

  bool isMultipleOf(unsigned idx, int64_t factor) const;

  void reset(AffineMap map);

  //-simplify-affine-expr pass).
  void simplify();
  AffineMap getAffineMap() const;

private:
  // Same meaning as AffineMap's fields.
  SmallVector<AffineExpr, 8> results;
  unsigned numDims;
  unsigned numSymbols;
  MLIRContext *context;
};

struct MutableIntegerSet {
public:
  MutableIntegerSet(IntegerSet set, MLIRContext *context);

  MutableIntegerSet(unsigned numDims, unsigned numSymbols,
                    MLIRContext *context);

  unsigned getNumDims() const { return numDims; }
  unsigned getNumSymbols() const { return numSymbols; }
  unsigned getNumConstraints() const { return constraints.size(); }

  void clear() {
    constraints.clear();
    eqFlags.clear();
  }

private:
  unsigned numDims;
  unsigned numSymbols;

  SmallVector<AffineExpr, 8> constraints;
  SmallVector<bool, 8> eqFlags;
};

// An affine value map can readily be constructed from an AffineApplyOp, or an
// AffineBound of a AffineForOp. It can be further transformed, substituted
// into, or simplified. Unlike AffineMap's, AffineValueMap's are created and
// destroyed during analysis. Only the AffineMap expressions that are pointed by
// them are unique'd. An affine value map, and the operations on it, maintain
// the invariant that operands are always positionally aligned with the
// AffineDimExpr and AffineSymbolExpr in the underlying AffineMap.
// TODO(bondhugula): Some of these classes could go into separate files.
class AffineValueMap {
public:
  // Creates an empty AffineValueMap (users should call 'reset' to reset map
  // and operands).
  AffineValueMap() {}
  AffineValueMap(AffineMap map, ArrayRef<Value> operands,
                 ArrayRef<Value> results = llvm::None);

  explicit AffineValueMap(AffineApplyOp applyOp);
  explicit AffineValueMap(AffineBound bound);

  ~AffineValueMap();

  // Resets this AffineValueMap with 'map', 'operands', and 'results'.
  void reset(AffineMap map, ArrayRef<Value> operands,
             ArrayRef<Value> results = llvm::None);

  static void difference(const AffineValueMap &a, const AffineValueMap &b,
                         AffineValueMap *res);

  inline bool isMultipleOf(unsigned idx, int64_t factor) const;

  bool isFunctionOf(unsigned idx, Value value) const;

  bool isConstant(unsigned idx) const;

  bool isIdentity() const;

  void setResult(unsigned i, AffineExpr e) { map.setResult(i, e); }
  AffineExpr getResult(unsigned i) { return map.getResult(i); }
  inline unsigned getNumOperands() const { return operands.size(); }
  inline unsigned getNumDims() const { return map.getNumDims(); }
  inline unsigned getNumSymbols() const { return map.getNumSymbols(); }
  inline unsigned getNumResults() const { return map.getNumResults(); }

  Value getOperand(unsigned i) const;
  ArrayRef<Value> getOperands() const;
  AffineMap getAffineMap() const;

private:
  // A mutable affine map.
  MutableAffineMap map;

  // TODO: make these trailing objects?
  SmallVector<Value, 4> operands;
  SmallVector<Value, 4> results;
};

// Both, the integer set being pointed to and the operands can change during
// analysis, simplification, and transformation.
class IntegerValueSet {
  // This will lead to a single equality in 'set'.
  explicit IntegerValueSet(const AffineValueMap &avm);

  // constructed ones), this method returns false conservatively.
  bool isEmpty() const;

  bool getNumDims() const { return set.getNumDims(); }
  bool getNumSymbols() const { return set.getNumSymbols(); }

private:
  // The set pointed to may itself change unlike in IR structures like
  // 'AffineCondition'.
  MutableIntegerSet set;
  SmallVector<Value, 4> operands;
};

class FlatAffineConstraints {
public:
  enum IdKind { Dimension, Symbol, Local };

  FlatAffineConstraints(unsigned numReservedInequalities,
                        unsigned numReservedEqualities,
                        unsigned numReservedCols, unsigned numDims = 0,
                        unsigned numSymbols = 0, unsigned numLocals = 0,
                        ArrayRef<Optional<Value>> idArgs = {})
      : numReservedCols(numReservedCols), numDims(numDims),
        numSymbols(numSymbols) {
    assert(numReservedCols >= numDims + numSymbols + 1);
    assert(idArgs.empty() || idArgs.size() == numDims + numSymbols + numLocals);
    equalities.reserve(numReservedCols * numReservedEqualities);
    inequalities.reserve(numReservedCols * numReservedInequalities);
    numIds = numDims + numSymbols + numLocals;
    ids.reserve(numReservedCols);
    if (idArgs.empty())
      ids.resize(numIds, None);
    else
      ids.append(idArgs.begin(), idArgs.end());
  }

  FlatAffineConstraints(unsigned numDims = 0, unsigned numSymbols = 0,
                        unsigned numLocals = 0,
                        ArrayRef<Optional<Value>> idArgs = {})
      : numReservedCols(numDims + numSymbols + numLocals + 1), numDims(numDims),
        numSymbols(numSymbols) {
    assert(numReservedCols >= numDims + numSymbols + 1);
    assert(idArgs.empty() || idArgs.size() == numDims + numSymbols + numLocals);
    numIds = numDims + numSymbols + numLocals;
    ids.reserve(numIds);
    if (idArgs.empty())
      ids.resize(numIds, None);
    else
      ids.append(idArgs.begin(), idArgs.end());
  }

  explicit FlatAffineConstraints(const HyperRectangularSet &set);

  // TODO(bondhugula)
  explicit FlatAffineConstraints(const AffineValueMap &avm);
  explicit FlatAffineConstraints(ArrayRef<const AffineValueMap *> avmRef);

  explicit FlatAffineConstraints(IntegerSet set);

  // TODO(bondhugula)
  explicit FlatAffineConstraints(const IntegerValueSet &set);

  FlatAffineConstraints(const FlatAffineConstraints &other);

  FlatAffineConstraints(ArrayRef<const AffineValueMap *> avmRef,
                        IntegerSet set);

  FlatAffineConstraints(const MutableAffineMap &map);

  ~FlatAffineConstraints() {}

  // Clears any existing data and reserves memory for the specified constraints.
  void reset(unsigned numReservedInequalities, unsigned numReservedEqualities,
             unsigned numReservedCols, unsigned numDims, unsigned numSymbols,
             unsigned numLocals = 0, ArrayRef<Value> idArgs = {});

  void reset(unsigned numDims = 0, unsigned numSymbols = 0,
             unsigned numLocals = 0, ArrayRef<Value> idArgs = {});

  void append(const FlatAffineConstraints &other);

  // Checks for emptiness by performing variable elimination on all identifiers,
  // running the GCD test on each equality constraint, and checking for invalid
  // constraints.
  // Returns true if the GCD test fails for any equality, or if any invalid
  // constraints are discovered on any row. Returns false otherwise.
  bool isEmpty() const;

  // Runs the GCD test on all equality constraints. Returns 'true' if this test
  // fails on any equality. Returns 'false' otherwise.
  // This test can be used to disprove the existence of a solution. If it
  // returns true, no integer solution to the equality constraints can exist.
  bool isEmptyByGCDTest() const;

  // Clones this object.
  std::unique_ptr<FlatAffineConstraints> clone() const;

  inline int64_t atEq(unsigned i, unsigned j) const {
    return equalities[i * numReservedCols + j];
  }
  inline int64_t &atEq(unsigned i, unsigned j) {
    return equalities[i * numReservedCols + j];
  }

  inline int64_t atIneq(unsigned i, unsigned j) const {
    return inequalities[i * numReservedCols + j];
  }

  inline int64_t &atIneq(unsigned i, unsigned j) {
    return inequalities[i * numReservedCols + j];
  }

  inline unsigned getNumCols() const { return numIds + 1; }

  inline unsigned getNumEqualities() const {
    assert(equalities.size() % numReservedCols == 0 &&
           "inconsistent equality buffer size");
    return equalities.size() / numReservedCols;
  }

  inline unsigned getNumInequalities() const {
    assert(inequalities.size() % numReservedCols == 0 &&
           "inconsistent inequality buffer size");
    return inequalities.size() / numReservedCols;
  }

  inline unsigned getNumReservedEqualities() const {
    return equalities.capacity() / numReservedCols;
  }

  inline unsigned getNumReservedInequalities() const {
    return inequalities.capacity() / numReservedCols;
  }

  inline ArrayRef<int64_t> getEquality(unsigned idx) const {
    return ArrayRef<int64_t>(&equalities[idx * numReservedCols], getNumCols());
  }

  inline ArrayRef<int64_t> getInequality(unsigned idx) const {
    return ArrayRef<int64_t>(&inequalities[idx * numReservedCols],
                             getNumCols());
  }

  AffineExpr toAffineExpr(unsigned idx, MLIRContext *context);

  //  TODO(bondhugula): add support for non-unit strides.
  LogicalResult addAffineForOpDomain(AffineForOp forOp);

  LogicalResult addLowerOrUpperBound(unsigned pos, AffineMap boundMap,
                                     ArrayRef<Value> operands, bool eq,
                                     bool lower = true);

  void getSliceBounds(unsigned offset, unsigned num, MLIRContext *context,
                      SmallVectorImpl<AffineMap> *lbMaps,
                      SmallVectorImpl<AffineMap> *ubMaps);

  LogicalResult addSliceBounds(ArrayRef<Value> values,
                               ArrayRef<AffineMap> lbMaps,
                               ArrayRef<AffineMap> ubMaps,
                               ArrayRef<Value> operands);

  // Adds an inequality (>= 0) from the coefficients specified in inEq.
  void addInequality(ArrayRef<int64_t> inEq);
  // Adds an equality from the coefficients specified in eq.
  void addEquality(ArrayRef<int64_t> eq);

  void addConstantLowerBound(unsigned pos, int64_t lb);
  void addConstantUpperBound(unsigned pos, int64_t ub);

  void addLocalFloorDiv(ArrayRef<int64_t> dividend, int64_t divisor);

  void addConstantLowerBound(ArrayRef<int64_t> expr, int64_t lb);
  void addConstantUpperBound(ArrayRef<int64_t> expr, int64_t ub);

  void setIdToConstant(unsigned pos, int64_t val);

  void setIdToConstant(Value id, int64_t val);

  bool findId(Value id, unsigned *pos) const;

  bool containsId(Value id) const;

  // Add identifiers of the specified kind - specified positions are relative to
  // the kind of identifier. The coefficient column corresponding to the added
  // identifier is initialized to zero. 'id' is the Value corresponding to the
  // identifier that can optionally be provided.
  void addDimId(unsigned pos, Value id = nullptr);
  void addSymbolId(unsigned pos, Value id = nullptr);
  void addLocalId(unsigned pos);
  void addId(IdKind kind, unsigned pos, Value id = nullptr);

  void addInductionVarOrTerminalSymbol(Value id);

  LogicalResult composeMap(const AffineValueMap *vMap);

  LogicalResult composeMatchingMap(AffineMap other);

  // TODO(bondhugula): deal with integer exactness when necessary - can return a
  // value to mark exactness for example.
  void projectOut(unsigned pos, unsigned num);
  inline void projectOut(unsigned pos) { return projectOut(pos, 1); }

  void projectOut(Value id);

  void removeId(IdKind idKind, unsigned pos);
  void removeId(unsigned pos);

  void removeDim(unsigned pos);

  void removeEquality(unsigned pos);
  void removeInequality(unsigned pos);

  void setDimSymbolSeparation(unsigned newSymbolCount);

  void convertLoopIVSymbolsToDims();

  void setAndEliminate(unsigned pos, int64_t constVal);

  LogicalResult constantFoldId(unsigned pos);

  void constantFoldIdRange(unsigned pos, unsigned num);

  LogicalResult unionBoundingBox(const FlatAffineConstraints &other);

  bool areIdsAlignedWithOther(const FlatAffineConstraints &other);

  //  Eg: Input: 'this'  has ((%i %j) [%M %N])
  //             'other' has (%k, %j) [%P, %N, %M])
  //      Output: both 'this', 'other' have (%i, %j, %k) [%M, %N, %P]
  //
  void mergeAndAlignIdsWithOther(unsigned offset, FlatAffineConstraints *other);

  unsigned getNumConstraints() const {
    return getNumInequalities() + getNumEqualities();
  }
  inline unsigned getNumIds() const { return numIds; }
  inline unsigned getNumDimIds() const { return numDims; }
  inline unsigned getNumSymbolIds() const { return numSymbols; }
  inline unsigned getNumDimAndSymbolIds() const { return numDims + numSymbols; }
  inline unsigned getNumLocalIds() const {
    return numIds - numDims - numSymbols;
  }

  inline ArrayRef<Optional<Value>> getIds() const {
    return {ids.data(), ids.size()};
  }
  inline MutableArrayRef<Optional<Value>> getIds() {
    return {ids.data(), ids.size()};
  }

  inline Optional<Value> getId(unsigned pos) const { return ids[pos]; }
  inline Optional<Value> &getId(unsigned pos) { return ids[pos]; }

  inline Value getIdValue(unsigned pos) const {
    assert(ids[pos].hasValue() && "identifier's Value not set");
    return ids[pos].getValue();
  }

  void getIdValues(unsigned start, unsigned end,
                   SmallVectorImpl<Value> *values) const {
    assert((start < numIds || start == end) && "invalid start position");
    assert(end <= numIds && "invalid end position");
    values->clear();
    values->reserve(end - start);
    for (unsigned i = start; i < end; i++) {
      values->push_back(getIdValue(i));
    }
  }
  inline void getAllIdValues(SmallVectorImpl<Value> *values) const {
    getIdValues(0, numIds, values);
  }

  inline void setIdValue(unsigned pos, Value val) {
    assert(pos < numIds && "invalid id position");
    ids[pos] = val;
  }
  void setIdValues(unsigned start, unsigned end, ArrayRef<Value> values) {
    assert((start < numIds || end == start) && "invalid start position");
    assert(end <= numIds && "invalid end position");
    assert(values.size() == end - start);
    for (unsigned i = start; i < end; ++i)
      ids[i] = values[i - start];
  }

  void clearAndCopyFrom(const FlatAffineConstraints &other);

  Optional<int64_t>
  getConstantBoundOnDimSize(unsigned pos,
                            SmallVectorImpl<int64_t> *lb = nullptr,
                            int64_t *lbFloorDivisor = nullptr,
                            SmallVectorImpl<int64_t> *ub = nullptr) const;

  Optional<int64_t> getConstantLowerBound(unsigned pos) const;

  Optional<int64_t> getConstantUpperBound(unsigned pos) const;

  std::pair<AffineMap, AffineMap>
  getLowerAndUpperBound(unsigned pos, unsigned offset, unsigned num,
                        unsigned symStartPos, ArrayRef<AffineExpr> localExprs,
                        MLIRContext *context) const;

  bool isHyperRectangular(unsigned pos, unsigned num) const;

  void removeTrivialRedundancy();

  void removeRedundantInequalities();

  // Removes all equalities and inequalities.
  void clearConstraints();

  void print(raw_ostream &os) const;
  void dump() const;

private:
  bool hasConsistentState() const;

  bool hasInvalidConstraint() const;

  template <bool isLower>
  Optional<int64_t> computeConstantLowerOrUpperBound(unsigned pos);

  // Eliminates a single identifier at 'position' from equality and inequality
  // constraints. Returns 'success' if the identifier was eliminated, and
  // 'failure' otherwise.
  inline LogicalResult gaussianEliminateId(unsigned position) {
    return success(gaussianEliminateIds(position, position + 1) == 1);
  }

  // Eliminates identifiers from equality and inequality constraints
  // in column range [posStart, posLimit).
  // Returns the number of variables eliminated.
  unsigned gaussianEliminateIds(unsigned posStart, unsigned posLimit);

  // See implementation comments for more details.
  void FourierMotzkinEliminate(unsigned pos, bool darkShadow = false,
                               bool *isResultIntegerExact = nullptr);

  void GCDTightenInequalities();

  void normalizeConstraintsByGCD();

  void removeIdRange(unsigned idStart, unsigned idLimit);

  SmallVector<int64_t, 64> equalities;

  SmallVector<int64_t, 64> inequalities;

  unsigned numReservedCols;

  unsigned numIds;

  unsigned numDims;

  unsigned numSymbols;

  SmallVector<Optional<Value>, 8> ids;

  // The rationale for 32 is that in the typical simplest of cases, an
  // identifier is expected to have one lower bound and one upper bound
  // constraint. With a level of tiling or a connection to another identifier
  // through a div or mod, an extra pair of bounds gets added. As a limit, we
  // don't expect an identifier to have more than 32 lower/upper/equality
  // constraints. This is conservatively set low and can be raised if needed.
  constexpr static unsigned kExplosionFactor = 32;
};

AffineExpr simplifyAffineExpr(AffineExpr expr, unsigned numDims,
                              unsigned numSymbols);

LogicalResult getFlattenedAffineExpr(AffineExpr expr, unsigned numDims,
                                     unsigned numSymbols,
                                     SmallVectorImpl<int64_t> *flattenedExpr,
                                     FlatAffineConstraints *cst = nullptr);

LogicalResult
getFlattenedAffineExprs(AffineMap map,
                        std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
                        FlatAffineConstraints *cst = nullptr);
LogicalResult
getFlattenedAffineExprs(IntegerSet set,
                        std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
                        FlatAffineConstraints *cst = nullptr);

} // end namespace mlir.

#endif // MLIR_ANALYSIS_AFFINE_STRUCTURES_H