[mlir][Affine] Revisit and simplify composeAffineMapAndOperands.

[nicolasvasilache]

Summary

In prehistorical times, AffineApplyOp was allowed to produce multiple values.
This allowed the creation of intricate SSA use-def chains.
AffineApplyNormalizer was originally introduced as a means of reusing the AffineMap::compose method to write SSA use-def chains.
Unfortunately, symbols that were produced by an AffineApplyOp needed to be promoted to dims and reordered for the mathematical composition to be valid.

Since then, single result AffineApplyOp became the law of the land but the original assumptions were not revisited.

This revision revisits these assumptions and retires AffineApplyNormalizer.

[nicolasvasilache]

clang-tidy

[ThomasRaoux]

Thanks @nicolasvasilache! I'm don't know this code in details but the changes look like a great clean up to me.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

1771

nit: are all those formatting changes intended? It looks like the previous code was fitting in 80 col?

[ftynse]

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

578

LogicalResult

623–630

Nit: some top-level doc?

624

Nit: why pointer rather than reference?

664

This compares actual underlying vectors. Does not seem to be a correctness problems because dim/symbol expressions are unique, but probably less efficient. How about comparing pointers instead?

[nicolasvasilache]

Address.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

624

prehistoric API, not looking to change it in this CL ..

1771

def not intentional .. sigh internal tooling..

mlir/include/mlir/Dialect/Affine/IR/AffineOps.h

private:
   AffineMap map;
 
   AffineBound(AffineForOp op, unsigned opStart, unsigned opEnd, AffineMap map)
       : op(op), opStart(opStart), opEnd(opEnd), map(map) {}
 
   friend class AffineForOp;
 };
 
-/// An `AffineApplyNormalizer` is a helper class that supports renumbering
-/// operands of AffineApplyOp. This acts as a reindexing map of Value to
-/// positional dims or symbols and allows simplifications such as:
-///
-/// ```mlir
-///    %1 = affine.apply (d0, d1) -> (d0 - d1) (%0, %0)
-/// ```
-///
-/// into:
-///
-/// ```mlir
-///    %1 = affine.apply () -> (0)
-/// ```
-struct AffineApplyNormalizer {
-  AffineApplyNormalizer(AffineMap map, ArrayRef<Value> operands);
-
-  /// Returns the AffineMap resulting from normalization.
-  AffineMap getAffineMap() { return affineMap; }
-
-  SmallVector<Value, 8> getOperands() {
-    SmallVector<Value, 8> res(reorderedDims);
-    res.append(concatenatedSymbols.begin(), concatenatedSymbols.end());
-    return res;
-  }
-
-  unsigned getNumSymbols() { return concatenatedSymbols.size(); }
-  unsigned getNumDims() { return reorderedDims.size(); }
-
-  /// Normalizes 'otherMap' and its operands 'otherOperands' to map to this
-  /// normalizer's coordinate space.
-  void normalize(AffineMap *otherMap, SmallVectorImpl<Value> *otherOperands);
-
-private:
-  /// Helper function to insert `v` into the coordinate system of the current
-  /// AffineApplyNormalizer. Returns the AffineDimExpr with the corresponding
-  /// renumbered position.
-  AffineDimExpr renumberOneDim(Value v);
-
-  /// Given an `other` normalizer, this rewrites `other.affineMap` in the
-  /// coordinate system of the current AffineApplyNormalizer.
-  /// Returns the rewritten AffineMap and updates the dims and symbols of
-  /// `this`.
-  AffineMap renumber(const AffineApplyNormalizer &other);
-
-  /// Maps of Value to position in `affineMap`.
-  DenseMap<Value, unsigned> dimValueToPosition;
-
-  /// Ordered dims and symbols matching positional dims and symbols in
-  /// `affineMap`.
-  SmallVector<Value, 8> reorderedDims;
-  SmallVector<Value, 8> concatenatedSymbols;
-
-  /// The number of symbols in concatenated symbols that belong to the original
-  /// map as opposed to those concatendated during map composition.
-  unsigned numProperSymbols;
-
-  AffineMap affineMap;
-
-  /// Used with RAII to control the depth at which AffineApply are composed
-  /// recursively. Only accepts depth 1 for now to allow a behavior where a
-  /// newly composed AffineApplyOp does not increase the length of the chain of
-  /// AffineApplyOps. Full composition is implemented iteratively on top of
-  /// this behavior.
-  static unsigned &affineApplyDepth() {
-    static thread_local unsigned depth = 0;
-    return depth;
-  }
-  static constexpr unsigned kMaxAffineApplyDepth = 1;
-
-  AffineApplyNormalizer() : numProperSymbols(0) { affineApplyDepth()++; }
-
-public:
-  ~AffineApplyNormalizer() { affineApplyDepth()--; }
-};
-
 } // end namespace mlir
 
 #endif

mlir/include/mlir/IR/AffineExpr.h

public:
   int64_t getLargestKnownDivisor() const;
 
   /// Return true if the affine expression is a multiple of 'factor'.
   bool isMultipleOf(int64_t factor) const;
 
   /// Return true if the affine expression involves AffineDimExpr `position`.
   bool isFunctionOfDim(unsigned position) const;
 
+  /// Return true if the affine expression involves AffineSymbolExpr `position`.
+  bool isFunctionOfSymbol(unsigned position) const;
+
   /// Walk all of the AffineExpr's in this expression in postorder.
   void walk(std::function<void(AffineExpr)> callback) const;
 
   /// This method substitutes any uses of dimensions and symbols (e.g.
   /// dim#0 with dimReplacements[0]) and returns the modified expression tree.
   /// This is a dense replacement method: a replacement must be specified for
   /// every single dim and symbol.
   AffineExpr replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
                                    ArrayRef<AffineExpr> symReplacements) const;
 
   /// Sparse replace method. Replace `expr` by `replacement` and return the
   /// modified expression tree.
   AffineExpr replace(AffineExpr expr, AffineExpr replacement) const;
 
   /// Sparse replace method. If `*this` appears in `map` replaces it by
   /// `map[*this]` and return the modified expression tree. Otherwise traverse
   /// `*this` and apply replace with `map` on its subexpressions.
   AffineExpr replace(const DenseMap<AffineExpr, AffineExpr> &map) const;
 
-  /// Replace symbols[0 .. numDims - 1] by
-  /// symbols[shift .. shift + numDims - 1].
+  /// Replace dims[0 .. numDims - 1] by dims[shift .. shift + numDims - 1].
+  AffineExpr shiftDims(unsigned numDims, unsigned shift) const;
+
+  /// Replace symbols[0 .. numSymbols - 1] by
+  ///         symbols[shift .. shift + numSymbols - 1].
   AffineExpr shiftSymbols(unsigned numSymbols, unsigned shift) const;
 
   AffineExpr operator+(int64_t v) const;
   AffineExpr operator+(AffineExpr other) const;
   AffineExpr operator-() const;
   AffineExpr operator-(int64_t v) const;
   AffineExpr operator-(AffineExpr other) const;
   AffineExpr operator*(int64_t v) const;

mlir/include/mlir/IR/AffineMap.h

public:
 
   ArrayRef<AffineExpr> getResults() const;
   AffineExpr getResult(unsigned idx) const;
 
   /// Extracts the position of the dimensional expression at the given result,
   /// when the caller knows it is safe to do so.
   unsigned getDimPosition(unsigned idx) const;
 
+  /// Return true if any affine expression involves AffineDimExpr `position`.
+  bool isFunctionOfDim(unsigned position) const {
+    return llvm::any_of(getResults(), [&](AffineExpr e) {
+      return e.isFunctionOfDim(position);
+    });
+  }
+
+  /// Return true if any affine expression involves AffineSymbolExpr `position`.
+  bool isFunctionOfSymbol(unsigned position) const {
+    return llvm::any_of(getResults(), [&](AffineExpr e) {
+      return e.isFunctionOfSymbol(position);
+    });
+  }
+
   /// Walk all of the AffineExpr's in this mapping. Each node in an expression
   /// tree is visited in postorder.
   void walkExprs(std::function<void(AffineExpr)> callback) const;
 
   /// This method substitutes any uses of dimensions and symbols (e.g.
   /// dim#0 with dimReplacements[0]) in subexpressions and returns the modified
   /// expression mapping.  Because this can be used to eliminate dims and
   /// symbols, the client needs to specify the number of dims and symbols in
   /// the result.  The returned map always has the same number of results.
   AffineMap replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
                                   ArrayRef<AffineExpr> symReplacements,
                                   unsigned numResultDims,
                                   unsigned numResultSyms) const;
 
+  /// Sparse replace method. Apply AffineExpr::replace(`expr`, `replacement`) to
+  /// each of the results and return a new AffineMap with the new results and
+  /// with the specified number of dims and symbols.
+  AffineMap replace(AffineExpr expr, AffineExpr replacement,
+                    unsigned numResultDims, unsigned numResultSyms) const;
+
+  /// Sparse replace method. Apply AffineExpr::replace(`map`) to each of the
+  /// results and return a new AffineMap with the new results and with the
+  /// specified number of dims and symbols.
+  AffineMap replace(const DenseMap<AffineExpr, AffineExpr> &map,
+                    unsigned numResultDims, unsigned numResultSyms) const;
+
+  /// Replace dims[0 .. numDims - 1] by dims[shift .. shift + numDims - 1].
+  AffineMap shiftDims(unsigned shift) const {
+    return AffineMap::get(
+        getNumDims() + shift, getNumSymbols(),
+        llvm::to_vector<4>(llvm::map_range(
+            getResults(),
+            [&](AffineExpr e) { return e.shiftDims(getNumDims(), shift); })),
+        getContext());
+  }
+
+  /// Replace symbols[0 .. numSymbols - 1] by
+  ///         symbols[shift .. shift + numSymbols - 1].
+  AffineMap shiftSymbols(unsigned shift) const {
+    return AffineMap::get(getNumDims(), getNumSymbols() + shift,
+                          llvm::to_vector<4>(llvm::map_range(
+                              getResults(),
+                              [&](AffineExpr e) {
+                                return e.shiftSymbols(getNumSymbols(), shift);
+                              })),
+                          getContext());
+  }
+
   /// Folds the results of the application of an affine map on the provided
   /// operands to a constant if possible.
   LogicalResult constantFold(ArrayRef<Attribute> operandConstants,
                              SmallVectorImpl<Attribute> &results) const;
 
   /// Propagates the constant operands into this affine map. Operands are
   /// allowed to be null, at which point they are treated as non-constant. This
   /// does not change the number of symbols and dimensions. Returns a new map,

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

OpFoldResult AffineApplyOp::fold(ArrayRef<Attribute> operands) {
 
   // Otherwise, default to folding the map.
   SmallVector<Attribute, 1> result;
   if (failed(map.constantFold(operands, result)))
     return {};
   return result[0];
 }
 
