[mlir][AffineMap] NFC - Refactor getProjectedMap and split into projectDims and projectSymbols

[nicolasvasilache]

Summary

The default behavior of getProjectedMap may be surprising as it implicitly compresses the dims and
the unused symbols.

Make these explicit in the API and refactor to more idiomatic implementations with better reuse.

[nicolasvasilache]

Update

[bondhugula]

mlir/include/mlir/IR/AffineMap.h

575–576

Please avoid Flag suffixes; it's redundant and inconsistent with the style here.

[qcolombet]

mlir/include/mlir/IR/AffineMap.h

417–419

Was there a typo in the original comment: not listed -> listed?

[ftynse]

mlir/lib/IR/AffineMap.cpp

574–580

Add documentation plz.

721–722

Nit: llvm::is_one_of<> makes this slightly more concise.

724

Nit: camelCase.

731–738

Optional: this is now possible without lambdas, along the lines of

for (...) {
  ...
  if constexpr (isDim) {
    replacements.push_back(getAffineDimExpr(...));
  } else {
    replacements.push_back(getAffineSymbolExpr(...));
  }
}}

which may or may not be more readable.

802

Nit: expand auto (presumably to AffineMap that we don't use by refernece)

[nicolasvasilache]

mlir/include/mlir/IR/AffineMap.h

575–576

The issue is that the natural name collides with the compressDims / compressSymbols functions in the .cpp
Do you have a preferred way to spell these names ?

575–576

Please feel free to just commit an NFC on top of this with your preferred spelling.

mlir/lib/IR/AffineMap.cpp

731–738

nice, didn't know about this yet.
not more readable here IMO so not applying but good to know for the future.

802

thanks missed that copypastaing

[nicolasvasilache]

Address comments.

[ftynse]

mlir/lib/IR/AffineMap.cpp

731–738

Yeah, that's what I expected. It is more useful to replace things like enable_if<foo> / enable_if<!foo> pairs.

mlir/include/mlir/IR/AffineMap.h

private:
   /// A pointer to the IR's context to store all newly created
   /// AffineExprStorage's.
   MLIRContext *context = nullptr;
 };
 
 /// Simplifies an affine map by simplifying its underlying AffineExpr results.
 AffineMap simplifyAffineMap(AffineMap map);
 
+/// Drop the dims that are listed in `unusedDims`.
+AffineMap compressDims(AffineMap map, const llvm::SmallBitVector &unusedDims);
+
 /// Drop the dims that are not used.
 AffineMap compressUnusedDims(AffineMap map);
 
 /// Drop the dims that are not used by any of the individual maps in `maps`.
 /// Asserts that all maps in `maps` are normalized to the same number of
 /// dims and symbols.
 SmallVector<AffineMap> compressUnusedDims(ArrayRef<AffineMap> maps);
 
-/// Drop the dims that are not listed in `unusedDims`.
-AffineMap compressDims(AffineMap map, const llvm::SmallBitVector &unusedDims);
+/// Drop the symbols that are listed in `unusedSymbols`.
+AffineMap compressSymbols(AffineMap map,
+                          const llvm::SmallBitVector &unusedSymbols);
 
 /// Drop the symbols that are not used.
 AffineMap compressUnusedSymbols(AffineMap map);
 
 /// Drop the symbols that are not used by any of the individual maps in `maps`.
 /// Asserts that all maps in `maps` are normalized to the same number of
 /// dims and symbols.
 SmallVector<AffineMap> compressUnusedSymbols(ArrayRef<AffineMap> maps);
 
-/// Drop the symbols that are not listed in `unusedSymbols`.
-AffineMap compressSymbols(AffineMap map,
-                          const llvm::SmallBitVector &unusedSymbols);
-
 /// Returns a map with the same dimension and symbol count as `map`, but whose
 /// results are the unique affine expressions of `map`.
 AffineMap removeDuplicateExprs(AffineMap map);
 
 /// Returns a map of codomain to domain dimensions such that the first codomain
 /// dimension for a particular domain dimension is selected.
 /// Returns an empty map if the input map is empty.
 /// Returns null map (not empty map) if `map` is not invertible (i.e. `map` does

 /// ```mlir
 ///    (d0, d1, d2, d3, d4, d5, d6, d7) -> (d2, d0, d3)
 /// ```
 AffineMap inversePermutation(AffineMap map);
 
 /// Return the reverse map of a projected permutation where the projected
 /// dimensions are transformed into 0s.
 ///
