Introduce AffineMap::compose(AffineMap)

This CL is the 2nd on the path to simplifying AffineMap composition.
This CL uses the now accepted `AffineExpr::compose(AffineMap)` to
implement `AffineMap::compose(AffineMap)`.

Implications of keeping the simplification function in
Analysis are documented where relevant.

PiperOrigin-RevId: 228276646

mlir/include/mlir/Analysis/AffineAnalysis.h

@@ -41,30 +41,16 @@ class OperationInst;
 class Instruction;
 class Value;
 
-/// Simplify an affine expression through flattening and some amount of
+/// Simplify an affine expression by flattening and some amount of
 /// simple analysis. This has complexity linear in the number of nodes in
 /// 'expr'. Returns the simplified expression, which is the same as the input
 ///  expression if it can't be simplified.
 AffineExpr simplifyAffineExpr(AffineExpr expr, unsigned numDims,
                               unsigned numSymbols);
 
-/// Given 2 unbounded AffineMaps `f` and `g`, returns the AffineMap f o g
-/// (i.e. f(g), i.e. f composed with g).
-/// The resulting AffineMap has as many AffineDimExpr as `g` and as many
-/// AffineSymbolExpr as the max of either `f` or `g`.
-/// In particular, this does not try and compress the unused AffineDimExpr and
-/// AffineSymbolExpr. This could be implemented later as a canonicalization if
-/// needed.
-///
-/// Prerequisites:
-/// The maps are composable, i.e. that the number of AffineDimExpr of `f`
-/// matches the number of results of `g`.
-///
-/// Examples:
-/// compose((d0, d1) -> (d0 + 1, d1 - 1),
-///         (d0)[s0] -> (d0 + s0, d0 - s0))
-/// Returns (d0)[s0] -> (d0 + s0 + 1, d0 - s0 - 1)
-AffineMap composeUnboundedMaps(AffineMap f, AffineMap g);
+/// Simplify an affine map through simplifying its underlying AffineExpr results
+/// and sizes.
+AffineMap simplifyAffineMap(AffineMap map);
 
 /// Returns the sequence of AffineApplyOp OperationInsts operation in
 /// 'affineApplyOps', which are reachable via a search starting from 'operands',

mlir/include/mlir/IR/AffineMap.h

@@ -117,6 +117,25 @@ public:
   bool constantFold(ArrayRef<Attribute> operandConstants,
                     SmallVectorImpl<Attribute> &results) const;
 
+  /// Returns the AffineMap resulting from composing `this` with `map`.
+  /// The resulting AffineMap has as many AffineDimExpr as `map` and as many
+  /// AffineSymbolExpr as the concatenation of `this` and `map` (in which case
+  /// the symbols of `this` map come first).
+  ///
+  /// Prerequisites:
+  /// The maps are composable, i.e. that the number of AffineDimExpr of `this`
+  /// matches the number of results of `map`.
+  /// At this time, composition of bounded AffineMap is not supported. Both
+  /// `this` and `map` must be unbounded.
+  ///
+  /// Example:
+  ///   map1: `(d0, d1)[s0, s1] -> (d0 + 1 + s1, d1 - 1 - s0)`
+  ///   map2: `(d0)[s0] -> (d0 + s0, d0 - s0))`
+  ///   map1.compose(map2):
+  ///     `(d0)[s0, s1, s2] -> (d0 + s1 + s2 + 1, d0 - s0 - s2 - 1)`
+  // TODO(ntv): support composition of bounded maps when we have a need for it.
+  AffineMap compose(AffineMap map);
+
   friend ::llvm::hash_code hash_value(AffineMap arg);
 
 private:

mlir/lib/Analysis/AffineAnalysis.cpp

@@ -77,6 +77,20 @@ static AffineExpr toAffineExpr(ArrayRef<int64_t> eq, unsigned numDims,
   return expr;
 }
 
+AffineMap mlir::simplifyAffineMap(AffineMap map) {
+  auto exprs = functional::map(
+      [map](AffineExpr e) {
+        return simplifyAffineExpr(e, map.getNumDims(), map.getNumSymbols());
+      },
+      map.getResults());
+  auto sizes = functional::map(
+      [map](AffineExpr e) {
+        return simplifyAffineExpr(e, map.getNumDims(), map.getNumSymbols());
+      },
+      map.getRangeSizes());
+  return AffineMap::get(map.getNumDims(), map.getNumSymbols(), exprs, sizes);
+}
+
 namespace {
 
 // This class is used to flatten a pure affine expression (AffineExpr,
@@ -392,24 +406,6 @@ AffineExpr mlir::simplifyAffineExpr(AffineExpr expr, unsigned numDims,
   return simplifiedExpr;
 }
 
-AffineMap mlir::composeUnboundedMaps(AffineMap f, AffineMap g) {
-  assert(f.getNumDims() == g.getNumResults() &&
-         "Num dims of f must be the same as num results of g for maps to be "
-         "composable");
-  assert(g.getRangeSizes().empty() && "Expected unbounded AffineMap");
-  assert(f.getRangeSizes().empty() && "Expected unbounded AffineMap");
-  auto exprs = functional::map(
-      [g](AffineExpr expr) {
-        return simplifyAffineExpr(expr.compose(g), g.getNumDims(),
-                                  g.getNumSymbols());
-      },
-      f.getResults());
-  auto composed =
-      AffineMap::get(g.getNumDims(),
-                     std::max(f.getNumSymbols(), g.getNumSymbols()), exprs, {});
-  return composed;
-}
-
 // Flattens the expressions in map. Returns true on success or false
 // if 'expr' was unable to be flattened (i.e., semi-affine expressions not
 // handled yet).

mlir/lib/IR/AffineMap.cpp

@@ -20,6 +20,7 @@
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Types.h"
+#include "mlir/Support/Functional.h"
 #include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/StringRef.h"
 
@@ -201,3 +202,29 @@ AffineMap AffineMap::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
 
   return get(numResultDims, numResultSyms, results, resultRanges);
 }