-AffineDimExpr AffineApplyNormalizer::renumberOneDim(Value v) {
-  DenseMap<Value, unsigned>::iterator iterPos;
-  bool inserted = false;
-  std::tie(iterPos, inserted) =
-      dimValueToPosition.insert(std::make_pair(v, dimValueToPosition.size()));
-  if (inserted) {
-    reorderedDims.push_back(v);
-  }
-  return getAffineDimExpr(iterPos->second, v.getContext())
-      .cast<AffineDimExpr>();
-}
-
-AffineMap AffineApplyNormalizer::renumber(const AffineApplyNormalizer &other) {
-  SmallVector<AffineExpr, 8> dimRemapping;
-  for (auto v : other.reorderedDims) {
-    auto kvp = other.dimValueToPosition.find(v);
-    if (dimRemapping.size() <= kvp->second)
-      dimRemapping.resize(kvp->second + 1);
-    dimRemapping[kvp->second] = renumberOneDim(kvp->first);
-  }
-  unsigned numSymbols = concatenatedSymbols.size();
-  unsigned numOtherSymbols = other.concatenatedSymbols.size();
-  SmallVector<AffineExpr, 8> symRemapping(numOtherSymbols);
-  for (unsigned idx = 0; idx < numOtherSymbols; ++idx) {
-    symRemapping[idx] =
-        getAffineSymbolExpr(idx + numSymbols, other.affineMap.getContext());
-  }
-  concatenatedSymbols.insert(concatenatedSymbols.end(),
-                             other.concatenatedSymbols.begin(),
-                             other.concatenatedSymbols.end());
-  auto map = other.affineMap;
-  return map.replaceDimsAndSymbols(dimRemapping, symRemapping,
-                                   reorderedDims.size(),
-                                   concatenatedSymbols.size());
-}
-
-// Gather the positions of the operands that are produced by an AffineApplyOp.
-static llvm::SetVector<unsigned>
-indicesFromAffineApplyOp(ArrayRef<Value> operands) {
-  llvm::SetVector<unsigned> res;
-  for (auto en : llvm::enumerate(operands))
-    if (isa_and_nonnull<AffineApplyOp>(en.value().getDefiningOp()))
-      res.insert(en.index());
-  return res;
-}
-
-// Support the special case of a symbol coming from an AffineApplyOp that needs
-// to be composed into the current AffineApplyOp.
-// This case is handled by rewriting all such symbols into dims for the purpose
-// of allowing mathematical AffineMap composition.
-// Returns an AffineMap where symbols that come from an AffineApplyOp have been
-// rewritten as dims and are ordered after the original dims.
-// TODO: This promotion makes AffineMap lose track of which
-// symbols are represented as dims. This loss is static but can still be
-// recovered dynamically (with `isValidSymbol`). Still this is annoying for the
-// semi-affine map case. A dynamic canonicalization of all dims that are valid
-// symbols (a.k.a `canonicalizePromotedSymbols`) into symbols helps and even
-// results in better simplifications and foldings. But we should evaluate
-// whether this behavior is what we really want after using more.
-static AffineMap promoteComposedSymbolsAsDims(AffineMap map,
-                                              ArrayRef<Value> symbols) {
-  if (symbols.empty()) {
-    return map;
-  }
-
-  // Extract the symbol positions that come from an AffineApplyOp and
-  // needs to be rewritten as dims.
-  auto symPositions = indicesFromAffineApplyOp(symbols);
-  if (symPositions.empty()) {
-    return map;
-  }
-
-  // Create the new map by replacing each symbol at pos by the next new dim.
-  unsigned numDims = map.getNumDims();
-  unsigned numSymbols = map.getNumSymbols();
-  unsigned numNewDims = 0;
-  unsigned numNewSymbols = 0;
-  SmallVector<AffineExpr, 8> symReplacements(numSymbols);
-  for (unsigned i = 0; i < numSymbols; ++i) {
-    symReplacements[i] =
-        symPositions.count(i) > 0
-            ? getAffineDimExpr(numDims + numNewDims++, map.getContext())
-            : getAffineSymbolExpr(numNewSymbols++, map.getContext());
-  }
-  assert(numSymbols >= numNewDims);
-  AffineMap newMap = map.replaceDimsAndSymbols(
-      {}, symReplacements, numDims + numNewDims, numNewSymbols);
-
-  return newMap;
-}
-
-/// The AffineNormalizer composes AffineApplyOp recursively. Its purpose is to
-/// keep a correspondence between the mathematical `map` and the `operands` of
-/// a given AffineApplyOp. This correspondence is maintained by iterating over
-/// the operands and forming an `auxiliaryMap` that can be composed
-/// mathematically with `map`. To keep this correspondence in cases where
-/// symbols are produced by affine.apply operations, we perform a local rewrite
-/// of symbols as dims.
-///
-/// Rationale for locally rewriting symbols as dims:
-/// ================================================
-/// The mathematical composition of AffineMap must always concatenate symbols
-/// because it does not have enough information to do otherwise. For example,
-/// composing `(d0)[s0] -> (d0 + s0)` with itself must produce
-/// `(d0)[s0, s1] -> (d0 + s0 + s1)`.
-///
-/// The result is only equivalent to `(d0)[s0] -> (d0 + 2 * s0)` when
-/// applied to the same mlir::Value for both s0 and s1.
-/// As a consequence mathematical composition of AffineMap always concatenates
-/// symbols.
-///
-/// When AffineMaps are used in AffineApplyOp however, they may specify
-/// composition via symbols, which is ambiguous mathematically. This corner case
-/// is handled by locally rewriting such symbols that come from AffineApplyOp
-/// into dims and composing through dims.
-/// TODO: Composition via symbols comes at a significant code
-/// complexity. Alternatively we should investigate whether we want to
-/// explicitly disallow symbols coming from affine.apply and instead force the
-/// user to compose symbols beforehand. The annoyances may be small (i.e. 1 or 2
-/// extra API calls for such uses, which haven't popped up until now) and the
-/// benefit potentially big: simpler and more maintainable code for a
-/// non-trivial, recursive, procedure.
-AffineApplyNormalizer::AffineApplyNormalizer(AffineMap map,
-                                             ArrayRef<Value> operands)
-    : AffineApplyNormalizer() {
-  static_assert(kMaxAffineApplyDepth > 0, "kMaxAffineApplyDepth must be > 0");
-  assert(map.getNumInputs() == operands.size() &&
-         "number of operands does not match the number of map inputs");
-
-  LLVM_DEBUG(map.print(dbgs() << "\nInput map: "));
-
-  // Promote symbols that come from an AffineApplyOp to dims by rewriting the
-  // map to always refer to:
-  //   (dims, symbols coming from AffineApplyOp, other symbols).
-  // The order of operands can remain unchanged.
-  // This is a simplification that relies on 2 ordering properties:
-  //   1. rewritten symbols always appear after the original dims in the map;
-  //   2. operands are traversed in order and either dispatched to:
-  //      a. auxiliaryExprs (dims and symbols rewritten as dims);
-  //      b. concatenatedSymbols (all other symbols)
-  // This allows operand order to remain unchanged.
-  unsigned numDimsBeforeRewrite = map.getNumDims();
-  map = promoteComposedSymbolsAsDims(map,
-                                     operands.take_back(map.getNumSymbols()));
-
-  LLVM_DEBUG(map.print(dbgs() << "\nRewritten map: "));
-
-  SmallVector<AffineExpr, 8> auxiliaryExprs;
-  bool furtherCompose = (affineApplyDepth() <= kMaxAffineApplyDepth);
-  // We fully spell out the 2 cases below. In this particular instance a little
-  // code duplication greatly improves readability.
-  // Note that the first branch would disappear if we only supported full
-  // composition (i.e. infinite kMaxAffineApplyDepth).
-  if (!furtherCompose) {
-    // 1. Only dispatch dims or symbols.
-    for (auto en : llvm::enumerate(operands)) {
-      auto t = en.value();
-      assert(t.getType().isIndex());
-      bool isDim = (en.index() < map.getNumDims());
-      if (isDim) {
-        // a. The mathematical composition of AffineMap composes dims.
-        auxiliaryExprs.push_back(renumberOneDim(t));
-      } else {
-        // b. The mathematical composition of AffineMap concatenates symbols.
-        //    We do the same for symbol operands.
-        concatenatedSymbols.push_back(t);
-      }
-    }
-  } else {
-    assert(numDimsBeforeRewrite <= operands.size());
-    // 2. Compose AffineApplyOps and dispatch dims or symbols.
-    for (unsigned i = 0, e = operands.size(); i < e; ++i) {
-      auto t = operands[i];
-      auto affineApply = t.getDefiningOp<AffineApplyOp>();
-      if (affineApply) {
-        // a. Compose affine.apply operations.
-        LLVM_DEBUG(affineApply->print(
-            dbgs() << "\nCompose AffineApplyOp recursively: "));
-        AffineMap affineApplyMap = affineApply.getAffineMap();
-        SmallVector<Value, 8> affineApplyOperands(
-            affineApply.getOperands().begin(), affineApply.getOperands().end());
-        AffineApplyNormalizer normalizer(affineApplyMap, affineApplyOperands);
-
-        LLVM_DEBUG(normalizer.affineMap.print(
-            dbgs() << "\nRenumber into current normalizer: "));
-
-        auto renumberedMap = renumber(normalizer);
-
-        LLVM_DEBUG(
-            renumberedMap.print(dbgs() << "\nRecursive composition yields: "));
+/// Replace all occurrences of AffineExpr at position `pos` in `map` by the
+/// defining AffineApplyOp expression and operands.
+/// When `dimOrSymbolPosition < dims.size()`, AffineDimExpr@[pos] is replaced.
+/// When `dimOrSymbolPosition >= dims.size()`,
+/// AffineSymbolExpr@[pos - dims.size()] is replaced.
+/// Mutate `map`,`dims` and `syms` in place as follows:
+///   1. `dims` and `syms` are only appended to.
+///   2. `map` dim and symbols are gradually shifted to higer positions.
+///   3. Old `dim` and `sym` entries are replaced by nullptr
+/// This avoids the need for any bookkeeping.
+static LogicalResult replaceDimOrSym(AffineMap *map,
+                                     unsigned dimOrSymbolPosition,
+                                     SmallVectorImpl<Value> &dims,
+                                     SmallVectorImpl<Value> &syms) {
+  bool isDimReplacement = (dimOrSymbolPosition < dims.size());
+  unsigned pos = isDimReplacement ? dimOrSymbolPosition
+                                  : dimOrSymbolPosition - dims.size();
+  Value &v = isDimReplacement ? dims[pos] : syms[pos];
+  if (!v)
+    return failure();
+
+  auto affineApply = v.getDefiningOp<AffineApplyOp>();
+  if (!affineApply)
+    return failure();
+
+  // At this point we will perform a replacement of `v`, set the entry in `dim`
+  // or `sym` to nullptr immediately.
+  v = nullptr;
+
+  // Compute the map, dims and symbols coming from the AffineApplyOp.
+  AffineMap composeMap = affineApply.getAffineMap();
+  assert(composeMap.getNumResults() == 1 && "affine.apply with >1 results");
+  AffineExpr composeExpr =
+      composeMap.shiftDims(dims.size()).shiftSymbols(syms.size()).getResult(0);
+  ValueRange composeDims =
+      affineApply.getMapOperands().take_front(composeMap.getNumDims());
+  ValueRange composeSyms =
+      affineApply.getMapOperands().take_back(composeMap.getNumSymbols());
+
+  // Perform the replacement and append the dims and symbols where relevant.
+  MLIRContext *ctx = map->getContext();
+  AffineExpr toReplace = isDimReplacement ? getAffineDimExpr(pos, ctx)
+                                          : getAffineSymbolExpr(pos, ctx);
+  *map = map->replace(toReplace, composeExpr, dims.size(), syms.size());
+  dims.append(composeDims.begin(), composeDims.end());
+  syms.append(composeSyms.begin(), composeSyms.end());
 
-        auxiliaryExprs.push_back(renumberedMap.getResult(0));
-      } else {
-        if (i < numDimsBeforeRewrite) {
-          // b. The mathematical composition of AffineMap composes dims.
-          auxiliaryExprs.push_back(renumberOneDim(t));
-        } else {
-          // c. The mathematical composition of AffineMap concatenates symbols.
-          //    Note that the map composition will put symbols already present
-          //    in the map before any symbols coming from the auxiliary map, so
-          //    we insert them before any symbols that are due to renumbering,
-          //    and after the proper symbols we have seen already.
-          concatenatedSymbols.insert(
-              std::next(concatenatedSymbols.begin(), numProperSymbols++), t);
-        }
-      }
+  return success();
+
+/// Iterate over `operands` and fold away all those produced by an AffineApplyOp
+/// iteratively. Perform canonicalization of map and operands as well as
+/// AffineMap simplification. `map` and `operands` are mutated in place.
+static void composeAffineMapAndOperands(AffineMap *map,
+                                        SmallVectorImpl<Value> *operands) {
+  if (map->getNumResults() == 0) {
+    canonicalizeMapAndOperands(map, operands);
+    *map = simplifyAffineMap(*map);
+    return;
+
+  MLIRContext *ctx = map->getContext();
+  SmallVector<Value, 4> dims(operands->begin(),
+                             operands->begin() + map->getNumDims());
+  SmallVector<Value, 4> syms(operands->begin() + map->getNumDims(),
+                             operands->end());
+
+  // Iterate over dims and symbols coming from AffineApplyOp and replace until
+  // exhaustion. This iteratively mutates `map`, `dims` and `syms`. Both `dims`
+  // and `syms` can only increase by construction.
+  // The implementation uses a `while` loop to support the case of symbols
+  // that may be constructed from dims ;this may be overkill.
+  while (true) {
+    bool changed = false;
+    for (unsigned pos = 0; pos != dims.size() + syms.size(); ++pos)
+      if ((changed |= succeeded(replaceDimOrSym(map, pos, dims, syms))))
+        break;
+    if (!changed)
+      break;
+  }
+
+  // Clear operands so we can fill them anew.
+  operands->clear();
+
+  // At this point we may have introduced null operands, prune them out before
+  // canonicalizing map and operands.
+  unsigned nDims = 0, nSyms = 0;
+  SmallVector<AffineExpr, 4> dimReplacements, symReplacements;
+  dimReplacements.reserve(dims.size());
+  symReplacements.reserve(syms.size());
+  for (auto *container : {&dims, &syms}) {
+    bool isDim = (container == &dims);
+    auto &repls = isDim ? dimReplacements : symReplacements;
+    for (auto en : llvm::enumerate(*container)) {
+      Value v = en.value();
+      if (!v) {
+        assert(isDim ? !map->isFunctionOfDim(en.index())
+                     : !map->isFunctionOfSymbol(en.index()) &&
+                           "map is function of unexpected expr@pos");
+        repls.push_back(getAffineConstantExpr(0, ctx));
+        continue;
       }
+      repls.push_back(isDim ? getAffineDimExpr(nDims++, ctx)
+                            : getAffineSymbolExpr(nSyms++, ctx));
+      operands->push_back(v);
     }
-
-  // Early exit if `map` is already composed.
-  if (auxiliaryExprs.empty()) {
-    affineMap = map;
-    return;
   }
+  *map = map->replaceDimsAndSymbols(dimReplacements, symReplacements, nDims,
+                                    nSyms);
 
-  assert(concatenatedSymbols.size() >= map.getNumSymbols() &&
-         "Unexpected number of concatenated symbols");
-  auto numDims = dimValueToPosition.size();
-  auto numSymbols = concatenatedSymbols.size() - map.getNumSymbols();
-  auto auxiliaryMap =
-      AffineMap::get(numDims, numSymbols, auxiliaryExprs, map.getContext());
-
-  LLVM_DEBUG(map.print(dbgs() << "\nCompose map: "));
-  LLVM_DEBUG(auxiliaryMap.print(dbgs() << "\nWith map: "));
-  LLVM_DEBUG(map.compose(auxiliaryMap).print(dbgs() << "\nResult: "));
-
-  // TODO: Disabling simplification results in major speed gains.
-  // Another option is to cache the results as it is expected a lot of redundant
-  // work is performed in practice.
-  affineMap = simplifyAffineMap(map.compose(auxiliaryMap));
-
-  LLVM_DEBUG(affineMap.print(dbgs() << "\nSimplified result: "));
-  LLVM_DEBUG(dbgs() << "\n");
-}
-
-void AffineApplyNormalizer::normalize(AffineMap *otherMap,
-                                      SmallVectorImpl<Value> *otherOperands) {
-  AffineApplyNormalizer other(*otherMap, *otherOperands);
-  *otherMap = renumber(other);
-
-  otherOperands->reserve(reorderedDims.size() + concatenatedSymbols.size());
-  otherOperands->assign(reorderedDims.begin(), reorderedDims.end());
-  otherOperands->append(concatenatedSymbols.begin(), concatenatedSymbols.end());
-}
-
-/// Implements `map` and `operands` composition and simplification to support
-/// `makeComposedAffineApply`. This can be called to achieve the same effects
-/// on `map` and `operands` without creating an AffineApplyOp that needs to be
-/// immediately deleted.
-static void composeAffineMapAndOperands(AffineMap *map,
-                                        SmallVectorImpl<Value> *operands) {
-  AffineApplyNormalizer normalizer(*map, *operands);
-  auto normalizedMap = normalizer.getAffineMap();
-  auto normalizedOperands = normalizer.getOperands();
-  canonicalizeMapAndOperands(&normalizedMap, &normalizedOperands);
-  *map = normalizedMap;
-  *operands = normalizedOperands;
-  assert(*map);
+  // Canonicalize and simplify before returning.
+  canonicalizeMapAndOperands(map, operands);
+  *map = simplifyAffineMap(*map);
 }
 
 void mlir::fullyComposeAffineMapAndOperands(AffineMap *map,
                                             SmallVectorImpl<Value> *operands) {
   while (llvm::any_of(*operands, [](Value v) {
     return isa_and_nonnull<AffineApplyOp>(v.getDefiningOp());
   })) {
     composeAffineMapAndOperands(map, operands);
void mlir::extractForInductionVars(ArrayRef<AffineForOp> forInsts,
                                    SmallVectorImpl<Value> *ivs) {
   ivs->reserve(forInsts.size());
   for (auto forInst : forInsts)
     ivs->push_back(forInst.getInductionVar());
 }
 
 /// Builds an affine loop nest, using "loopCreatorFn" to create individual loop
 /// operations.
 template <typename BoundListTy, typename LoopCreatorTy>
 static void buildAffineLoopNestImpl(
     OpBuilder &builder, Location loc, BoundListTy lbs, BoundListTy ubs,
     ArrayRef<int64_t> steps,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn,
     LoopCreatorTy &&loopCreatorFn) {
   assert(lbs.size() == ubs.size() && "Mismatch in number of arguments");
   assert(lbs.size() == steps.size() && "Mismatch in number of arguments");
 

mlir/lib/Dialect/Affine/IR/AffineValueMap.cpp

void AffineValueMap::reset(AffineMap map, ValueRange operands,
   this->operands.assign(operands.begin(), operands.end());
   this->results.assign(results.begin(), results.end());
 }
 
 void AffineValueMap::difference(const AffineValueMap &a,
                                 const AffineValueMap &b, AffineValueMap *res) {
   assert(a.getNumResults() == b.getNumResults() && "invalid inputs");
 
-  // Fully compose A's map + operands.
-  auto aMap = a.getAffineMap();
-  SmallVector<Value, 4> aOperands(a.getOperands().begin(),
-                                  a.getOperands().end());
-  fullyComposeAffineMapAndOperands(&aMap, &aOperands);
-
-  // Use the affine apply normalizer to get B's map into A's coordinate space.
-  AffineApplyNormalizer normalizer(aMap, aOperands);
-  SmallVector<Value, 4> bOperands(b.getOperands().begin(),
-                                  b.getOperands().end());
-  auto bMap = b.getAffineMap();
-  normalizer.normalize(&bMap, &bOperands);
-
-  assert(std::equal(bOperands.begin(), bOperands.end(),
-                    normalizer.getOperands().begin()) &&
-         "operands are expected to be the same after normalization");
+  SmallVector<Value, 4> allOperands;
+  allOperands.reserve(a.getNumOperands() + b.getNumOperands());
+  auto aDims = a.getOperands().take_front(a.getNumDims());
+  auto bDims = b.getOperands().take_front(b.getNumDims());
+  auto aSyms = a.getOperands().take_back(a.getNumSymbols());
+  auto bSyms = b.getOperands().take_back(b.getNumSymbols());
+  allOperands.append(aDims.begin(), aDims.end());
+  allOperands.append(bDims.begin(), bDims.end());
+  allOperands.append(aSyms.begin(), aSyms.end());
+  allOperands.append(bSyms.begin(), bSyms.end());
+
+  // Shift dims and symbols of b's map.
+  auto bMap = b.getAffineMap()
+                  .shiftDims(a.getNumDims())
+                  .shiftSymbols(a.getNumSymbols());
 
   // Construct the difference expressions.
+  auto aMap = a.getAffineMap();
   SmallVector<AffineExpr, 4> diffExprs;
   diffExprs.reserve(a.getNumResults());
   for (unsigned i = 0, e = bMap.getNumResults(); i < e; ++i)
-    diffExprs.push_back(normalizer.getAffineMap().getResult(i) -
-                        bMap.getResult(i));
+    diffExprs.push_back(aMap.getResult(i) - bMap.getResult(i));
 
-  auto diffMap =
-      AffineMap::get(normalizer.getNumDims(), normalizer.getNumSymbols(),
-                     diffExprs, aMap.getContext());
-  canonicalizeMapAndOperands(&diffMap, &bOperands);
+  auto diffMap = AffineMap::get(bMap.getNumDims(), bMap.getNumSymbols(),
+                                diffExprs, bMap.getContext());
+  fullyComposeAffineMapAndOperands(&diffMap, &allOperands);
+  canonicalizeMapAndOperands(&diffMap, &allOperands);
   diffMap = simplifyAffineMap(diffMap);
-  res->reset(diffMap, bOperands);
+  res->reset(diffMap, allOperands);
 }
 
 // Returns true and sets 'indexOfMatch' if 'valueToMatch' is found in
 // 'valuesToSearch' beginning at 'indexStart'. Returns false otherwise.
 static bool findIndex(Value valueToMatch, ArrayRef<Value> valuesToSearch,
                       unsigned indexStart, unsigned *indexOfMatch) {
   unsigned size = valuesToSearch.size();
   for (unsigned i = indexStart; i < size; ++i) {

mlir/lib/IR/AffineExpr.cpp

case AffineExprKind::Mod:
     if (newLHS == lhs && newRHS == rhs)
       return *this;
     return getAffineBinaryOpExpr(getKind(), newLHS, newRHS);
   }
   llvm_unreachable("Unknown AffineExpr");
 }
 
 /// Replace symbols[0 .. numDims - 1] by symbols[shift .. shift + numDims - 1].
+AffineExpr AffineExpr::shiftDims(unsigned numDims, unsigned shift) const {
+  SmallVector<AffineExpr, 4> dims;
+  for (unsigned idx = 0; idx < numDims; ++idx)
+    dims.push_back(getAffineDimExpr(idx + shift, getContext()));
+  return replaceDimsAndSymbols(dims, {});
+}
+
+/// Replace symbols[0 .. numSymbols - 1] by
+/// symbols[shift .. shift + numSymbols - 1].
 AffineExpr AffineExpr::shiftSymbols(unsigned numSymbols, unsigned shift) const {
   SmallVector<AffineExpr, 4> symbols;
   for (unsigned idx = 0; idx < numSymbols; ++idx)
     symbols.push_back(getAffineSymbolExpr(idx + shift, getContext()));
   return replaceDimsAndSymbols({}, symbols);
 }
 
 /// Sparse replace method. Return the modified expression tree.
bool AffineExpr::isFunctionOfDim(unsigned position) const {
   }
   if (auto expr = this->dyn_cast<AffineBinaryOpExpr>()) {
     return expr.getLHS().isFunctionOfDim(position) ||
            expr.getRHS().isFunctionOfDim(position);
   }
   return false;
 }
 
+bool AffineExpr::isFunctionOfSymbol(unsigned position) const {
+  if (getKind() == AffineExprKind::SymbolId) {
+    return *this == mlir::getAffineSymbolExpr(position, getContext());
+  }
+  if (auto expr = this->dyn_cast<AffineBinaryOpExpr>()) {
+    return expr.getLHS().isFunctionOfSymbol(position) ||
+           expr.getRHS().isFunctionOfSymbol(position);
+  }
+  return false;
+}
+
 AffineBinaryOpExpr::AffineBinaryOpExpr(AffineExpr::ImplType *ptr)
     : AffineExpr(ptr) {}
 AffineExpr AffineBinaryOpExpr::getLHS() const {
   return static_cast<ImplType *>(expr)->lhs;
 }
 AffineExpr AffineBinaryOpExpr::getRHS() const {
   return static_cast<ImplType *>(expr)->rhs;
 }

mlir/lib/IR/AffineMap.cpp

AffineMap AffineMap::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
                                            ArrayRef<AffineExpr> symReplacements,
                                            unsigned numResultDims,
                                            unsigned numResultSyms) const {
   SmallVector<AffineExpr, 8> results;
   results.reserve(getNumResults());
   for (auto expr : getResults())
     results.push_back(
         expr.replaceDimsAndSymbols(dimReplacements, symReplacements));
-
   return get(numResultDims, numResultSyms, results, getContext());
 }
 
+/// Sparse replace method. Apply AffineExpr::replace(`expr`, `replacement`) to
+/// each of the results and return a new AffineMap with the new results and
+/// with the specified number of dims and symbols.
+AffineMap AffineMap::replace(AffineExpr expr, AffineExpr replacement,
+                             unsigned numResultDims,
+                             unsigned numResultSyms) const {
+  SmallVector<AffineExpr, 4> newResults;
+  newResults.reserve(getNumResults());
+  for (AffineExpr e : getResults())
+    newResults.push_back(e.replace(expr, replacement));
+  return AffineMap::get(numResultDims, numResultSyms, newResults, getContext());
+}
+
+/// Sparse replace method. Apply AffineExpr::replace(`map`) to each of the
+/// results and return a new AffineMap with the new results and with the
+/// specified number of dims and symbols.
+AffineMap AffineMap::replace(const DenseMap<AffineExpr, AffineExpr> &map,
+                             unsigned numResultDims,
+                             unsigned numResultSyms) const {
+  SmallVector<AffineExpr, 4> newResults;
+  newResults.reserve(getNumResults());
+  for (AffineExpr e : getResults())
+    newResults.push_back(e.replace(map));
+  return AffineMap::get(numResultDims, numResultSyms, newResults, getContext());
+}
+
 AffineMap AffineMap::compose(AffineMap map) {
   assert(getNumDims() == map.getNumResults() && "Number of results mismatch");
   // Prepare `map` by concatenating the symbols and rewriting its exprs.
   unsigned numDims = map.getNumDims();
   unsigned numSymbolsThisMap = getNumSymbols();
   unsigned numSymbols = numSymbolsThisMap + map.getNumSymbols();
   SmallVector<AffineExpr, 8> newDims(numDims);
   for (unsigned idx = 0; idx < numDims; ++idx) {

mlir/test/Dialect/Affine/SuperVectorize/normalize_maps.mlir

-// RUN: mlir-opt %s -affine-super-vectorizer-test -normalize-maps |  FileCheck %s
-
-// CHECK-DAG: #[[$ZERO:[a-zA-Z0-9]+]] = affine_map<() -> (0)>
-// CHECK-DAG: #[[$ID1:[a-zA-Z0-9]+]] = affine_map<(d0) -> (d0)>
-// CHECK-DAG: #[[$D0TIMES2:[a-zA-Z0-9]+]] = affine_map<(d0) -> (d0 * 2)>
-// CHECK-DAG: #[[$D0PLUSD1:[a-zA-Z0-9]+]] = affine_map<(d0, d1) -> (d0 + d1)>
-// CHECK-DAG: #[[$MINSD0PLUSD1:[a-zA-Z0-9]+]] = affine_map<(d0, d1) -> (-d0 + d1)>
-// CHECK-DAG: #[[$D0MINUSD1:[a-zA-Z0-9]+]] = affine_map<(d0, d1) -> (d0 - d1)>
-
-// CHECK-LABEL: func @simple()
-func @simple() {
-  affine.for %i0 = 0 to 7 {
-    %0 = affine.apply affine_map<(d0) -> (d0)> (%i0)
-    %1 = affine.apply affine_map<(d0) -> (d0)> (%0)
-    %2 = affine.apply affine_map<(d0, d1) -> (d0 + d1)> (%0, %0)
-    %3 = affine.apply affine_map<(d0, d1) -> (d0 - d1)> (%0, %0)
-  }
-  // CHECK-NEXT: affine.for %{{.*}} = 0 to 7
-  // CHECK-NEXT:   {{.*}} affine.apply #[[$ID1]](%{{.*}})
-  // CHECK-NEXT:   {{.*}} affine.apply #[[$D0TIMES2]](%{{.*}})
-  // CHECK-NEXT:   {{.*}} affine.apply #[[$ZERO]]()
-
-  affine.for %i1 = 0 to 7 {
-    affine.for %i2 = 0 to 42 {
-      %20 = affine.apply affine_map<(d0, d1) -> (d1)> (%i1, %i2)
-      %21 = affine.apply affine_map<(d0, d1) -> (d0)> (%i1, %i2)
-      %22 = affine.apply affine_map<(d0, d1) -> (d0 + d1)> (%20, %21)
-      %23 = affine.apply affine_map<(d0, d1) -> (d0 - d1)> (%20, %21)
-      %24 = affine.apply affine_map<(d0, d1) -> (-d0 + d1)> (%20, %21)
-    }
-  }
-  //      CHECK: affine.for %{{.*}} = 0 to 7
-  // CHECK-NEXT:   affine.for %{{.*}} = 0 to 42
-  // CHECK-NEXT:     {{.*}} affine.apply #[[$D0PLUSD1]](%{{.*}}, %{{.*}})
-  // CHECK-NEXT:     {{.*}} affine.apply #[[$MINSD0PLUSD1]](%{{.*}}, %{{.*}})
-  // CHECK-NEXT:     {{.*}} affine.apply #[[$D0MINUSD1]](%{{.*}}, %{{.*}})
-
-  affine.for %i3 = 0 to 16 {
-    affine.for %i4 = 0 to 47 step 2 {
-      affine.for %i5 = 0 to 78 step 16 {
-        %50 = affine.apply affine_map<(d0) -> (d0)> (%i3)
-        %51 = affine.apply affine_map<(d0) -> (d0)> (%i4)
-        %52 = affine.apply affine_map<(d0) -> (d0)> (%i5)
-        %53 = affine.apply affine_map<(d0, d1, d2) -> (d0)> (%50, %51, %52)
-        %54 = affine.apply affine_map<(d0, d1, d2) -> (d1)> (%50, %51, %52)
-        %55 = affine.apply affine_map<(d0, d1, d2) -> (d2)> (%50, %51, %52)
-      }
-    }
-  }
-  // CHECK:      affine.for %{{.*}} = 0 to 16
-  // CHECK-NEXT:   affine.for %{{.*}} = 0 to 47 step 2
-  // CHECK-NEXT:     affine.for %{{.*}} = 0 to 78 step 16
-  // CHECK-NEXT:       {{.*}} affine.apply #[[$ID1]](%{{.*}})
-  // CHECK-NEXT:       {{.*}} affine.apply #[[$ID1]](%{{.*}})
-  // CHECK-NEXT:       {{.*}} affine.apply #[[$ID1]](%{{.*}})
-
-  return
-}

mlir/test/Dialect/Affine/affine-data-copy.mlir

 // CHECK: affine.for %[[I:.*]] = 0 to 4096 step 128 {
 // CHECK:   affine.for %[[J:.*]] = 0 to 4096 step 128 {
 // CHECK:     [[BUFC:%[0-9]+]] = alloc() : memref<128x128xf32>
 // The result matrix's copy gets hoisted out.
 // Result matrix copy-in.
 // CHECK:     affine.for %[[II:.*]] = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {
 // CHECK:       affine.for %[[JJ:.*]] = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {
 // CHECK:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<4096x4096xf32>
-// CHECK:         affine.store %{{.*}}, [[BUFC]][-%[[I]] + %[[II]], -%[[J]] + %[[JJ]]] : memref<128x128xf32>
+// CHECK:         affine.store %{{.*}}, [[BUFC]][%[[II]] - %[[I]], %[[JJ]] - %[[J]]] : memref<128x128xf32>
 // CHECK:       }
 // CHECK:     }
 
 // LHS matrix copy-in.
 // CHECK:     affine.for %[[K:.*]] = 0 to 4096 step 128 {
 // CHECK:      [[BUFA:%[0-9]+]] = alloc() : memref<128x128xf32>
 // CHECK:       affine.for %[[II:.*]] = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {
 // CHECK:         affine.for %[[KK:.*]] = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {
 // CHECK:           affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<4096x4096xf32>
-// CHECK:           affine.store %{{.*}}, [[BUFA]][-%[[I]] + %[[II]], -%[[K]] + %[[KK]]] : memref<128x128xf32>
+// CHECK:           affine.store %{{.*}}, [[BUFA]][%[[II]] - %[[I]], %[[KK]] - %[[K]]] : memref<128x128xf32>
 // CHECK:         }
 // CHECK:       }
 
 // RHS matrix copy-in.
 // CHECK:       [[BUFB:%[0-9]+]] = alloc() : memref<128x128xf32>
 // CHECK:       affine.for %[[KK:.*]] = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {
 // CHECK:         affine.for %[[JJ:.*]] = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {
 // CHECK:           affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<4096x4096xf32>
-// CHECK:           affine.store %{{.*}}, [[BUFB]][-%[[K]] + %[[KK]], -%[[J]] + %[[JJ]]] : memref<128x128xf32>
+// CHECK:           affine.store %{{.*}}, [[BUFB]][%[[KK]] - %[[K]], %[[JJ]] - %[[J]]] : memref<128x128xf32>
 // CHECK:         }
 // CHECK:       }
 
 // Computation on the fast buffers.
 // CHECK:       affine.for %{{.*}} = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {
 // CHECK:         affine.for %{{.*}} = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {
 // CHECK:           affine.for %{{.*}} = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {
 // CHECK:             affine.load [[BUFA]][-%{{.*}} + %{{.*}}, -%{{.*}} + %{{.*}}] : memref<128x128xf32>
 // CHECK:             affine.load [[BUFB]][-%{{.*}} + %{{.*}}, -%{{.*}} + %{{.*}}] : memref<128x128xf32>
 // CHECK:             affine.load [[BUFC]][-%{{.*}} + %{{.*}}, -%{{.*}} + %{{.*}}] : memref<128x128xf32>
 // CHECK:             mulf %{{.*}}, %{{.*}} : f32
 // CHECK:             addf %{{.*}}, %{{.*}} : f32
 // CHECK:             affine.store %{{.*}}, [[BUFC]][-%{{.*}} + %{{.*}}, -%{{.*}} + %{{.*}}] : memref<128x128xf32>
 // CHECK:           }
 // CHECK:         }
 // CHECK:       }
 // CHECK:       dealloc [[BUFB]] : memref<128x128xf32>
 // CHECK:       dealloc [[BUFA]] : memref<128x128xf32>
 // CHECK:     }
 
 // Result matrix copy out.
 // CHECK:     affine.for %{{.*}} = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {
 // CHECK:       affine.for %{{.*}} = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {
-// CHECK:         affine.load [[BUFC]][-%{{.*}} + %{{.*}}, -%{{.*}} + %{{.*}}] : memref<128x128xf32>
+// CHECK:         affine.load [[BUFC]][%{{.*}} - %{{.*}}, %{{.*}} - %{{.*}}] : memref<128x128xf32>
 // CHECK:         store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<4096x4096xf32>
 // CHECK:       }
 // CHECK:     }
 // CHECK:     dealloc [[BUFC]] : memref<128x128xf32>
 // CHECK:   }
 // CHECK: }
 
 // Check that only one memref is copied when memref filter is used.
affine.for %i = 0 to 4096 step 100 {
     }
   }
   return %A : memref<4096xf32>
 }
 // CHECK:      affine.for %[[IV1:.*]] = 0 to 4096 step 100
 // CHECK:        %[[BUF:.*]] = alloc() : memref<100xf32>
 // CHECK-NEXT:   affine.for %[[IV2:.*]] = #[[$MAP_IDENTITY]](%[[IV1]]) to min #[[$MAP_MIN_UB1]](%[[IV1]]) {
 // CHECK-NEXT:     affine.load %{{.*}}[%[[IV2]]] : memref<4096xf32>
-// CHECK-NEXT:     affine.store %{{.*}}, %[[BUF]][-%[[IV1]] + %[[IV2]]] : memref<100xf32>
+// CHECK-NEXT:     affine.store %{{.*}}, %[[BUF]][%[[IV2]] - %[[IV1]]] : memref<100xf32>
 // CHECK-NEXT:   }
 // CHECK-NEXT:   affine.for %[[IV2:.*]] = #[[$MAP_IDENTITY]](%[[IV1]]) to min #[[$MAP_MIN_UB2]](%[[IV1]]) {
 // CHECK-NEXT:     affine.load %[[BUF]][-%[[IV1]] + %[[IV2]]] : memref<100xf32>
 // CHECK-NEXT:     mulf
 // CHECK-NEXT:     affine.store %{{.*}}, %[[BUF]][-%[[IV1]] + %[[IV2]]] : memref<100xf32>
 // CHECK-NEXT:   }
 // CHECK:        affine.for %[[IV2:.*]] = #[[$MAP_IDENTITY]](%[[IV1]]) to min #[[$MAP_MIN_UB1]](%[[IV1]]) {
-// CHECK-NEXT:     affine.load %[[BUF]][-%[[IV1]] + %[[IV2]]] : memref<100xf32>
+// CHECK-NEXT:     affine.load %[[BUF]][%[[IV2]] - %[[IV1]]] : memref<100xf32>
 // CHECK-NEXT:     affine.store %{{.*}}, %{{.*}}[%[[IV2]]] : memref<4096xf32>
 // CHECK-NEXT:   }
 // CHECK-NEXT:   dealloc %[[BUF]] : memref<100xf32>
 // CHECK-NEXT: }
 
 // -----
 
 // Lower bound is a max; upper bound is a min. This pattern typically appears

mlir/test/Dialect/Affine/canonicalize.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -pass-pipeline='func(canonicalize)' | FileCheck %s
 
-// Affine maps for test case: compose_affine_maps_1dto2d_no_symbols
-// CHECK-DAG: [[$MAP0:#map[0-9]+]] = affine_map<(d0) -> (d0 - 1)>
-// CHECK-DAG: [[$MAP1:#map[0-9]+]] = affine_map<(d0) -> (d0 + 1)>
-
-// Affine maps for test case: compose_affine_maps_1dto2d_with_symbols
-// CHECK-DAG: [[$MAP4:#map[0-9]+]] = affine_map<(d0) -> (d0 - 4)>
-// CHECK-DAG: [[$MAP4b:#map[0-9]+]] = affine_map<(d0) -> (d0 - 7)>
-// CHECK-DAG: [[$MAP7:#map[0-9]+]] = affine_map<(d0) -> (d0 * 2 - 3)>
-// CHECK-DAG: [[$MAP7a:#map[0-9]+]] = affine_map<(d0) -> (d0 * 2 + 1)>
-
-// Affine map for test case: compose_affine_maps_d2_tile
-// CHECK-DAG: [[$MAP8:#map[0-9]+]] = affine_map<(d0, d1) -> (d1 + (d0 ceildiv 4) * 4 - (d1 floordiv 4) * 4)>
-// CHECK-DAG: [[$MAP8a:#map[0-9]+]] = affine_map<(d0, d1) -> (d1 + (d0 ceildiv 8) * 8 - (d1 floordiv 8) * 8)>
-
-// Affine maps for test case: compose_affine_maps_dependent_loads
-// CHECK-DAG: [[$MAP9:#map[0-9]+]] = affine_map<(d0) -> (d0 + 3)>
-// CHECK-DAG: [[$MAP10:#map[0-9]+]] = affine_map<(d0) -> (d0 * 3)>
-// CHECK-DAG: [[$MAP11:#map[0-9]+]] = affine_map<(d0) -> ((d0 + 3) ceildiv 3)>
-// CHECK-DAG: [[$MAP12:#map[0-9]+]] = affine_map<(d0) -> (d0 * 7 - 49)>
-
-// Affine maps for test case: compose_affine_maps_diamond_dependency
-// CHECK-DAG: [[$MAP13A:#map[0-9]+]] = affine_map<(d0) -> ((d0 + 6) ceildiv 8)>
-// CHECK-DAG: [[$MAP13B:#map[0-9]+]] = affine_map<(d0) -> ((d0 * 4 - 4) floordiv 3)>
-
-// Affine maps for test case: compose_affine_maps_multiple_symbols
-// CHECK-DAG: [[$MAP14:#map[0-9]+]] = affine_map<()[s0, s1] -> (((s1 + s0) * 4) floordiv s0)>
-
-// Affine maps for test case: partial_fold_map
-// CHECK-DAG: [[$MAP15:#map[0-9]+]] = affine_map<()[s0] -> (s0 - 42)>
-
-// Affine maps for test cases: symbolic_composition_*
-// CHECK-DAG: [[$MAP_symbolic_composition_a:#map[0-9]+]] = affine_map<()[s0] -> (s0 * 512)>
-// CHECK-DAG: [[$MAP_symbolic_composition_b:#map[0-9]+]] = affine_map<()[s0] -> (s0 * 4)>
-// CHECK-DAG: [[$MAP_symbolic_composition_c:#map[0-9]+]] = affine_map<()[s0, s1] -> (s0 * 3 + s1)>
-// CHECK-DAG: [[$MAP_symbolic_composition_d:#map[0-9]+]] = affine_map<()[s0, s1] -> (s1 * 3 + s0)>
-
-// Affine maps for test cases: map_mix_dims_and_symbols_*
-// CHECK-DAG: [[$MAP_mix_dims_and_symbols_b:#map[0-9]+]] = affine_map<()[s0, s1] -> (s1 + s0 * 42 + 6)>
-// CHECK-DAG: [[$MAP_mix_dims_and_symbols_c:#map[0-9]+]] = affine_map<()[s0, s1] -> (s1 * 4 + s0 * 168 - 4)>
-// CHECK-DAG: [[$MAP_mix_dims_and_symbols_d:#map[0-9]+]] = affine_map<()[s0, s1] -> ((s1 + s0 * 42 + 6) ceildiv 8)>
-// CHECK-DAG: [[$MAP_mix_dims_and_symbols_e:#map[0-9]+]] = affine_map<()[s0, s1] -> ((s1 * 4 + s0 * 168 - 4) floordiv 3)>
+// -----
 
-// Affine maps for test case: $symbolic_semi_affine
-// CHECK-DAG: [[$symbolic_semi_affine:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 floordiv (s0 + 1))>
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> (d0 - 1)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (d0 + 1)>
 
 // CHECK-LABEL: func @compose_affine_maps_1dto2d_no_symbols() {
 func @compose_affine_maps_1dto2d_no_symbols() {
   %0 = alloc() : memref<4x4xf32>
 
   affine.for %i0 = 0 to 15 {
     // Test load[%x, %x]
 
     %x0 = affine.apply affine_map<(d0) -> (d0 - 1)> (%i0)
     %x1_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%x0, %x0)
     %x1_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%x0, %x0)
 
-    // CHECK: [[I0A:%[0-9]+]] = affine.apply [[$MAP0]](%{{.*}})
-    // CHECK-NEXT: [[V0:%[0-9]+]] = load %0{{\[}}[[I0A]], [[I0A]]{{\]}}
+    // CHECK: %[[I0A:.*]] = affine.apply #[[$MAP0]](%{{.*}})
+    // CHECK-NEXT: %[[V0:.*]] = load %0[%[[I0A]], %[[I0A]]]
     %v0 = load %0[%x1_0, %x1_1] : memref<4x4xf32>
 
     // Test store[%y, %y]
     %y0 = affine.apply affine_map<(d0) -> (d0 + 1)> (%i0)
     %y1_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%y0, %y0)
     %y1_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%y0, %y0)
 
-    // CHECK-NEXT: [[I1A:%[0-9]+]] = affine.apply [[$MAP1]](%{{.*}})
-    // CHECK-NEXT: store [[V0]], %0{{\[}}[[I1A]], [[I1A]]{{\]}}
+    // CHECK-NEXT: %[[I1A:.*]] = affine.apply #[[$MAP1]](%{{.*}})
+    // CHECK-NEXT: store %[[V0]], %0[%[[I1A]], %[[I1A]]]
     store %v0, %0[%y1_0, %y1_1] : memref<4x4xf32>
 
     // Test store[%x, %y]
     %xy_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%x0, %y0)
     %xy_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%x0, %y0)
 
-    // CHECK-NEXT: store [[V0]], %0{{\[}}[[I0A]], [[I1A]]{{\]}}
+    // CHECK-NEXT: store %[[V0]], %0[%[[I0A]], %[[I1A]]]
     store %v0, %0[%xy_0, %xy_1] : memref<4x4xf32>
 
     // Test store[%y, %x]
     %yx_0 = affine.apply affine_map<(d0, d1) -> (d0)> (%y0, %x0)
     %yx_1 = affine.apply affine_map<(d0, d1) -> (d1)> (%y0, %x0)
-    // CHECK-NEXT: store [[V0]], %0{{\[}}[[I1A]], [[I0A]]{{\]}}
+    // CHECK-NEXT: store %[[V0]], %0[%[[I1A]], %[[I0A]]]
     store %v0, %0[%yx_0, %yx_1] : memref<4x4xf32>
   }
   return
 }
 
+// -----
+
+// CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0) -> (d0 - 4)>
+// CHECK-DAG: #[[$MAP7:.*]] = affine_map<(d0) -> (d0 * 2 - 3)>
+// CHECK-DAG: #[[$MAP7a:.*]] = affine_map<(d0) -> (d0 * 2 + 1)>
+
 // CHECK-LABEL: func @compose_affine_maps_1dto2d_with_symbols() {
 func @compose_affine_maps_1dto2d_with_symbols() {
   %0 = alloc() : memref<4x4xf32>
 
   affine.for %i0 = 0 to 15 {
     // Test load[%x0, %x0] with symbol %c4
     %c4 = constant 4 : index
     %x0 = affine.apply affine_map<(d0)[s0] -> (d0 - s0)> (%i0)[%c4]
 
-    // CHECK: [[I0:%[0-9]+]] = affine.apply [[$MAP4]](%{{.*}})
-    // CHECK-NEXT: [[V0:%[0-9]+]] = load %{{[0-9]+}}{{\[}}[[I0]], [[I0]]{{\]}}
+    // CHECK: %[[I0:.*]] = affine.apply #[[$MAP4]](%{{.*}})
+    // CHECK-NEXT: %[[V0:.*]] = load %{{.*}}[%[[I0]], %[[I0]]]
     %v0 = load %0[%x0, %x0] : memref<4x4xf32>
 
     // Test load[%x0, %x1] with symbol %c4 captured by '%x0' map.
     %x1 = affine.apply affine_map<(d0) -> (d0 + 1)> (%i0)
     %y1 = affine.apply affine_map<(d0, d1) -> (d0+d1)> (%x0, %x1)
-    // CHECK-NEXT: [[I1:%[0-9]+]] = affine.apply [[$MAP7]](%{{.*}})
-    // CHECK-NEXT: store [[V0]], %{{[0-9]+}}{{\[}}[[I1]], [[I1]]{{\]}}
+    // CHECK-NEXT: %[[I1:.*]] = affine.apply #[[$MAP7]](%{{.*}})
+    // CHECK-NEXT: store %[[V0]], %{{.*}}[%[[I1]], %[[I1]]]
     store %v0, %0[%y1, %y1] : memref<4x4xf32>
 
     // Test store[%x1, %x0] with symbol %c4 captured by '%x0' map.
     %y2 = affine.apply affine_map<(d0, d1) -> (d0 + d1)> (%x1, %x0)
-    // CHECK-NEXT: [[I2:%[0-9]+]] = affine.apply [[$MAP7]](%{{.*}})
-    // CHECK-NEXT: store [[V0]], %{{[0-9]+}}{{\[}}[[I2]], [[I2]]{{\]}}
+    // CHECK-NEXT: %[[I2:.*]] = affine.apply #[[$MAP7]](%{{.*}})
+    // CHECK-NEXT: store %[[V0]], %{{.*}}[%[[I2]], %[[I2]]]
     store %v0, %0[%y2, %y2] : memref<4x4xf32>
 
     // Test store[%x2, %x0] with symbol %c4 from '%x0' and %c5 from '%x2'
     %c5 = constant 5 : index
     %x2 = affine.apply affine_map<(d0)[s0] -> (d0 + s0)> (%i0)[%c5]
     %y3 = affine.apply affine_map<(d0, d1) -> (d0 + d1)> (%x2, %x0)
-    // CHECK: [[I3:%[0-9]+]] = affine.apply [[$MAP7a]](%{{.*}})
-    // CHECK-NEXT: store [[V0]], %{{[0-9]+}}{{\[}}[[I3]], [[I3]]{{\]}}
+    // CHECK: %[[I3:.*]] = affine.apply #[[$MAP7a]](%{{.*}})
+    // CHECK-NEXT: store %[[V0]], %{{.*}}[%[[I3]], %[[I3]]]
     store %v0, %0[%y3, %y3] : memref<4x4xf32>
   }
   return
 }
 
+// -----
+
+// CHECK-DAG: #[[$MAP8:.*]] = affine_map<(d0, d1) -> (d1 + (d0 ceildiv 4) * 4 - (d1 floordiv 4) * 4)>
+// CHECK-DAG: #[[$MAP8a:.*]] = affine_map<(d0, d1) -> (d1 + (d0 ceildiv 8) * 8 - (d1 floordiv 8) * 8)>
+
 // CHECK-LABEL: func @compose_affine_maps_2d_tile() {
 func @compose_affine_maps_2d_tile() {
   %0 = alloc() : memref<16x32xf32>
   %1 = alloc() : memref<16x32xf32>
 
   %c4 = constant 4 : index
   %c8 = constant 8 : index
 
   affine.for %i0 = 0 to 3 {
     %x0 = affine.apply affine_map<(d0)[s0] -> (d0 ceildiv s0)> (%i0)[%c4]
     affine.for %i1 = 0 to 3 {
       %x1 = affine.apply affine_map<(d0)[s0] -> (d0 ceildiv s0)> (%i1)[%c8]
       affine.for %i2 = 0 to 3 {
         %x2 = affine.apply affine_map<(d0)[s0] -> (d0 mod s0)> (%i2)[%c4]
         affine.for %i3 = 0 to 3 {
           %x3 = affine.apply affine_map<(d0)[s0] -> (d0 mod s0)> (%i3)[%c8]
 
           %x40 = affine.apply affine_map<(d0, d1, d2, d3)[s0, s1] ->
             ((d0 * s0) + d2)> (%x0, %x1, %x2, %x3)[%c4, %c8]
           %x41 = affine.apply affine_map<(d0, d1, d2, d3)[s0, s1] ->
             ((d1 * s1) + d3)> (%x0, %x1, %x2, %x3)[%c4, %c8]
-          // CHECK: [[I0:%[0-9]+]] = affine.apply [[$MAP8]](%{{.*}}, %{{.*}})
-          // CHECK: [[I1:%[0-9]+]] = affine.apply [[$MAP8a]](%{{.*}}, %{{.*}})
-          // CHECK-NEXT: [[L0:%[0-9]+]] = load %{{[0-9]+}}{{\[}}[[I0]], [[I1]]{{\]}}
+          // CHECK: %[[I0:.*]] = affine.apply #[[$MAP8]](%{{.*}}, %{{.*}})
+          // CHECK: %[[I1:.*]] = affine.apply #[[$MAP8a]](%{{.*}}, %{{.*}})
+          // CHECK-NEXT: %[[L0:.*]] = load %{{.*}}[%[[I0]], %[[I1]]]
           %v0 = load %0[%x40, %x41] : memref<16x32xf32>
 
-          // CHECK-NEXT: store [[L0]], %{{[0-9]+}}{{\[}}[[I0]], [[I1]]{{\]}}
+          // CHECK-NEXT: store %[[L0]], %{{.*}}[%[[I0]], %[[I1]]]
           store %v0, %1[%x40, %x41] : memref<16x32xf32>
         }
       }
     }
   }
   return
 }
 
+// -----
+
+// CHECK-DAG: #[[$MAP4b:.*]] = affine_map<(d0) -> (d0 - 7)>
+// CHECK-DAG: #[[$MAP9:.*]] = affine_map<(d0) -> (d0 + 3)>
+// CHECK-DAG: #[[$MAP10:.*]] = affine_map<(d0) -> (d0 * 3)>
+// CHECK-DAG: #[[$MAP11:.*]] = affine_map<(d0) -> ((d0 + 3) ceildiv 3)>
+// CHECK-DAG: #[[$MAP12:.*]] = affine_map<(d0) -> (d0 * 7 - 49)>
+
 // CHECK-LABEL: func @compose_affine_maps_dependent_loads() {
 func @compose_affine_maps_dependent_loads() {
   %0 = alloc() : memref<16x32xf32>
   %1 = alloc() : memref<16x32xf32>
 
   affine.for %i0 = 0 to 3 {
     affine.for %i1 = 0 to 3 {
       affine.for %i2 = 0 to 3 {
         %c3 = constant 3 : index
         %c7 = constant 7 : index
 
         %x00 = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d0 + s0)>
             (%i0, %i1, %i2)[%c3, %c7]
         %x01 = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d1 - s1)>
             (%i0, %i1, %i2)[%c3, %c7]
         %x02 = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d2 * s0)>
             (%i0, %i1, %i2)[%c3, %c7]
 
-        // CHECK: [[I0:%[0-9]+]] = affine.apply [[$MAP9]](%{{.*}})
-        // CHECK: [[I1:%[0-9]+]] = affine.apply [[$MAP4b]](%{{.*}})
-        // CHECK: [[I2:%[0-9]+]] = affine.apply [[$MAP10]](%{{.*}})
-        // CHECK-NEXT: [[V0:%[0-9]+]] = load %{{[0-9]+}}{{\[}}[[I0]], [[I1]]{{\]}}
+        // CHECK: %[[I0:.*]] = affine.apply #[[$MAP9]](%{{.*}})
+        // CHECK: %[[I1:.*]] = affine.apply #[[$MAP4b]](%{{.*}})
+        // CHECK: %[[I2:.*]] = affine.apply #[[$MAP10]](%{{.*}})
+        // CHECK-NEXT: %[[V0:.*]] = load %{{.*}}[%[[I0]], %[[I1]]]
         %v0 = load %0[%x00, %x01] : memref<16x32xf32>
 
-        // CHECK-NEXT: store [[V0]], %{{[0-9]+}}{{\[}}[[I0]], [[I2]]{{\]}}
+        // CHECK-NEXT: store %[[V0]], %{{.*}}[%[[I0]], %[[I2]]]
         store %v0, %0[%x00, %x02] : memref<16x32xf32>
 
         // Swizzle %i0, %i1
-        // CHECK-NEXT: store [[V0]], %{{[0-9]+}}{{\[}}[[I1]], [[I0]]{{\]}}
+        // CHECK-NEXT: store %[[V0]], %{{.*}}[%[[I1]], %[[I0]]]
         store %v0, %0[%x01, %x00] : memref<16x32xf32>
 
         // Swizzle %x00, %x01 and %c3, %c7
         %x10 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d0 * s1)>
            (%x01, %x00)[%c3, %c7]
         %x11 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d1 ceildiv s0)>
            (%x01, %x00)[%c3, %c7]
 
