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 @@ -223,10 +223,10 @@ } static SmallVector makeTiledViews(OpBuilder &b, Location loc, - LinalgOp linalgOp, + LinalgOp linalgOp, AffineMap map, ArrayRef ivs, ArrayRef tileSizes, - ArrayRef viewSizes) { + ArrayRef allViewSizes) { assert(linalgOp.hasBufferSemantics() && "expected linalg op with buffer semantics"); assert(ivs.size() == static_cast(llvm::count_if( @@ -236,6 +236,7 @@ using namespace edsc::op; + auto viewSizes = applyMapToValues(b, loc, map, allViewSizes); // Construct (potentially temporary) mins and maxes on which to apply maps // that define tile subviews. SmallVector lbs, subViewSizes; @@ -356,7 +357,7 @@ return llvm::None; // 2. Build the tiled loop ranges. - auto viewSizes = getViewSizes(b, op); + auto allViewSizes = getViewSizes(b, op); // The flattened loopToOperandRangesMaps is expected to be an invertible // permutation map (asserted in the inverse calculation). auto mapsRange = op.indexing_maps().getAsRange(); @@ -369,7 +370,7 @@ SmallVector loopRanges; LoopIndexToRangeIndexMap loopIndexToRangeIndex; std::tie(loopRanges, loopIndexToRangeIndex) = makeTiledLoopRanges( - b, scope.getLocation(), viewSizesToLoopsMap, viewSizes, tileSizes); + b, scope.getLocation(), viewSizesToLoopsMap, allViewSizes, tileSizes); if (!options.interchangeVector.empty()) applyPermutationToVector(loopRanges, options.interchangeVector); @@ -395,7 +396,8 @@ if (!options.interchangeVector.empty()) ivValues = applyMapToValues(b, loc, invPermutationMap, ivValues); - auto views = makeTiledViews(b, loc, op, ivValues, tileSizes, viewSizes); + auto views = makeTiledViews(b, loc, op, viewSizesToLoopsMap, ivValues, + tileSizes, allViewSizes); auto operands = getAssumedNonViewOperands(op); views.append(operands.begin(), operands.end()); res = op.clone(b, loc, views); 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 @@ -9,36 +9,38 @@ linalg.conv(%arg0, %arg1, %arg2) {dilations = [10, 20], strides = [30, 40]} : memref, memref, memref return } -// TILE-23004-LABEL: func @conv( -// TILE-23004: %{{.*}}: memref, %{{.*}}: memref, %{{.*}}: memref) { -// TILE-23004-DAG: %[[C0:.*]] = constant 0 : index -// TILE-23004-DAG: %[[C2:.*]] = constant 2 : index -// TILE-23004-DAG: %[[C3:.*]] = constant 3 : index -// TILE-23004-DAG: %[[C4:.*]] = constant 4 : index -// TILE-23004: %[[Q:.*]] = dim %{{.*}}, %c2 : memref -// TILE-23004: %[[B:.*]] = dim %{{.*}}, %c0 : memref -// TILE-23004: %[[PaddedInput0:.*]] = dim %{{.*}}, %c1 : memref -// TILE-23004: %[[X0:.*]] = dim %{{.*}}, %c1 : memref +// TILE-23004: func @conv( +// TILE-23004-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref +// 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: %[[C2:.*]] = constant 2 : index +// TILE-23004-DAG: %[[C3:.*]] = constant 3 : index +// TILE-23004-DAG: %[[C4:.*]] = constant 4 : index +// TILE-23004: %[[Z0:.*]] = dim %[[ARG0]], %c0 : memref +// TILE-23004: %[[Q:.*]] = dim %[[ARG0]], %c2 : memref +// TILE-23004: %[[B:.*]] = dim %[[ARG1]], %c0 : memref +// TILE-23004: %[[X0:.*]] = dim %[[ARG2]], %c1 : memref // 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:.*]] = dim %{{.*}}, %c0 : memref -// TILE-23004: %[[Z1:.*]] = dim %{{.*}}, %c1 : memref -// TILE-23004: %[[Z2:.*]] = dim %{{.*}}, %c2 : memref +// 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: %[[K:.*]] = dim %{{.*}}, %c3 : memref -// TILE-23004: %[[FilterView:.*]] = subview %{{.*}}[0, 0, %[[ivK]], 0] [%[[Z0]], %[[Z1]], %[[szK]], %[[K]]] [1, 1, 1, 1] : memref to memref +// 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 %{{.*}}, %c1 : memref -// TILE-23004: %[[I1pStep:.*]] = affine.min #[[$S0x10p90D0x30pS1]](%[[ivJ]])[%[[PaddedInput0]], %[[PaddedInput0b]]] -// TILE-23004: %[[SZ2:.*]] = dim %{{.*}}, %c2 : memref -// TILE-23004: %[[dim3:.*]] = dim %{{.*}}, %c3 +// TILE-23004: %[[PaddedInput0b:.*]] = dim %[[ARG1]], %c1 : memref +// TILE-23004: %[[I1pStep:.*]] = affine.min #[[$S0x10p90D0x30pS1]](%[[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: %[[InputView:.*]] = subview %{{.*}}[%[[ivI]], %[[J1]], 0, %[[ivK]]] [%{{.*}}, %{{.*}}, %[[SZ2]], %[[sz3]]] [1, 1, 1, 1] : memref to memref // -// TILE-23004: %[[X0:.*]] = dim %{{.*}}, %c2 : memref -// TILE-23004: %[[X1:.*]] = dim %{{.*}}, %c3 : memref +// TILE-23004: %[[X0:.*]] = dim %[[ARG2]], %c2 : memref +// TILE-23004: %[[X1:.*]] = dim %[[ARG2]], %c3 : memref // TILE-23004: %[[OutputView:.*]] = subview %{{.*}}[%[[ivI]], %[[ivJ]], 0, 0] [%{{.*}}, %{{.*}}, %[[X0]], %[[X1]]] [1, 1, 1, 1] : memref to memref // // TILE-23004: linalg.conv(%[[FilterView]], %[[InputView]], %[[OutputView]]) {dilations = [10, 20], strides = [30, 40]} : memref, memref, memref diff --git a/mlir/test/Dialect/Linalg/tile_simple_conv.mlir b/mlir/test/Dialect/Linalg/tile_simple_conv.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/Linalg/tile_simple_conv.mlir @@ -0,0 +1,49 @@ +// 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)> + +func @conv(%arg0 : memref, %arg1 : memref, %arg2 : memref) { + linalg.conv(%arg0, %arg1, %arg2) : memref, memref, memref + return +} + +// CHECK: func @conv +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref +// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: memref +// CHECK-DAG: %[[C0:.*]] = constant 0 : index +// CHECK-DAG: %[[C1:.*]] = constant 1 : index +// 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: %[[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: %[[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