+
+AffineMap AffineMap::compose(AffineMap map) {
+  assert(getNumDims() == map.getNumResults() && "Number of results mismatch");
+  assert(getRangeSizes().empty() && "TODO: support bounded AffineMap");
+  assert(map.getRangeSizes().empty() && "TODO: support bounded AffineMap");
+  // 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) {
+    newDims[idx] = getAffineDimExpr(idx, getContext());
+  }
+  SmallVector<AffineExpr, 8> newSymbols(numSymbols);
+  for (unsigned idx = numSymbolsThisMap; idx < numSymbols; ++idx) {
+    newSymbols[idx - numSymbolsThisMap] =
+        getAffineSymbolExpr(idx, getContext());
+  }
+  auto newMap =
+      map.replaceDimsAndSymbols(newDims, newSymbols, numDims, numSymbols);
+  SmallVector<AffineExpr, 8> exprs;
+  exprs.reserve(getResults().size());
+  for (auto expr : getResults())
+    exprs.push_back(expr.compose(newMap));
+  return AffineMap::get(numDims, numSymbols, exprs, {});
+}

mlir/lib/Transforms/MaterializeVectors.cpp

@@ -470,7 +470,7 @@ static AffineMap projectedPermutationMap(VectorTransferOpTy *transfer,
   }
   auto projectionMap = AffineMap::get(optionalRatio->size(), 0, keep, {});
   LLVM_DEBUG(projectionMap.print(dbgs() << "\nprojectionMap: "));
-  return composeUnboundedMaps(projectionMap, permutationMap);
+  return simplifyAffineMap(projectionMap.compose(permutationMap));
 }
 
 /// Creates an instantiated version of `read` for the instance of

mlir/lib/Transforms/Vectorization/VectorizerTestPass.cpp

@@ -222,9 +222,9 @@ void VectorizerTestPass::testComposeMaps(Function *f) {
   }
   AffineMap res;
   for (auto m : maps) {
-    res = res ? composeUnboundedMaps(res, m) : m;
+    res = res ? res.compose(m) : m;
   }
-  res.print(outs() << "\nComposed map: ");
+  simplifyAffineMap(res).print(outs() << "\nComposed map: ");
 }
 
 bool affineApplyOp(const Instruction &inst) {

mlir/test/Transforms/Vectorize/compose_maps.mlir

@@ -17,14 +17,14 @@ func @simple1() {
 }
 
 func @simple2() {
-  // CHECK: Composed map: (d0)[s0] -> (d0)
+  // CHECK: Composed map: (d0)[s0, s1] -> (d0 - s0 + s1)
   "test_affine_map"() { affine_map: (d0)[s0] -> (d0 + s0 - 1) } : () -> ()
   "test_affine_map"() { affine_map: (d0)[s0] -> (d0 - s0 + 1) } : () -> ()
   return
 }
 
 func @simple3a() {
-  // CHECK: Composed map: (d0, d1)[s0, s1] -> ((d0 ceildiv s0) * s0, (d1 ceildiv s1) * s1)
+  // CHECK: Composed map: (d0, d1)[s0, s1, s2, s3] -> ((d0 ceildiv s2) * s0, (d1 ceildiv s3) * s1)
   "test_affine_map"() { affine_map: (d0, d1)[s0, s1] -> (d0 ceildiv s0, d1 ceildiv s1) } : () -> ()
   "test_affine_map"() { affine_map: (d0, d1)[s0, s1] -> (d0 * s0, d1 * s1) } : () -> ()
   return
@@ -37,7 +37,7 @@ func @simple3b() {
 }
 
 func @simple3c() {
-  // CHECK: Composed map: (d0, d1)[s0, s1, s2, s3] -> ((d0 ceildiv s0) * s0 + d0 mod s2, (d1 ceildiv s1) * s1 + d1 mod s3)
+  // CHECK: Composed map: (d0, d1)[s0, s1, s2, s3, s4, s5] -> ((d0 ceildiv s4) * s4 + d0 mod s2, (d1 ceildiv s5) * s5 + d1 mod s3)
   "test_affine_map"() { affine_map: (d0, d1)[s0, s1] -> ((d0 ceildiv s0) * s0, (d1 ceildiv s1) * s1, d0, d1) } : () -> ()
   "test_affine_map"() { affine_map: (d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0 + d2 mod s2, d1 + d3 mod s3) } : () -> ()
   return
@@ -121,4 +121,11 @@ func @drop() {
   "test_affine_map"() { affine_map: (d0, d1, d2)[s0, s1] -> (d0 + s1, d1 + s0, d0 + d1 + d2) } : () -> ()
   "test_affine_map"() { affine_map: (d0, d1, d2) -> (d0 + d2) } : () -> ()
   return
+}
+
+func @multi_symbols() {
+  // CHECK: Composed map: (d0)[s0, s1, s2] -> (d0 + s1 + s2 + 1, d0 - s0 - s2 - 1)
+  "test_affine_map"() { affine_map: (d0)[s0] -> (d0 + s0, d0 - s0) } : () -> ()
+  "test_affine_map"() { affine_map: (d0, d1)[s0, s1] -> (d0 + 1 + s1, d1 - 1 - s0) } : () -> ()
+  return
 }