-        // CHECK-NEXT: [[I2A:%[0-9]+]] = affine.apply [[$MAP12]](%{{.*}})
-        // CHECK-NEXT: [[I2B:%[0-9]+]] = affine.apply [[$MAP11]](%{{.*}})
-        // CHECK-NEXT: store [[V0]], %{{[0-9]+}}{{\[}}[[I2A]], [[I2B]]{{\]}}
+        // CHECK-NEXT: %[[I2A:.*]] = affine.apply #[[$MAP12]](%{{.*}})
+        // CHECK-NEXT: %[[I2B:.*]] = affine.apply #[[$MAP11]](%{{.*}})
+        // CHECK-NEXT: store %[[V0]], %{{.*}}[%[[I2A]], %[[I2B]]]
         store %v0, %0[%x10, %x11] : memref<16x32xf32>
       }
     }
   }
   return
 }
 
+// -----
+
+// CHECK-DAG: #[[$MAP13A:.*]] = affine_map<(d0) -> ((d0 + 6) ceildiv 8)>
+// CHECK-DAG: #[[$MAP13B:.*]] = affine_map<(d0) -> ((d0 * 4 - 4) floordiv 3)>
+
 // CHECK-LABEL: func @compose_affine_maps_diamond_dependency
 func @compose_affine_maps_diamond_dependency(%arg0: f32, %arg1: memref<4x4xf32>) {
   affine.for %i0 = 0 to 15 {
     %a = affine.apply affine_map<(d0) -> (d0 - 1)> (%i0)
     %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
     %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
     %d0 = affine.apply affine_map<(d0, d1) -> (d0 ceildiv 8)> (%b, %c)
     %d1 = affine.apply affine_map<(d0, d1) -> (d1 floordiv 3)> (%b, %c)
-    // CHECK: [[I0:%[0-9]+]] = affine.apply [[$MAP13A]](%{{.*}})
-    // CHECK: [[I1:%[0-9]+]] = affine.apply [[$MAP13B]](%{{.*}})
-    // CHECK-NEXT: store %arg0, %arg1{{\[}}[[I0]], [[I1]]{{\]}}
+    // CHECK: %[[I0:.*]] = affine.apply #[[$MAP13A]](%{{.*}})
+    // CHECK: %[[I1:.*]] = affine.apply #[[$MAP13B]](%{{.*}})
+    // CHECK-NEXT: store %arg0, %arg1[%[[I0]], %[[I1]]]
     store %arg0, %arg1[%d0, %d1] : memref<4x4xf32>
   }
 
   return
 }
 
