This is an archive of the discontinued LLVM Phabricator instance.

Differential D100805

Fix FoldReshapeOpWithUnitExtent generating illegal reshape
ClosedPublic

Authored by asaadaldien on Apr 19 2021, 4:56 PM.

Details

Reviewers
mravishankar
nicolasvasilache
Commits
rG10d7924581f8: Fix FoldReshapeOpWithUnitExtent generating illegal reshape
Summary

This will prevent fusion that spans all dims and generates
(d0, d1, ...) -> () reshape that isn't legal

Diff Detail

Event Timeline

asaadaldien created this revision.Apr 19 2021, 4:56 PM
Herald added subscribers: dcaballe, cota, teijeong and 16 others. · View Herald TranscriptApr 19 2021, 4:56 PM
asaadaldien requested review of this revision.Apr 19 2021, 4:56 PM
Herald added a reviewer: nicolasvasilache. · View Herald TranscriptApr 19 2021, 4:56 PM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: stephenneuendorffer, nicolasvasilache. · View Herald Transcript
asaadaldien edited the summary of this revision. (Show Details)Apr 19 2021, 4:59 PM
Harbormaster completed remote builds in B99589: Diff 338672.Apr 19 2021, 5:25 PM
mravishankar requested changes to this revision.Apr 20 2021, 12:38 PM
mravishankar added inline comments.
mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
517

Could we update the logic above to not create such reassociationExprs to begin with?

This revision now requires changes to proceed.Apr 20 2021, 12:38 PM
asaadaldien updated this revision to Diff 339013.Apr 20 2021, 3:25 PM

Move validation within folding loop

asaadaldien marked an inline comment as done.Apr 20 2021, 3:26 PM
asaadaldien added inline comments.
mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
517

I moved the condition within the folding loop and added a comment to explain why this should fail the pattern. PTAL

Harbormaster completed remote builds in B99820: Diff 339013.Apr 20 2021, 4:13 PM
mravishankar accepted this revision.Apr 20 2021, 9:53 PM
This revision is now accepted and ready to land.Apr 20 2021, 9:53 PM
Closed by commit rG10d7924581f8: Fix FoldReshapeOpWithUnitExtent generating illegal reshape (authored by asaadaldien). · Explain WhyApr 21 2021, 11:31 AM
This revision was automatically updated to reflect the committed changes.
asaadaldien marked an inline comment as done.
asaadaldien added a commit: rG10d7924581f8: Fix FoldReshapeOpWithUnitExtent generating illegal reshape.

Revision Contents

PathSize
mlir/
lib/
Dialect/
Linalg/
Transforms/
9 lines
test/
Dialect/
Linalg/
18 lines

Diff 339013

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

Show First 20 Lines • Show All 494 Lines • ▼ Show 20 Lines while (expandedDim < expandedShape.size() &&
expandedShape[expandedDim] == 1) { expandedShape[expandedDim] == 1) {
reassociationExprs[foldedDim].push_back( reassociationExprs[foldedDim].push_back(
rewriter.getAffineDimExpr(expandedDim++)); rewriter.getAffineDimExpr(expandedDim++));
} }
} }
} else { } else {
return failure(); return failure();
} }
foldedDim++; foldedDim++;
// If inner most dims are folded there shouldn't be any leading 1 dims.
// otherwise these dims are not mapped and will lead into an illegal
// reshape.
if (expandedDim == expandedShape.size()) {
if (foldedDim < foldedShape.size() && foldedShape[foldedDim] == 1) {
return failure();
}
}
} }
if (expandedDim != expandedShape.size()) if (expandedDim != expandedShape.size())
return failure(); return failure();
SmallVector<AffineMap, 4> reassociationMaps = SmallVector<AffineMap, 4> reassociationMaps =
mravishankarUnsubmitted
Done
ReplyInline Actions

Could we update the logic above to not create such reassociationExprs to begin with?

mravishankar: Could we update the logic above to not create such reassociationExprs to begin with?
asaadaldienAuthorUnsubmitted
Done
ReplyInline Actions

I moved the condition within the folding loop and added a comment to explain why this should fail the pattern. PTAL

asaadaldien: I moved the condition within the folding loop and added a comment to explain why this should…
llvm::to_vector<4>(llvm::map_range( llvm::to_vector<4>(llvm::map_range(
reassociationExprs, [&](ArrayRef<AffineExpr> exprs) -> AffineMap { reassociationExprs, [&](ArrayRef<AffineExpr> exprs) -> AffineMap {
return AffineMap::get(expandedShape.size(), 0, exprs, return AffineMap::get(expandedShape.size(), 0, exprs,
rewriter.getContext()); rewriter.getContext());
})); }));
rewriter.replaceOpWithNewOp<TensorReshapeOp>( rewriter.replaceOpWithNewOp<TensorReshapeOp>(
reshapeOp, dstType, parentReshapeOp.src(), reshapeOp, dstType, parentReshapeOp.src(),
rewriter.getAffineMapArrayAttr(reassociationMaps)); rewriter.getAffineMapArrayAttr(reassociationMaps));
▲ Show 20 Lines • Show All 111 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

Show First 20 Lines • Show All 576 Lines • ▼ Show 20 Lines linalg.generic {
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
outs(%arg0 : memref<1x8xindex>) { outs(%arg0 : memref<1x8xindex>) {
^bb0(%arg1: index): // no predecessors ^bb0(%arg1: index): // no predecessors
%0 = linalg.index 1 : index %0 = linalg.index 1 : index
linalg.yield %0 : index linalg.yield %0 : index
} }
return return
} }
// -----
func @no_fold_reshape_empty_expr(%arg0: tensor<3x2x2xf32>) -> tensor<12x1xf32> {
%0 = linalg.tensor_reshape %arg0 [affine_map<(d0, d1, d2, d3) -> (d0)>, affine_map<(d0, d1, d2, d3) -> (d1)>, affine_map<(d0, d1, d2, d3) -> (d2, d3)>] : tensor<3x2x2xf32> into tensor<3x2x2x1xf32>
%1 = linalg.tensor_reshape %0 [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>, affine_map<(d0, d1, d2, d3) -> (d3)>] : tensor<3x2x2x1xf32> into tensor<12x1xf32>
return %1 : tensor<12x1xf32>
}
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d1)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d2, d3)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
// CHECK: func @no_fold_reshape_empty_expr
// CHECK-SAME: %[[ARG0:.+]]: tensor<3x2x2xf32>
// CHECK: %[[RARG0:.+]] = linalg.tensor_reshape %[[ARG0:.+]] [#[[MAP0]], #[[MAP1]], #[[MAP2]]
// CHECK: %[[RES:.+]] = linalg.tensor_reshape %[[RARG0:.+]] [#[[MAP3]], #[[MAP4]]]
// CHECK: return %[[RES:.+]] : tensor<12x1xf32>