maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

q[mlir][Vector] Add a ShapeCastOp(BroadcastOp) canonicalization pattern 

This pattern can kick in when the source of the broadcast has a shape
that is a prefix/suffix of the result of the shape_cast.

Differential Revision: https://reviews.llvm.org/D128734
This commit is contained in:
Nicolas Vasilache 2022-06-28 08:29:46 -07:00
parent 2b747241a6
commit a48bdee686
2 changed files with 50 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
mlir/lib/Dialect/Vector/IR/VectorOps.cpp
View File

@ -4188,12 +4188,46 @@ public:
}
};
/// Pattern to rewrite a ShapeCast(Broadcast) -> Broadcast.
/// This only applies when the shape of the broadcast source is a suffix of the
/// shape of the result (i.e. when broadcast without reshape is expressive
/// enough to capture the result in a single op).
class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
public:
using OpRewritePattern<ShapeCastOp>::OpRewritePattern;
LogicalResult matchAndRewrite(ShapeCastOp shapeCastOp,
PatternRewriter &rewriter) const override {
auto broadcastOp =
shapeCastOp.getSource().getDefiningOp<vector::BroadcastOp>();
if (!broadcastOp)
return failure();
auto broadcastSourceVectorType =
broadcastOp.getSourceType().dyn_cast<VectorType>();
auto broadcastSourceShape = broadcastSourceVectorType
? broadcastSourceVectorType.getShape()
: ArrayRef<int64_t>{};
auto shapeCastTargetShape = shapeCastOp.getResultVectorType().getShape();
// Bail if `broadcastSourceShape` is not a suffix of the result.
bool isSuffix = (broadcastSourceShape == shapeCastTargetShape.take_back(
broadcastSourceShape.size()));
if (!isSuffix)
return failure();
rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
shapeCastOp, shapeCastOp.getResultVectorType(),
broadcastOp.getSource());
return success();
}
};
} // namespace
void ShapeCastOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
// Pattern to rewrite a ShapeCastOp(ConstantOp) -> ConstantOp.
results.add<ShapeCastConstantFolder>(context);
results.add<ShapeCastConstantFolder, ShapeCastBroadcastFolder>(context);
}
//===----------------------------------------------------------------------===//

17
mlir/test/Dialect/Vector/canonicalize.mlir
View File

@ -646,10 +646,10 @@ func.func @fold_broadcast_shapecast(%arg0: vector<4xf32>) -> vector<4xf32> {
// -----
// CHECK-LABEL: func @dont_fold_broadcast_shapecast_scalar
// CHECK-LABEL: func @canonicalize_broadcast_shapecast_scalar
// CHECK: vector.broadcast
// CHECK: vector.shape_cast
func.func @dont_fold_broadcast_shapecast_scalar(%arg0: f32) -> vector<1xf32> {
// CHECK-NOT: vector.shape_cast
func.func @canonicalize_broadcast_shapecast_scalar(%arg0: f32) -> vector<1xf32> {
%0 = vector.broadcast %arg0 : f32 to vector<1x1x1xf32>
%1 = vector.shape_cast %0 : vector<1x1x1xf32> to vector<1xf32>
return %1 : vector<1xf32>
@ -668,6 +668,17 @@ func.func @dont_fold_broadcast_shapecast_diff_shape(%arg0: vector<4xf32>) -> vec
// -----
// CHECK-LABEL: func @canonicalize_broadcast_shapecast
// CHECK: vector.broadcast
// CHECK-NOT: vector.shape_cast
func.func @canonicalize_broadcast_shapecast(%arg0: vector<3xf32>) -> vector<8x3xf32> {
%0 = vector.broadcast %arg0 : vector<3xf32> to vector<2x4x3xf32>
%1 = vector.shape_cast %0 : vector<2x4x3xf32> to vector<8x3xf32>
return %1 : vector<8x3xf32>
}
// -----
// CHECK-LABEL: fold_vector_transfers
func.func @fold_vector_transfers(%A: memref<?x8xf32>) -> (vector<4x8xf32>, vector<4x9xf32>) {
%c0 = arith.constant 0 : index