+// -----
+
+// CHECK-DAG: #[[$MAP14:.*]] = affine_map<()[s0, s1] -> (((s1 + s0) * 4) floordiv s0)>
+
 // CHECK-LABEL: func @compose_affine_maps_multiple_symbols
 func @compose_affine_maps_multiple_symbols(%arg0: index, %arg1: index) -> index {
   %a = affine.apply affine_map<(d0)[s0] -> (s0 + d0)> (%arg0)[%arg1]
   %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
   %e = affine.apply affine_map<(d0)[s0] -> (d0 floordiv s0)> (%c)[%arg1]
-  // CHECK: [[I0:%[0-9]+]] = affine.apply [[$MAP14]]()[%{{.*}}, %{{.*}}]
+  // CHECK: [[I0:.*]] = affine.apply #[[$MAP14]]()[%{{.*}}, %{{.*}}]
   return %e : index
 }
 
+// -----
+
 // CHECK-LABEL: func @arg_used_as_dim_and_symbol
 func @arg_used_as_dim_and_symbol(%arg0: memref<100x100xf32>, %arg1: index, %arg2: f32) {
   %c9 = constant 9 : index
   %1 = alloc() : memref<100x100xf32, 1>
   %2 = alloc() : memref<1xi32>
   affine.for %i0 = 0 to 100 {
     affine.for %i1 = 0 to 100 {
       %3 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d1 + s0 + s1)>
         (%i0, %i1)[%arg1, %c9]
       %4 = affine.apply affine_map<(d0, d1, d3) -> (d3 - (d0 + d1))>
         (%arg1, %c9, %3)
-      // CHECK: store %arg2, %{{[0-9]+}}{{\[}}%{{.*}}, %{{.*}}{{\]}}
+      // CHECK: store %arg2, %{{.*}}[%{{.*}}, %{{.*}}]
       store %arg2, %1[%4, %arg1] : memref<100x100xf32, 1>
     }
   }
   return
 }
 
