Skip to content

Commit 0a0fa15

Browse files
d-smirnovd-smirnov
committed
bug fixes
1 parent c262828 commit 0a0fa15

File tree

3 files changed

+43
-45
lines changed

3 files changed

+43
-45
lines changed

‎mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp‎

Lines changed: 20 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -3160,11 +3160,6 @@ WinogradInputTransformOp::getTiledImplementation(OpBuilder &builder,
31603160
ArrayRef<OpFoldResult> offsets,
31613161
ArrayRef<OpFoldResult> sizes) {
31623162
IntegerAttr oneAttr = builder.getI64IntegerAttr(1);
3163-
IntegerAttr zeroAttr = builder.getI64IntegerAttr(0);
3164-
ShapedType inputType = getInputOperandType();
3165-
ArrayRef<int64_t> inputShape = inputType.getShape();
3166-
int64_t inputH = inputShape[getInputHDim()];
3167-
int64_t inputW = inputShape[getInputWDim()];
31683163
int64_t m = getM();
31693164
int64_t r = getR();
31703165

@@ -3175,12 +3170,16 @@ WinogradInputTransformOp::getTiledImplementation(OpBuilder &builder,
31753170

31763171
Location loc = getLoc();
31773172
MLIRContext *context = builder.getContext();
3173+
auto identityAffineMap =
3174+
AffineMap::get(1, 0, {builder.getAffineDimExpr(0)}, context);
31783175
auto offsetAffineMap =
31793176
AffineMap::get(1, 0, {builder.getAffineDimExpr(0) * m}, context);
31803177
Value mappedOffsetH = affine::makeComposedAffineApply(
3181-
builder, loc, offsetAffineMap, offsets[getOutputTileHDim()]);
3178+
builder, loc, (alphaH != 1 ? offsetAffineMap : identityAffineMap),
3179+
offsets[getOutputTileHDim()]);
31823180
Value mappedOffsetW = affine::makeComposedAffineApply(
3183-
builder, loc, offsetAffineMap, offsets[getOutputTileWDim()]);
3181+
builder, loc, (alphaW != 1 ? offsetAffineMap : identityAffineMap),
3182+
offsets[getOutputTileWDim()]);
31843183
auto sizeAffineMap = AffineMap::get(
31853184
1, 0, {builder.getAffineDimExpr(0) * m + (r - 1)}, context);
31863185
Value mappedSizeH = affine::makeComposedAffineApply(
@@ -3191,10 +3190,8 @@ WinogradInputTransformOp::getTiledImplementation(OpBuilder &builder,
31913190
SmallVector<Value> tiledOperands;
31923191
SmallVector<OpFoldResult> sliceOffsets, sliceSizes;
31933192

3194-
OpFoldResult offsetH =
3195-
inputH != 1 ? OpFoldResult(mappedOffsetH) : OpFoldResult(zeroAttr);
3196-
OpFoldResult offsetW =
3197-
inputW != 1 ? OpFoldResult(mappedOffsetW) : OpFoldResult(zeroAttr);
3193+
OpFoldResult offsetH = OpFoldResult(mappedOffsetH);
3194+
OpFoldResult offsetW = OpFoldResult(mappedOffsetW);
31983195
sliceOffsets.append(
31993196
{offsets[getOutputNDim()], offsetH, offsetW, offsets[getOutputCDim()]});
32003197
OpFoldResult sizeH =
@@ -3308,28 +3305,29 @@ LogicalResult WinogradOutputTransformOp::getResultTilePosition(
33083305

33093306
Location loc = getLoc();
33103307
MLIRContext *context = builder.getContext();
3308+
auto identityAffineMap =
3309+
AffineMap::get(1, 0, {builder.getAffineDimExpr(0)}, context);
33113310
auto affineMap =
33123311
AffineMap::get(1, 0, {builder.getAffineDimExpr(0) * m}, context);
33133312

3313+
ShapedType valueType = getValueOperandType();
3314+
ArrayRef<int64_t> valueShape = valueType.getShape();
3315+
int64_t valueH = valueShape[0];
3316+
int64_t valueW = valueShape[1];
33143317
Value mappedOffsetH = affine::makeComposedAffineApply(
3315-
builder, loc, affineMap, offsets[getValueTileHDim()]);
3318+
builder, loc, (valueH != 1 ? affineMap : identityAffineMap),
3319+
offsets[getValueTileHDim()]);
33163320
Value mappedOffsetW = affine::makeComposedAffineApply(
3317-
builder, loc, affineMap, offsets[getValueTileWDim()]);
3321+
builder, loc, (valueW != 1 ? affineMap : identityAffineMap),
3322+
offsets[getValueTileWDim()]);
33183323
Value mappedSizeH = affine::makeComposedAffineApply(
33193324
builder, loc, affineMap, sizes[getValueTileHDim()]);
33203325
Value mappedSizeW = affine::makeComposedAffineApply(
33213326
builder, loc, affineMap, sizes[getValueTileWDim()]);
33223327

3323-
ShapedType valueType = getValueOperandType();
3324-
ArrayRef<int64_t> valueShape = valueType.getShape();
3325-
int64_t valueH = valueShape[0];
3326-
int64_t valueW = valueShape[1];
33273328
IntegerAttr oneAttr = builder.getI64IntegerAttr(1);
3328-
IntegerAttr zeroAttr = builder.getI64IntegerAttr(0);
3329-
OpFoldResult offsetH =
3330-
valueH != 1 ? OpFoldResult(mappedOffsetH) : OpFoldResult(zeroAttr);
3331-
OpFoldResult offsetW =
3332-
valueW != 1 ? OpFoldResult(mappedOffsetW) : OpFoldResult(zeroAttr);
3329+
OpFoldResult offsetH = OpFoldResult(mappedOffsetH);
3330+
OpFoldResult offsetW = OpFoldResult(mappedOffsetW);
33333331
OpFoldResult sizeH =
33343332
valueH != 1 ? OpFoldResult(mappedSizeH) : OpFoldResult(oneAttr);
33353333
OpFoldResult sizeW =

‎mlir/test/Dialect/Linalg/transform-tile-and-winograd-rewrite.mlir‎

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -379,10 +379,9 @@ module attributes {transform.with_named_sequence} {
379379
}
380380
}
381381

