[Mlir-commits] [mlir] 7b47de7 - [mlir] NFC - Add AffineMap::replace variant with dim/symbol inference

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2021-07-14T20:29:12Z
New Revision: 7b47de774fd43eb5be7aee86c6f8c15999295b36

URL: https://github.com/llvm/llvm-project/commit/7b47de774fd43eb5be7aee86c6f8c15999295b36
DIFF: https://github.com/llvm/llvm-project/commit/7b47de774fd43eb5be7aee86c6f8c15999295b36.diff

LOG: [mlir] NFC - Add AffineMap::replace variant with dim/symbol inference

Added: 
    

Modified: 
    mlir/include/mlir/IR/AffineMap.h
    mlir/lib/IR/AffineMap.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h
index 336e6188c69ee..b308bea6ec67b 100644
--- a/mlir/include/mlir/IR/AffineMap.h
+++ b/mlir/include/mlir/IR/AffineMap.h
@@ -196,6 +196,11 @@ class AffineMap {
   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 inferred
+  /// number of dims and symbols.
+  AffineMap replace(const DenseMap<AffineExpr, AffineExpr> &map) 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.

diff  --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp
index 2713426492dcf..2257617bd903e 100644
--- a/mlir/lib/IR/AffineMap.cpp
+++ b/mlir/lib/IR/AffineMap.cpp
@@ -439,6 +439,15 @@ AffineMap AffineMap::replace(const DenseMap<AffineExpr, AffineExpr> &map,
   return AffineMap::get(numResultDims, numResultSyms, newResults, getContext());
 }
 
+AffineMap
+AffineMap::replace(const DenseMap<AffineExpr, AffineExpr> &map) const {
+  SmallVector<AffineExpr, 4> newResults;
+  newResults.reserve(getNumResults());
+  for (AffineExpr e : getResults())
+    newResults.push_back(e.replace(map));
+  return AffineMap::inferFromExprList(newResults).front();
+}
+
 AffineMap AffineMap::compose(AffineMap map) const {
   assert(getNumDims() == map.getNumResults() && "Number of results mismatch");
   // Prepare `map` by concatenating the symbols and rewriting its exprs.