+// -----
+
 // CHECK-LABEL: func @trivial_maps
 func @trivial_maps() {
   // CHECK-NOT: affine.apply
 
   %0 = alloc() : memref<10xf32>
   %c0 = constant 0 : index
   %cst = constant 0.000000e+00 : f32
   affine.for %i1 = 0 to 10 {
     %1 = affine.apply affine_map<()[s0] -> (s0)>()[%c0]
     store %cst, %0[%1] : memref<10xf32>
     %2 = load %0[%c0] : memref<10xf32>
 
     %3 = affine.apply affine_map<()[] -> (0)>()[]
     store %cst, %0[%3] : memref<10xf32>
     store %2, %0[%c0] : memref<10xf32>
   }
   return
 }
 
+// -----
+
+// CHECK-DAG: #[[$MAP15:.*]] = affine_map<()[s0] -> (s0 - 42)>
+
 // CHECK-LABEL: func @partial_fold_map
 func @partial_fold_map(%arg1: index, %arg2: index) -> index {
   // TODO: Constant fold one index into affine.apply
   %c42 = constant 42 : index
   %2 = affine.apply affine_map<(d0, d1) -> (d0 - d1)> (%arg1, %c42)
-  // CHECK: [[X:%[0-9]+]] = affine.apply [[$MAP15]]()[%{{.*}}]
+  // CHECK: [[X:.*]] = affine.apply #[[$MAP15]]()[%{{.*}}]
   return %2 : index
 }
 
