From 98f661dfb5f02a014394d7b7ea68206fbb7c8ed9 Mon Sep 17 00:00:00 2001 From: Quinn Dawkins Date: Mon, 9 Sep 2024 11:08:00 -0400 Subject: [PATCH] Drop upstreamed canonicalization pattern --- .../Dialect/Flow/Transforms/Canonicalizer.cpp | 79 ----------------- .../Transforms/test/flow_canonicalize.mlir | 85 ------------------- 2 files changed, 164 deletions(-) diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Canonicalizer.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Canonicalizer.cpp index e7da87363c9a..2c780988579e 100644 --- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Canonicalizer.cpp +++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Canonicalizer.cpp @@ -19,84 +19,6 @@ namespace mlir::iree_compiler::IREE::Flow { namespace { -/// Folds a chain of `tensor.pad` ops with the same constant padding value. -/// -/// Example: -/// -/// ```mlir -/// %1 = tensor.pad %0 low[0, 1] high[0, 2] { -/// tensor.yield %val -/// } : tensor<1x2xf32> to tensor<2x5xf32> -/// %res = tensor.pad %1 low[0, 2] high[3, 0] { -/// tensor.yield %val -/// } : tensor<1x5xf32> to tensor<5x7xf32> -/// ``` -/// -/// folds into: -/// -/// ```mlir -/// %res = tensor.pad %0 low[0, 3] high[3, 2] { -/// tensor.yield %val -/// } : tensor<1x2xf32> to tensor<5x7xf32> -/// ``` -/// -/// NOTE: This wasn't sent upstream as a canonicalization due to the use of -/// the Affine dialect. -struct FoldConsecutiveConstantPadding : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(tensor::PadOp padOp, - PatternRewriter &rewriter) const override { - if (padOp.getNofold()) { - return failure(); - } - auto producerPad = padOp.getSource().getDefiningOp(); - if (!producerPad || producerPad.getNofold()) { - return rewriter.notifyMatchFailure( - padOp, "producer is not a foldable tensor.pad op"); - } - - // Fail if the tensor::PadOps padding values do not match. - Value consumerPadValue = padOp.getConstantPaddingValue(); - Value producerPadValue = producerPad.getConstantPaddingValue(); - if (!consumerPadValue || !producerPadValue || - consumerPadValue != producerPadValue) { - return rewriter.notifyMatchFailure( - padOp, "cannot fold PadOps with different padding values"); - } - - Location loc = padOp.getLoc(); - AffineExpr d0, d1; - bindDims(rewriter.getContext(), d0, d1); - - // Combine the low/high paddings of the two tensor::PadOps. - auto addPaddings = [&](ArrayRef consumerPaddings, - ArrayRef producerPaddings) { - SmallVector sumPaddings; - for (auto [consumerIndex, producerIndex] : - llvm::zip_equal(consumerPaddings, producerPaddings)) { - sumPaddings.push_back(affine::makeComposedFoldedAffineApply( - rewriter, loc, d0 + d1, {consumerIndex, producerIndex})); - } - return sumPaddings; - }; - - SmallVector newHighPad = - addPaddings(padOp.getMixedHighPad(), producerPad.getMixedHighPad()); - SmallVector newLowPad = - addPaddings(padOp.getMixedLowPad(), producerPad.getMixedLowPad()); - - auto newPadOp = rewriter.create( - padOp.getLoc(), padOp.getResultType(), producerPad.getSource(), - newLowPad, newHighPad, padOp.getNofold(), - getPrunedAttributeList(padOp, tensor::PadOp::getAttributeNames())); - rewriter.inlineRegionBefore(padOp.getRegion(), newPadOp.getRegion(), - newPadOp.getRegion().begin()); - rewriter.replaceOp(padOp, newPadOp.getResult()); - return success(); - } -}; - /// Canonicalize operations in nested regions. struct CanonicalizerPass : public impl::CanonicalizerPassBase { @@ -124,7 +46,6 @@ struct CanonicalizerPass // compilation phase. tensor::populateMergeConsecutiveInsertExtractSlicePatterns(owningPatterns); IREE::Flow::populateTensorDialectCastOpPattern(context, owningPatterns); - owningPatterns.add(context); patterns = std::make_shared(std::move(owningPatterns)); diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/flow_canonicalize.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/flow_canonicalize.mlir index effadc47da5d..a8b0d126d43e 100644 --- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/flow_canonicalize.mlir +++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/flow_canonicalize.mlir @@ -1,90 +1,5 @@ // RUN: iree-opt --iree-flow-canonicalize %s --split-input-file --mlir-print-local-scope | FileCheck %s -util.func public @merge_constant_padding(%arg0: tensor<2x3xf32>, %pad_value: f32) -> tensor<7x8xf32> { - %pad0 = tensor.pad %arg0 low[1, 1] high[1, 0] { - ^bb0(%b0: index, %b1 : index): - tensor.yield %pad_value : f32 - } : tensor<2x3xf32> to tensor<4x4xf32> - %pad1 = tensor.pad %pad0 low[0, 2] high[3, 2] { - ^bb0(%b2: index, %b3 : index): - tensor.yield %pad_value : f32 - } : tensor<4x4xf32> to tensor<7x8xf32> - util.return %pad1 : tensor<7x8xf32> -} -// CHECK-LABEL: util.func public @merge_constant_padding -// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]: tensor<2x3xf32> -// CHECK-SAME: %[[PADVAL:[A-Za-z0-9]+]]: f32 -// CHECK: %[[PAD:.+]] = tensor.pad %[[ARG0]] low[1, 3] high[4, 2] -// CHECK: tensor.yield %[[PADVAL]] -// CHECK: util.return %[[PAD]] - -// ----- - -util.func public @merge_constant_padding_dynamic(%arg0: tensor, %idx: index, %pad_value: f32) -> tensor { - %pad0 = tensor.pad %arg0 low[%idx, 1] high[1, 0] { - ^bb0(%b0: index, %b1 : index): - tensor.yield %pad_value : f32 - } : tensor to tensor - %pad1 = tensor.pad %pad0 low[0, 2] high[%idx, 2] { - ^bb0(%b2: index, %b3 : index): - tensor.yield %pad_value : f32 - } : tensor to tensor - util.return %pad1 : tensor -} -// CHECK-LABEL: util.func public @merge_constant_padding_dynamic -// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]: tensor -// CHECK-SAME: %[[IDX:[A-Za-z0-9]+]]: index -// CHECK-SAME: %[[PADVAL:[A-Za-z0-9]+]]: f32 -// CHECK: %[[HIGH:.+]] = affine.apply affine_map<()[s0] -> (s0 + 1)>()[%[[IDX]]] -// CHECK: %[[PAD:.+]] = tensor.pad %[[ARG0]] low[%[[IDX]], 3] high[%[[HIGH]], 2] -// CHECK: tensor.yield %[[PADVAL]] -// CHECK: util.return %[[PAD]] - -// ----- - -util.func public @dont_merge_constant_padding_nofold(%arg0: tensor<2x3xf32>, %pad_value: f32) -> tensor<7x8xf32> { - %pad0 = tensor.pad %arg0 low[1, 1] high[1, 0] { - ^bb0(%b0: index, %b1 : index): - tensor.yield %pad_value : f32 - } : tensor<2x3xf32> to tensor<4x4xf32> - %pad1 = tensor.pad %pad0 nofold low[0, 2] high[3, 2] { - ^bb0(%b2: index, %b3 : index): - tensor.yield %pad_value : f32 - } : tensor<4x4xf32> to tensor<7x8xf32> - util.return %pad1 : tensor<7x8xf32> -} - -// Verify that folding does not happen if it would drop a nofold attribute - -// CHECK-LABEL: util.func public @dont_merge_constant_padding_nofold -// CHECK: tensor.pad -// CHECK: tensor.pad {{.*}} nofold - -// ----- - -util.func public @dont_merge_constant_padding_different_vals( - %arg0: tensor<2x3xf32>, - %pad_value0: f32, - %pad_value1: f32) -> tensor<7x8xf32> { - %pad0 = tensor.pad %arg0 low[1, 1] high[1, 0] { - ^bb0(%b0: index, %b1 : index): - tensor.yield %pad_value0 : f32 - } : tensor<2x3xf32> to tensor<4x4xf32> - %pad1 = tensor.pad %pad0 nofold low[0, 2] high[3, 2] { - ^bb0(%b2: index, %b3 : index): - tensor.yield %pad_value1 : f32 - } : tensor<4x4xf32> to tensor<7x8xf32> - util.return %pad1 : tensor<7x8xf32> -} - -// Verify that folding does not happen if it would drop a nofold attribute - -// CHECK-LABEL: util.func public @dont_merge_constant_padding_different_vals -// CHECK: tensor.pad -// CHECK: tensor.pad - -// ----- - util.func public @tensor_cast_to_reshape(%reshape_17 : tensor, %65 : tensor, %0 : index, %1 : index) -> tensor { %cast = tensor.cast %reshape_17 : tensor to tensor %66 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,