382-
// CHECK: #[[$MAP0:.+]] = affine_map<(d0) -> (d0 * 4)>
383-
// CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1) -> ()>
384-
// CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1) -> (d0, d1)>
385-
// CHECK-LABEL: func.func @conv2d_mx1_rx1
382+
// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1) -> ()>
383+
// CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)>
384+
// CHECK-LABEL: func.func @conv2d_mx1_rx1_2
386385
// CHECK-SAME: (%[[ARG0:.*]]: tensor<2x6x2x5xf32>, %[[ARG1:.*]]: tensor<2x3x1x5xf32>, %[[ARG2:.*]]: tensor<2x4x2x2xf32>) -> tensor<2x4x2x2xf32> {
387386
// CHECK: %[[CST:.*]] = arith.constant 3.200000e+01 : f32
388387
// CHECK: %[[CST_0:.*]] = arith.constant dense<{{.*}}> : tensor<4x6xf32>
@@ -405,8 +404,7 @@ module attributes {transform.with_named_sequence} {
405404
// CHECK: scf.yield %[[S7]]
406405
// CHECK: %[[S2:.*]] = tensor.empty() : tensor<6x1x1x2x2x5xf32>
407406
// CHECK: %[[S3:.*]] = scf.for %[[ARG3:.*]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG4:.*]] = %[[S2]])
408-
// CHECK: %[[S8:.*]] = affine.apply #[[$MAP0]](%[[ARG3]])
409-
// CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[ARG0]][0, 0, %8, 0] [2, 6, 1, 5] [1, 1, 1, 1]
407+
// CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[ARG0]][0, 0, %[[ARG3]], 0] [2, 6, 1, 5] [1, 1, 1, 1]
410408
// CHECK: %[[EXTRACTED_SLICE_5:.*]] = tensor.extract_slice %[[ARG4]][0, 0, 0, %[[ARG3]], 0, 0] [6, 1, 1, 1, 2, 5] [1, 1, 1, 1, 1, 1]
411409
// CHECK: %[[S9:.*]] = scf.for %[[ARG5:.*]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG6:.*]] = %[[EXTRACTED_SLICE_5]])
412410
// CHECK: %[[S10:.*]] = scf.for %[[ARG7:.*]] = %[[C0]] to %[[C5]] step %[[C1]] iter_args(%[[ARG8:.*]] = %[[ARG6]])
@@ -427,15 +425,15 @@ module attributes {transform.with_named_sequence} {
427425
// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[S6]] {{\[}}[0, 1], [2, 3, 4], [5]] output_shape [6, 1, 1, 2, 2, 2]
428426
// CHECK: %[[S7:.*]] = scf.for %[[ARG3:.*]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG4:.*]] = %[[ARG2]])
429427
// CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[EXPANDED]][0, 0, 0, %[[ARG3]], 0, 0] [6, 1, 1, 1, 2, 2] [1, 1, 1, 1, 1, 1]
430-
// CHECK: %[[EXTRACTED_SLICE_5:.*]] = tensor.extract_slice %[[ARG4]][0, 0, 0, 0] [2, 4, 1, 2] [1, 1, 1, 1]
428+
// CHECK: %[[EXTRACTED_SLICE_5:.*]] = tensor.extract_slice %[[ARG4]][0, 0, %[[ARG3]], 0] [2, 4, 1, 2] [1, 1, 1, 1]
431429
// CHECK: %[[S8:.*]] = scf.for %[[ARG5:.*]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG6:.*]] = %[[EXTRACTED_SLICE_5]])
432430
// CHECK: %[[S9:.*]] = scf.for %[[ARG7:.*]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG8:.*]] = %[[ARG6]])
433431
// CHECK: %[[EXTRACTED_SLICE_6:.*]] = tensor.extract_slice %[[EXTRACTED_SLICE]][0, 0, 0, 0, %[[ARG5]], %[[ARG7]]] [6, 1, 1, 1, 1, 1] [1, 1, 1, 1, 1, 1]
434432
// CHECK: %[[EXTRACTED_SLICE_7:.*]] = tensor.extract_slice %[[ARG8]][%[[ARG5]], 0, 0, %[[ARG7]]] [1, 4, 1, 1] [1, 1, 1, 1]
435433
// CHECK: %[[S10:.*]] = tensor.empty() : tensor<4x1xf32>
436434
// CHECK: %[[S11:.*]] = linalg.fill ins(%[[CST_3]] : f32) outs(%[[S10]] : tensor<4x1xf32>) -> tensor<4x1xf32>
437435
// CHECK: %[[S12:.*]] = linalg.matmul ins(%[[CST_0]], %[[EXTRACTED_SLICE_6]] : tensor<4x6xf32>, tensor<6x1xf32>) outs(%[[S11]] : tensor<4x1xf32>) -> tensor<4x1xf32>
438-
// CHECK: %[[S13:.*]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel"]} ins(%[[CST]], %[[S12]] : f32, tensor<4x1xf32>) outs(%[[EXTRACTED_SLICE_7]] : tensor<4x1xf32>) {
436+
// CHECK: %[[S13:.*]] = linalg.generic {indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel"]} ins(%[[CST]], %[[S12]] : f32, tensor<4x1xf32>) outs(%[[EXTRACTED_SLICE_7]] : tensor<4x1xf32>) {
439437
// CHECK: ^bb0(%[[IN1:.*]]: f32, %[[IN2:.*]]: f32, %[[OUT:.*]]: f32):
440438
// CHECK: %[[VAL_57:.*]] = arith.mulf %[[IN1]], %[[IN2]] : f32
441439
// CHECK: %[[VAL_58:.*]] = arith.addf %[[VAL_57]], %[[OUT]] : f32
@@ -444,6 +442,6 @@ module attributes {transform.with_named_sequence} {
444442
// CHECK: %[[INSERTED_SLICE_8:.*]] = tensor.insert_slice %[[S13]] into %[[ARG8]][%[[ARG5]], 0, 0, %[[ARG7]]] [1, 4, 1, 1] [1, 1, 1, 1]
445443
// CHECK: scf.yield %[[INSERTED_SLICE_8]]
446444
// CHECK: scf.yield %[[S9]]
447-
// CHECK: %[[INSERTED_SLICE:.*]] = tensor.insert_slice %[[S8]] into %[[ARG4]][0, 0, 0, 0] [2, 4, 1, 2] [1, 1, 1, 1]
445+
// CHECK: %[[INSERTED_SLICE:.*]] = tensor.insert_slice %[[S8]] into %[[ARG4]][0, 0, %[[ARG3]], 0] [2, 4, 1, 2] [1, 1, 1, 1]
448446
// CHECK: scf.yield %[[INSERTED_SLICE]]
449447
// CHECK: return %[[S7]]

‎mlir/test/Dialect/Linalg/transform-tile-winograd.mlir‎

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -231,8 +231,9 @@ module attributes {transform.with_named_sequence} {
231231
}
232232
}
233233

234-
// CHECK: #[[$MAP0:.+]] = affine_map<(d0) -> (d0 * 4)>
235-
// CHECK: #[[$MAP1:.+]] = affine_map<() -> (6)>
234+
// CHECK: #[[$MAP:.+]] = affine_map<(d0) -> (d0)>
235+
// CHECK: #[[$MAP1:.+]] = affine_map<(d0) -> (d0 * 4)>
236+
// CHECK: #[[$MAP2:.+]] = affine_map<() -> (6)>
236237
// CHECK-LABEL: func.func @tile_winograd_input(
237238
// CHECK-SAME: %[[ARG0:.*]]: tensor<2x1x10x5xf32>, %[[ARG1:.*]]: tensor<1x6x1x2x2x5xf32>) -> tensor<1x6x1x2x2x5xf32> {
238239
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
@@ -251,11 +252,11 @@ module attributes {transform.with_named_sequence} {
251252
// CHECK: %[[S2:.*]] = scf.for %[[ARG4:.*]] = %[[C0_1]] to %[[C2]] step %[[C1_2]] iter_args(%[[ARG5:.*]] = %[[ARG3]])
252253
// CHECK: %[[S3:.*]] = scf.for %[[ARG6:.*]] = %[[C0_3]] to %[[C2_4]] step %[[C1_5]] iter_args(%[[ARG7:.*]] = %[[ARG5]])
253254
// CHECK: %[[S4:.*]] = scf.for %[[ARG8:.*]] = %[[C0_6]] to %[[C5]] step %[[C1_7]] iter_args(%[[ARG9:.*]] = %[[ARG7]])
254-
// CHECK: %[[S5:.*]] = affine.apply #[[$MAP0]](%[[ARG2]])
255-
// CHECK: %[[S6:.*]] = affine.apply #[[$MAP0]](%[[ARG4]])
256-
// CHECK: %[[S7:.*]] = affine.apply #[[$MAP1]]()
257-
// CHECK: %[[S8:.*]] = affine.apply #[[$MAP1]]()
258-
// CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[ARG0]][%[[ARG6]], 0, %[[S6]], %[[ARG8]]] [1, 1, %[[S8]], 1] [1, 1, 1, 1] : tensor<2x1x10x5xf32> to tensor<1x1x?x1xf32>
255+
// CHECK: %[[S5:.*]] = affine.apply #[[$MAP]](%[[ARG2]])
256+
// CHECK: %[[S6:.*]] = affine.apply #[[$MAP1]](%[[ARG4]])
257+
// CHECK: %[[S7:.*]] = affine.apply #[[$MAP2]]()
258+
// CHECK: %[[S8:.*]] = affine.apply #[[$MAP2]]()
259+
// CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[ARG0]][%[[ARG6]], %[[S5]], %[[S6]], %[[ARG8]]] [1, 1, %[[S8]], 1] [1, 1, 1, 1] : tensor<2x1x10x5xf32> to tensor<1x1x?x1xf32>
259260
// CHECK: %[[EXTRACTED_SLICE_10:.*]] = tensor.extract_slice %[[ARG9]][0, 0, %[[ARG2]], %[[ARG4]], %[[ARG6]], %[[ARG8]]] [1, 6, 1, 1, 1, 1] [1, 1, 1, 1, 1, 1] : tensor<1x6x1x2x2x5xf32> to tensor<1x6x1x1x1x1xf32>
260261
// CHECK: %[[S9:.*]] = linalg.winograd_input_transform m(4) r(3) ins(%[[EXTRACTED_SLICE]] : tensor<1x1x?x1xf32>) outs(%[[EXTRACTED_SLICE_10]] : tensor<1x6x1x1x1x1xf32>) -> tensor<1x6x1x1x1x1xf32>
261262
// CHECK: %[[INSERTED_SLICE:.*]] = tensor.insert_slice %[[S9]] into %[[ARG9]][0, 0, %[[ARG2]], %[[ARG4]], %[[ARG6]], %[[ARG8]]] [1, 6, 1, 1, 1, 1] [1, 1, 1, 1, 1, 1] : tensor<1x6x1x1x1x1xf32> into tensor<1x6x1x2x2x5xf32>
@@ -357,8 +358,9 @@ module attributes {transform.with_named_sequence} {
357358
}
358359
}
359360

360-
// CHECK: #[[$MAP0:.+]] = affine_map<(d0) -> (d0 * 4)>
361-
// CHECK: #[[$MAP1:.+]] = affine_map<() -> (4)>
361+
// CHECK: #[[$MAP:.+]] = affine_map<(d0) -> (d0 * 4)>
362+
// CHECK: #[[$MAP1:.+]] = affine_map<(d0) -> (d0)>
363+
// CHECK: #[[$MAP2:.+]] = affine_map<() -> (4)>
362364
// CHECK-LABEL: func.func @tile_winograd_output(
363365
// CHECK-SAME: %[[ARG0:.*]]: tensor<6x1x2x1x3x5xf32>, %[[ARG1:.*]]: tensor<3x8x1x5xf32>) -> tensor<3x8x1x5xf32> {
364366
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
@@ -378,9 +380,9 @@ module attributes {transform.with_named_sequence} {
378380
// CHECK: %[[S3:.*]] = scf.for %[[ARG6:.*]] = %[[C0_3]] to %[[C3]] step %[[C1_4]] iter_args(%[[ARG11:.*]] = %[[ARG10]]) -> (tensor<3x8x1x5xf32>)
379381
// CHECK: %[[S4:.*]] = scf.for %[[ARG8:.*]] = %[[C0_5]] to %[[C5]] step %[[C1_6]] iter_args(%[[ARG12:.*]] = %[[ARG11]]) -> (tensor<3x8x1x5xf32>)
380382
// CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[ARG0]][0, 0, %[[ARG2]], %[[ARG4]], %[[ARG6]], %[[ARG8]]] [6, 1, 1, 1, 1, 1] [1, 1, 1, 1, 1, 1] : tensor<6x1x2x1x3x5xf32> to tensor<6x1x1x1x1x1xf32>
381-
// CHECK: %[[S5:.*]] = affine.apply #[[$MAP0]](%[[ARG2]])
382-
// CHECK: %[[S6:.*]] = affine.apply #[[$MAP0]](%[[ARG4]])
383-
// CHECK: %[[S7:.*]] = affine.apply #[[$MAP1]]()
384-
// CHECK: %[[S8:.*]] = affine.apply #[[$MAP1]]()
385-
// CHECK: %[[EXTRACTED_SLICE_9:.*]] = tensor.extract_slice %[[ARG12]][%[[ARG6]], %[[S5]], 0, %[[ARG8]]] [1, %[[S7]], 1, 1] [1, 1, 1, 1] : tensor<3x8x1x5xf32> to tensor<1x?x1x1xf32>
383+
// CHECK: %[[S5:.*]] = affine.apply #[[$MAP]](%[[ARG2]])
384+
// CHECK: %[[S6:.*]] = affine.apply #[[$MAP1]](%[[ARG4]])
385+
// CHECK: %[[S7:.*]] = affine.apply #[[$MAP2]]()
386+
// CHECK: %[[S8:.*]] = affine.apply #[[$MAP2]]()
387+
// CHECK: %[[EXTRACTED_SLICE_9:.*]] = tensor.extract_slice %[[ARG12]][%[[ARG6]], %[[S5]], %[[S6]], %[[ARG8]]] [1, %[[S7]], 1, 1] [1, 1, 1, 1] : tensor<3x8x1x5xf32> to tensor<1x?x1x1xf32>
386388
// CHECK: %[[S9:.*]] = linalg.winograd_output_transform m(4) r(3) ins(%[[EXTRACTED_SLICE]] : tensor<6x1x1x1x1x1xf32>) outs(%[[EXTRACTED_SLICE_9]] : tensor<1x?x1x1xf32>) -> tensor<1x?x1x1xf32>

0 commit comments

Comments
 (0)