+// -----
+
+// CHECK-DAG: #[[$MAP_symbolic_composition_a:.*]] = affine_map<()[s0] -> (s0 * 512)>
+
 // CHECK-LABEL: func @symbolic_composition_a(%{{.*}}: index, %{{.*}}: index) -> index {
 func @symbolic_composition_a(%arg0: index, %arg1: index) -> index {
   %0 = affine.apply affine_map<(d0) -> (d0 * 4)>(%arg0)
   %1 = affine.apply affine_map<()[s0, s1] -> (8 * s0)>()[%0, %arg0]
   %2 = affine.apply affine_map<()[s0, s1] -> (16 * s1)>()[%arg1, %1]
-  // CHECK: %{{.*}} = affine.apply [[$MAP_symbolic_composition_a]]()[%{{.*}}]
+  // CHECK: %{{.*}} = affine.apply #[[$MAP_symbolic_composition_a]]()[%{{.*}}]
   return %2 : index
 }
 
+// -----
+
+// CHECK-DAG: #[[$MAP_symbolic_composition_b:.*]] = affine_map<()[s0] -> (s0 * 4)>
+
 // CHECK-LABEL: func @symbolic_composition_b(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
 func @symbolic_composition_b(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
   %0 = affine.apply affine_map<(d0) -> (d0)>(%arg0)
   %1 = affine.apply affine_map<()[s0, s1, s2, s3] -> (s0 + s1 + s2 + s3)>()[%0, %0, %0, %0]
-  // CHECK: %{{.*}} = affine.apply [[$MAP_symbolic_composition_b]]()[%{{.*}}]
+  // CHECK: %{{.*}} = affine.apply #[[$MAP_symbolic_composition_b]]()[%{{.*}}]
   return %1 : index
 }
 
+// -----
+
+// CHECK-DAG: #[[$MAP_symbolic_composition_c:.*]] = affine_map<()[s0, s1] -> (s0 * 3 + s1)>
+
 // CHECK-LABEL: func @symbolic_composition_c(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
 func @symbolic_composition_c(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
   %0 = affine.apply affine_map<(d0) -> (d0)>(%arg0)
   %1 = affine.apply affine_map<(d0) -> (d0)>(%arg1)
   %2 = affine.apply affine_map<()[s0, s1, s2, s3] -> (s0 + s1 + s2 + s3)>()[%0, %0, %0, %1]
-  // CHECK: %{{.*}} = affine.apply [[$MAP_symbolic_composition_c]]()[%{{.*}}, %{{.*}}]
+  // CHECK: %{{.*}} = affine.apply #[[$MAP_symbolic_composition_c]]()[%{{.*}}, %{{.*}}]
   return %2 : index
 }
 
-// CHECK-LABEL: func @symbolic_composition_d(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
+// -----
+
+// CHECK-DAG: #[[$MAP_symbolic_composition_d:.*]] = affine_map<()[s0, s1] -> (s0 + s1 * 3)>
+
+// CHECK-LABEL: func @symbolic_composition_d(
+//  CHECK-SAME:   %[[ARG0:[0-9a-zA-Z]+]]: index
+//  CHECK-SAME:   %[[ARG1:[0-9a-zA-Z]+]]: index
 func @symbolic_composition_d(%arg0: index, %arg1: index, %arg2: index, %arg3: index) -> index {
   %0 = affine.apply affine_map<(d0) -> (d0)>(%arg0)
   %1 = affine.apply affine_map<()[s0] -> (s0)>()[%arg1]
   %2 = affine.apply affine_map<()[s0, s1, s2, s3] -> (s0 + s1 + s2 + s3)>()[%0, %0, %0, %1]
-  // CHECK: %{{.*}} = affine.apply [[$MAP_symbolic_composition_d]]()[%{{.*}}, %{{.*}}]
+  // CHECK: %{{.*}} = affine.apply #[[$MAP_symbolic_composition_d]]()[%[[ARG1]], %[[ARG0]]]
   return %2 : index
 }
 
+// -----
+
+// CHECK-DAG: #[[$MAP_mix_dims_and_symbols_b:.*]] = affine_map<()[s0, s1] -> (s0 * 42 + s1 + 6)>
 
 // CHECK-LABEL: func @mix_dims_and_symbols_b(%arg0: index, %arg1: index) -> index {
 func @mix_dims_and_symbols_b(%arg0: index, %arg1: index) -> index {
   %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
   %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
-  // CHECK: {{.*}} = affine.apply [[$MAP_mix_dims_and_symbols_b]]()[%{{.*}}, %{{.*}}]
+  // CHECK: {{.*}} = affine.apply #[[$MAP_mix_dims_and_symbols_b]]()[%{{.*}}, %{{.*}}]
 
   return %b : index
 }
 
+// -----
+
+// CHECK-DAG: #[[$MAP_mix_dims_and_symbols_c:.*]] = affine_map<()[s0, s1] -> (s0 * 168 + s1 * 4 - 4)>
+
 // CHECK-LABEL: func @mix_dims_and_symbols_c(%arg0: index, %arg1: index) -> index {
 func @mix_dims_and_symbols_c(%arg0: index, %arg1: index) -> index {
   %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
   %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
   %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
-  // CHECK: {{.*}} = affine.apply [[$MAP_mix_dims_and_symbols_c]]()[%{{.*}}, %{{.*}}]
+  // CHECK: {{.*}} = affine.apply #[[$MAP_mix_dims_and_symbols_c]]()[%{{.*}}, %{{.*}}]
   return %c : index
 }
 
+// -----
+
+// CHECK-DAG: #[[$MAP_mix_dims_and_symbols_d:.*]] = affine_map<()[s0, s1] -> ((s0 * 42 + s1 + 6) ceildiv 8)>
+
 // CHECK-LABEL: func @mix_dims_and_symbols_d(%arg0: index, %arg1: index) -> index {
 func @mix_dims_and_symbols_d(%arg0: index, %arg1: index) -> index {
   %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
   %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
   %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
   %d = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)> ()[%b]
-  // CHECK: {{.*}} = affine.apply [[$MAP_mix_dims_and_symbols_d]]()[%{{.*}}, %{{.*}}]
+  // CHECK: {{.*}} = affine.apply #[[$MAP_mix_dims_and_symbols_d]]()[%{{.*}}, %{{.*}}]
   return %d : index
 }
 
+// -----
+
+// CHECK-DAG: #[[$MAP_mix_dims_and_symbols_e:.*]] = affine_map<()[s0, s1] -> ((s0 * 168 + s1 * 4 - 4) floordiv 3)>
+
 // CHECK-LABEL: func @mix_dims_and_symbols_e(%arg0: index, %arg1: index) -> index {
 func @mix_dims_and_symbols_e(%arg0: index, %arg1: index) -> index {
   %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
   %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
   %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
   %d = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)> ()[%b]
   %e = affine.apply affine_map<(d0) -> (d0 floordiv 3)> (%c)
-  // CHECK: {{.*}} = affine.apply [[$MAP_mix_dims_and_symbols_e]]()[%{{.*}}, %{{.*}}]
+  // CHECK: {{.*}} = affine.apply #[[$MAP_mix_dims_and_symbols_e]]()[%{{.*}}, %{{.*}}]
   return %e : index
 }
 
+// -----
+
 // CHECK-LABEL: func @mix_dims_and_symbols_f(%arg0: index, %arg1: index) -> index {
 func @mix_dims_and_symbols_f(%arg0: index, %arg1: index) -> index {
   %a = affine.apply affine_map<(d0)[s0] -> (d0 - 1 + 42 * s0)> (%arg0)[%arg1]
   %b = affine.apply affine_map<(d0) -> (d0 + 7)> (%a)
   %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
   %d = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)> ()[%b]
   %e = affine.apply affine_map<(d0) -> (d0 floordiv 3)> (%c)
   %f = affine.apply affine_map<(d0, d1)[s0, s1] -> (d0 - s1 +  d1 - s0)> (%d, %e)[%e, %d]
   // CHECK: {{.*}} = constant 0 : index
 
   return %f : index
 }
 
+// -----
+
+// CHECK-DAG: #[[$MAP_symbolic_composition_b:.*]] = affine_map<()[s0] -> (s0 * 4)>
+
 // CHECK-LABEL: func @mix_dims_and_symbols_g(%arg0: index, %arg1: index) -> (index, index, index) {
 func @mix_dims_and_symbols_g(%M: index, %N: index) -> (index, index, index) {
   %K = affine.apply affine_map<(d0) -> (4*d0)> (%M)
   %res1 = affine.apply affine_map<()[s0, s1] -> (4 * s0)>()[%N, %K]
   %res2 = affine.apply affine_map<()[s0, s1] -> (s1)>()[%N, %K]
   %res3 = affine.apply affine_map<()[s0, s1] -> (1024)>()[%N, %K]
   // CHECK-DAG: {{.*}} = constant 1024 : index
-  // CHECK-DAG: {{.*}} = affine.apply [[$MAP_symbolic_composition_b]]()[%{{.*}}]
-  // CHECK-DAG: {{.*}} = affine.apply [[$MAP_symbolic_composition_b]]()[%{{.*}}]
+  // CHECK-DAG: {{.*}} = affine.apply #[[$MAP_symbolic_composition_b]]()[%{{.*}}]
+  // CHECK-DAG: {{.*}} = affine.apply #[[$MAP_symbolic_composition_b]]()[%{{.*}}]
   return %res1, %res2, %res3 : index, index, index
 }
 
