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 @@ -243,7 +243,9 @@ 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]); + // Before composing, we need to make range a closed interval. + Value size = isTiled ? tileSizes[idx] : viewSizes[idx]; + subViewSizes.push_back(std_subi(size, std_constant_index(1))); } auto *op = linalgOp.getOperation(); @@ -282,7 +284,9 @@ auto m = map.getSubMap({r}); auto offset = applyMapToValues(b, loc, m, lbs).front(); offsets.push_back(offset); - auto size = applyMapToValues(b, loc, m, subViewSizes).front(); + auto closed_int_size = applyMapToValues(b, loc, m, subViewSizes).front(); + // Resulting size needs to be made half open interval again. + auto size = (Value)std_addi(closed_int_size, std_constant_index(1)); // The size of the subview should be trimmed to avoid out-of-bounds // accesses, unless we statically know the subview size divides the view 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,10 @@ // 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) -> (d0 * 10 + 60)> // 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: #[[$step_map:.*]] = affine_map<(d0, d1)[s0] -> (d0, d1 * -30 + s0)> func @conv(%arg0: memref, %arg1: memref, %arg2: memref) { linalg.conv(%arg0, %arg1, %arg2) {dilations = [10, 20], strides = [30, 40]} : memref, memref, memref @@ -14,6 +15,7 @@ // TILE-23004-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref // TILE-23004-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: memref) // TILE-23004-DAG: %[[C0:.*]] = constant 0 : index +// TILE-23004-DAG: %[[C1:.*]] = constant 1 : index // TILE-23004-DAG: %[[C2:.*]] = constant 2 : index // TILE-23004-DAG: %[[C3:.*]] = constant 3 : index // TILE-23004-DAG: %[[C4:.*]] = constant 4 : index @@ -24,6 +26,7 @@ // 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: %[[Z0_minus_1:.*]] = subi %[[Z0]], %[[C1]] : index // TILE-23004: %[[Z0_1:.*]] = dim %[[ARG0]], %c0 : memref // TILE-23004: %[[Z1:.*]] = dim %[[ARG0]], %c1 : memref // TILE-23004: %[[Z2:.*]] = dim %[[ARG0]], %c2 : memref @@ -32,8 +35,10 @@ // 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: %[[closed_int_size:.*]] = affine.apply #[[$S0x10p90D0x30pS1]](%[[Z0_minus_1]]) +// TILE-23004: %[[size:.*]] = addi %[[closed_int_size]], %[[C1]] : index // TILE-23004: %[[PaddedInput0b:.*]] = dim %[[ARG1]], %c1 : memref -// TILE-23004: %[[I1pStep:.*]] = affine.min #[[$S0x10p90D0x30pS1]](%[[ivJ]])[%[[Z0]], %[[PaddedInput0b]]] +// TILE-23004: %[[I1pStep:.*]] = affine.min #[[$step_map]](%[[size]], %[[ivJ]])[%[[PaddedInput0b]]] // TILE-23004: %[[SZ2:.*]] = dim %[[ARG1]], %c2 : memref // TILE-23004: %[[dim3:.*]] = dim %[[ARG1]], %c3 // TILE-23004: %[[sz3:.*]] = affine.min #[[$bound_map_4]](%[[ivK]])[%[[dim3]]] 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,12 @@ // 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: #[[MAP1:.*]] = affine_map<(d0) -> (d0 + 2)> +// CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0, d1)[s0] -> (d0, -d1 + s0)> +// CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0) -> (d0 + 3)> +// CHECK-DAG: #[[MAP4:.*]] = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3)> +// CHECK-DAG: #[[MAP5:.*]] = affine_map<(d0)[s0] -> (3, -d0 + s0)> +// CHECK-DAG: #[[MAP6:.*]] = affine_map<(d0)[s0] -> (4, -d0 + s0)> func @conv(%arg0 : memref, %arg1 : memref, %arg2 : memref) { linalg.conv(%arg0, %arg1, %arg2) : memref, memref, memref @@ -28,22 +30,28 @@ // 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: %[[T5:.*]] = subi %[[T0]], %[[C1]] : index +// CHECK: %[[T6:.*]] = subi %[[T1]], %[[C1]] : index +// CHECK: %[[T7:.*]] = dim %[[ARG1]], %[[C0]] +// CHECK: %[[T8:.*]] = affine.min #[[MAP0]](%[[ARG3]])[%[[T7]]] +// CHECK: %[[T9:.*]] = affine.apply #[[MAP1]](%[[T5]]) +// CHECK: %[[T10:.*]] = addi %[[T9]], %[[C1]] : index +// CHECK: %[[T11:.*]] = dim %[[ARG1]], %[[C1]] +// CHECK: %[[T12:.*]] = affine.min #[[MAP2]](%[[T10]], %[[ARG4]])[%[[T11]]] +// CHECK: %[[T13:.*]] = affine.apply #[[MAP3]](%[[T6]]) +// CHECK: %[[T14:.*]] = addi %[[T13]], %[[C1]] : index +// CHECK: %[[T15:.*]] = dim %[[ARG1]], %[[C2]] +// CHECK: %[[T16:.*]] = affine.min #[[MAP2]](%[[T14]], %[[ARG5]])[%[[T15]]] +// CHECK: %[[T17:.*]] = 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: [%[[T8]], %[[T12]], %[[T16]], %[[T17]]] +// CHECK: %[[T18:.*]] = dim %[[ARG2]], %[[C0]] +// CHECK: %[[T19:.*]] = affine.min #[[MAP0]](%[[ARG3]])[%[[T18]]] +// CHECK: %[[T20:.*]] = dim %[[ARG2]], %[[C1]] +// CHECK: %[[T21:.*]] = affine.min #[[MAP5]](%[[ARG4]])[%[[T20]]] +// CHECK: %[[T22:.*]] = dim %[[ARG2]], %[[C2]] +// CHECK: %[[T23:.*]] = affine.min #[[MAP6]](%[[ARG5]])[%[[T22]]] +// CHECK: %[[T24:.*]] = dim %[[ARG2]], %[[C3]] // CHECK: %[[SV2:.*]] = subview %[[ARG2]][%[[ARG3]], %[[ARG4]], %[[ARG5]], 0] -// CHECK-SAME: [%[[T14]], %[[T16]], %[[T18]], %[[T19]]] -// CHECK: linalg.conv(%[[ARG0]], %[[SV1]], %[[SV2]]) \ No newline at end of file +// CHECK-SAME: [%[[T19]], %[[T21]], %[[T23]], %[[T24]]] +// CHECK: linalg.conv(%[[ARG0]], %[[SV1]], %[[SV2]])