diff --git a/mlir/include/mlir/Dialect/Affine/Utils.h b/mlir/include/mlir/Dialect/Affine/Utils.h --- a/mlir/include/mlir/Dialect/Affine/Utils.h +++ b/mlir/include/mlir/Dialect/Affine/Utils.h @@ -20,14 +20,22 @@ class AffineForOp; class AffineIfOp; class AffineParallelOp; +class AffineMinOp; struct LogicalResult; class Operation; +class Value; /// Replaces parallel affine.for op with 1-d affine.parallel op. /// mlir::isLoopParallel detect the parallel affine.for ops. /// There is no cost model currently used to drive this parallelization. void affineParallelize(AffineForOp forOp); +/// Computes a constant upper bound for a given value. Returns the constant as +/// an int64_t value if found, or None if the min computation failed. This is +/// mostly targeted for values that are result of affine.min operations (which +/// is why it lives here). +Optional getUpperBound(Value value); + /// Hoists out affine.if/else to as high as possible, i.e., past all invariant /// affine.fors/parallel's. Returns success if any hoisting happened; folded` is /// set to true if the op was folded or erased. This hoisting could lead to diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp --- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp @@ -144,6 +144,70 @@ forOp.erase(); } +/// Forward declaration of method to compute upper-bound of a given value. +static Optional getUpperBoundImpl(Value value); + +/// Returns the bound of an affine expression within the map of `minOp`. It +/// recursively walks the expression tree to compute the bound if possible. +static Optional getUpperBoundImpl(AffineMinOp minOp, AffineExpr expr) { + if (auto cst = expr.dyn_cast()) + return cst.getValue(); + if (auto dim = expr.dyn_cast()) + return getUpperBoundImpl(minOp.getOperand(dim.getPosition())); + if (auto symbol = expr.dyn_cast()) + return getUpperBoundImpl(minOp.getOperand( + minOp.getAffineMap().getNumDims() + symbol.getPosition())); + if (expr.getKind() == AffineExprKind::Add) { + AffineBinaryOpExpr binaryExpr = expr.cast(); + AffineExpr lhs = binaryExpr.getLHS(); + AffineExpr rhs = binaryExpr.getRHS(); + AffineConstantExpr constExpr = lhs.dyn_cast(); + AffineExpr nonConstExpr = rhs; + if (!constExpr) { + constExpr = rhs.dyn_cast(); + if (!constExpr) + return {}; + nonConstExpr = lhs; + } + if (constExpr.getValue() < 0) + return {}; + if (Optional recValue = getUpperBoundImpl(minOp, nonConstExpr)) + return *recValue + constExpr.getValue(); + return {}; + } + return {}; +} + +static Optional getUpperBoundImpl(Value value) { + Operation *definingOp = value.getDefiningOp(); + if (!definingOp) + return {}; + + if (auto affineMinOp = dyn_cast(definingOp)) { + Optional boundingConstant = {}; + AffineMap map = affineMinOp.getAffineMap(); + for (auto e : map.getResults()) { + if (Optional exprBound = getUpperBoundImpl(affineMinOp, e)) { + boundingConstant = boundingConstant + ? std::min(*boundingConstant, *exprBound) + : *exprBound; + } + } + return boundingConstant; + } + if (auto constIndexOp = dyn_cast(definingOp)) { + if (constIndexOp.getType().isa()) + return constIndexOp.value().cast().getInt(); + } + return {}; +} + +/// Computes a constant upper bound for a affine.min op. Returns the constant as +/// an int64_t value if found, or None if the min computation failed. +Optional mlir::getUpperBound(Value value) { + return getUpperBoundImpl(value); +} + // Returns success if any hoisting happened. LogicalResult mlir::hoistAffineIfOp(AffineIfOp ifOp, bool *folded) { // Bail out early if the ifOp returns a result. TODO: Consider how to diff --git a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt --- a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt +++ b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt @@ -18,6 +18,7 @@ LINK_LIBS PUBLIC MLIRAffineOps + MLIRAffineUtils MLIRAnalysis MLIREDSC MLIRIR diff --git a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp @@ -12,6 +12,7 @@ #include "PassDetail.h" #include "mlir/Dialect/Affine/EDSC/Intrinsics.h" +#include "mlir/Dialect/Affine/Utils.h" #include "mlir/Dialect/Linalg/EDSC/FoldedIntrinsics.h" #include "mlir/Dialect/Linalg/IR/LinalgOps.h" #include "mlir/Dialect/Linalg/IR/LinalgTypes.h" @@ -49,17 +50,7 @@ /// Otherwise return size. static Value extractSmallestConstantBoundingSize(OpBuilder &b, Location loc, Value size) { - Optional boundingConst = {}; - if (auto affineMinOp = size.getDefiningOp()) { - for (auto e : affineMinOp.getAffineMap().getResults()) - if (auto cst = e.dyn_cast()) - boundingConst = boundingConst - ? std::min(boundingConst.getValue(), cst.getValue()) - : cst.getValue(); - } else if (auto constIndexOp = size.getDefiningOp()) { - if (constIndexOp.getType().isa()) - boundingConst = constIndexOp.value().cast().getInt(); - } + Optional boundingConst = getUpperBound(size); return boundingConst && *boundingConst >= 0 ? b.create(loc, *boundingConst) : size; diff --git a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp @@ -240,10 +240,15 @@ // Construct (potentially temporary) mins and maxes on which to apply maps // that define tile subviews. SmallVector lbs, subViewSizes; + AffineMap subviewSizeMap = + AffineMap::get(0, 2, {b.getAffineSymbolExpr(0), b.getAffineSymbolExpr(1)}, + b.getContext()); for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) { bool isTiled = !isZero(tileSizes[idx]); lbs.push_back(isTiled ? ivs[idxIvs++] : (Value)std_constant_index(0)); - subViewSizes.push_back(isTiled ? tileSizes[idx] : viewSizes[idx]); + Value subviewSize = + affine_min(subviewSizeMap, ValueRange{tileSizes[idx], viewSizes[idx]}); + subViewSizes.push_back(isTiled ? subviewSize : viewSizes[idx]); } auto *op = linalgOp.getOperation(); diff --git a/mlir/test/Dialect/Linalg/tile.mlir b/mlir/test/Dialect/Linalg/tile.mlir --- a/mlir/test/Dialect/Linalg/tile.mlir +++ b/mlir/test/Dialect/Linalg/tile.mlir @@ -8,17 +8,22 @@ // TILE-002-DAG: #[[$strided1D:.*]] = affine_map<(d0)[s0] -> (d0 + s0)> // TILE-234-DAG: #[[$strided1D:.*]] = affine_map<(d0)[s0] -> (d0 + s0)> +// TILE-2-DAG: #[[$minTileSize:.*]] = affine_map<()[s0] -> (2, s0)> +// TILE-02-DAG: #[[$minTileSize:.*]] = affine_map<()[s0] -> (2, s0)> +// TILE-002-DAG: #[[$minTileSize:.*]] = affine_map<()[s0] -> (2, s0)> +// TILE-234-DAG: #[[$minTileSize0:.*]] = affine_map<()[s0] -> (2, s0)> +// TILE-234-DAG: #[[$minTileSize1:.*]] = affine_map<()[s0] -> (3, s0)> +// TILE-234-DAG: #[[$minTileSize2:.*]] = affine_map<()[s0] -> (4, s0)> + // TILE-2-DAG: #[[$strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)> // TILE-02-DAG: #[[$strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)> // TILE-002-DAG: #[[$strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)> // TILE-234-DAG: #[[$strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)> -// TILE-2-DAG: #[[$bound_map:.*]] = affine_map<(d0)[s0] -> (2, -d0 + s0)> -// TILE-02-DAG: #[[$bound_map:.*]] = affine_map<(d0)[s0] -> (2, -d0 + s0)> -// TILE-002-DAG: #[[$bound_map:.*]] = affine_map<(d0)[s0] -> (2, -d0 + s0)> -// TILE-234-DAG: #[[$bound_map_2:.*]] = affine_map<(d0)[s0] -> (2, -d0 + s0)> -// TILE-234-DAG: #[[$bound_map_3:.*]] = affine_map<(d0)[s0] -> (3, -d0 + s0)> -// TILE-234-DAG: #[[$bound_map_4:.*]] = affine_map<(d0)[s0] -> (4, -d0 + s0)> +// TILE-2-DAG: #[[$bound_map:.*]] = affine_map<(d0, d1)[s0] -> (d0, -d1 + s0)> +// TILE-02-DAG: #[[$bound_map:.*]] = affine_map<(d0, d1)[s0] -> (d0, -d1 + s0)> +// TILE-002-DAG: #[[$bound_map:.*]] = affine_map<(d0, d1)[s0] -> (d0, -d1 + s0)> +// TILE-234-DAG: #[[$bound_map:.*]] = affine_map<(d0, d1)[s0] -> (d0, -d1 + s0)> // TILE-2-DAG: #[[$bound_map_static:.*]] = affine_map<(d0) -> (2, -d0 + 10)> // TILE-02-DAG: #[[$bound_map_static:.*]] = affine_map<(d0) -> (2, -d0 + 12)> @@ -42,12 +47,13 @@ // TILE-2-DAG: %[[C2:.*]] = constant 2 : index // TILE-2: %[[M:.*]] = dim %{{.*}}, %c0 : memref // TILE-2: scf.for %[[I:.*]] = %{{.*}}{{.*}} to %[[M]] step %{{.*}} { +// TILE-2: %[[dim0:.*]] = affine.min #[[$minTileSize]]()[%[[M]]] // TILE-2: %[[localM:.*]] = dim %{{.*}}, %c0 -// TILE-2: %[[szM:.*]] = affine.min #[[$bound_map]](%[[I]])[%[[localM]]] +// TILE-2: %[[szM:.*]] = affine.min #[[$bound_map]](%[[dim0]], %[[I]])[%[[localM]]] // TILE-2: %[[K:.*]] = dim %{{.*}}, %c1 : memref // TILE-2: %[[sAi:.*]] = subview %{{.*}}[%[[I]], 0] [%[[szM]], %[[K]]] [1, 1] : memref to memref // TILE-2: %[[localK:.*]] = dim %{{.*}}, %c0 -// TILE-2: %[[szK:.*]] = affine.min #[[$bound_map]](%[[I]])[%[[localK]]] +// TILE-2: %[[szK:.*]] = affine.min #[[$bound_map]](%[[dim0]], %[[I]])[%[[localK]]] // TILE-2: %[[N:.*]] = dim %{{.*}}, %c1 : memref // TILE-2: %[[sCi:.*]] = subview %{{.*}}[%[[I]], 0] [%[[szK]], %[[N]]] [1, 1] : memref to memref // TILE-2: linalg.matmul %[[sAi]], %{{.*}}, %[[sCi]] : @@ -60,13 +66,14 @@ // TILE-02-DAG: %[[C2:.*]] = constant 2 : index // TILE-02: %[[N:.*]] = dim %arg1, %c1 : memref // TILE-02: scf.for %[[J:.*]] = %{{.*}} to %[[N]] step %{{.*}} { +// TILE-02: %[[dim0:.]] = affine.min #[[$minTileSize]]()[%[[N]]] // TILE-02: %[[K:.*]] = dim %{{.*}}, %c0 : memref // TILE-02: %[[localN:.*]] = dim %{{.*}}, %c1 -// TILE-02: %[[szN:.*]] = affine.min #[[$bound_map]](%[[J]])[%[[localN]]] +// TILE-02: %[[szN:.*]] = affine.min #[[$bound_map]](%[[dim0]], %[[J]])[%[[localN]]] // TILE-02: %[[sBj:.*]] = subview %{{.*}}[0, %[[J]]] [%[[K]], %[[szN]]] [1, 1] : memref to memref // TILE-02: %[[M:.*]] = dim %{{.*}}, %c0 : memref // TILE-02: %[[localK:.*]] = dim %{{.*}}, %c1 -// TILE-02: %[[szK:.*]] = affine.min #[[$bound_map]](%[[J]])[%[[localK]]] +// TILE-02: %[[szK:.*]] = affine.min #[[$bound_map]](%[[dim0]], %[[J]])[%[[localK]]] // TILE-02: %[[sCj:.*]] = subview %{{.*}}[0, %[[J]]] [%[[M]], %[[szK]]] [1, 1] : memref to memref // TILE-02: linalg.matmul %{{.*}}, %[[sBj]], %[[sCj]] : // TILE-02: (memref, @@ -78,12 +85,13 @@ // TILE-002-DAG: %[[C2:.*]] = constant 2 : index // TILE-002: %[[ubK:.*]] = dim %{{.*}}, %c1 : memref // TILE-002: scf.for %[[K:.*]] = %{{.*}}{{.*}} to %[[ubK]] step %{{.*}} { +// TILE-002: %[[dim0:.]] = affine.min #[[$minTileSize]]()[%[[ubK]] // TILE-002: %[[M:.*]] = dim %{{.*}}, %c0 : memref // TILE-002: %[[localK:.*]] = dim %{{.*}}, %c1 -// TILE-002: %[[szK:.*]] = affine.min #[[$bound_map]](%[[K]])[%[[localK]]] +// TILE-002: %[[szK:.*]] = affine.min #[[$bound_map]](%[[dim0]], %[[K]])[%[[localK]]] // TILE-002: %[[sAj:.*]] = subview %{{.*}}[0, %[[K]]] [%[[M]], %[[szK]]] [1, 1] : memref to memref // TILE-002: %[[localK:.*]] = dim %{{.*}}, %c0 -// TILE-002: %[[szK:.*]] = affine.min #[[$bound_map]](%[[K]])[%[[localK]]] +// TILE-002: %[[szK:.*]] = affine.min #[[$bound_map]](%[[dim0]], %[[K]])[%[[localK]]] // TILE-002: %[[N:.*]] = dim %{{.*}}, %c1 : memref // TILE-002: %[[sBj:.*]] = subview %{{.*}}[%[[K]], 0] [%[[szK]], %[[N]]] [1, 1] : memref to memref // TILE-002: linalg.matmul %[[sAj]], %[[sBj]], %{{.*}} : @@ -102,20 +110,23 @@ // TILE-234: scf.for %[[I:.*]] = %{{.*}}{{.*}} to %[[ubM]] step %{{.*}} { // TILE-234: scf.for %[[J:.*]] = %{{.*}}{{.*}} to %[[ubN]] step %{{.*}} { // TILE-234: scf.for %[[K:.*]] = %{{.*}}{{.*}} to %[[ubK]] step %{{.*}} { +// TILE-234: %[[dimM:.]] = affine.min #[[$minTileSize0]]()[%[[ubM]]] +// TILE-234: %[[dimN:.]] = affine.min #[[$minTileSize1]]()[%[[ubN]]] +// TILE-234: %[[dimK:.]] = affine.min #[[$minTileSize2]]()[%[[ubK]]] // TILE-234: %[[localM:.*]] = dim %{{.*}}, %c0 -// TILE-234: %[[szM:.*]] = affine.min #[[$bound_map_2]](%[[I]])[%[[localM]]] +// TILE-234: %[[szM:.*]] = affine.min #[[$bound_map]](%[[dimM]], %[[I]])[%[[localM]]] // TILE-234: %[[localK:.*]] = dim %{{.*}}, %c1 -// TILE-234: %[[szK:.*]] = affine.min #[[$bound_map_4]](%[[K]])[%[[localK]]] +// TILE-234: %[[szK:.*]] = affine.min #[[$bound_map]](%[[dimK]], %[[K]])[%[[localK]]] // TILE-234: %[[sAik:.*]] = subview %{{.*}}[%[[I]], %[[K]]] [%[[szM]], %[[szK]]] [1, 1] : memref to memref // TILE-234: %[[localK:.*]] = dim %{{.*}}, %c0 -// TILE-234: %[[szK:.*]] = affine.min #[[$bound_map_4]](%[[K]])[%[[localK]]] +// TILE-234: %[[szK:.*]] = affine.min #[[$bound_map]](%[[dimK]], %[[K]])[%[[localK]]] // TILE-234: %[[localN:.*]] = dim %{{.*}}, %c1 -// TILE-234: %[[szN:.*]] = affine.min #[[$bound_map_3]](%[[J]])[%[[localN]]] +// TILE-234: %[[szN:.*]] = affine.min #[[$bound_map]](%[[dimN]], %[[J]])[%[[localN]]] // TILE-234: %[[sBkj:.*]] = subview %{{.*}}[%[[K]], %[[J]]] [%[[szK]], %[[szN]]] [1, 1] : memref to memref // TILE-234: %[[localM:.*]] = dim %{{.*}}, %c0 -// TILE-234: %[[szM:.*]] = affine.min #[[$bound_map_2]](%[[I]])[%[[localM]]] +// TILE-234: %[[szM:.*]] = affine.min #[[$bound_map]](%[[dimM]], %[[I]])[%[[localM]]] // TILE-234: %[[localN:.*]] = dim %{{.*}}, %c1 -// TILE-234: %[[szN:.*]] = affine.min #[[$bound_map_3]](%[[J]])[%[[localN]]] +// TILE-234: %[[szN:.*]] = affine.min #[[$bound_map]](%[[dimN]], %[[J]])[%[[localN]]] // TILE-234: %[[sCij:.*]] = subview %{{.*}}[%[[I]], %[[J]]] [%[[szM]], %[[szN]]] [1, 1] : memref to memref // // TILE-234: linalg.matmul %[[sAik]], %[[sBkj]], %[[sCij]] : @@ -213,12 +224,13 @@ // TILE-2-DAG: %[[C2:.*]] = constant 2 : index // TILE-2: %[[M:.*]] = dim %{{.*}}, %c0 : memref // TILE-2: scf.for %[[I:.*]] = %{{.*}}{{.*}} to %[[M]] step %{{.*}} { +// TILE-2: %[[dim0:.]] = affine.min #[[$minTileSize]]()[%[[M]]] // TILE-2: %[[localM:.*]] = dim %[[ARG0]], %c0 -// TILE-2: %[[szM:.*]] = affine.min #[[$bound_map]](%[[I]])[%[[localM]]] +// TILE-2: %[[szM:.*]] = affine.min #[[$bound_map]](%[[dim0]], %[[I]])[%[[localM]]] // TILE-2: %[[N:.*]] = dim %{{.*}}, %c1 : memref // TILE-2: %[[sAi:.*]] = subview %{{.*}}[%[[I]], 0] [%[[szM]], %[[N]]] [1, 1] : memref to memref // TILE-2: %[[localN:.*]] = dim %{{.*}}, %c0 -// TILE-2: %[[szN:.*]] = affine.min #[[$bound_map]](%[[I]])[%[[localN]]] +// TILE-2: %[[szN:.*]] = affine.min #[[$bound_map]](%[[dim0]], %[[I]])[%[[localN]]] // TILE-2: %[[sCi:.*]] = subview %{{.*}}[%[[I]]] [%[[szN]]] [1] : memref to memref // TILE-2: linalg.matvec %[[sAi]], %{{.*}}, %[[sCi]] : (memref, memref, memref) @@ -230,12 +242,13 @@ // TILE-02-DAG: %[[C2:.*]] = constant 2 : index // TILE-02: %[[K:.*]] = dim %{{.*}}, %c1 : memref // TILE-02: scf.for %[[J:.*]] = %{{.*}}{{.*}} to %[[K]] step %{{.*}} { +// TILE-02: %[[dim0:.]] = affine.min #[[$minTileSize]]()[%[[K]]] // TILE-02: %[[M:.*]] = dim %{{.*}}, %c0 : memref // TILE-02: %[[localN:.*]] = dim %{{.*}}, %c1 -// TILE-02: %[[szN:.*]] = affine.min #[[$bound_map]](%[[J]])[%[[localN]]] +// TILE-02: %[[szN:.*]] = affine.min #[[$bound_map]](%[[dim0]], %[[J]])[%[[localN]]] // TILE-02: %[[sAj:.*]] = subview %{{.*}}[0, %[[J]]] [%[[M]], %[[szN]]] [1, 1] : memref to memref // TILE-02: %[[localN:.*]] = dim %{{.*}}, %c0 -// TILE-02: %[[szN:.*]] = affine.min #[[$bound_map]](%[[J]])[%[[localN]]] +// TILE-02: %[[szN:.*]] = affine.min #[[$bound_map]](%[[dim0]], %[[J]])[%[[localN]]] // TILE-02: %[[sBj:.*]] = subview %{{.*}}[%[[J]]] [%[[szN]]] [1] : memref to memref // TILE-02: linalg.matvec %[[sAj]], %[[sBj]], %{{.*}} : (memref, memref, memref) @@ -256,16 +269,18 @@ // TILE-234: %[[K:.*]] = dim %{{.*}}, %c1 : memref // TILE-234: scf.for %[[I:.*]] = %{{.*}}{{.*}} to %[[M]] step %{{.*}} { // TILE-234: scf.for %[[J:.*]] = %{{.*}}{{.*}} to %[[K]] step %{{.*}} { +// TILE-234: %[[dimM:.]] = affine.min #[[$minTileSize0]]()[%[[M]]] +// TILE-234: %[[dimK:.]] = affine.min #[[$minTileSize1]]()[%[[K]]] // TILE-234: %[[localM:.*]] = dim %{{.*}}, %c0 -// TILE-234: %[[szM:.*]] = affine.min #[[$bound_map_2]](%[[I]])[%[[localM]]] +// TILE-234: %[[szM:.*]] = affine.min #[[$bound_map]](%[[dimM]], %[[I]])[%[[localM]]] // TILE-234: %[[localN:.*]] = dim %{{.*}}, %c1 -// TILE-234: %[[szN:.*]] = affine.min #[[$bound_map_3]](%[[J]])[%[[localN]]] +// TILE-234: %[[szN:.*]] = affine.min #[[$bound_map]](%[[dimK]], %[[J]])[%[[localN]]] // TILE-234: %[[sAij:.*]] = subview %{{.*}}[%[[I]], %[[J]]] [%[[szM]], %[[szN]]] [1, 1] : memref to memref // TILE-234: %[[localN:.*]] = dim %{{.*}}, %c0 -// TILE-234: %[[szN:.*]] = affine.min #[[$bound_map_3]](%[[J]])[%[[localN]]] +// TILE-234: %[[szN:.*]] = affine.min #[[$bound_map]](%[[dimK]], %[[J]])[%[[localN]]] // TILE-234: %[[sBj:.*]] = subview %{{.*}}[%[[J]]] [%[[szN]]] [1] : memref to memref // TILE-234: %[[localM:.*]] = dim %{{.*}}, %c0 -// TILE-234: %[[szM:.*]] = affine.min #[[$bound_map_2]](%[[I]])[%[[localM]]] +// TILE-234: %[[szM:.*]] = affine.min #[[$bound_map]](%[[dimM]], %[[I]])[%[[localM]]] // TILE-234: %[[sCi:.*]] = subview %{{.*}}[%[[I]]] [%[[szM]]] [1] : memref to memref // // TILE-234: linalg.matvec %[[sAij]], %[[sBj]], %[[sCi]] : (memref, memref, memref) @@ -281,11 +296,12 @@ // TILE-2-DAG: %[[C2:.*]] = constant 2 : index // TILE-2: %[[M:.*]] = dim %{{.*}}, %c0 : memref // TILE-2: scf.for %[[I:.*]] = %{{.*}}{{.*}} to %[[M]] step %{{.*}} { +// TILE-2: %[[dim0:.]] = affine.min #[[$minTileSize]]()[%[[M]]] // TILE-2: %[[localM:.*]] = dim %{{.*}}, %c0 -// TILE-2: %[[szM:.*]] = affine.min #[[$bound_map]](%[[I]])[%[[localM]]] +// TILE-2: %[[szM:.*]] = affine.min #[[$bound_map]](%[[dim0]], %[[I]])[%[[localM]]] // TILE-2: %[[sAi:.*]] = subview %{{.*}}[%[[I]]] [%[[szM]]] [1] : memref to memref // TILE-2: %[[localM:.*]] = dim %{{.*}}, %c0 -// TILE-2: %[[szM:.*]] = affine.min #[[$bound_map]](%[[I]])[%[[localM]]] +// TILE-2: %[[szM:.*]] = affine.min #[[$bound_map]](%[[dim0]], %[[I]])[%[[localM]]] // TILE-2: %[[sBi:.*]] = subview %{{.*}}[%[[I]]] [%[[szM]]] [1] : memref to memref // TILE-2: linalg.dot %[[sAi]], %[[sBi]], {{.*}} : (memref, memref, memref) @@ -300,11 +316,12 @@ // TILE-234-DAG: %[[C2:.*]] = constant 2 : index // TILE-234: %[[ubK:.*]] = dim %{{.*}}, %c0 : memref // TILE-234: scf.for %[[I:.*]] = %{{.*}} to %[[ubK]] step %{{.*}} { +// TILE-234: %[[dimK:.]] = affine.min #[[$minTileSize0]]()[%[[ubK]]] // TILE-234: %[[localM:.*]] = dim %{{.*}}, %c0 -// TILE-234: %[[szM:.*]] = affine.min #[[$bound_map_2]](%[[I]])[%[[localM]]] +// TILE-234: %[[szM:.*]] = affine.min #[[$bound_map]](%[[dimK]], %[[I]])[%[[localM]]] // TILE-234: %[[sAi:.*]] = subview %{{.*}}[%[[I]]] [%[[szM]]] [1] : memref to memref // TILE-234: %[[localM:.*]] = dim %{{.*}}, %c0 -// TILE-234: %[[szM:.*]] = affine.min #[[$bound_map_2]](%[[I]])[%[[localM]]] +// TILE-234: %[[szM:.*]] = affine.min #[[$bound_map]](%[[dimK]], %[[I]])[%[[localM]]] // TILE-234: %[[sBi:.*]] = subview %{{.*}}[%[[I]]] [%[[szM]]] [1] : memref to memref // TILE-234: linalg.dot %[[sAi]], %[[sBi]], %{{.*}} : (memref, memref, memref) diff --git a/mlir/test/Dialect/Linalg/tile_conv.mlir b/mlir/test/Dialect/Linalg/tile_conv.mlir --- a/mlir/test/Dialect/Linalg/tile_conv.mlir +++ b/mlir/test/Dialect/Linalg/tile_conv.mlir @@ -1,9 +1,11 @@ // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3,0,0,4" | FileCheck %s -check-prefix=TILE-23004 // TILE-23004-DAG: #[[$D0x30pS0x10:.*]] = affine_map<(d0) -> (d0 * 30)> -// TILE-23004-DAG: #[[$S0x10p90D0x30pS1:.*]] = affine_map<(d0)[s0, s1] -> (s0 * 10 + 90, d0 * -30 + s1)> +// TILE-23004-DAG: #[[$S0x10p90D0x30pS1:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * 30 + s0 * 10, d1 * -30 + s1)> // TILE-23004-DAG: #[[$strided4D:.*]] = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3)> -// TILE-23004-DAG: #[[$bound_map_4:.*]] = affine_map<(d0)[s0] -> (4, -d0 + s0)> +// TILE-23004-DAG: #[[$minTileSize1:.*]] = affine_map<()[s0] -> (4, s0)> +// TILE-23004-DAG: #[[$minTileSize2:.*]] = affine_map<()[s0] -> (3, s0)> +// TILE-23004-DAG: #[[$bound_map_4:.*]] = affine_map<(d0, d1)[s0] -> (d0, -d1 + s0)> func @conv(%arg0: memref, %arg1: memref, %arg2: memref) { linalg.conv(%arg0, %arg1, %arg2) {dilations = [10, 20], strides = [30, 40]} : memref, memref, memref @@ -24,19 +26,21 @@ // TILE-23004: scf.for %[[ivI:.*]] = %{{.*}} to %[[B]] step %{{.*}} { // TILE-23004: scf.for %[[ivJ:.*]] = %{{.*}} to %[[X0]] step %{{.*}} { // TILE-23004: scf.for %[[ivK:.*]] = %{{.*}} to %[[Q]] step %{{.*}} { +// TILE-23004: %[[dimX0:.*]] = affine.min #[[$minTileSize2]]()[%[[X0]]] +// TILE-23004: %[[dimQ:.*]] = affine.min #[[$minTileSize1]]()[%[[Q]]] // TILE-23004: %[[Z0_1:.*]] = dim %[[ARG0]], %c0 : memref // TILE-23004: %[[Z1:.*]] = dim %[[ARG0]], %c1 : memref // TILE-23004: %[[Z2:.*]] = dim %[[ARG0]], %c2 : memref -// TILE-23004: %[[szK:.*]] = affine.min #[[$bound_map_4]](%[[ivK]])[%[[Z2]]] +// TILE-23004: %[[szK:.*]] = affine.min #[[$bound_map_4]](%[[dimQ]], %[[ivK]])[%[[Z2]]] // TILE-23004: %[[K:.*]] = dim %[[ARG0]], %c3 : memref // TILE-23004: %[[FilterView:.*]] = subview %{{.*}}[0, 0, %[[ivK]], 0] [%[[Z0_1]], %[[Z1]], %[[szK]], %[[K]]] [1, 1, 1, 1] : memref to memref // // TILE-23004: %[[J1:.*]] = affine.apply #[[$D0x30pS0x10]](%[[ivJ]]) // TILE-23004: %[[PaddedInput0b:.*]] = dim %[[ARG1]], %c1 : memref -// TILE-23004: %[[I1pStep:.*]] = affine.min #[[$S0x10p90D0x30pS1]](%[[ivJ]])[%[[Z0]], %[[PaddedInput0b]]] +// TILE-23004: %[[I1pStep:.*]] = affine.min #[[$S0x10p90D0x30pS1]](%[[dimX0]], %[[ivJ]])[%[[Z0]], %[[PaddedInput0b]]] // TILE-23004: %[[SZ2:.*]] = dim %[[ARG1]], %c2 : memref // TILE-23004: %[[dim3:.*]] = dim %[[ARG1]], %c3 -// TILE-23004: %[[sz3:.*]] = affine.min #[[$bound_map_4]](%[[ivK]])[%[[dim3]]] +// TILE-23004: %[[sz3:.*]] = affine.min #[[$bound_map_4]](%[[dimQ]], %[[ivK]])[%[[dim3]]] // TILE-23004: %[[InputView:.*]] = subview %{{.*}}[%[[ivI]], %[[J1]], 0, %[[ivK]]] [%{{.*}}, %{{.*}}, %[[SZ2]], %[[sz3]]] [1, 1, 1, 1] : memref to memref // // TILE-23004: %[[X0:.*]] = dim %[[ARG2]], %c2 : memref diff --git a/mlir/test/Dialect/Linalg/tile_conv_padding.mlir b/mlir/test/Dialect/Linalg/tile_conv_padding.mlir --- a/mlir/test/Dialect/Linalg/tile_conv_padding.mlir +++ b/mlir/test/Dialect/Linalg/tile_conv_padding.mlir @@ -3,7 +3,8 @@ // TILE-23004-DAG: #[[$strided4D:.*]] = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3)> // TILE-20000-DAG: #[[$strided4D:.*]] = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3)> -// TILE-20000-DAG: #[[$minmap:.*]] = affine_map<(d0)[s0] -> (2, -d0 + s0)> +// TILE-20000-DAG: #[[$minTileSize:.*]] = affine_map<()[s0] -> (2, s0)> +// TILE-20000-DAG: #[[$minmap:.*]] = affine_map<(d0, d1)[s0] -> (d0, -d1 + s0)> func @conv_padding(%arg0: memref, %arg1: memref, %arg2: memref) { linalg.conv(%arg0, %arg1, %arg2) {dilations = [10, 20], padding = dense<[[1, 1], [0, 1]]> : tensor<2x2xi64>, strides = [30, 40]} : memref, memref, memref @@ -23,14 +24,15 @@ // TILE-20000-DAG: %[[C2:.*]] = constant 2 : index // TILE-20000: %[[B:.*]] = dim %[[ARG1]], %c0 // TILE-20000: scf.for %[[ivI:.*]] = %[[C0]] to %[[B]] step %[[C2]] { +// TILE-20000: %[[DIMB:.*]] = affine.min #[[$minTileSize]]()[%[[B]]] // TILE-20000: %[[DIM10:.*]] = dim %[[ARG1]], %c0 -// TILE-20000: %[[EXTENT:.*]] = affine.min #[[$minmap]](%[[ivI]])[%[[DIM10]]] +// TILE-20000: %[[EXTENT:.*]] = affine.min #[[$minmap]](%[[DIMB]], %[[ivI]])[%[[DIM10]]] // TILE-20000: %[[DIM11:.*]] = dim %[[ARG1]], %c1 // TILE-20000: %[[DIM12:.*]] = dim %[[ARG1]], %c2 // TILE-20000: %[[DIM13:.*]] = dim %[[ARG1]], %c3 // TILE-20000: %[[SUBVIEW1:.*]] = subview %[[ARG1]][%[[ivI]], 0, 0, 0] [%[[EXTENT]], %[[DIM11]], %[[DIM12]], %[[DIM13]]] // TILE-20000: %[[DIM20:.*]] = dim %[[ARG2]], %c0 -// TILE-20000: %[[EXTENT:.*]] = affine.min #[[$minmap]](%[[ivI]])[%[[DIM20]]] +// TILE-20000: %[[EXTENT:.*]] = affine.min #[[$minmap]](%[[DIMB]], %[[ivI]])[%[[DIM20]]] // TILE-20000: %[[DIM21:.*]] = dim %[[ARG2]], %c1 // TILE-20000: %[[DIM22:.*]] = dim %[[ARG2]], %c2 // TILE-20000: %[[DIM23:.*]] = dim %[[ARG2]], %c3 diff --git a/mlir/test/Dialect/Linalg/tile_simple_conv.mlir b/mlir/test/Dialect/Linalg/tile_simple_conv.mlir --- a/mlir/test/Dialect/Linalg/tile_simple_conv.mlir +++ b/mlir/test/Dialect/Linalg/tile_simple_conv.mlir @@ -1,10 +1,10 @@ // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3,4" | FileCheck %s -// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0] -> (2, -d0 + s0)> -// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0, s1] -> (s0 + 3, -d0 + s1)> -// CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0)[s0, s1] -> (s0 + 4, -d0 + s1)> -// CHECK-DAG: #[[MAP4:.*]] = affine_map<(d0)[s0] -> (3, -d0 + s0)> -// CHECK-DAG: #[[MAP5:.*]] = affine_map<(d0)[s0] -> (4, -d0 + s0)> +// CHECK-DAG: #[[$minTileSize0:.*]] = affine_map<()[s0] -> (2, s0)> +// CHECK-DAG: #[[$minTileSize1:.*]] = affine_map<()[s0] -> (3, s0)> +// CHECK-DAG: #[[$minTileSize2:.*]] = affine_map<()[s0] -> (4, s0)> +// CHECK-DAG: #[[$bound1:.*]] = affine_map<(d0, d1)[s0] -> (d0, -d1 + s0)> +// CHECK-DAG: #[[$bound2:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 + s0, -d1 + s1)> func @conv(%arg0 : memref, %arg1 : memref, %arg2 : memref) { linalg.conv(%arg0, %arg1, %arg2) : memref, memref, memref @@ -20,30 +20,33 @@ // CHECK-DAG: %[[C2:.*]] = constant 2 : index // CHECK-DAG: %[[C3:.*]] = constant 3 : index // CHECK-DAG: %[[C4:.*]] = constant 4 : index -// CHECK: %[[T0:.*]] = dim %[[ARG0]], %[[C0]] -// CHECK: %[[T1:.*]] = dim %[[ARG0]], %[[C1]] -// CHECK: %[[T2:.*]] = dim %[[ARG1]], %[[C0]] -// CHECK: %[[T3:.*]] = dim %[[ARG2]], %[[C1]] -// CHECK: %[[T4:.*]] = dim %[[ARG2]], %[[C2]] -// CHECK: scf.for %[[ARG3:.*]] = %[[C0]] to %[[T2]] step %[[C2]] -// CHECK: scf.for %[[ARG4:.*]] = %[[C0]] to %[[T3]] step %[[C3]] -// CHECK: scf.for %[[ARG5:.*]] = %[[C0]] to %[[T4]] step %[[C4]] -// CHECK: %[[T5:.*]] = dim %[[ARG1]], %[[C0]] -// CHECK: %[[T6:.*]] = affine.min #[[MAP0]](%[[ARG3]])[%[[T5]]] -// CHECK: %[[T7:.*]] = dim %[[ARG1]], %[[C1]] -// CHECK: %[[T8:.*]] = affine.min #[[MAP1]](%[[ARG4]])[%[[T0]], %[[T7]]] -// CHECK: %[[T9:.*]] = dim %[[ARG1]], %[[C2]] -// CHECK: %[[T10:.*]] = affine.min #[[MAP2]](%[[ARG5]])[%[[T1]], %[[T9]]] -// CHECK: %[[T11:.*]] = dim %[[ARG1]], %[[C3]] +// CHECK: %[[F0:.*]] = dim %[[ARG0]], %[[C0]] +// CHECK: %[[F1:.*]] = dim %[[ARG0]], %[[C1]] +// CHECK: %[[N:.*]] = dim %[[ARG1]], %[[C0]] +// CHECK: %[[Y0:.*]] = dim %[[ARG2]], %[[C1]] +// CHECK: %[[Y1:.*]] = dim %[[ARG2]], %[[C2]] +// CHECK: scf.for %[[ARG3:.*]] = %[[C0]] to %[[N]] step %[[C2]] +// CHECK: scf.for %[[ARG4:.*]] = %[[C0]] to %[[Y0]] step %[[C3]] +// CHECK: scf.for %[[ARG5:.*]] = %[[C0]] to %[[Y1]] step %[[C4]] +// CHECK: %[[dimN:.*]] = affine.min #[[$minTileSize0]]()[%[[N]]] +// CHECK: %[[dimY0:.*]] = affine.min #[[$minTileSize1]]()[%[[Y0]]] +// CHECK: %[[dimY1:.*]] = affine.min #[[$minTileSize2]]()[%[[Y1]]] +// CHECK: %[[N_2:.*]] = dim %[[ARG1]], %[[C0]] +// CHECK: %[[SVSIZEN:.*]] = affine.min #[[$bound1]](%[[dimN]], %[[ARG3]])[%[[N_2]]] +// CHECK: %[[X0:.*]] = dim %[[ARG1]], %[[C1]] +// CHECK: %[[SVSIZEX0:.*]] = affine.min #[[$bound2]](%[[dimY0]], %[[ARG4]])[%[[F0]], %[[X0]]] +// CHECK: %[[X1:.*]] = dim %[[ARG1]], %[[C2]] +// CHECK: %[[SVSIZEX1:.*]] = affine.min #[[$bound2]](%[[dimY1]], %[[ARG5]])[%[[F1]], %[[X1]]] +// CHECK: %[[SVSIZEC:.*]] = dim %[[ARG1]], %[[C3]] // CHECK: %[[SV1:.*]] = subview %[[ARG1]][%[[ARG3]], %[[ARG4]], %[[ARG5]], 0] -// CHECK-SAME: [%[[T6]], %[[T8]], %[[T10]], %[[T11]]] -// CHECK: %[[T13:.*]] = dim %[[ARG2]], %[[C0]] -// CHECK: %[[T14:.*]] = affine.min #[[MAP0]](%[[ARG3]])[%[[T13]]] -// CHECK: %[[T15:.*]] = dim %[[ARG2]], %[[C1]] -// CHECK: %[[T16:.*]] = affine.min #[[MAP4]](%[[ARG4]])[%[[T15]]] -// CHECK: %[[T17:.*]] = dim %[[ARG2]], %[[C2]] -// CHECK: %[[T18:.*]] = affine.min #[[MAP5]](%[[ARG5]])[%[[T17]]] -// CHECK: %[[T19:.*]] = dim %[[ARG2]], %[[C3]] +// CHECK-SAME: [%[[SVSIZEN]], %[[SVSIZEX0]], %[[SVSIZEX1]], %[[SVSIZEC]]] +// CHECK: %[[N_3:.*]] = dim %[[ARG2]], %[[C0]] +// CHECK: %[[SVSIZEN_2:.*]] = affine.min #[[$bound1]](%[[dimN]], %[[ARG3]])[%[[N_3]]] +// CHECK: %[[Y0_2:.*]] = dim %[[ARG2]], %[[C1]] +// CHECK: %[[SVSIZEY0:.*]] = affine.min #[[$bound1]](%[[dimY0]], %[[ARG4]])[%[[Y0_2]]] +// CHECK: %[[Y1_2:.*]] = dim %[[ARG2]], %[[C2]] +// CHECK: %[[SVSIZEY1:.*]] = affine.min #[[$bound1]](%[[dimY1]], %[[ARG5]])[%[[Y1_2]]] +// CHECK: %[[SVSIZEK:.*]] = dim %[[ARG2]], %[[C3]] // CHECK: %[[SV2:.*]] = subview %[[ARG2]][%[[ARG3]], %[[ARG4]], %[[ARG5]], 0] -// CHECK-SAME: [%[[T14]], %[[T16]], %[[T18]], %[[T19]]] +// CHECK-SAME: [%[[SVSIZEN_2]], %[[SVSIZEY0]], %[[SVSIZEY1]], %[[SVSIZEK]]] // CHECK: linalg.conv(%[[ARG0]], %[[SV1]], %[[SV2]]) \ No newline at end of file