+// -----
+
+// CHECK-DAG: #[[$symbolic_semi_affine:.*]] = affine_map<(d0)[s0] -> (d0 floordiv (s0 + 1))>
+
 // CHECK-LABEL: func @symbolic_semi_affine(%arg0: index, %arg1: index, %arg2: memref<?xf32>) {
 func @symbolic_semi_affine(%M: index, %N: index, %A: memref<?xf32>) {
   %f1 = constant 1.0 : f32
   affine.for %i0 = 1 to 100 {
     %1 = affine.apply affine_map<()[s0] -> (s0 + 1)> ()[%M]
     %2 = affine.apply affine_map<(d0)[s0] -> (d0 floordiv s0)> (%i0)[%1]
-    // CHECK-DAG: {{.*}} = affine.apply [[$symbolic_semi_affine]](%{{.*}})[%{{.*}}]
+    // CHECK-DAG: {{.*}} = affine.apply #[[$symbolic_semi_affine]](%{{.*}})[%{{.*}}]
     store %f1, %A[%2] : memref<?xf32>
   }
   return
 }
 
 // -----
 
-// CHECK: [[$MAP0:#map[0-9]+]] = affine_map<()[s0] -> (0, s0)>
-// CHECK: [[$MAP1:#map[0-9]+]] = affine_map<()[s0] -> (100, s0)>
+// CHECK: #[[$MAP0:.*]] = affine_map<()[s0] -> (0, s0)>
+// CHECK: #[[$MAP1:.*]] = affine_map<()[s0] -> (100, s0)>
 
 // CHECK-LABEL:  func @constant_fold_bounds(%arg0: index) {
 func @constant_fold_bounds(%N : index) {
   // CHECK:      constant 3 : index
   // CHECK-NEXT: "foo"() : () -> index
   %c9 = constant 9 : index
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = affine.apply affine_map<(d0, d1) -> (d0 + d1)> (%c1, %c2)
   %l = "foo"() : () -> index
 
   // CHECK:  affine.for %{{.*}} = 5 to 7 {
   affine.for %i = max affine_map<(d0, d1) -> (0, d0 + d1)> (%c2, %c3) to min affine_map<(d0, d1) -> (d0 - 2, 32*d1)> (%c9, %c1) {
     "foo"(%i, %c3) : (index, index) -> ()
   }
 
   // Bound takes a non-constant argument but can still be folded.
   // CHECK:  affine.for %{{.*}} = 1 to 7 {
   affine.for %j = max affine_map<(d0) -> (0, 1)> (%N) to min affine_map<(d0, d1) -> (7, 9)> (%N, %l) {
     "foo"(%j, %c3) : (index, index) -> ()
   }
 
   // None of the bounds can be folded.
-  // CHECK: affine.for %{{.*}} = max [[$MAP0]]()[%{{.*}}] to min [[$MAP1]]()[%{{.*}}] {
+  // CHECK: affine.for %{{.*}} = max #[[$MAP0]]()[%{{.*}}] to min #[[$MAP1]]()[%{{.*}}] {
   affine.for %k = max affine_map<()[s0] -> (0, s0)> ()[%l] to min affine_map<()[s0] -> (100, s0)> ()[%N] {
     "foo"(%k, %c3) : (index, index) -> ()
   }
   return
 }
 
 // -----
 
 // CHECK-LABEL:  func @fold_empty_loop() {
 func @fold_empty_loop() {
   // CHECK-NOT: affine.for
   affine.for %i = 0 to 10 {
   }
   return
 }
 // CHECK: return
 
+
 // -----
 
-// CHECK-DAG: [[$SET:#set[0-9]*]] = affine_set<(d0, d1)[s0] : (d0 >= 0, -d0 + 1022 >= 0, d1 >= 0, -d1 + s0 - 2 >= 0)>
+// CHECK-DAG: #[[$SET:.*]] = affine_set<(d0, d1)[s0] : (d0 >= 0, -d0 + 1022 >= 0, d1 >= 0, -d1 + s0 - 2 >= 0)>
 
 // CHECK-LABEL: func @canonicalize_affine_if
-// CHECK-SAME: [[M:%.*]]: index,
-// CHECK-SAME: [[N:%.*]]: index)
+//  CHECK-SAME:   %[[M:[0-9a-zA-Z]*]]: index,
+//  CHECK-SAME:   %[[N:[0-9a-zA-Z]*]]: index)
 func @canonicalize_affine_if(%M : index, %N : index) {
   %c1022 = constant 1022 : index
   // Drop unused operand %M, propagate %c1022, and promote %N to symbolic.
   affine.for %i = 0 to 1024 {
     affine.for %j = 0 to %N {
-      // CHECK: affine.if [[$SET]](%{{.*}}, %{{.*}}){{\[}}[[N]]{{\]}}
-      affine.if affine_set<(d0, d1, d2, d3)[s0] : (d1 >= 0, d0 - d1 >= 0, d2 >= 0, d3 - d2 - 2 >= 0)> (%c1022, %i, %j, %N)[%M] {
+      // CHECK: affine.if #[[$SET]](%{{.*}}, %{{.*}})[%[[N]]]
+      affine.if affine_set<(d0, d1, d2, d3)[s0] : (d1 >= 0, d0 - d1 >= 0, d2 >= 0, d3 - d2 - 2 >= 0)>
+          (%c1022, %i, %j, %N)[%M] {
         "foo"() : () -> ()
       }
       "bar"() : () -> ()
     }
   }
   return
 }
 
 // -----
 
-// CHECK-DAG: [[$LBMAP:#map[0-9]+]] = affine_map<()[s0] -> (0, s0)>
-// CHECK-DAG: [[$UBMAP:#map[0-9]+]] = affine_map<()[s0] -> (1024, s0 * 2)>
+// CHECK-DAG: #[[$LBMAP:.*]] = affine_map<()[s0] -> (0, s0)>
+// CHECK-DAG: #[[$UBMAP:.*]] = affine_map<()[s0] -> (1024, s0 * 2)>
 
 // CHECK-LABEL: func @canonicalize_bounds
-// CHECK-SAME: [[M:%.*]]: index,
-// CHECK-SAME: [[N:%.*]]: index)
+// CHECK-SAME: %[[M:.*]]: index,
+// CHECK-SAME: %[[N:.*]]: index)
 func @canonicalize_bounds(%M : index, %N : index) {
   %c0 = constant 0 : index
   %c1024 = constant 1024 : index
   // Drop unused operand %N, drop duplicate operand %M, propagate %c1024, and
   // promote %M to a symbolic one.
-  // CHECK: affine.for %{{.*}} = 0 to min [[$UBMAP]](){{\[}}[[M]]{{\]}}
+  // CHECK: affine.for %{{.*}} = 0 to min #[[$UBMAP]]()[%[[M]]]
   affine.for %i = 0 to min affine_map<(d0, d1, d2, d3) -> (d0, d1 + d2)> (%c1024, %M, %M, %N) {
     "foo"() : () -> ()
   }
   // Promote %M to symbolic position.
-  // CHECK: affine.for %{{.*}} = 0 to #map{{[0-9]+}}(){{\[}}[[M]]{{\]}}
+  // CHECK: affine.for %{{.*}} = 0 to #{{.*}}()[%[[M]]]
   affine.for %i = 0 to affine_map<(d0) -> (4 * d0)> (%M) {
     "foo"() : () -> ()
   }
   // Lower bound canonicalize.
-  // CHECK: affine.for %{{.*}} = max [[$LBMAP]](){{\[}}[[N]]{{\]}} to [[M]]
+  // CHECK: affine.for %{{.*}} = max #[[$LBMAP]]()[%[[N]]] to %[[M]]
   affine.for %i = max affine_map<(d0, d1) -> (d0, d1)> (%c0, %N) to %M {
     "foo"() : () -> ()
   }
   return
 }
 
 // -----
 

 // symbols were concatenated in the wrong order during map folding. Map
 // composition places the symbols of the original map before those of the map
 // it is composed with, e.g. A.compose(B) will first have all symbols of A,
 // then all symbols of B.
 
 #map1 = affine_map<(d0)[s0, s1] -> (d0 * s0 + s1)>
 #map2 = affine_map<(d0)[s0] -> (1024, -d0 + s0)>
 
-// CHECK: #[[$MAP:.*]] = affine_map<()[s0, s1] -> (1024, s1 * -1024 + s0)>
+// CHECK: #[[$MAP:.*]] = affine_map<()[s0, s1] -> (1024, s0 - s1 * 1024)>
 
 // CHECK: func @rep(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index)
 func @rep(%arg0 : index, %arg1 : index) -> index {
   // CHECK-NOT: constant
   %c0 = constant 0 : index
   %c1024 = constant 1024 : index
   // CHECK-NOT: affine.apply
   %0 = affine.apply #map1(%arg0)[%c1024, %c0]
affine.parallel (%i) = (%c0) to (%c0 + %c4) {
     %1 = affine.apply #map3(%i)
     // CHECK: affine.parallel (%{{.*}}) = (0) to (%{{.*}} * 5)
     affine.parallel (%j) = (%c0) to (%1) {
       affine.store %cst, %0[%j] : memref<4xf32>
     }
   }
   return
 }
+
+// -----
+
+func @compose_affine_maps_div_symbol(%A : memref<i64>, %i0 : index, %i1 : index) {
+  %0 = affine.apply affine_map<()[s0] -> (2 * s0)> ()[%i0]
+  %1 = affine.apply affine_map<()[s0] -> (3 * s0)> ()[%i0]
+  %2 = affine.apply affine_map<(d0)[s0, s1] -> (d0 mod s1 + s0 * s1 + s0 * 4)> (%i1)[%0, %1]
+  %3 = index_cast %2: index to i64
+  store %3, %A[]: memref<i64>
+  affine.for %i2 = 0 to 3 {
+    %4 = affine.apply affine_map<(d0)[s0, s1] -> (d0 ceildiv s1 + s0 + s0 * 3)> (%i2)[%0, %1]
+    %5 = index_cast %4: index to i64
+    store %5, %A[]: memref<i64>
+  }
+  return
+}

mlir/test/Dialect/Linalg/reshape_fusion.mlir

^bb0(%arg3 : index, %arg4 : index, %arg5 : index, %arg6: i32, %arg7: i32, %s: i32):
       %7 = addi %5, %6 : i32
       linalg.yield %7 : i32
   } -> tensor<?x?x?xi32>
   return %1 : tensor<?x?x?xi32>
 }
 
 // The generic op version of the test check for the op structure. Only
 // checking the op body here.
-//       CHECK: #[[MAP:.+]] =  affine_map<(d0, d1) -> (d0 * 4 + d1)>
+//       CHECK: #[[MAP:.+]] =  affine_map<(d0, d1) -> (d0 + d1 * 4)>
 //       CHECK: func @indexed_generic_op_reshape_producer_fusion
 //       CHECK:   linalg.indexed_generic
 //       CHECK:   ^{{.*}}(
 //  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index, %[[ARG3:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: index, %[[ARG5:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG6:[a-zA-Z0-9]+]]: i32, %[[ARG7:[a-zA-Z0-9]+]]: i32,
 //  CHECK-SAME:     %[[ARG8:[a-zA-Z0-9]+]]: i32)
-//       CHECK:     %[[T3:.+]] = affine.apply #[[MAP]](%[[ARG2]], %[[ARG3]])
+//       CHECK:     %[[T3:.+]] = affine.apply #[[MAP]](%[[ARG3]], %[[ARG2]])
 //       CHECK:     %[[T4:.+]] = muli %[[ARG6]], %[[ARG7]]
 //       CHECK:     %[[T5:.+]] = index_cast %[[T3]]
 //       CHECK:     %[[T6:.+]] = addi %[[T4]], %[[T5]]
 //       CHECK:     %[[T7:.+]] = index_cast %[[ARG4]]
 //       CHECK:     %[[T8:.+]] = addi %[[T6]], %[[T7]]
 //       CHECK:     %[[T9:.+]] = index_cast %[[ARG5]]
 //       CHECK:     %[[T10:.+]] = addi %[[T8]], %[[T9]]
 //       CHECK:     linalg.yield %[[T10]]
func @indexed_generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?xi32>,
   } -> tensor<?x?xi32>
   %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
                                  affine_map<(i, j, k, l) -> (j, k, l)>] :
     tensor<?x?xi32> into tensor<?x?x4x5xi32>
   return %1 : tensor<?x?x4x5xi32>
 }
 // The generic op version of the test check for the op structure. Only
 // checking the op body here.
-//       CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 * 20 + d1 * 5 + d2)>
+//       CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 + d1 * 5 + d2 * 20)>
 //       CHECK: func @indexed_generic_op_reshape_consumer_fusion
 //       CHECK:   linalg.indexed_generic
 //       CHECK:   ^{{.*}}(
 //  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index, %[[ARG3:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: index, %[[ARG5:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG6:[a-zA-Z0-9]+]]: i32, %[[ARG7:[a-zA-Z0-9]+]]: i32,
 //  CHECK-SAME:     %[[ARG8:[a-zA-Z0-9]+]]: i32)
