diff options
| -rw-r--r-- | mlir/include/mlir/Analysis/AffineAnalysis.h | 17 | ||||
| -rw-r--r-- | mlir/lib/Analysis/AffineAnalysis.cpp | 17 |
2 files changed, 22 insertions, 12 deletions
diff --git a/mlir/include/mlir/Analysis/AffineAnalysis.h b/mlir/include/mlir/Analysis/AffineAnalysis.h index 8f7fbd6ec25..84c17a76600 100644 --- a/mlir/include/mlir/Analysis/AffineAnalysis.h +++ b/mlir/include/mlir/Analysis/AffineAnalysis.h @@ -48,22 +48,31 @@ 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). This assumes (and enforces) that the -/// maps are composable, i.e. that the number of AffineDimExpr of `f` matches -/// the number of results of `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); +/// 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 OperationStmts operation in /// 'affineApplyOps', which are reachable via a search starting from 'operands', /// and ending at operands which are not defined by AffineApplyOps. diff --git a/mlir/lib/Analysis/AffineAnalysis.cpp b/mlir/lib/Analysis/AffineAnalysis.cpp index aa88eb2348d..7bd800cecd7 100644 --- a/mlir/lib/Analysis/AffineAnalysis.cpp +++ b/mlir/lib/Analysis/AffineAnalysis.cpp @@ -336,18 +336,18 @@ AffineExpr mlir::simplifyAffineExpr(AffineExpr expr, unsigned numDims, /// 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.template dyn_cast<AffineBinaryOpExpr>()) { + if (auto binExpr = e.dyn_cast<AffineBinaryOpExpr>()) { AffineExpr lhs, rhs; AffineExprBinaryOp binOp; std::tie(lhs, rhs, binOp) = matchBinaryOpExpr(binExpr); return binOp(substExprs(lhs, exprs), substExprs(rhs, exprs)); } - if (auto dim = e.template dyn_cast<AffineDimExpr>()) { + 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.template dyn_cast<AffineSymbolExpr>()) { + if (auto sym = e.dyn_cast<AffineSymbolExpr>()) { return sym; } return e.template cast<AffineConstantExpr>(); @@ -360,11 +360,7 @@ 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) { - auto e = simplifyAffineExpr(substExprs(expr, g.getResults()), - g.getNumDims(), g.getNumSymbols()); - return e; - }, + [g](AffineExpr expr) { return mlir::composeWithUnboundedMap(expr, g); }, f.getResults()); auto composed = AffineMap::get(g.getNumDims(), @@ -372,6 +368,11 @@ 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 'expr' into 'flattenedExpr'. Returns true on success or false // if 'expr' was unable to be flattened (i.e., semi-affine expression has not // been implemented yet). |