-/// Prerequisites: `map` must be a projected permuation.
+/// Prerequisites: `map` must be a projected permutation.
 ///
 /// Example 1:
 ///
 /// ```mlir
 ///    affine_map<(d0, d1, d2, d3) -> (d2, d0)>
 /// ```
 ///
 /// returns:

 /// 2) map                  : affine_map<(d0, d1) -> (d0 + d1)>
 ///    projected_dimensions : {1}
 ///    result               : affine_map<(d0) -> (d0)>
 ///
 /// 3) map                  : affine_map<(d0, d1, d2) -> (d0, d1)>
 ///    projected_dimensions : {1}
 ///    result               : affine_map<(d0, d1) -> (d0, 0)>
 ///
-/// This function also compresses unused symbols away.
+/// This function also compresses the dims when the boolean flag is true.
+AffineMap projectDims(AffineMap map,
+                      const llvm::SmallBitVector &projectedDimensions,
+                      bool compressDimsFlag = false);
+/// Symbol counterpart of `projectDims`.
+/// This function also compresses the symbols when the boolean flag is true.
+AffineMap projectSymbols(AffineMap map,
+                         const llvm::SmallBitVector &projectedSymbols,
+                         bool compressSymbolsFlag = false);
+/// Calls `projectDims(map, projectedDimensions, compressDimsFlag)`.
+/// If `compressSymbolsFlag` is true, additionally call `compressUnusedSymbols`.
 AffineMap getProjectedMap(AffineMap map,
-                          const llvm::SmallBitVector &projectedDimensions);
+                          const llvm::SmallBitVector &projectedDimensions,
+                          bool compressDimsFlag = true,
+                          bool compressSymbolsFlag = true);
+
+// Return a bitvector where each bit set indicates a dimension that is not used
+// by any of the maps in the input array `maps`.
+llvm::SmallBitVector getUnusedDimsBitVector(ArrayRef<AffineMap> maps);
+
+// Return a bitvector where each bit set indicates a symbol that is not used
+// by any of the maps in the input array `maps`.
+llvm::SmallBitVector getUnusedSymbolsBitVector(ArrayRef<AffineMap> maps);
+
+inline raw_ostream &operator<<(raw_ostream &os, AffineMap map) {
+  map.print(os);
+  return os;
+}
+
+//===----------------------------------------------------------------------===//
+// Templated helper functions.
+//===----------------------------------------------------------------------===//
 
 /// Apply a permutation from `map` to `source` and return the result.
 template <typename T>
 SmallVector<T> applyPermutationMap(AffineMap map, llvm::ArrayRef<T> source) {
   assert(map.isProjectedPermutation());
   assert(map.getNumInputs() == source.size());
   SmallVector<T> result;
   result.reserve(map.getNumResults());
   for (AffineExpr expr : map.getResults()) {
     if (auto dimExpr = expr.dyn_cast<AffineDimExpr>()) {
       result.push_back(source[dimExpr.getPosition()]);
     } else if (auto constExpr = expr.dyn_cast<AffineConstantExpr>()) {
       assert(constExpr.getValue() == 0 &&
              "Unexpected constant in projected permutation map");
       result.push_back(0);
     } else {
       llvm_unreachable("Unexpected result in projected permutation map");
     }
   }
   return result;
 }
 
-/// Calculates maxmimum dimension and symbol positions from the expressions
+/// Calculates maximum dimension and symbol positions from the expressions
 /// in `exprsLists` and stores them in `maxDim` and `maxSym` respectively.
 template <typename AffineExprContainer>
 static void getMaxDimAndSymbol(ArrayRef<AffineExprContainer> exprsList,
                                int64_t &maxDim, int64_t &maxSym) {
   for (const auto &exprs : exprsList) {
     for (auto expr : exprs) {
       expr.walk([&maxDim, &maxSym](AffineExpr e) {
         if (auto d = e.dyn_cast<AffineDimExpr>())
           maxDim = std::max(maxDim, static_cast<int64_t>(d.getPosition()));
         if (auto s = e.dyn_cast<AffineSymbolExpr>())
           maxSym = std::max(maxSym, static_cast<int64_t>(s.getPosition()));
       });
     }
   }
 }
 
