Skip to content

Commit 0fa5401

Browse files
committed
Address review requests
- Renamed allowInsertSliceLowering to lowerPadLikeWithInsertSlice - Renamed allowExtractSliceLowering to lowerUnpadLikeWithExtractSlice - Removed the redundant unit test since this is NFC change This reverts commit 46b7202.
1 parent 46b7202 commit 0fa5401

File tree

5 files changed

+12
-68
lines changed

5 files changed

+12
-68
lines changed

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -549,7 +549,7 @@ def LowerPackOp : Op<Transform_Dialect, "structured.lower_pack", [
549549
}];
550550

551551
let arguments = (ins Transform_ConcreteOpType<"tensor.pack">:$target,
552-
DefaultValuedAttr<BoolAttr, "true">:$allowInsertSliceLowering);
552+
DefaultValuedAttr<BoolAttr, "true">:$lowerPadLikeWithInsertSlice);
553553
let results = (outs Transform_ConcreteOpType<"tensor.pad">:$pad_op,
554554
Transform_ConcreteOpType<"tensor.expand_shape">:$expand_shape_op,
555555
Transform_ConcreteOpType<"linalg.transpose">:$transpose_op);
@@ -590,7 +590,7 @@ def LowerUnPackOp : Op<Transform_Dialect, "structured.lower_unpack", [
590590
}];
591591

