diff --git a/mlir/include/mlir/EDSC/Intrinsics.h b/mlir/include/mlir/EDSC/Intrinsics.h --- a/mlir/include/mlir/EDSC/Intrinsics.h +++ b/mlir/include/mlir/EDSC/Intrinsics.h @@ -195,6 +195,7 @@ using affine_apply = ValueBuilder; using affine_if = OperationBuilder; using affine_load = ValueBuilder; +using affine_min = ValueBuilder; using affine_store = OperationBuilder; using alloc = ValueBuilder; using call = OperationBuilder; 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 @@ -260,7 +260,8 @@ for (unsigned viewIndex = 0; viewIndex < linalgOp.getNumInputsAndOutputs(); ++viewIndex) { Value view = *(viewIteratorBegin + viewIndex); - unsigned rank = view.getType().cast().getRank(); + auto viewType = view.getType().cast(); + unsigned rank = viewType.getRank(); auto map = loopToOperandRangesMaps(linalgOp)[viewIndex]; // If the view is not tiled, we can use it as is. if (!isTiled(map, tileSizes)) { @@ -287,12 +288,29 @@ auto offset = applyMapToValues(b, loc, m, lbs, folder).front(); offsets.push_back(offset); auto size = applyMapToValues(b, loc, m, subViewSizes, folder).front(); + + // The size of the subview should be trimmed to avoid out-of-bounds + // accesses, unless we statically know the subview size divides the view + // size evenly. + int64_t viewSize = viewType.getDimSize(r); + auto sizeCst = dyn_cast_or_null(size.getDefiningOp()); + if (ShapedType::isDynamic(viewSize) || !sizeCst || + (viewSize % sizeCst.getValue()) != 0) { + // Compute min(size, dim - offset) to avoid out-of-bounds accesses. + auto minMap = AffineMap::get( + /*dimCount=*/3, /*symbolCount=*/0, + {getAffineDimExpr(/*position=*/0, b.getContext()), + getAffineDimExpr(/*position=*/1, b.getContext()) - + getAffineDimExpr(/*position=*/2, b.getContext())}); + auto d = dim(folder, view, r); + size = affine_min(folder, b.getIndexType(), minMap, + ValueRange{size, d, offset}); + } + sizes.push_back(size); strides.push_back(constant_index(folder, 1)); } - // TODO(b/144419024) Atm std.subview is not guaranteed in-bounds. Depending - // on the semantics we attach to it, we may need to use min(size, dim) here - // and canonicalize later. + res.push_back(b.create(loc, view, offsets, sizes, strides)); } 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 @@ -13,6 +13,11 @@ // 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, d1, d2) -> (d0, d1 - d2)> +// TILE-02-DAG: #[[bound_map:.*]] = affine_map<(d0, d1, d2) -> (d0, d1 - d2)> +// TILE-002-DAG: #[[bound_map:.*]] = affine_map<(d0, d1, d2) -> (d0, d1 - d2)> +// TILE-234-DAG: #[[bound_map:.*]] = affine_map<(d0, d1, d2) -> (d0, d1 - d2)> + // TILE-2-DAG: #[[strided1D_dynamic:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> // TILE-02-DAG: #[[strided1D_dynamic:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> // T_ILE-002-DAG: #[[strided1D_dynamic:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> @@ -37,10 +42,14 @@ // TILE-2-DAG: %[[C2:.*]] = constant 2 : index // TILE-2: %[[M:.*]] = dim %{{.*}}, 0 : memref // TILE-2: loop.for %[[I:.*]] = %{{.*}}{{.*}} to %[[M]] step %{{.*}} { +// TILE-2: %[[localM:.*]] = dim %{{.*}}, 0 +// TILE-2: %[[szM:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localM]], %[[I]]) // TILE-2: %[[K:.*]] = dim %{{.*}}, 1 : memref -// TILE-2: %[[sAi:.*]] = std.subview %{{.*}}[%[[I]], %[[C0]]][%[[C2]], %[[K]]][%[[C1]], %[[C1]]] : memref to memref +// TILE-2: %[[sAi:.*]] = std.subview %{{.*}}[%[[I]], %[[C0]]][%[[szM]], %[[K]]][%[[C1]], %[[C1]]] : memref to memref +// TILE-2: %[[localK:.*]] = dim %{{.*}}, 0 +// TILE-2: %[[szK:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localK]], %[[I]]) // TILE-2: %[[N:.*]] = dim %{{.*}}, 1 : memref -// TILE-2: %[[sCi:.*]] = std.subview %{{.*}}[%[[I]], %[[C0]]][%[[C2]], %[[N]]][%[[C1]], %[[C1]]] : memref to memref +// TILE-2: %[[sCi:.*]] = std.subview %{{.*}}[%[[I]], %[[C0]]][%[[szK]], %[[N]]][%[[C1]], %[[C1]]] : memref to memref // TILE-2: linalg.matmul(%[[sAi]], %{{.*}}, %[[sCi]]) : memref, memref, memref // TILE-02-LABEL: func @matmul( @@ -50,9 +59,13 @@ // TILE-02: %[[N:.*]] = dim %arg1, 1 : memref // TILE-02: loop.for %[[J:.*]] = %{{.*}} to %[[N]] step %{{.*}} { // TILE-02: %[[K:.*]] = dim %{{.*}}, 0 : memref -// TILE-02: %[[sBj:.*]] = std.subview %{{.*}}[%[[C0]], %[[J]]][%[[K]], %[[C2]]][%[[C1]], %[[C1]]] : memref to memref +// TILE-02: %[[localN:.*]] = dim %{{.*}}, 1 +// TILE-02: %[[szN:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localN]], %[[J]]) +// TILE-02: %[[sBj:.*]] = std.subview %{{.*}}[%[[C0]], %[[J]]][%[[K]], %[[szN]]][%[[C1]], %[[C1]]] : memref to memref // TILE-02: %[[M:.*]] = dim %{{.*}}, 0 : memref -// TILE-02: %[[sCj:.*]] = std.subview %{{.*}}[%[[C0]], %[[J]]][%[[M]], %[[C2]]][%[[C1]], %[[C1]]] : memref to memref +// TILE-02: %[[localK:.*]] = dim %{{.*}}, 1 +// TILE-02: %[[szK:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localK]], %[[J]]) +// TILE-02: %[[sCj:.*]] = std.subview %{{.*}}[%[[C0]], %[[J]]][%[[M]], %[[szK]]][%[[C1]], %[[C1]]] : memref to memref // TILE-02: linalg.matmul(%{{.*}}, %[[sBj]], %[[sCj]]) : memref, memref, memref // TILE-002-LABEL: func @matmul( @@ -62,9 +75,13 @@ // TILE-002: %[[ubK:.*]] = dim %{{.*}}, 1 : memref // TILE-002: loop.for %[[K:.*]] = %{{.*}}{{.*}} to %[[ubK]] step %{{.*}} { // TILE-002: %[[M:.*]] = dim %{{.*}}, 0 : memref -// TILE-002: %[[sAj:.*]] = std.subview %{{.*}}[%[[C0]], %[[K]]][%[[M]], %[[C2]]][%[[C1]], %[[C1]]] : memref to memref +// TILE-002: %[[localK:.*]] = dim %{{.*}}, 1 +// TILE-002: %[[szK:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localK]], %[[K]]) +// TILE-002: %[[sAj:.*]] = std.subview %{{.*}}[%[[C0]], %[[K]]][%[[M]], %[[szK]]][%[[C1]], %[[C1]]] : memref to memref +// TILE-002: %[[localK:.*]] = dim %{{.*}}, 0 +// TILE-002: %[[szK:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localK]], %[[K]]) // TILE-002: %[[N:.*]] = dim %{{.*}}, 1 : memref -// TILE-002: %[[sBj:.*]] = std.subview %{{.*}}[%[[K]], %[[C0]]][%[[C2]], %[[N]]][%[[C1]], %[[C1]]] : memref to memref +// TILE-002: %[[sBj:.*]] = std.subview %{{.*}}[%[[K]], %[[C0]]][%[[szK]], %[[N]]][%[[C1]], %[[C1]]] : memref to memref // TILE-002: linalg.matmul(%[[sAj]], %[[sBj]], %{{.*}}) : memref, memref, memref // TILE-234-LABEL: func @matmul( @@ -79,9 +96,90 @@ // TILE-234: loop.for %[[I:.*]] = %{{.*}}{{.*}} to %[[ubM]] step %{{.*}} { // TILE-234: loop.for %[[J:.*]] = %{{.*}}{{.*}} to %[[ubN]] step %{{.*}} { // TILE-234: loop.for %[[K:.*]] = %{{.*}}{{.*}} to %[[ubK]] step %{{.*}} { -// TILE-234: %[[sAik:.*]] = std.subview %{{.*}}[%[[I]], %[[K]]][%[[C2]], %[[C4]]][%[[C1]], %[[C1]]] : memref to memref -// TILE-234: %[[sBkj:.*]] = std.subview %{{.*}}[%[[K]], %[[J]]][%[[C4]], %[[C3]]][%[[C1]], %[[C1]]] : memref to memref -// TILE-234: %[[sCij:.*]] = std.subview %{{.*}}[%[[I]], %[[J]]][%[[C2]], %[[C3]]][%[[C1]], %[[C1]]] : memref to memref +// TILE-234: %[[localM:.*]] = dim %{{.*}}, 0 +// TILE-234: %[[szM:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localM]], %[[I]]) +// TILE-234: %[[localK:.*]] = dim %{{.*}}, 1 +// TILE-234: %[[szK:.*]] = affine.min #[[bound_map]](%[[C4]], %[[localK]], %[[K]]) +// TILE-234: %[[sAik:.*]] = std.subview %{{.*}}[%[[I]], %[[K]]][%[[szM]], %[[szK]]][%[[C1]], %[[C1]]] : memref to memref +// TILE-234: %[[localK:.*]] = dim %{{.*}}, 0 +// TILE-234: %[[szK:.*]] = affine.min #[[bound_map]](%[[C4]], %[[localK]], %[[K]]) +// TILE-234: %[[localN:.*]] = dim %{{.*}}, 1 +// TILE-234: %[[szN:.*]] = affine.min #[[bound_map]](%[[C3]], %[[localN]], %[[J]]) +// TILE-234: %[[sBkj:.*]] = std.subview %{{.*}}[%[[K]], %[[J]]][%[[szK]], %[[szN]]][%[[C1]], %[[C1]]] : memref to memref +// TILE-234: %[[localM:.*]] = dim %{{.*}}, 0 +// TILE-234: %[[szM:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localM]], %[[I]]) +// TILE-234: %[[localN:.*]] = dim %{{.*}}, 1 +// TILE-234: %[[szN:.*]] = affine.min #[[bound_map]](%[[C3]], %[[localN]], %[[J]]) +// TILE-234: %[[sCij:.*]] = std.subview %{{.*}}[%[[I]], %[[J]]][%[[szM]], %[[szN]]][%[[C1]], %[[C1]]] : memref to memref +// +// TILE-234: linalg.matmul(%[[sAik]], %[[sBkj]], %[[sCij]]) : memref, memref, memref + +// When the buffer shapes are known at compile time, it is possible to avoid +// the "min" in subview size computation. This test uses buffer sizes divisible +// by respective tile sizes (M=10 divisble by 2, N=12 divisible by 2 and 3, +// K=16 divisble by 2 and 4). +func @matmul_static(%arg0: memref<10x16xf32, offset: ?, strides: [?, 1]>, %arg1: memref<16x12xf32, offset: ?, strides: [?, 1]>, %arg2: memref<10x12xf32, offset: ?, strides: [?, 1]>) { + linalg.matmul(%arg0, %arg1, %arg2) : memref<10x16xf32, offset: ?, strides: [?, 1]>, memref<16x12xf32, offset: ?, strides: [?, 1]>, memref<10x12xf32, offset: ?, strides: [?, 1]> + return +} +// TILE-2-LABEL: func @matmul_static( +// TILE-2-DAG: %[[C0:.*]] = constant 0 : index +// TILE-2-DAG: %[[C1:.*]] = constant 1 : index +// TILE-2-DAG: %[[C2:.*]] = constant 2 : index +// TILE-2: %[[M:.*]] = dim %{{.*}}, 0 : memref<10x16xf32, #[[strided2D]]> +// TILE-2: loop.for %[[I:.*]] = %{{.*}}{{.*}} to %[[M]] step %{{.*}} { +// TILE-2: %[[K:.*]] = dim %{{.*}}, 1 : memref<10x16xf32, #[[strided2D]]> +// TILE-2: %[[sAi:.*]] = std.subview %{{.*}}[%[[I]], %[[C0]]][%[[C2]], %[[K]]][%[[C1]], %[[C1]]] : memref<10x16xf32, #[[strided2D]]> to memref +// TILE-2: %[[N:.*]] = dim %{{.*}}, 1 : memref<10x12xf32, #[[strided2D]]> +// TILE-2: %[[sCi:.*]] = std.subview %{{.*}}[%[[I]], %[[C0]]][%[[C2]], %[[N]]][%[[C1]], %[[C1]]] : memref<10x12xf32, #[[strided2D]]> to memref +// TILE-2: linalg.matmul(%[[sAi]], %{{.*}}, %[[sCi]]) + +// TILE-02-LABEL: func @matmul_static( +// TILE-02-DAG: %[[C0:.*]] = constant 0 : index +// TILE-02-DAG: %[[C1:.*]] = constant 1 : index +// TILE-02-DAG: %[[C2:.*]] = constant 2 : index +// TILE-02: %[[N:.*]] = dim %arg1, 1 : memref<16x12xf32, #[[strided2D]]> +// TILE-02: loop.for %[[J:.*]] = %{{.*}} to %[[N]] step %{{.*}} { +// TILE-02: %[[K:.*]] = dim %{{.*}}, 0 : memref<16x12xf32, #[[strided2D]]> +// TILE-02-NOT: affine.min +// TILE-02: %[[sBj:.*]] = std.subview %{{.*}}[%[[C0]], %[[J]]][%[[K]], %[[C2]]][%[[C1]], %[[C1]]] : memref<16x12xf32, #[[strided2D]]> to memref +// TILE-02: %[[M:.*]] = dim %{{.*}}, 0 : memref<10x12xf32, #[[strided2D]]> +// TILE-02-NOT: affine.min +// TILE-02: %[[sCj:.*]] = std.subview %{{.*}}[%[[C0]], %[[J]]][%[[M]], %[[C2]]][%[[C1]], %[[C1]]] : memref<10x12xf32, #[[strided2D]]> to memref +// TILE-02: linalg.matmul(%{{.*}}, %[[sBj]], %[[sCj]]) : memref<10x16xf32, #[[strided2D]]>, memref, memref + +// TILE-002-LABEL: func @matmul_static( +// TILE-002-DAG: %[[C0:.*]] = constant 0 : index +// TILE-002-DAG: %[[C1:.*]] = constant 1 : index +// TILE-002-DAG: %[[C2:.*]] = constant 2 : index +// TILE-002: %[[ubK:.*]] = dim %{{.*}}, 1 : memref<10x16xf32, #[[strided2D]]> +// TILE-002: loop.for %[[K:.*]] = %{{.*}}{{.*}} to %[[ubK]] step %{{.*}} { +// TILE-002: %[[M:.*]] = dim %{{.*}}, 0 : memref<10x16xf32, #[[strided2D]]> +// TILE-002-NOT: affine.min +// TILE-002: %[[sAj:.*]] = std.subview %{{.*}}[%[[C0]], %[[K]]][%[[M]], %[[C2]]][%[[C1]], %[[C1]]] : memref<10x16xf32, #[[strided2D]]> to memref +// TILE-002: %[[N:.*]] = dim %{{.*}}, 1 : memref<16x12xf32, #[[strided2D]]> +// TILE-002-NOT: affine.min +// TILE-002: %[[sBj:.*]] = std.subview %{{.*}}[%[[K]], %[[C0]]][%[[C2]], %[[N]]][%[[C1]], %[[C1]]] : memref<16x12xf32, #[[strided2D]]> to memref +// TILE-002: linalg.matmul(%[[sAj]], %[[sBj]], %{{.*}}) : memref, memref, memref<10x12xf32, #[[strided2D]]> + +// TILE-234-LABEL: func @matmul_static( +// TILE-234-DAG: %[[C0:.*]] = constant 0 : index +// TILE-234-DAG: %[[C1:.*]] = constant 1 : index +// TILE-234-DAG: %[[C2:.*]] = constant 2 : index +// TILE-234-DAG: %[[C3:.*]] = constant 3 : index +// TILE-234-DAG: %[[C4:.*]] = constant 4 : index +// TILE-234: %[[ubM:.*]] = dim %{{.*}}, 0 : memref<10x16xf32, #[[strided2D]]> +// TILE-234: %[[ubK:.*]] = dim %{{.*}}, 1 : memref<10x16xf32, #[[strided2D]]> +// TILE-234: %[[ubN:.*]] = dim %{{.*}}, 1 : memref<16x12xf32, #[[strided2D]]> +// TILE-234: loop.for %[[I:.*]] = %{{.*}}{{.*}} to %[[ubM]] step %{{.*}} { +// TILE-234: loop.for %[[J:.*]] = %{{.*}}{{.*}} to %[[ubN]] step %{{.*}} { +// TILE-234: loop.for %[[K:.*]] = %{{.*}}{{.*}} to %[[ubK]] step %{{.*}} { +// TILE-234-NOT: affine.min +// TILE-234: %[[sAik:.*]] = std.subview %{{.*}}[%[[I]], %[[K]]][%[[C2]], %[[C4]]][%[[C1]], %[[C1]]] : memref<10x16xf32, #[[strided2D]]> to memref +// TILE-234-NOT: affine.min +// TILE-234: %[[sBkj:.*]] = std.subview %{{.*}}[%[[K]], %[[J]]][%[[C4]], %[[C3]]][%[[C1]], %[[C1]]] : memref<16x12xf32, #[[strided2D]]> to memref +// TILE-234-NOT: affine.min +// TILE-234: %[[sCij:.*]] = std.subview %{{.*}}[%[[I]], %[[J]]][%[[C2]], %[[C3]]][%[[C1]], %[[C1]]] : memref<10x12xf32, #[[strided2D]]> to memref // // TILE-234: linalg.matmul(%[[sAik]], %[[sBkj]], %[[sCij]]) : memref, memref, memref @@ -94,10 +192,14 @@ // TILE-2-DAG: %[[C1:.*]] = constant 1 : index // TILE-2-DAG: %[[C2:.*]] = constant 2 : index // TILE-2: %[[M:.*]] = dim %{{.*}}, 0 : memref -// TILE-2: loop.for %[[I]] = %{{.*}}{{.*}} to %[[M]] step %{{.*}} { +// TILE-2: loop.for %[[I:.*]] = %{{.*}}{{.*}} to %[[M]] step %{{.*}} { +// TILE-2: %[[localM:.*]] = dim %{{.*}}, 0 +// TILE-2: %[[szM:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localM]], %[[I]]) // TILE-2: %[[N:.*]] = dim %{{.*}}, 1 : memref -// TILE-2: %[[sAi:.*]] = std.subview %{{.*}}[%[[I]], %[[C0]]][%[[C2]], %[[N]]][%[[C1]], %[[C1]]] : memref to memref -// TILE-2: %[[sCi:.*]] = std.subview %{{.*}}[%[[I]]][%[[C2]]][%[[C1]]] : memref to memref +// TILE-2: %[[sAi:.*]] = std.subview %{{.*}}[%[[I]], %[[C0]]][%[[szM]], %[[N]]][%[[C1]], %[[C1]]] : memref to memref +// TILE-2: %[[localN:.*]] = dim %{{.*}}, 0 +// TILE-2: %[[szN:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localN]], %[[I]]) +// TILE-2: %[[sCi:.*]] = std.subview %{{.*}}[%[[I]]][%[[szN]]][%[[C1]]] : memref to memref // TILE-2: linalg.matvec(%[[sAi]], %{{.*}}, %[[sCi]]) : memref, memref, memref // TILE-02-LABEL: func @matvec( @@ -107,8 +209,12 @@ // TILE-02: %[[K:.*]] = dim %{{.*}}, 1 : memref // TILE-02: loop.for %[[J]] = %{{.*}}{{.*}} to %[[K]] step %{{.*}} { // TILE-02: %[[M:.*]] = dim %{{.*}}, 0 : memref -// TILE-02: %[[sAj:.*]] = std.subview %{{.*}}[%[[C0]], %[[J]]][%[[M]], %[[C2]]][%[[C1]], %[[C1]]] : memref to memref -// TILE-02: %[[sBj:.*]] = std.subview %{{.*}}[%[[J]]][%[[C2]]][%[[C1]]] : memref to memref +// TILE-02: %[[localN:.*]] = dim %{{.*}}, 1 +// TILE-02: %[[szN:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localN]], %[[J]]) +// TILE-02: %[[sAj:.*]] = std.subview %{{.*}}[%[[C0]], %[[J]]][%[[M]], %[[szN]]][%[[C1]], %[[C1]]] : memref to memref +// TILE-02: %[[localN:.*]] = dim %{{.*}}, 0 +// TILE-02: %[[szN:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localN]], %[[J]]) +// TILE-02: %[[sBj:.*]] = std.subview %{{.*}}[%[[J]]][%[[szN]]][%[[C1]]] : memref to memref // TILE-02: linalg.matvec(%[[sAj]], %[[sBj]], %{{.*}}) : memref, memref, memref // TILE-002-LABEL: func @matvec( @@ -123,9 +229,17 @@ // TILE-234: %[[K:.*]] = dim %{{.*}}, 1 : memref // TILE-234: loop.for %[[I:.*]] = %{{.*}}{{.*}} to %[[M]] step %{{.*}} { // TILE-234: loop.for %[[J:.*]] = %{{.*}}{{.*}} to %[[K]] step %{{.*}} { -// TILE-234: %[[sAij:.*]] = std.subview %{{.*}}[%[[I]], %[[J]]][%[[C2]], %[[C3]]][%[[C1]], %[[C1]]] : memref to memref -// TILE-234: %[[sBj:.*]] = std.subview %{{.*}}[%[[J]]][%[[C3]]][%[[C1]]] : memref to memref -// TILE-234: %[[sCi:.*]] = std.subview %{{.*}}[%[[I]]][%[[C2]]][%[[C1]]] : memref to memref +// TILE-234: %[[localM:.*]] = dim %{{.*}}, 0 +// TILE-234: %[[szM:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localM]], %[[I]]) +// TILE-234: %[[localN:.*]] = dim %{{.*}}, 1 +// TILE-234: %[[szN:.*]] = affine.min #[[bound_map]](%[[C3]], %[[localN]], %[[J]]) +// TILE-234: %[[sAij:.*]] = std.subview %{{.*}}[%[[I]], %[[J]]][%[[szM]], %[[szN]]][%[[C1]], %[[C1]]] : memref to memref +// TILE-234: %[[localN:.*]] = dim %{{.*}}, 0 +// TILE-234: %[[szN:.*]] = affine.min #[[bound_map]](%[[C3]], %[[localN]], %[[J]]) +// TILE-234: %[[sBj:.*]] = std.subview %{{.*}}[%[[J]]][%[[szN]]][%[[C1]]] : memref to memref +// TILE-234: %[[localM:.*]] = dim %{{.*}}, 0 +// TILE-234: %[[szM:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localM]], %[[I]]) +// TILE-234: %[[sCi:.*]] = std.subview %{{.*}}[%[[I]]][%[[szM]]][%[[C1]]] : memref to memref // // TILE-234: linalg.matvec(%[[sAij]], %[[sBj]], %[[sCi]]) : memref, memref, memref @@ -138,9 +252,13 @@ // TILE-2-DAG: %[[C1:.*]] = constant 1 : index // TILE-2-DAG: %[[C2:.*]] = constant 2 : index // TILE-2: %[[M:.*]] = dim %{{.*}}, 0 : memref -// TILE-2: loop.for %[[I]] = %{{.*}}{{.*}} to %[[M]] step %{{.*}} { -// TILE-2: %[[sAi:.*]] = std.subview %{{.*}}[%[[I]]][%[[C2]]][%[[C1]]] : memref to memref -// TILE-2: %[[sBi:.*]] = std.subview %{{.*}}[%[[I]]][%[[C2]]][%[[C1]]] : memref to memref +// TILE-2: loop.for %[[I:.*]] = %{{.*}}{{.*}} to %[[M]] step %{{.*}} { +// TILE-2: %[[localM:.*]] = dim %{{.*}}, 0 +// TILE-2: %[[szM:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localM]], %[[I]]) +// TILE-2: %[[sAi:.*]] = std.subview %{{.*}}[%[[I]]][%[[szM]]][%[[C1]]] : memref to memref +// TILE-2: %[[localM:.*]] = dim %{{.*}}, 0 +// TILE-2: %[[szM:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localM]], %[[I]]) +// TILE-2: %[[sBi:.*]] = std.subview %{{.*}}[%[[I]]][%[[szM]]][%[[C1]]] : memref to memref // TILE-2: linalg.dot(%[[sAi]], %[[sBi]], {{.*}}) : memref, memref, memref // TILE-02-LABEL: func @dot( @@ -155,8 +273,12 @@ // TILE-234-DAG: %[[C2:.*]] = constant 2 : index // TILE-234: %[[ubK:.*]] = dim %{{.*}}, 0 : memref // TILE-234: loop.for %[[I:.*]] = %{{.*}} to %[[ubK]] step %{{.*}} { -// TILE-234: %[[sAi:.*]] = std.subview %{{.*}}[%[[I]]][%[[C2]]][%[[C1]]] : memref to memref -// TILE-234: %[[sBi:.*]] = std.subview %{{.*}}[%[[I]]][%[[C2]]][%[[C1]]] : memref to memref +// TILE-234: %[[localM:.*]] = dim %{{.*}}, 0 +// TILE-234: %[[szM:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localM]], %[[I]]) +// TILE-234: %[[sAi:.*]] = std.subview %{{.*}}[%[[I]]][%[[szM]]][%[[C1]]] : memref to memref +// TILE-234: %[[localM:.*]] = dim %{{.*}}, 0 +// TILE-234: %[[szM:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localM]], %[[I]]) +// TILE-234: %[[sBi:.*]] = std.subview %{{.*}}[%[[I]]][%[[szM]]][%[[C1]]] : memref to memref // TILE-234: linalg.dot(%[[sAi]], %[[sBi]], %{{.*}}) : memref, memref, memref func @fill_static(%arg0: memref<127x99xf32>, %arg1: f32) { 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 @@ -4,6 +4,7 @@ // TILE-23004-DAG: #[[S0x10p90:.*]] = affine_map<()[s0] -> (s0 * 10 + 90)> // 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: #[[strided4D_dynamic:.*]] = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3, s4] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3 * s4)> +// TILE-23004-DAG: #[[bound_map:.*]] = affine_map<(d0, d1, d2) -> (d0, d1 - d2)> func @conv(%arg0: memref, %arg1: memref, %arg2: memref) { linalg.conv(%arg0, %arg1, %arg2) {dilations = [10, 20], strides = [30, 40]} : memref, memref, memref @@ -25,16 +26,20 @@ // TILE-23004: loop.for %[[ivK:.*]] = %{{.*}} to %[[Q]] step %{{.*}} { // TILE-23004: %[[Z0:.*]] = dim %{{.*}}, 0 : memref // TILE-23004: %[[Z1:.*]] = dim %{{.*}}, 1 : memref +// TILE-23004: %[[Z2:.*]] = dim %{{.*}}, 2 : memref +// TILE-23004: %[[szK:.*]] = affine.min #[[bound_map]](%[[C4]], %[[Z2]], %[[ivK]]) // TILE-23004: %[[K:.*]] = dim %{{.*}}, 3 : memref -// TILE-23004: %[[FilterView:.*]] = std.subview %{{.*}}[%[[C0]], %[[C0]], %[[ivK]], %[[C0]]][%[[Z0]], %[[Z1]], %[[C4]], %[[K]]][%[[C1]], %[[C1]], %[[C1]], %[[C1]]] : memref to memref +// TILE-23004: %[[FilterView:.*]] = std.subview %{{.*}}[%[[C0]], %[[C0]], %[[ivK]], %[[C0]]][%[[Z0]], %[[Z1]], %[[szK]], %[[K]]][%[[C1]], %[[C1]], %[[C1]], %[[C1]]] : memref to memref // // TILE-23004: %[[J1:.*]] = affine.apply #[[D0x30pS0x10]](%[[ivJ]]) // T__ILE-23004: %[[I1pStep:.*]] = affine.apply #[[S0x10p90]]()[%[[I1]]] // TILE-23004: %[[SZ2:.*]] = dim %{{.*}}, 2 : memref -// TILE-23004: %[[InputView:.*]] = std.subview %{{.*}}[%[[ivI]], %[[J1]], %[[C0]], %[[ivK]]][%[[C2]], %{{.*}}, %[[SZ2]], %[[C4]]][%[[C1]], %[[C1]], %[[C1]], %[[C1]]] : memref to memref +// TILE-23004: %[[dim3:.*]] = dim %{{.*}}, 3 +// TILE-23004: %[[sz3:.*]] = affine.min #[[bound_map]](%[[C4]], %[[dim3]], %[[ivK]] +// TILE-23004: %[[InputView:.*]] = std.subview %{{.*}}[%[[ivI]], %[[J1]], %[[C0]], %[[ivK]]][%{{.*}}, %{{.*}}, %[[SZ2]], %[[sz3]]][%[[C1]], %[[C1]], %[[C1]], %[[C1]]] : memref to memref // // TILE-23004: %[[X0:.*]] = dim %{{.*}}, 2 : memref // TILE-23004: %[[X1:.*]] = dim %{{.*}}, 3 : memref -// TILE-23004: %[[OutputView:.*]] = std.subview %{{.*}}[%[[ivI]], %[[ivJ]], %[[C0]], %[[C0]]][%[[C2]], %[[C3]], %[[X0]], %[[X1]]][%[[C1]], %[[C1]], %[[C1]], %[[C1]]] : memref to memref +// TILE-23004: %[[OutputView:.*]] = std.subview %{{.*}}[%[[ivI]], %[[ivJ]], %[[C0]], %[[C0]]][%{{.*}}, %{{.*}}, %[[X0]], %[[X1]]][%[[C1]], %[[C1]], %[[C1]], %[[C1]]] : memref to memref // // TILE-23004: linalg.conv(%[[FilterView]], %[[InputView]], %[[OutputView]]) {dilations = [10, 20], strides = [30, 40]} : memref, memref, memref