Skip to content

[mlir][Linalg] Move linalg.fill -> linalg.pack pattern into fill canonicalization patterns. #66002

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Sep 11, 2023
Merged

[mlir][Linalg] Move linalg.fill -> linalg.pack pattern into fill canonicalization patterns. #66002

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a3475fa
[mlir][Linalg] Move `linalg.fill` -> `linalg.pack` pattern into `fill…
MaheshRavishankar Sep 11, 2023
5f387a2
Commit to retrigger.
MaheshRavishankar Sep 11, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 57 additions & 6 deletions mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -737,7 +737,8 @@ struct FoldFillWithTensorExtract : public OpRewritePattern<tensor::ExtractOp> {

LogicalResult matchAndRewrite(tensor::ExtractOp extractOp,
PatternRewriter &rewriter) const override {
// See if tensor input of tensor.extract op is the result of a linalg.fill op.
// See if tensor input of tensor.extract op is the result of a linalg.fill
// op.
auto fillOp = extractOp.getTensor().getDefiningOp<linalg::FillOp>();
if (!fillOp)
return failure();
Expand All @@ -751,15 +752,65 @@ struct FoldFillWithTensorExtract : public OpRewritePattern<tensor::ExtractOp> {
}
};

/// Folds pack(fill) into a single fill op if
/// 1. The pack op does not have padding value, or
/// 2. The filled value and padding value are the same.
static FailureOr<FillOp> foldFillPackIntoFillOp(RewriterBase &rewriter,
tensor::PackOp packOp) {
auto fillOp = packOp.getSource().getDefiningOp<FillOp>();
if (!fillOp)
return failure();

if (auto paddingValue = packOp.getPaddingValue())
if (!isEqualConstantIntOrValue(paddingValue, fillOp.value()))
return failure();

OpBuilder::InsertionGuard guard(rewriter);
rewriter.setInsertionPoint(fillOp);

Value packOpDest = packOp.getDest();
if (!packOpDest.hasOneUse())
return failure();
if (auto emptyOp = packOpDest.getDefiningOp<tensor::EmptyOp>()) {
packOpDest = tensor::PackOp::createDestinationTensor(
rewriter, fillOp.getLoc(), fillOp.getDpsInitOperand(0)->get(),
packOp.getMixedTiles(), packOp.getInnerDimsPos(),
packOp.getOuterDimsPerm());
} else {
DominanceInfo dom(fillOp);
if (!dom.properlyDominates(packOpDest, fillOp))
return failure();
}

Value fillDest = packOpDest;
return clone(rewriter, fillOp, packOpDest.getType(),
{fillOp.value(), fillDest});
}

/// Wrapper pattern that applies foldFillPackIntoFillOp method.
struct FoldFillWithPack : public OpRewritePattern<tensor::PackOp> {
public:
FoldFillWithPack(MLIRContext *context)
: OpRewritePattern<tensor::PackOp>(context) {}

LogicalResult matchAndRewrite(tensor::PackOp packOp,
PatternRewriter &rewriter) const override {
auto fillOp = foldFillPackIntoFillOp(rewriter, packOp);
if (failed(fillOp))
return failure();
rewriter.replaceOp(packOp, fillOp.value().result());
return success();
}
};

} // namespace

void FillOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results
.add<FoldFillWithTensorExtract,
FoldFillWithPad, FoldFillWithTensorReshape<tensor::CollapseShapeOp>,
FoldFillWithTensorReshape<tensor::ExpandShapeOp>,
FoldInsertPadIntoFill>(context);
results.add<FoldFillWithTensorExtract, FoldFillWithPack, FoldFillWithPad,
FoldFillWithTensorReshape<tensor::CollapseShapeOp>,
FoldFillWithTensorReshape<tensor::ExpandShapeOp>,
FoldInsertPadIntoFill>(context);
}

//===----------------------------------------------------------------------===//
Expand Down
60 changes: 0 additions & 60 deletions mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -448,46 +448,6 @@ bubbleUpPackOpThroughGenericOp(RewriterBase &rewriter, tensor::PackOp packOp,
*packInfo);
}

/// Folds pack(fill) into a single fill op if
/// 1. The pack op does not have padding value, or
/// 2. The filled value and padding value are the same.
static FailureOr<FillOp>
foldFillPackIntoFillOp(RewriterBase &rewriter, tensor::PackOp packOp,
ControlPropagationFn controlFn) {
auto fillOp = packOp.getSource().getDefiningOp<FillOp>();
if (!fillOp)
return failure();

// User controlled propagation function.
if (!controlFn(fillOp))
return failure();

if (auto paddingValue = packOp.getPaddingValue())
if (!isEqualConstantIntOrValue(paddingValue, fillOp.value()))
return failure();

OpBuilder::InsertionGuard guard(rewriter);
rewriter.setInsertionPoint(fillOp);

Value packOpDest = packOp.getDest();
if (!packOpDest.hasOneUse())
return failure();
if (auto emptyOp = packOpDest.getDefiningOp<tensor::EmptyOp>()) {
packOpDest = tensor::PackOp::createDestinationTensor(
rewriter, fillOp.getLoc(), fillOp.getDpsInitOperand(0)->get(),
packOp.getMixedTiles(), packOp.getInnerDimsPos(),
packOp.getOuterDimsPerm());
} else {
DominanceInfo dom(fillOp);
if (!dom.properlyDominates(packOpDest, fillOp))
return failure();
}

Value fillDest = packOpDest;
return clone(rewriter, fillOp, packOpDest.getType(),
{fillOp.value(), fillDest});
}

/// Wrapper pattern that applies bubbleUpPackOpThroughGenericOp method.
struct BubbleUpPackOpThroughGenericOpPattern
: public OpRewritePattern<tensor::PackOp> {
Expand All @@ -510,25 +470,6 @@ struct BubbleUpPackOpThroughGenericOpPattern
ControlPropagationFn controlFn;
};

/// Wrapper pattern that applies foldFillPackIntoFillOp method.
struct FoldFillPackIntoFillOpPattern : public OpRewritePattern<tensor::PackOp> {
public:
FoldFillPackIntoFillOpPattern(MLIRContext *context, ControlPropagationFn fun)
: OpRewritePattern<tensor::PackOp>(context), controlFn(std::move(fun)) {}

LogicalResult matchAndRewrite(tensor::PackOp packOp,
PatternRewriter &rewriter) const override {
auto fillOp = foldFillPackIntoFillOp(rewriter, packOp, controlFn);
if (failed(fillOp))
return failure();
rewriter.replaceOp(packOp, fillOp.value().result());
return success();
}

private:
ControlPropagationFn controlFn;
};

// TODO: Relax this restriction. We should unpack a generic op also
// in the presence of multiple unpack ops as producers.
/// Return the unpacked operand, if present, for the current generic op.
Expand Down Expand Up @@ -750,7 +691,6 @@ void mlir::linalg::populateDataLayoutPropagationPatterns(
RewritePatternSet &patterns,
const ControlPropagationFn &controlPackUnPackPropagation) {
patterns.insert<BubbleUpPackOpThroughGenericOpPattern,
FoldFillPackIntoFillOpPattern,
PushDownUnPackOpThroughGenericOp, PushDownUnPackThroughPadOp>(
patterns.getContext(), controlPackUnPackPropagation);
}
44 changes: 44 additions & 0 deletions mlir/test/Dialect/Linalg/canonicalize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -353,6 +353,50 @@ func.func @fold_fill_extract(%arg0 : i1) -> i1 {

// -----

func.func @fill_pack() -> tensor<24x32x16x16xf32> {
%dest = tensor.empty() : tensor<384x512xf32>
%cst = arith.constant 0.000000e+00 : f32
%0 = tensor.empty() : tensor<24x32x16x16xf32>
%1 = linalg.fill ins(%cst : f32) outs(%dest : tensor<384x512xf32>) -> tensor<384x512xf32>
%pack = tensor.pack %1 inner_dims_pos = [0, 1] inner_tiles = [16, 16] into %0 : tensor<384x512xf32> -> tensor<24x32x16x16xf32>
return %pack : tensor<24x32x16x16xf32>
}
// CHECK-LABEL: func.func @fill_pack
// CHECK: %[[PACKED_EMPTY:.+]] = tensor.empty() : tensor<24x32x16x16xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}) outs(%[[PACKED_EMPTY]]
// CHECK: return %[[FILL]]

// -----

#map = affine_map<()[s0] -> (s0 ceildiv 16)>
func.func @dynamic_fill_pack(%arg0: tensor<?x?xf32>) -> tensor<?x?x16x16xf32> {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%0 = linalg.fill ins(%cst : f32) outs(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32>
%dim = tensor.dim %0, %c0 : tensor<?x?xf32>
%dim_0 = tensor.dim %0, %c1 : tensor<?x?xf32>
%1 = affine.apply #map()[%dim]
%2 = affine.apply #map()[%dim_0]
%3 = tensor.empty(%1, %2) : tensor<?x?x16x16xf32>
%pack = tensor.pack %0 padding_value(%cst : f32) inner_dims_pos = [0, 1] inner_tiles = [16, 16] into %3 : tensor<?x?xf32> -> tensor<?x?x16x16xf32>
return %pack : tensor<?x?x16x16xf32>
}
// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0] -> (s0 ceildiv 16)>
// CHECK: func.func @dynamic_fill_pack
// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK: %[[D0:.+]] = tensor.dim %[[DEST]], %[[C0]]
// CHECK: %[[D1:.+]] = tensor.dim %[[DEST]], %[[C1]]
// CHECK: %[[PACKED_D0:.+]] = affine.apply #[[MAP]]()[%[[D0]]]
// CHECK: %[[PACKED_D1:.+]] = affine.apply #[[MAP]]()[%[[D1]]]
// CHECK: %[[PACKED_EMPTY:.+]] = tensor.empty(%[[PACKED_D0]], %[[PACKED_D1]]) : tensor<?x?x16x16xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}) outs(%[[PACKED_EMPTY]]
// CHECK: return %[[FILL]]

// -----

// CHECK: func @fold_self_copy
func.func @fold_self_copy(%0 : memref<4x16xf32>) {
// CHECK-NEXT: return
Expand Down
44 changes: 0 additions & 44 deletions mlir/test/Dialect/Linalg/data-layout-propagation.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -839,47 +839,3 @@ func.func @unpack_different_destination_shape(%arg0: tensor<1x1x1080x1920x16xi32
// CHECK-SAME: inner_dims_pos = [0] inner_tiles = [16]
// CHECK-SAME: into %[[UNPACK_NEW_DEST]]
// CHECK: return %[[UNPACK]] : tensor<16x540x960xi32>

// -----

func.func @fill_pack() -> tensor<24x32x16x16xf32> {
%dest = tensor.empty() : tensor<384x512xf32>
%cst = arith.constant 0.000000e+00 : f32
%0 = tensor.empty() : tensor<24x32x16x16xf32>
%1 = linalg.fill ins(%cst : f32) outs(%dest : tensor<384x512xf32>) -> tensor<384x512xf32>
%pack = tensor.pack %1 inner_dims_pos = [0, 1] inner_tiles = [16, 16] into %0 : tensor<384x512xf32> -> tensor<24x32x16x16xf32>
return %pack : tensor<24x32x16x16xf32>
}
// CHECK-LABEL: func.func @fill_pack
// CHECK: %[[PACKED_EMPTY:.+]] = tensor.empty() : tensor<24x32x16x16xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}) outs(%[[PACKED_EMPTY]]
// CHECK: return %[[FILL]]

// -----

#map = affine_map<()[s0] -> (s0 ceildiv 16)>
func.func @dynamic_fill_pack(%arg0: tensor<?x?xf32>) -> tensor<?x?x16x16xf32> {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%0 = linalg.fill ins(%cst : f32) outs(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32>
%dim = tensor.dim %0, %c0 : tensor<?x?xf32>
%dim_0 = tensor.dim %0, %c1 : tensor<?x?xf32>
%1 = affine.apply #map()[%dim]
%2 = affine.apply #map()[%dim_0]
%3 = tensor.empty(%1, %2) : tensor<?x?x16x16xf32>
%pack = tensor.pack %0 padding_value(%cst : f32) inner_dims_pos = [0, 1] inner_tiles = [16, 16] into %3 : tensor<?x?xf32> -> tensor<?x?x16x16xf32>
return %pack : tensor<?x?x16x16xf32>
}
// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0] -> (s0 ceildiv 16)>
// CHECK: func.func @dynamic_fill_pack
// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK: %[[D0:.+]] = tensor.dim %[[DEST]], %[[C0]]
// CHECK: %[[D1:.+]] = tensor.dim %[[DEST]], %[[C1]]
// CHECK: %[[PACKED_D0:.+]] = affine.apply #[[MAP]]()[%[[D0]]]
// CHECK: %[[PACKED_D1:.+]] = affine.apply #[[MAP]]()[%[[D1]]]
// CHECK: %[[PACKED_EMPTY:.+]] = tensor.empty(%[[PACKED_D0]], %[[PACKED_D1]]) : tensor<?x?x16x16xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}) outs(%[[PACKED_EMPTY]]
// CHECK: return %[[FILL]]