592592
let arguments = (ins Transform_ConcreteOpType<"tensor.unpack">:$target,
593-
DefaultValuedAttr<BoolAttr, "true">:$allowExtractSliceLowering);
593+
DefaultValuedAttr<BoolAttr, "true">:$lowerUnpadLikeWithExtractSlice);
594594
let results = (outs Transform_ConcreteOpType<"tensor.empty">:$empty_op,
595595
Transform_ConcreteOpType<"linalg.transpose">:$transpose_op,
596596
Transform_ConcreteOpType<"tensor.collapse_shape">:$collapse_shape_op,

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1122,7 +1122,7 @@ struct LowerPackResult {
11221122
/// Rewrite pack as pad + reshape + transpose.
11231123
FailureOr<LowerPackResult> lowerPack(RewriterBase &rewriter,
11241124
tensor::PackOp packOp,
1125-
bool allowInsertSliceLowering = true);
1125+
bool lowerPadLikeWithInsertSlice = true);
11261126

11271127
struct LowerUnPackOpResult {
11281128
tensor::EmptyOp emptyOp;
@@ -1134,7 +1134,7 @@ struct LowerUnPackOpResult {
11341134
/// Rewrite pack as empty + transpose + reshape + extract_slice.
11351135
FailureOr<LowerUnPackOpResult>
11361136
lowerUnPack(RewriterBase &rewriter, tensor::UnPackOp unPackOp,
1137-
bool allowExtractSliceLowering = true);
1137+
bool lowerUnpadLikeWithExtractSlice = true);
11381138

11391139
/// Struct to hold the result of a `pack` call.
11401140
struct PackResult {

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1171,9 +1171,9 @@ DiagnosedSilenceableFailure transform::LowerPackOp::applyToOne(
11711171
transform::ApplyToEachResultList &transformResults,
11721172
transform::TransformState &state) {
11731173
rewriter.setInsertionPoint(target);
1174-
bool allowInsertSliceLowering = getAllowInsertSliceLowering();
1174+
bool lowerPadLikeWithInsertSlice = getLowerPadLikeWithInsertSlice();
11751175
FailureOr<LowerPackResult> res =
1176-
lowerPack(rewriter, target, allowInsertSliceLowering);
1176+
lowerPack(rewriter, target, lowerPadLikeWithInsertSlice);
11771177
if (failed(res)) {
11781178
return mlir::emitSilenceableFailure(target->getLoc())
11791179
<< "cannot lower to pad + expand + transpose";
@@ -1193,9 +1193,9 @@ DiagnosedSilenceableFailure transform::LowerUnPackOp::applyToOne(
11931193
transform::ApplyToEachResultList &transformResults,
11941194
transform::TransformState &state) {
11951195
rewriter.setInsertionPoint(target);
1196-
bool allowExtractSliceLowering = getAllowExtractSliceLowering();
1196+
bool lowerUnpadLikeWithExtractSlice = getLowerUnpadLikeWithExtractSlice();
11971197
FailureOr<LowerUnPackOpResult> res =
1198-
lowerUnPack(rewriter, target, allowExtractSliceLowering);
1198+
lowerUnPack(rewriter, target, lowerUnpadLikeWithExtractSlice);
11991199
if (failed(res)) {
12001200
DiagnosedSilenceableFailure diag =
12011201
emitSilenceableError()

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -218,7 +218,7 @@ struct PackedOperandsDimList {
218218

219219
FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
220220
tensor::PackOp packOp,
221-
bool allowInsertSliceLowering) {
221+
bool lowerPadLikeWithInsertSlice) {
222222
// 1. Filter out NYI cases.
223223
auto packedTensorType =
224224
cast<RankedTensorType>(packOp->getResultTypes().front());
@@ -296,7 +296,7 @@ FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
296296
llvm::interleaveComma(stripMinedShape, DBGS() << "stripMinedShape: ");
297297
DBGSNL(); DBGS() << "collapsed type: " << collapsed; DBGSNL(););
298298

299-
if (allowInsertSliceLowering && packOp.isLikePad()) {
299+
if (lowerPadLikeWithInsertSlice && packOp.isLikePad()) {
300300
// Pack ops which operate as simple pads may not produce legal
301301
// tensor.insert_slice operations when the packed type does not rank reduce
302302
// to the padded type.
@@ -354,7 +354,7 @@ FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
354354

355355
FailureOr<LowerUnPackOpResult>
356356
linalg::lowerUnPack(RewriterBase &rewriter, tensor::UnPackOp unPackOp,
357-
bool allowExtractSliceLowering) {
357+
bool lowerUnpadLikeWithExtractSlice) {
358358
Location loc = unPackOp->getLoc();
359359
OpBuilder::InsertionGuard g(rewriter);
360360
rewriter.setInsertionPoint(unPackOp);
@@ -364,7 +364,7 @@ linalg::lowerUnPack(RewriterBase &rewriter, tensor::UnPackOp unPackOp,
364364

365365
OpFoldResult zero = rewriter.getIndexAttr(0), one = rewriter.getIndexAttr(1);
366366
auto destTensorType = cast<RankedTensorType>(unPackOp.getDest().getType());
367-
if (allowExtractSliceLowering && unPackOp.isLikeUnPad()) {
367+
if (lowerUnpadLikeWithExtractSlice && unPackOp.isLikeUnPad()) {
368368
// This unpack is just a plain unpad.
369369
// Just extract the slice from the higher ranked tensor.
370370
ArrayRef<int64_t> destShape = destTensorType.getShape();

mlir/test/Dialect/Linalg/transform-lower-pack.mlir

Lines changed: 0 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -96,34 +96,6 @@ module attributes {transform.with_named_sequence} {
9696

9797
// -----
9898

99-
// This is same as pack_as_pad but since we explicitly added {allowInsertSliceLowering = false}, it should not
100-
// be lowered to insert_slice.
101-
// CHECK-LABEL: func.func @pack_disallowed_as_pad(
102-
// CHECK: %[[SRC:.+]]: tensor<129x47x16x16xf32>,
103-
// CHECK: %[[OUT:.+]]: tensor<1x1x1x1x136x64x16x16xf32>)
104-
func.func @pack_disallowed_as_pad(%arg0: tensor<129x47x16x16xf32>, %arg1: tensor<1x1x1x1x136x64x16x16xf32>) -> tensor<1x1x1x1x136x64x16x16xf32> {
105-
%cst_0 = arith.constant 0.0 : f32
106-
// tensor.pack is lowered to tensor.pad + tensor.expand_shape + tensor.insert_slice
107-
// CHECK: %[[PAD:.*]] = tensor.pad %[[SRC]] low[0, 0, 0, 0] high[7, 17, 0, 0]
108-
// CHECK: : tensor<129x47x16x16xf32> to tensor<136x64x16x16xf32>
109-
// CHECK-NOT: %[[RES:.*]] = tensor.insert_slice %[[PAD]] into %[[OUT]]
110-
%pack = tensor.pack %arg0 padding_value(%cst_0 : f32) inner_dims_pos = [0, 1, 2, 3] inner_tiles = [136, 64, 16, 16] into %arg1
111-
: tensor<129x47x16x16xf32> -> tensor<1x1x1x1x136x64x16x16xf32>
112-
return %pack : tensor<1x1x1x1x136x64x16x16xf32>
113-
}
114-
115-
module attributes {transform.with_named_sequence} {
116-
transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
117-
%pack = transform.structured.match ops{["tensor.pack"]} in %module_op
118-
: (!transform.any_op) -> !transform.op<"tensor.pack">
119-
transform.structured.lower_pack %pack {allowInsertSliceLowering = false}: (!transform.op<"tensor.pack">)
120-
-> (!transform.op<"tensor.pad">, !transform.op<"tensor.expand_shape">, !transform.op<"linalg.transpose">)
121-
transform.yield
122-
}
123-
}
124-
125-
// -----
126-
12799
// Check that we don't lower the following pack as a pad.
128100
// Although all the outer most dimensions in the resulting shape are 1s,
129101
// some of the original dimensions are not part of the inner_dims_pos, hence
@@ -261,34 +233,6 @@ module attributes {transform.with_named_sequence} {
261233

262234
// -----
263235

264-
// This is same as upack_as_pad but since we explicitly added {allowExtractSlicelowering = false}, it should not
265-
// be lowered to extract_slice.
266-
// CHECK-LABEL: func.func @unpack_disallowed_as_pad(
267-
func.func @unpack_disallowed_as_pad(%arg0: tensor<1x1x1x1x136x64x16x16xf32>, %arg1: tensor<129x47x16x16xf32>) -> tensor<129x47x16x16xf32> {
268-
%cst_0 = arith.constant 0.0 : f32
269-
270-
// CHECK-SAME: %[[ARG0:[^:]*]]: tensor<1x1x1x1x136x64x16x16xf32>
271-
// CHECK-NOT: %[[RES:.*]] = tensor.extract_slice %[[ARG0]]
272-
%pack = tensor.unpack %arg0 inner_dims_pos = [0, 1, 2, 3] inner_tiles = [136, 64, 16, 16] into %arg1
273-
: tensor<1x1x1x1x136x64x16x16xf32> -> tensor<129x47x16x16xf32>
274-
return %pack : tensor<129x47x16x16xf32>
275-
}
276-
277-
module attributes {transform.with_named_sequence} {
278-
transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
279-
%unpack = transform.structured.match ops{["tensor.unpack"]} in %module_op
280-
: (!transform.any_op) -> !transform.op<"tensor.unpack">
281-
transform.structured.lower_unpack %unpack {allowExtractSliceLowering = false}: (!transform.op<"tensor.unpack">)
282-
-> (!transform.op<"tensor.empty">,
283-
!transform.op<"linalg.transpose">,
284-
!transform.op<"tensor.collapse_shape">,
285-
!transform.op<"tensor.extract_slice">)
286-
transform.yield
287-
}
288-
}
289-
290-
// -----
291-
292236
// CHECK-LABEL: func.func @pack_with_outer_dims_perm(
293237
func.func @pack_with_outer_dims_perm(%src: tensor<100x200x128x256xi32>,
294238
%dest: tensor<200x4x16x100x16x32xi32>)

0 commit comments

Comments
 (0)