From 4ff48db68dbb209ccdb190d1661832bb2a2d47fd Mon Sep 17 00:00:00 2001 From: MaheshRavishankar Date: Thu, 23 Jul 2020 10:20:54 -0700 Subject: [PATCH] [mlir][Linalg] Fixing bug in subview size computation in Linalg tiling. The `makeTiledViews` did not use the sizes of the tiled views based on the result of the loop bound inference computation. This manifested as an error in computing tile sizes with convolution where not all the result expression of concatenated affine maps are simple AffineDimExpr. Differential Revision: https://reviews.llvm.org/D84366 --- mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp | 12 +++-- mlir/test/Dialect/Linalg/tile_conv.mlir | 44 +++++++++-------- .../test/Dialect/Linalg/tile_simple_conv.mlir | 49 +++++++++++++++++++ 3 files changed, 79 insertions(+), 26 deletions(-) create mode 100644 mlir/test/Dialect/Linalg/tile_simple_conv.mlir diff --git a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp index 6ffc181fed67..2c327276610a 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp @@ -223,10 +223,10 @@ static bool isTiled(AffineMap map, ArrayRef tileSizes) { } 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 @@ static SmallVector makeTiledViews(OpBuilder &b, Location loc, 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 @@ Optional static tileLinalgOpImpl( 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 @@ Optional static tileLinalgOpImpl( 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 @@ Optional static tileLinalgOpImpl( 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 index 1bbb8b60382e..a08a2f1e585c 100644 --- a/mlir/test/Dialect/Linalg/tile_conv.mlir +++ b/mlir/test/Dialect/Linalg/tile_conv.mlir @@ -9,36 +9,38 @@ func @conv(%arg0: memref, %arg1: 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 index 000000000000..f854f7570fef --- /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