-inline raw_ostream &operator<<(raw_ostream &os, AffineMap map) {
-  map.print(os);
-  return os;
-}
-
-// Return a bitvector where each bit set indicates a dimension that is not used
-// by any of the maps in the input array `maps`.
-llvm::SmallBitVector getUnusedDimsBitVector(ArrayRef<AffineMap> maps);
-
 } // namespace mlir
 
 namespace llvm {
 
 // AffineExpr hash just like pointers
 template <>
 struct DenseMapInfo<mlir::AffineMap> {
   static mlir::AffineMap getEmptyKey() {

mlir/lib/IR/AffineMap.cpp

 //===- AffineMap.cpp - MLIR Affine Map Classes ----------------------------===//
 //
 // Part of the LLVM 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/AffineMap.h"
 #include "AffineMapDetail.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/MathExtras.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/raw_ostream.h"
 #include <numeric>
 #include <optional>
+#include <type_traits>
 
 using namespace mlir;
 
 namespace {
 
 // AffineExprConstantFolder evaluates an affine expression using constant
 // operands passed in 'operandConsts'. Returns an IntegerAttr attribute
 // representing the constant value of the affine expression evaluated on

 AffineMap AffineMap::getMinorSubMap(unsigned numResults) const {
   if (numResults == 0)
     return AffineMap();
   if (numResults > getNumResults())
     return *this;
   return getSliceMap(getNumResults() - numResults, numResults);
 }
 
-AffineMap mlir::compressDims(AffineMap map,
-                             const llvm::SmallBitVector &unusedDims) {
-  unsigned numDims = 0;
-  SmallVector<AffineExpr> dimReplacements;
-  dimReplacements.reserve(map.getNumDims());
-  MLIRContext *context = map.getContext();
-  for (unsigned dim = 0, e = map.getNumDims(); dim < e; ++dim) {
-    if (unusedDims.test(dim))
-      dimReplacements.push_back(getAffineConstantExpr(0, context));
-    else
-      dimReplacements.push_back(getAffineDimExpr(numDims++, context));
-  }
-  SmallVector<AffineExpr> resultExprs;
-  resultExprs.reserve(map.getNumResults());
-  for (auto e : map.getResults())
-    resultExprs.push_back(e.replaceDims(dimReplacements));
-  return AffineMap::get(numDims, map.getNumSymbols(), resultExprs, context);
-}
-
-AffineMap mlir::compressUnusedDims(AffineMap map) {
-  return compressDims(map, getUnusedDimsBitVector({map}));
-}
-
-static SmallVector<AffineMap>
-compressUnusedImpl(ArrayRef<AffineMap> maps,
-                   llvm::function_ref<AffineMap(AffineMap)> compressionFun) {
+/// Implementation detail to compress multiple affine maps with a compressionFun
+/// that is expected to be either compressUnusedDims or compressUnusedSymbols.
+/// The implementation keeps track of num dims and symbols across the different
+/// affine maps.
+static SmallVector<AffineMap> compressUnusedListImpl(
+    ArrayRef<AffineMap> maps,
   if (maps.empty())
     return SmallVector<AffineMap>();
   SmallVector<AffineExpr> allExprs;
   allExprs.reserve(maps.size() * maps.front().getNumResults());
   unsigned numDims = maps.front().getNumDims(),
            numSymbols = maps.front().getNumSymbols();
   for (auto m : maps) {
     assert(numDims == m.getNumDims() && numSymbols == m.getNumSymbols() &&
for (auto m : maps) {
     res.push_back(AffineMap::get(unifiedNumDims, unifiedNumSymbols,
                                  unifiedResults.take_front(m.getNumResults()),
                                  m.getContext()));
     unifiedResults = unifiedResults.drop_front(m.getNumResults());
   }
   return res;
 }
 
+AffineMap mlir::compressDims(AffineMap map,
+                             const llvm::SmallBitVector &unusedDims) {
+  return projectDims(map, unusedDims, /*compressDimsFlag=*/true);
+}
+
+AffineMap mlir::compressUnusedDims(AffineMap map) {
+  return compressDims(map, getUnusedDimsBitVector({map}));
+}
+
 SmallVector<AffineMap> mlir::compressUnusedDims(ArrayRef<AffineMap> maps) {
-  return compressUnusedImpl(maps,
-                            [](AffineMap m) { return compressUnusedDims(m); });
+  return compressUnusedListImpl(
+      maps, [](AffineMap m) { return compressUnusedDims(m); });
 }
 
 AffineMap mlir::compressSymbols(AffineMap map,
                                 const llvm::SmallBitVector &unusedSymbols) {
-  unsigned numSymbols = 0;
-  SmallVector<AffineExpr> symReplacements;
-  symReplacements.reserve(map.getNumSymbols());
-  MLIRContext *context = map.getContext();
-  for (unsigned sym = 0, e = map.getNumSymbols(); sym < e; ++sym) {
-    if (unusedSymbols.test(sym))
-      symReplacements.push_back(getAffineConstantExpr(0, context));
-    else
-      symReplacements.push_back(getAffineSymbolExpr(numSymbols++, context));
-  }
-  SmallVector<AffineExpr> resultExprs;
-  resultExprs.reserve(map.getNumResults());
-  for (auto e : map.getResults())
-    resultExprs.push_back(e.replaceSymbols(symReplacements));
-  return AffineMap::get(map.getNumDims(), numSymbols, resultExprs, context);
+  return projectSymbols(map, unusedSymbols, /*compressSymbolsFlag=*/true);
 }
 
 AffineMap mlir::compressUnusedSymbols(AffineMap map) {
-  llvm::SmallBitVector unusedSymbols(map.getNumSymbols(), true);
-  map.walkExprs([&](AffineExpr expr) {
-    if (auto symExpr = expr.dyn_cast<AffineSymbolExpr>())
-      unusedSymbols.reset(symExpr.getPosition());
-  });
-  return compressSymbols(map, unusedSymbols);
+  return compressSymbols(map, getUnusedSymbolsBitVector({map}));
 }
 
 SmallVector<AffineMap> mlir::compressUnusedSymbols(ArrayRef<AffineMap> maps) {
-  return compressUnusedImpl(
+  return compressUnusedListImpl(
       maps, [](AffineMap m) { return compressUnusedSymbols(m); });
 }
 
 AffineMap mlir::simplifyAffineMap(AffineMap map) {
   SmallVector<AffineExpr, 8> exprs;
   for (auto e : map.getResults()) {
     exprs.push_back(
         simplifyAffineExpr(e, map.getNumDims(), map.getNumSymbols()));
for (auto m : maps) {
 
     numSymbols += m.getNumSymbols();
     numDims = std::max(m.getNumDims(), numDims);
   }
   return AffineMap::get(numDims, numSymbols, results,
                         maps.front().getContext());
 }
 
+/// Common implementation to project out dimensions or symbols from an affine
+/// map based on the template type.
+/// Additionally, if 'compress' is true, the projected out dimensions or symbols
+/// are also dropped from the resulting map.
+template <typename AffineDimOrSymExpr>
+static AffineMap projectCommonImpl(AffineMap map,
+                                   const llvm::SmallBitVector &toProject,
+                                   bool compress) {
+  static_assert(llvm::is_one_of<AffineDimOrSymExpr, AffineDimExpr,
+                                AffineSymbolExpr>::value,
+                "expected AffineDimExpr or AffineSymbolExpr");
+
+  constexpr bool isDim = std::is_same<AffineDimOrSymExpr, AffineDimExpr>::value;
+  int64_t numDimOrSym = (isDim) ? map.getNumDims() : map.getNumSymbols();
+  SmallVector<AffineExpr> replacements;
+  replacements.reserve(numDimOrSym);
+
+  auto createNewDimOrSym = (isDim) ? getAffineDimExpr : getAffineSymbolExpr;
+  auto replaceDims = [](AffineExpr e, ArrayRef<AffineExpr> replacements) {
+    return e.replaceDims(replacements);
+  };
+  auto replaceSymbols = [](AffineExpr e, ArrayRef<AffineExpr> replacements) {
+    return e.replaceSymbols(replacements);
+  };
+  auto replaceNewDimOrSym = (isDim) ? replaceDims : replaceSymbols;
+
+  MLIRContext *context = map.getContext();
+  int64_t newNumDimOrSym = 0;
+  for (unsigned dimOrSym = 0; dimOrSym < numDimOrSym; ++dimOrSym) {
+    if (toProject.test(dimOrSym)) {
+      replacements.push_back(getAffineConstantExpr(0, context));
+      continue;
+    }
+    int64_t newPos = compress ? newNumDimOrSym++ : dimOrSym;
+    replacements.push_back(createNewDimOrSym(newPos, context));
+  }
+  SmallVector<AffineExpr> resultExprs;
+  resultExprs.reserve(map.getNumResults());
+  for (auto e : map.getResults())
+    resultExprs.push_back(replaceNewDimOrSym(e, replacements));
+
+  int64_t numDims = (compress && isDim) ? newNumDimOrSym : map.getNumDims();
+  int64_t numSyms = (compress && !isDim) ? newNumDimOrSym : map.getNumSymbols();
+  return AffineMap::get(numDims, numSyms, resultExprs, context);
+}
+
+AffineMap mlir::projectDims(AffineMap map,
+                            const llvm::SmallBitVector &projectedDimensions,
+                            bool compressDimsFlag) {
+  return projectCommonImpl<AffineDimExpr>(map, projectedDimensions,
+                                          compressDimsFlag);
+}
+
+AffineMap mlir::projectSymbols(AffineMap map,
+                               const llvm::SmallBitVector &projectedSymbols,
+                               bool compressSymbolsFlag) {
+  return projectCommonImpl<AffineSymbolExpr>(map, projectedSymbols,
+                                             compressSymbolsFlag);
+}
+
 AffineMap mlir::getProjectedMap(AffineMap map,
-                                const llvm::SmallBitVector &unusedDims) {
-  return compressUnusedSymbols(compressDims(map, unusedDims));
+                                const llvm::SmallBitVector &projectedDimensions,
+                                bool compressDimsFlag,
+                                bool compressSymbolsFlag) {
+  map = projectDims(map, projectedDimensions, compressDimsFlag);
+  if (compressSymbolsFlag)
+    map = compressUnusedSymbols(map);
+  return map;
 }
 
 llvm::SmallBitVector mlir::getUnusedDimsBitVector(ArrayRef<AffineMap> maps) {
   unsigned numDims = maps[0].getNumDims();
   llvm::SmallBitVector numDimsBitVector(numDims, true);
-  for (const auto &m : maps) {
+  for (AffineMap m : maps) {
     for (unsigned i = 0; i < numDims; ++i) {
       if (m.isFunctionOfDim(i))
         numDimsBitVector.reset(i);
     }
   }
   return numDimsBitVector;
 }
 
+llvm::SmallBitVector mlir::getUnusedSymbolsBitVector(ArrayRef<AffineMap> maps) {
+  unsigned numSymbols = maps[0].getNumSymbols();
+  llvm::SmallBitVector numSymbolsBitVector(numSymbols, true);
+  for (AffineMap m : maps) {
+    for (unsigned i = 0; i < numSymbols; ++i) {
+      if (m.isFunctionOfSymbol(i))
+        numSymbolsBitVector.reset(i);
+    }
+  }
+  return numSymbolsBitVector;
+}
+
 //===----------------------------------------------------------------------===//
 // MutableAffineMap.
 //===----------------------------------------------------------------------===//
 
 MutableAffineMap::MutableAffineMap(AffineMap map)
     : results(map.getResults().begin(), map.getResults().end()),
       numDims(map.getNumDims()), numSymbols(map.getNumSymbols()),
       context(map.getContext()) {}
bool MutableAffineMap::isMultipleOf(unsigned idx, int64_t factor) const {
   if (results[idx].isMultipleOf(factor))
     return true;
 
   // TODO: use simplifyAffineExpr and FlatAffineConstraints to
   // complete this (for a more powerful analysis).
   return false;
 }
 
-// Simplifies the result affine expressions of this map. The expressions have to
-// be pure for the simplification implemented.
+// Simplifies the result affine expressions of this map. The expressions
+// have to be pure for the simplification implemented.
 void MutableAffineMap::simplify() {
   // Simplify each of the results if possible.
   // TODO: functional-style map
   for (unsigned i = 0, e = getNumResults(); i < e; i++) {
     results[i] = simplifyAffineExpr(getResult(i), numDims, numSymbols);
   }
 }
 
 AffineMap MutableAffineMap::getAffineMap() const {
   return AffineMap::get(numDims, numSymbols, results, context);
 }