-//       CHECK:     %[[T3:.+]] = affine.apply #[[MAP]](%[[ARG3]], %[[ARG4]], %[[ARG5]])
+//       CHECK:     %[[T3:.+]] = affine.apply #[[MAP]](%[[ARG5]], %[[ARG4]], %[[ARG3]])
 //       CHECK:     %[[T4:.+]] = muli %[[ARG6]], %[[ARG7]]
 //       CHECK:     %[[T5:.+]] = index_cast %[[ARG2]]
 //       CHECK:     %[[T6:.+]] = addi %[[T4]], %[[T5]]
 //       CHECK:     %[[T7:.+]] = index_cast %[[T3]]
 //       CHECK:     %[[T8:.+]] = addi %[[T6]], %[[T7]]
 //       CHECK:     linalg.yield %[[T8]]
 
 // -----
%c = linalg.indexed_generic {
        ^bb0(%arg0 : index, %arg1 : index, %arg2 : index, %arg3 : i32, %arg4: i32, %s: i32):
          %1 = addi %arg3, %arg4 : i32
          %2 = index_cast %arg0 : index to i32
          %3 = addi %1, %2 : i32
          %4 = index_cast %arg1 : index to i32
          %5 = addi %3, %4 : i32
          %6 = index_cast %arg2 : index to i32
          %7 = addi %5, %6 : i32
-   linalg.yield %7 : i32
+  linalg.yield %7 : i32
        } -> tensor<6x4x210xi32>
   %d = linalg.tensor_reshape %c
          [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1)>,
           affine_map<(d0, d1, d2, d3, d4, d5) -> (d2)>,
           affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4, d5)>]
        : tensor<6x4x210xi32> into tensor<2x3x4x5x6x7xi32>
   return %d : tensor<2x3x4x5x6x7xi32>
 }
 
 
 //   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2)>
 //   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4)>
 //   CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d5)>
 //   CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
 //   CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
 //   CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1)>
 //   CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2)>
 //   CHECK-DAG: #[[MAP7:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4, d5)>
 //   CHECK-DAG: #[[MAP8:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d0, d1, d5)>
 //   CHECK-DAG: #[[MAP9:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d5)>
 //   CHECK-DAG: #[[MAP10:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d2, d3, d4)>
-//   CHECK-DAG: #[[MAP11:.+]] = affine_map<(d0, d1) -> (d0 * 3 + d1)>
-//   CHECK-DAG: #[[MAP12:.+]] = affine_map<(d0, d1, d2) -> (d0 * 42 + d1 * 7 + d2)>
+//   CHECK-DAG: #[[MAP11:.+]] = affine_map<(d0, d1) -> (d0 + d1 * 3)>
+//   CHECK-DAG: #[[MAP12:.+]] = affine_map<(d0, d1, d2) -> (d0 + d1 * 7 + d2 * 42)>
 //       CHECK: func @reshape_as_consumer_permutation
 //  CHECK-SAME:   %[[ARG0:.+]]: tensor<210x6x4xi32>
 //  CHECK-SAME:   %[[ARG1:.+]]: tensor<210x4xi32>
 //   CHECK-DAG:   %[[T0:.+]] = linalg.init_tensor [6, 4, 210]
 //   CHECK-DAG:   %[[T1:.+]] = linalg.tensor_reshape %[[ARG0]]
 //  CHECK-SAME:     [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
 //   CHECK-DAG:   %[[T2:.+]] = linalg.tensor_reshape %[[ARG1]]
 //  CHECK-SAME:     [#[[MAP3]], #[[MAP4]]]
 //       CHECK:   %[[T3:.+]] = linalg.tensor_reshape %[[T0]]
 //  CHECK-SAME:     [#[[MAP5]], #[[MAP6]], #[[MAP7]]]
 //       CHECK:   %[[T4:.+]] = linalg.indexed_generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP8]], #[[MAP9]], #[[MAP10]]]
 //  CHECK-SAME:     ins(%[[T1]], %[[T2]] : tensor<5x6x7x2x3x4xi32>, tensor<5x6x7x4xi32>)
 //  CHECK-SAME:     outs(%[[T3]] : tensor<2x3x4x5x6x7xi32>)
 //       CHECK:   ^{{.+}}(
 //  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index, %[[ARG3:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: index, %[[ARG5:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG6:[a-zA-Z0-9]+]]: index, %[[ARG7:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG8:[a-zA-Z0-9]+]]: i32, %[[ARG9:[a-zA-Z0-9]+]]: i32,
 //  CHECK-SAME:     %[[ARG10:[a-zA-Z0-9]+]]: i32)
-//   CHECK-DAG:       %[[T5:.+]] = affine.apply #[[MAP11]](%[[ARG2]], %[[ARG3]])
-//   CHECK-DAG:       %[[T6:.+]] = affine.apply #[[MAP12]](%[[ARG4]], %[[ARG5]], %[[ARG6]])
+//   CHECK-DAG:       %[[T5:.+]] = affine.apply #[[MAP11]](%[[ARG3]], %[[ARG2]])
+//   CHECK-DAG:       %[[T6:.+]] = affine.apply #[[MAP12]](%[[ARG6]], %[[ARG5]], %[[ARG4]])
 //   CHECK-DAG:       %[[T7:.+]] = addi %[[ARG8]], %[[ARG9]]
 //       CHECK:       %[[T8:.+]] = index_cast %[[T5]]
 //       CHECK:       %[[T9:.+]] = addi %[[T7]], %[[T8]]
 //       CHECK:       %[[T10:.+]] = index_cast %[[T6]]
 //       CHECK:       %[[T11:.+]] = addi %[[T9]], %[[T10]]
 //       CHECK:       %[[T12:.+]] = index_cast %[[ARG7]]
 //       CHECK:       %[[T13:.+]] = addi %[[T11]], %[[T12]]
 
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: i32, %s: i32):  // no predecessors
     %7 = addi %5, %6 : i32
     linalg.yield %7 : i32
   } -> tensor<264x?x4xi32>
   return %1 : tensor<264x?x4xi32>
 }
 
 //   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
 //   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
-//   CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0 * 8 + d1)>
+//   CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0 + d1 * 8)>
 //   CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
 //   CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d2)>
 //   CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
 //       CHECK: @reshape_as_producer_projected_permutation
 //  CHECK-SAME:   %[[ARG0:.+]]: tensor<33x8x?xi32>
 //       CHECK:   %[[RES:.+]] = linalg.indexed_generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]]]
 //  CHECK-SAME:     ins(%[[ARG0]] : tensor<33x8x?xi32>)
 //       CHECK:   ^{{.+}}(
 //  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG3:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG5:[a-zA-Z0-9]+]]: i32,
 //  CHECK-SAME:     %[[ARG7:[a-zA-Z0-9]+]]: i32)
-//       CHECK:       %[[T0:.+]] = affine.apply #[[MAP2]](%[[ARG1]], %[[ARG2]])
+//       CHECK:       %[[T0:.+]] = affine.apply #[[MAP2]](%[[ARG2]], %[[ARG1]])
 //       CHECK:       %[[T1:.+]] = index_cast %[[T0]] : index to i32
 //       CHECK:       %[[T2:.+]] = addi %[[ARG5]], %[[T1]] : i32
 //       CHECK:       %[[T3:.+]] = index_cast %[[ARG3]] : index to i32
 //       CHECK:       %[[T4:.+]] = addi %[[T2]], %[[T3]] : i32
 //       CHECK:       %[[T5:.+]] = index_cast %[[ARG4]] : index to i32
 //       CHECK:       %[[T6:.+]] = addi %[[T4]], %[[T5]] : i32
 //       CHECK:       linalg.yield %[[T6]] : i32
 //       CHECK:    %[[RES2:.+]] = linalg.tensor_reshape %[[RES]]

mlir/test/EDSC/builder-api-test.cpp

affineLoopBuilder(lb, ub, 2, [&](Value j) {
       ((i7 + i13) / i7) % i13 - i7 *i13;
     });
   });
 
   // clang-format off
   // CHECK-LABEL: func @builder_dynamic_for_func_args(%{{.*}}: index, %{{.*}}: index) {
   //     CHECK:  affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%{{.*}}) to affine_map<(d0) -> (d0)>(%{{.*}}) step 3 {
   //     CHECK:  {{.*}} = affine.apply affine_map<()[s0] -> (s0 * 3)>()[%{{.*}}]
-  //     CHECK:  {{.*}} = affine.apply affine_map<()[s0, s1] -> (s1 + s0 * 3)>()[%{{.*}}, %{{.*}}]
+  //     CHECK:  {{.*}} = affine.apply affine_map<()[s0, s1] -> (s0 * 3 + s1)>()[%{{.*}}, %{{.*}}]
   //     CHECK:  {{.*}} = affine.apply affine_map<()[s0] -> (s0 + 3)>()[%{{.*}}]
   //     CHECK:  affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%{{.*}}) to affine_map<(d0) -> (d0)>(%{{.*}}) step 2 {
   //     CHECK:    {{.*}} = affine.apply affine_map<(d0, d1) -> ((d0 + d1 * 3) floordiv 32)>(%{{.*}}, %{{.*}})
   //     CHECK:    {{.*}} = affine.apply affine_map<(d0, d1) -> (((d0 + d1 * 3) floordiv 32) * 31)>(%{{.*}}, %{{.*}})
   //     CHECK:    {{.*}} = affine.apply affine_map<(d0, d1) -> ((((d0 + d1 * 3) floordiv 32) * 31) ceildiv 32)>(%{{.*}}, %{{.*}})
   // CHECK-DAG:    [[rf1:%[0-9]+]] = addf {{.*}}, {{.*}} : f32
   // CHECK-DAG:    [[rf2:%[0-9]+]] = divf [[rf1]], {{.*}} : f32
   // CHECK-DAG:    [[rf3:%[0-9]+]] = remf [[rf2]], {{.*}} : f32

mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp

static llvm::cl::opt<bool> clTestSlicingAnalysis(
     llvm::cl::cat(clOptionsCategory));
 static llvm::cl::opt<bool> clTestComposeMaps(
     "compose-maps",
     llvm::cl::desc(
         "Enable testing the composition of AffineMap where each "
         "AffineMap in the composition is specified as the affine_map attribute "
         "in a constant op."),
     llvm::cl::cat(clOptionsCategory));
-static llvm::cl::opt<bool> clTestNormalizeMaps(
-    "normalize-maps",
-    llvm::cl::desc(
-        "Enable testing the normalization of AffineAffineApplyOp "
-        "where each AffineAffineApplyOp in the composition is a single output "
-        "operation."),
-    llvm::cl::cat(clOptionsCategory));
 static llvm::cl::opt<bool> clTestVecAffineLoopNest(
     "vectorize-affine-loop-nest",
     llvm::cl::desc(
         "Enable testing for the 'vectorizeAffineLoopNest' utility by "
         "vectorizing the outermost loops found"),
     llvm::cl::cat(clOptionsCategory));
 
 namespace {
 struct VectorizerTestPass
     : public PassWrapper<VectorizerTestPass, FunctionPass> {
   static constexpr auto kTestAffineMapOpName = "test_affine_map";
   static constexpr auto kTestAffineMapAttrName = "affine_map";
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<vector::VectorDialect>();
   }
 
   void runOnFunction() override;
   void testVectorShapeRatio(llvm::raw_ostream &outs);
   void testForwardSlicing(llvm::raw_ostream &outs);
   void testBackwardSlicing(llvm::raw_ostream &outs);
   void testSlicing(llvm::raw_ostream &outs);
   void testComposeMaps(llvm::raw_ostream &outs);
-  void testNormalizeMaps();
 
   /// Test for 'vectorizeAffineLoopNest' utility.
   void testVecAffineLoopNest();
 };
 
 } // end anonymous namespace
 
 void VectorizerTestPass::testVectorShapeRatio(llvm::raw_ostream &outs) {

 
 static bool affineApplyOp(Operation &op) { return isa<AffineApplyOp>(op); }
 
 static bool singleResultAffineApplyOpWithoutUses(Operation &op) {
   auto app = dyn_cast<AffineApplyOp>(op);
   return app && app.use_empty();
 }
 
-void VectorizerTestPass::testNormalizeMaps() {
-  using matcher::Op;
-
-  auto f = getFunction();
-
-  // Save matched AffineApplyOp that all need to be erased in the end.
-  auto pattern = Op(affineApplyOp);
-  SmallVector<NestedMatch, 8> toErase;
-  pattern.match(f, &toErase);
-  {
-    // Compose maps.
-    auto pattern = Op(singleResultAffineApplyOpWithoutUses);
-    SmallVector<NestedMatch, 8> matches;
-    pattern.match(f, &matches);
-    for (auto m : matches) {
-      auto app = cast<AffineApplyOp>(m.getMatchedOperation());
-      OpBuilder b(m.getMatchedOperation());
-      SmallVector<Value, 8> operands(app.getOperands());
-      makeComposedAffineApply(b, app.getLoc(), app.getAffineMap(), operands);
-    }
-  }
-  // We should now be able to erase everything in reverse order in this test.
-  for (auto m : llvm::reverse(toErase)) {
-    m.getMatchedOperation()->erase();
-  }
-}
-
 /// Test for 'vectorizeAffineLoopNest' utility.
 void VectorizerTestPass::testVecAffineLoopNest() {
   std::vector<SmallVector<AffineForOp, 2>> loops;
   gatherLoops(getFunction(), loops);
 
   // Expected only one loop nest.
   if (loops.empty() || loops[0].size() != 1)
     return;
llvm::raw_string_ostream outs(str);
     if (clTestBackwardSlicingAnalysis)
       testBackwardSlicing(outs);
 
     if (clTestSlicingAnalysis)
       testSlicing(outs);
 
     if (clTestComposeMaps)
       testComposeMaps(outs);
-
-    if (clTestNormalizeMaps)
-      testNormalizeMaps();
   }
 
   if (clTestVecAffineLoopNest)
     testVecAffineLoopNest();
 
   if (!outs.str().empty()) {
     emitRemark(UnknownLoc::get(&getContext()), outs.str());
   }