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
535

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
535

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 339331

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

Show First 20 Lines • Show All 512 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 641 Lines • ▼ Show 20 Lines
// CHECK: %[[FILL:.+]] = linalg.fill(%[[INIT]], %{{.+}}) // CHECK: %[[FILL:.+]] = linalg.fill(%[[INIT]], %{{.+}})
// CHECK: %[[RESULT:.+]] = linalg.generic // CHECK: %[[RESULT:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP3]]] // CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP3]]]
// CHECK-SAME: iterator_types = ["parallel", "reduction"] // CHECK-SAME: iterator_types = ["parallel", "reduction"]
// CHECK-SAME: ins(%[[RESHAPE]] : tensor<?x?xf32>) // CHECK-SAME: ins(%[[RESHAPE]] : tensor<?x?xf32>)
// CHECK-SAME: outs(%[[FILL]] : tensor<?xf32>) // CHECK-SAME: outs(%[[FILL]] : tensor<?xf32>)
// CHECK: %[[RESULT_RESHAPE:.+]] = linalg.tensor_reshape %[[RESULT]] [#[[MAP2]]] // CHECK: %[[RESULT_RESHAPE:.+]] = linalg.tensor_reshape %[[RESULT]] [#[[MAP2]]]
// CHECK: return %[[RESULT_RESHAPE]] // CHECK: return %[[RESULT_RESHAPE]]
// -----
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>