Introduce AffineExpr::compose(AffineMap)

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

Arguably, `simplifyAffineExpr` should be part of IR and not Analysis but
this CL does not yet pull the trigger on that.

PiperOrigin-RevId: 228265845

mlir/include/mlir/Analysis/AffineAnalysis.h

@@ -66,14 +66,6 @@ AffineExpr simplifyAffineExpr(AffineExpr expr, unsigned numDims,
 /// Returns (d0)[s0] -> (d0 + s0 + 1, d0 - s0 - 1)
 AffineMap composeUnboundedMaps(AffineMap f, AffineMap g);
 
-/// Given an AffineExpr `e` and an unbounded AffineMap `g`, returns the
-/// AffineExpr e o g (i.e. e(g), i.e. e composed with g).
-///
-/// Prerequisites:
-/// `e` and `g` are composable, i.e. that the number of AffineDimExpr of `e` is
-/// smaller than the number of results of `g`.
-AffineExpr composeWithUnboundedMap(AffineExpr e, AffineMap g);
-
 /// Returns the sequence of AffineApplyOp OperationInsts operation in
 /// 'affineApplyOps', which are reachable via a search starting from 'operands',
 /// and ending at operands which are not defined by AffineApplyOps.

mlir/include/mlir/IR/AffineExpr.h

@@ -31,6 +31,7 @@
 namespace mlir {
 
 class MLIRContext;
+class AffineMap;
 
 namespace detail {
 
@@ -142,6 +143,21 @@ public:
   AffineExpr operator%(uint64_t v) const;
   AffineExpr operator%(AffineExpr other) const;
 
+  /// Compose with an AffineMap.
+  /// Returns the composition of this AffineExpr with `map`.
+  ///
+  /// Prerequisites:
+  /// `this` and `map` are composable, i.e. that the number of AffineDimExpr of
+  /// `this` is smaller than the number of results of `map`. If a result of a
+  /// map does not have a corresponding AffineDimExpr, that result simply does
+  /// not appear in the produced AffineExpr.
+  ///
+  /// Example:
+  ///   expr: `d0 + d2`
+  ///   map:  `(d0, d1, d2)[s0, s1] -> (d0 + s1, d1 + s0, d0 + d1 + d2)`
+  ///   returned expr: `d0 * 2 + d1 + d2 + s1`
+  AffineExpr compose(AffineMap map) const;
+
   friend ::llvm::hash_code hash_value(AffineExpr arg);
 
 protected:

mlir/lib/Analysis/AffineAnalysis.cpp

@@ -392,27 +392,6 @@ AffineExpr mlir::simplifyAffineExpr(AffineExpr expr, unsigned numDims,
   return simplifiedExpr;
 }
 
-/// Returns the AffineExpr that results from substituting `exprs[i]` into `e`
-/// for each AffineDimExpr of position i in `e`.
-/// Precondition: the maximal AffineDimExpr position in `e` is smaller than
-/// `exprs.size()`.
-static AffineExpr substExprs(AffineExpr e, llvm::ArrayRef<AffineExpr> exprs) {
-  if (auto binExpr = e.dyn_cast<AffineBinaryOpExpr>()) {
-    return getAffineBinaryOpExpr(binExpr.getKind(),
-                                 substExprs(binExpr.getLHS(), exprs),
-                                 substExprs(binExpr.getRHS(), exprs));
-  }
-  if (auto dim = e.dyn_cast<AffineDimExpr>()) {
-    assert(dim.getPosition() < exprs.size() &&
-           "Cannot compose due to dim mismatch");
-    return exprs[dim.getPosition()];
-  }
-  if (auto sym = e.dyn_cast<AffineSymbolExpr>()) {
-    return sym;
-  }
-  return e.template cast<AffineConstantExpr>();
-}
-
 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 "
@@ -420,7 +399,10 @@ AffineMap mlir::composeUnboundedMaps(AffineMap f, AffineMap g) {
   assert(g.getRangeSizes().empty() && "Expected unbounded AffineMap");
   assert(f.getRangeSizes().empty() && "Expected unbounded AffineMap");
   auto exprs = functional::map(
-      [g](AffineExpr expr) { return mlir::composeWithUnboundedMap(expr, g); },
+      [g](AffineExpr expr) {
+        return simplifyAffineExpr(expr.compose(g), g.getNumDims(),
+                                  g.getNumSymbols());
+      },
       f.getResults());
   auto composed =
       AffineMap::get(g.getNumDims(),
@@ -428,11 +410,6 @@ AffineMap mlir::composeUnboundedMaps(AffineMap f, AffineMap g) {
   return composed;
 }
 
-AffineExpr mlir::composeWithUnboundedMap(AffineExpr e, AffineMap g) {
-  return simplifyAffineExpr(substExprs(e, g.getResults()), g.getNumDims(),
-                            g.getNumSymbols());
-}
-
 // 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/Analysis/VectorAnalysis.cpp

@@ -429,7 +429,8 @@ SingleResultAffineNormalizer::SingleResultAffineNormalizer(
   auto exprsMap = AffineMap::get(dimValueToPosition.size(),
                                  symValueToPosition.size(), exprs, {});
 
-  expr = composeWithUnboundedMap(map.getResult(0), exprsMap);
+  expr = simplifyAffineExpr(map.getResult(0).compose(exprsMap),
+                            exprsMap.getNumDims(), exprsMap.getNumSymbols());
 
   LLVM_DEBUG(map.getResult(0).print(dbgs() << "\nCompose expr: "));
   LLVM_DEBUG(exprsMap.print(dbgs() << "\nWith map: "));

mlir/lib/IR/AffineExpr.cpp

@@ -18,6 +18,7 @@
 #include "mlir/IR/AffineExpr.h"
 #include "AffineExprDetail.h"
 #include "mlir/IR/AffineExprVisitor.h"
+#include "mlir/IR/AffineMap.h"
 #include "mlir/Support/STLExtras.h"
 #include "llvm/ADT/STLExtras.h"
 
@@ -283,7 +284,11 @@ AffineExpr AffineExpr::operator%(uint64_t v) const {
 AffineExpr AffineExpr::operator%(AffineExpr other) const {
   return AffineBinaryOpExprStorage::get(AffineExprKind::Mod, expr, other.expr);
 }
-
+AffineExpr AffineExpr::compose(AffineMap map) const {
+  SmallVector<AffineExpr, 8> dimReplacements(map.getResults().begin(),
+                                             map.getResults().end());
+  return replaceDimsAndSymbols(dimReplacements, {});
+}
 raw_ostream &operator<<(raw_ostream &os, AffineExpr &expr) {
   expr.print(os);
   return os;

mlir/test/Transforms/Vectorize/compose_maps.mlir

@@ -114,4 +114,11 @@ func @symbols1() {
   "test_affine_map"() { affine_map: (d0)[s0] -> (d0 + s0, d0 - s0) } : () -> ()
   "test_affine_map"() { affine_map: (d0, d1) -> (d0 + 1, d1 - 1) } : () -> ()
   return
+}
+
+func @drop() {
+  // CHECK: Composed map: (d0, d1, d2)[s0, s1] -> (d0 * 2 + d1 + d2 + s1)
+  "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
 }