[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
This commit is contained in:
MaheshRavishankar
parent
918ba93367
commit
4ff48db68d
@ -223,10 +223,10 @@ static bool isTiled(AffineMap map, ArrayRef<Value> tileSizes) {
|
||||
}
|
||||
|
||||
static SmallVector<Value, 4> makeTiledViews(OpBuilder &b, Location loc,
|
||||
LinalgOp linalgOp,
|
||||
LinalgOp linalgOp, AffineMap map,
|
||||
ArrayRef<Value> ivs,
|
||||
ArrayRef<Value> tileSizes,
|
||||
ArrayRef<Value> viewSizes) {
|
||||
ArrayRef<Value> allViewSizes) {
|
||||
assert(linalgOp.hasBufferSemantics() &&
|
||||
"expected linalg op with buffer semantics");
|
||||
assert(ivs.size() == static_cast<size_t>(llvm::count_if(
|
||||
@ -236,6 +236,7 @@ static SmallVector<Value, 4> 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<Value, 8> lbs, subViewSizes;
|
||||
@ -356,7 +357,7 @@ Optional<TiledLinalgOp> 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<AffineMapAttr>();
|
||||
@ -369,7 +370,7 @@ Optional<TiledLinalgOp> static tileLinalgOpImpl(
|
||||
SmallVector<SubViewOp::Range, 4> 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<TiledLinalgOp> 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);
|
||||
|
||||
@ -9,36 +9,38 @@ func @conv(%arg0: memref<?x?x?x?xf32, offset: ?, strides: [?, ?, ?, 1]>, %arg1:
|
||||
linalg.conv(%arg0, %arg1, %arg2) {dilations = [10, 20], strides = [30, 40]} : memref<?x?x?x?xf32, offset: ?, strides: [?, ?, ?, 1]>, memref<?x?x?x?xf32, offset: ?, strides: [?, ?, ?, 1]>, memref<?x?x?x?xf32, offset: ?, strides: [?, ?, ?, 1]>
|
||||
return
|
||||
}
|
||||
// TILE-23004-LABEL: func @conv(
|
||||
// TILE-23004: %{{.*}}: memref<?x?x?x?xf32, #[[$strided4D]]>, %{{.*}}: memref<?x?x?x?xf32, #[[$strided4D]]>, %{{.*}}: memref<?x?x?x?xf32, #[[$strided4D]]>) {
|
||||
// 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<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[B:.*]] = dim %{{.*}}, %c0 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[PaddedInput0:.*]] = dim %{{.*}}, %c1 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[X0:.*]] = dim %{{.*}}, %c1 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: func @conv(
|
||||
// TILE-23004-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: memref<?x?x?x?xf32, #[[$strided4D]]>)
|
||||
// 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<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[Q:.*]] = dim %[[ARG0]], %c2 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[B:.*]] = dim %[[ARG1]], %c0 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[X0:.*]] = dim %[[ARG2]], %c1 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// 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<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[Z1:.*]] = dim %{{.*}}, %c1 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[Z2:.*]] = dim %{{.*}}, %c2 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[Z0_1:.*]] = dim %[[ARG0]], %c0 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[Z1:.*]] = dim %[[ARG0]], %c1 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[Z2:.*]] = dim %[[ARG0]], %c2 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[szK:.*]] = affine.min #[[$bound_map_4]](%[[ivK]])[%[[Z2]]]
|
||||
// TILE-23004: %[[K:.*]] = dim %{{.*}}, %c3 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[FilterView:.*]] = subview %{{.*}}[0, 0, %[[ivK]], 0] [%[[Z0]], %[[Z1]], %[[szK]], %[[K]]] [1, 1, 1, 1] : memref<?x?x?x?xf32, #[[$strided4D]]> to memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[K:.*]] = dim %[[ARG0]], %c3 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[FilterView:.*]] = subview %{{.*}}[0, 0, %[[ivK]], 0] [%[[Z0_1]], %[[Z1]], %[[szK]], %[[K]]] [1, 1, 1, 1] : memref<?x?x?x?xf32, #[[$strided4D]]> to memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
//
|
||||
// TILE-23004: %[[J1:.*]] = affine.apply #[[$D0x30pS0x10]](%[[ivJ]])
|
||||
// TILE-23004: %[[PaddedInput0b:.*]] = dim %{{.*}}, %c1 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[I1pStep:.*]] = affine.min #[[$S0x10p90D0x30pS1]](%[[ivJ]])[%[[PaddedInput0]], %[[PaddedInput0b]]]
|
||||
// TILE-23004: %[[SZ2:.*]] = dim %{{.*}}, %c2 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[dim3:.*]] = dim %{{.*}}, %c3
|
||||
// TILE-23004: %[[PaddedInput0b:.*]] = dim %[[ARG1]], %c1 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[I1pStep:.*]] = affine.min #[[$S0x10p90D0x30pS1]](%[[ivJ]])[%[[Z0]], %[[PaddedInput0b]]]
|
||||
// TILE-23004: %[[SZ2:.*]] = dim %[[ARG1]], %c2 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// 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<?x?x?x?xf32, #[[$strided4D]]> to memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
//
|
||||
// TILE-23004: %[[X0:.*]] = dim %{{.*}}, %c2 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[X1:.*]] = dim %{{.*}}, %c3 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[X0:.*]] = dim %[[ARG2]], %c2 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[X1:.*]] = dim %[[ARG2]], %c3 : memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
// TILE-23004: %[[OutputView:.*]] = subview %{{.*}}[%[[ivI]], %[[ivJ]], 0, 0] [%{{.*}}, %{{.*}}, %[[X0]], %[[X1]]] [1, 1, 1, 1] : memref<?x?x?x?xf32, #[[$strided4D]]> to memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
//
|
||||
// TILE-23004: linalg.conv(%[[FilterView]], %[[InputView]], %[[OutputView]]) {dilations = [10, 20], strides = [30, 40]} : memref<?x?x?x?xf32, #[[$strided4D]]>, memref<?x?x?x?xf32, #[[$strided4D]]>, memref<?x?x?x?xf32, #[[$strided4D]]>
|
||||
|
||||
49
mlir/test/Dialect/Linalg/tile_simple_conv.mlir
Normal file
49
mlir/test/Dialect/Linalg/tile_simple_conv.mlir
Normal file
@ -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<?x?x?x?xf32>, %arg1 : memref<?x?x?x?xf32>, %arg2 : memref<?x?x?x?xf32>) {
|
||||
linalg.conv(%arg0, %arg1, %arg2) : memref<?x?x?x?xf32>, memref<?x?x?x?xf32>, memref<?x?x?x?xf32>
|
||||
return
|
||||
}
|
||||
|
||||
// CHECK: func @conv
|
||||
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<?x?x?x?xf32>
|
||||
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<?x?x?x?xf32>
|
||||
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: memref<?x?x?x?xf32>
|
||||
// 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]])
|
||||
Loading…
x
Reference in New Issue
Block a user