From ce1c8a1ebcc85e8446c05f7292d17ac2e6b7ed1d Mon Sep 17 00:00:00 2001 From: Ehsan Toosi Date: Mon, 14 Sep 2020 01:21:28 -0700 Subject: [PATCH] [MLIR][LHLO] Replace lhlo-copy-removal pass with mlir-copy-removal pass Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/43137 This PR removes lhlo-copy-removal pass entirely and replace its usages with ```mlir::createCopyRemovalPass()```. -- 7ce1a06f507c8db46c6d7b43c7870cf56002e18e by Ehsan Toosi : [mlir][lhlo] Replace lhlo-copy-removal pass with mlir-copy-removal pass COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/tensorflow/pull/43137 from dfki-ehna:using_mlir_copy_removal 7ce1a06f507c8db46c6d7b43c7870cf56002e18e PiperOrigin-RevId: 331498501 --- include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h | 1 + include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td | 8 +- .../Dialect/mhlo/transforms/lmhlo_passes.td | 6 - .../mlir-hlo/Dialect/mhlo/transforms/passes.h | 6 - lib/Dialect/mhlo/transforms/CMakeLists.txt | 1 - .../mhlo/transforms/lhlo_copy_removal.cc | 102 ---------------- tests/end2end/broadcast.mlir | 2 +- tests/end2end/reshape.mlir | 2 +- tests/lhlo-copy-removal.mlir | 115 ------------------ 9 files changed, 10 insertions(+), 233 deletions(-) delete mode 100644 lib/Dialect/mhlo/transforms/lhlo_copy_removal.cc delete mode 100644 tests/lhlo-copy-removal.mlir diff --git a/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h b/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h index bb9b290..6ad513d 100644 --- a/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h +++ b/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h @@ -27,6 +27,7 @@ limitations under the License. #include "mlir/IR/Operation.h" #include "mlir/IR/StandardTypes.h" #include "mlir/IR/Types.h" +#include "mlir/Interfaces/CopyOpInterface.h" #include "mlir/Interfaces/SideEffectInterfaces.h" #include "mlir/Interfaces/ViewLikeInterface.h" diff --git a/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td b/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td index 750cce6..1c9820e 100644 --- a/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td +++ b/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td @@ -34,6 +34,7 @@ limitations under the License. #define LHLO_OPS include "mlir/IR/OpBase.td" +include "mlir/Interfaces/CopyOpInterface.td" include "mlir/Interfaces/SideEffectInterfaces.td" include "mlir/Interfaces/ViewLikeInterface.td" include "mlir-hlo/Dialect/mhlo/IR/hlo_ops_base.td" @@ -616,11 +617,16 @@ def LHLO_ConvOp : LHLO_Op<"convolution", []>, BASE_HLO_ConvOp { ); } -def LHLO_CopyOp: LHLO_Op<"copy", []>, BASE_HLO_CopyOp { +def LHLO_CopyOp: LHLO_Op<"copy", [CopyOpInterface]>, BASE_HLO_CopyOp { let arguments = (ins Arg:$operand, Arg:$output ); + + let extraClassDeclaration = [{ + Value getSource() { return operand();} + Value getTarget() { return output(); } + }]; } def LHLO_DotOp: LHLO_Op<"dot", []>, BASE_HLO_DotOp { diff --git a/include/mlir-hlo/Dialect/mhlo/transforms/lmhlo_passes.td b/include/mlir-hlo/Dialect/mhlo/transforms/lmhlo_passes.td index 963ff5d..39b4ca6 100644 --- a/include/mlir-hlo/Dialect/mhlo/transforms/lmhlo_passes.td +++ b/include/mlir-hlo/Dialect/mhlo/transforms/lmhlo_passes.td @@ -15,12 +15,6 @@ limitations under the License. include "mlir/Pass/PassBase.td" -def LhloCopyRemovalPass : Pass<"lhlo-copy-removal", "FuncOp"> { - let summary = "Removes redundant LHLO copy operations."; - let constructor = "createLhloCopyRemovalPass()"; -} - - def LhloLegalizeToLinalgPass : Pass<"lhlo-legalize-to-linalg", "FuncOp"> { let summary = "Legalize from LHLO dialect to Linalg dialect."; let constructor = "createLegalizeLhloToLinalgPass()"; diff --git a/include/mlir-hlo/Dialect/mhlo/transforms/passes.h b/include/mlir-hlo/Dialect/mhlo/transforms/passes.h index 2f723f4..bc6e3fd 100644 --- a/include/mlir-hlo/Dialect/mhlo/transforms/passes.h +++ b/include/mlir-hlo/Dialect/mhlo/transforms/passes.h @@ -95,12 +95,6 @@ std::unique_ptr createLegalizeToGpuPass(); std::unique_ptr createLhloFuseLinalgPass( bool use_parallel_loops = false, llvm::ArrayRef tile_sizes = {}); -// Removes unnecessary LHLO copies which copy from the allocated buffers to the -// block arguments. The block arguments are used instead of all uses of these -// buffers. The buffers are freed. This pass only works in regions that contain -// a single block. -std::unique_ptr createLhloCopyRemovalPass(); - // Lowers from LHLO dialect to parallel loops. std::unique_ptr> createLegalizeLhloToParallelLoopsPass(); diff --git a/lib/Dialect/mhlo/transforms/CMakeLists.txt b/lib/Dialect/mhlo/transforms/CMakeLists.txt index 945fa0e..5ae0ec2 100644 --- a/lib/Dialect/mhlo/transforms/CMakeLists.txt +++ b/lib/Dialect/mhlo/transforms/CMakeLists.txt @@ -125,7 +125,6 @@ add_mlir_library(MhloLhloToLinalg ) add_mlir_library(LmhloPasses - lhlo_copy_removal.cc lhlo_fuse_linalg.cc lhlo_legalize_to_affine.cc lhlo_legalize_to_gpu.cc diff --git a/lib/Dialect/mhlo/transforms/lhlo_copy_removal.cc b/lib/Dialect/mhlo/transforms/lhlo_copy_removal.cc deleted file mode 100644 index 7a44184..0000000 --- a/lib/Dialect/mhlo/transforms/lhlo_copy_removal.cc +++ /dev/null @@ -1,102 +0,0 @@ -/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -// This file implements a pass to remove redundant LHLO copy operations. - -#include "mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h" -#include "mlir-hlo/Dialect/mhlo/transforms/passes.h" -#include "mlir/Dialect/StandardOps/IR/Ops.h" -#include "mlir/IR/Operation.h" -#include "mlir/Pass/Pass.h" - -namespace mlir { -namespace lmhlo { -namespace { - -// Removes LHLO copy operations that copy from allocated buffers to block -// arguments. All uses of each buffer are replaced with the corresponding block -// argument and the buffer is freed. Note that this pass only works in regions -// with a single block. -struct LhloCopyRemovalPass - : mlir::PassWrapper> { - void runOnOperation() override { - llvm::SmallVector eraseList; - auto operation = getOperation(); - operation->walk([&](mlir::lmhlo::CopyOp copyOp) { - // If this region contains more than one block, then ignore this copy - // operation. - if (copyOp.getParentRegion()->getBlocks().size() > 1) { - return; - } - - mlir::Value fromOperand = copyOp.operand(); - mlir::Value toOperand = copyOp.output(); - - // If the fromOperand value is a block argument or the toOperand - // value is not a block argument, then ignore this copy operation. - if (!fromOperand.getDefiningOp() || toOperand.getDefiningOp()) { - return; - } - - // The copy operation removal is illegal if there is at least a single use - // of toOperand value that lies between the first use of fromOperand value - // and the copy operation. - auto fromOperandUsers = fromOperand.getUsers(); - auto firstUser = *fromOperandUsers.begin(); - for (auto op : fromOperandUsers) { - if (op->isBeforeInBlock(firstUser)) firstUser = op; - } - for (auto op : toOperand.getUsers()) { - if (op->isBeforeInBlock(copyOp) && firstUser->isBeforeInBlock(op)) { - return; - } - } - - // TODO(DFKI): Use live variable analysis to solve aliasing issues among - // block arguments. - - // Remove the associated alloc operation. - auto allocOp = fromOperand.getDefiningOp(); - eraseList.push_back(allocOp); - - // Iterate over all uses of the fromOperand to find the associated - // deallocOp (if any). - for (auto op : fromOperandUsers) { - if (isa(op)) { - eraseList.push_back(op); - break; - } - } - - // Replace all uses of the fromOperand with the toOperand. This rewires - // all references pointing to the original alloc operation to the new - // target operation in order to safely remove the copy op. - fromOperand.replaceAllUsesWith(toOperand); - copyOp.erase(); - }); - for (auto op : eraseList) { - op->erase(); - } - }; -}; - -} // namespace - -std::unique_ptr createLhloCopyRemovalPass() { - return std::make_unique(); -} - -} // namespace lmhlo -} // namespace mlir diff --git a/tests/end2end/broadcast.mlir b/tests/end2end/broadcast.mlir index 43fb7c9..c50c4a0 100644 --- a/tests/end2end/broadcast.mlir +++ b/tests/end2end/broadcast.mlir @@ -1,4 +1,4 @@ -// RUN: mlir-hlo-opt %s -mhlo-test-chlo-legalize-to-hlo -hlo-legalize-to-lhlo=results-escape-function=true -buffer-placement -lhlo-copy-removal -canonicalize -cse -lhlo-legalize-to-linalg -lhlo-fuse-linalg -convert-linalg-to-loops -canonicalize -cse -convert-linalg-to-llvm -test-lhlo-legalize-to-llvm | mlir-cpu-runner -e main -entry-point-result=void -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext | FileCheck %s +// RUN: mlir-hlo-opt %s -mhlo-test-chlo-legalize-to-hlo -hlo-legalize-to-lhlo=results-escape-function=true -buffer-placement -copy-removal -canonicalize -cse -lhlo-legalize-to-linalg -lhlo-fuse-linalg -convert-linalg-to-loops -canonicalize -cse -convert-linalg-to-llvm -test-lhlo-legalize-to-llvm | mlir-cpu-runner -e main -entry-point-result=void -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext | FileCheck %s func @main() -> () { call @trivial_broadcast_wrapper() : () -> () diff --git a/tests/end2end/reshape.mlir b/tests/end2end/reshape.mlir index 3d98ce5..35f286c 100644 --- a/tests/end2end/reshape.mlir +++ b/tests/end2end/reshape.mlir @@ -1,4 +1,4 @@ -// RUN: mlir-hlo-opt %s -mhlo-test-chlo-legalize-to-hlo -hlo-legalize-to-lhlo=results-escape-function=true -buffer-placement -lhlo-copy-removal -canonicalize -cse -lhlo-legalize-to-linalg -lhlo-fuse-linalg -convert-linalg-to-loops -convert-scf-to-std -canonicalize -cse -test-lhlo-legalize-to-llvm | mlir-cpu-runner -e main -entry-point-result=void -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext | FileCheck %s +// RUN: mlir-hlo-opt %s -mhlo-test-chlo-legalize-to-hlo -hlo-legalize-to-lhlo=results-escape-function=true -buffer-placement -copy-removal -canonicalize -cse -lhlo-legalize-to-linalg -lhlo-fuse-linalg -convert-linalg-to-loops -convert-scf-to-std -canonicalize -cse -test-lhlo-legalize-to-llvm | mlir-cpu-runner -e main -entry-point-result=void -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext | FileCheck %s func @main() -> () { call @reshape_with_static_shape_size_matrix_to_1D() : () -> () diff --git a/tests/lhlo-copy-removal.mlir b/tests/lhlo-copy-removal.mlir deleted file mode 100644 index 3271595..0000000 --- a/tests/lhlo-copy-removal.mlir +++ /dev/null @@ -1,115 +0,0 @@ -// RUN: mlir-hlo-opt -lhlo-copy-removal %s -o - | FileCheck %s - -// CHECK-LABEL: func @remove_simple -func @remove_simple(%arg0: memref<2x2xf32>) { - %0 = alloc() {temp = true} : memref<2x2xf32> - "lmhlo.copy"(%0, %arg0) : (memref<2x2xf32>, memref<2x2xf32>) -> () - dealloc %0 : memref<2x2xf32> - // CHECK-NEXT: "lmhlo.terminator"() : () -> () - "lmhlo.terminator"() : () -> () -} - -// ----- - -// CHECK-LABEL: func @remove_without_dealloc -func @remove_without_dealloc(%arg0: memref<2x2xf32>) { - %0 = alloc() {temp = true} : memref<2x2xf32> - "lmhlo.copy"(%0, %arg0) : (memref<2x2xf32>, memref<2x2xf32>) -> () - // CHECK-NEXT: "lmhlo.terminator"() : () -> () - "lmhlo.terminator"() : () -> () -} - -// ----- - -// CHECK-LABEL: func @replace_dependency -func @replace_dependency(%arg0: memref<2x2xf32>, %arg1: memref<2x2xf32>) { - %0 = alloc() {temp = true} : memref<2x2xf32> - "lmhlo.exponential"(%arg0, %0) : (memref<2x2xf32>, memref<2x2xf32>) -> () - // CHECK-NEXT: "lmhlo.exponential"(%arg0, %arg1) : (memref<2x2xf32>, memref<2x2xf32>) -> () - "lmhlo.copy"(%0, %arg1) : (memref<2x2xf32>, memref<2x2xf32>) -> () - dealloc %0 : memref<2x2xf32> - // CHECK-NEXT: "lmhlo.terminator"() : () -> () - "lmhlo.terminator"() : () -> () -} - -// ----- - -// CHECK-LABEL: func @keep_copies -func @keep_copies(%arg0: memref<2x2xf32>, %arg1: memref<2x2xf32>) { - // CHECK-NEXT: "lmhlo.copy"(%arg0, %arg1) : (memref<2x2xf32>, memref<2x2xf32>) -> () - "lmhlo.copy"(%arg0, %arg1) : (memref<2x2xf32>, memref<2x2xf32>) -> () - // CHECK-NEXT: "lmhlo.terminator"() : () -> () - "lmhlo.terminator"() : () -> () -} - -// ----- - -// CHECK-LABEL: func @must_not_be_removed -func @must_not_be_removed(%arg0: memref<2x2xf32>, - %arg1: memref<2x2xf32>, - %arg2: memref<2x2xf32>) { - // CHECK-NEXT: %[[ALLOC:.*]] = alloc() {temp = true} : memref<2x2xf32> - %0 = alloc() {temp = true} : memref<2x2xf32> - // CHECK-NEXT: "lmhlo.exponential"(%arg0, %[[ALLOC]]) : (memref<2x2xf32>, memref<2x2xf32>) -> () - "lmhlo.exponential"(%arg0, %0) : (memref<2x2xf32>, memref<2x2xf32>) -> () - // CHECK-NEXT: "lmhlo.exponential"(%arg1, %arg2) : (memref<2x2xf32>, memref<2x2xf32>) -> () - "lmhlo.exponential"(%arg1, %arg2) : (memref<2x2xf32>, memref<2x2xf32>) -> () - // CHECK-NEXT: "lmhlo.copy"(%[[ALLOC]], %arg2) : (memref<2x2xf32>, memref<2x2xf32>) -> () - "lmhlo.copy"(%0, %arg2) : (memref<2x2xf32>, memref<2x2xf32>) -> () - dealloc %0 : memref<2x2xf32> - "lmhlo.terminator"() : () -> () -} - -// ----- - -// CHECK-LABEL: func @must_be_removed_first -func @must_be_removed_first(%arg0: memref<2x2xf32>, - %arg1: memref<2x2xf32>, - %arg2: memref<2x2xf32>) { - %0 = alloc() {temp = true} : memref<2x2xf32> - // CHECK-NEXT: "lmhlo.exponential"(%arg1, %arg2) : (memref<2x2xf32>, memref<2x2xf32>) -> () - "lmhlo.exponential"(%arg1, %arg2) : (memref<2x2xf32>, memref<2x2xf32>) -> () - // CHECK-NEXT: "lmhlo.exponential"(%arg0, %arg2) : (memref<2x2xf32>, memref<2x2xf32>) -> () - "lmhlo.exponential"(%arg0, %0) : (memref<2x2xf32>, memref<2x2xf32>) -> () - "lmhlo.copy"(%0, %arg2) : (memref<2x2xf32>, memref<2x2xf32>) -> () - dealloc %0 : memref<2x2xf32> - "lmhlo.terminator"() : () -> () -} - -// ----- - -// CHECK-LABEL: func @must_be_removed_second -func @must_be_removed_second(%arg0: memref<2x2xf32>, - %arg1: memref<2x2xf32>, - %arg2: memref<2x2xf32>) { - %0 = alloc() {temp = true} : memref<2x2xf32> - // CHECK-NEXT: "lmhlo.exponential"(%arg0, %arg2) : (memref<2x2xf32>, memref<2x2xf32>) -> () - "lmhlo.exponential"(%arg0, %0) : (memref<2x2xf32>, memref<2x2xf32>) -> () - "lmhlo.copy"(%0, %arg2) : (memref<2x2xf32>, memref<2x2xf32>) -> () - // CHECK-NEXT: "lmhlo.exponential"(%arg1, %arg2) : (memref<2x2xf32>, memref<2x2xf32>) -> () - "lmhlo.exponential"(%arg1, %arg2) : (memref<2x2xf32>, memref<2x2xf32>) -> () - dealloc %0 : memref<2x2xf32> - "lmhlo.terminator"() : () -> () -} - -// ----- - -// CHECK-LABEL: func @reduce -func @reduce(%arg0: memref<1x8xf32>, %arg1: memref, %arg2: memref<1xf32>) { - %0 = alloc() : memref<1xf32> - "lmhlo.reduce"(%arg0, %arg1, %0) ( { - // CHECK: ^bb0(%[[ARG0:.*]]: memref, %[[ARG1:.*]]: memref, - // CHECK-SAME: %[[ARG2:.*]]: memref) - ^bb0(%arg3: memref, %arg4: memref, %arg5: memref): - %1 = alloc() : memref - // CHECK: "lmhlo.add"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) - "lmhlo.add"(%arg3, %arg4, %1) - : (memref, memref, memref) -> () - // CHECK-NOT; lmhlo.copy - "lmhlo.copy"(%1, %arg5) : (memref, memref) -> () - "lmhlo.terminator"() : () -> () - }) {dimensions = dense<1> : tensor<1xi64>} - : (memref<1x8xf32>, memref, memref<1xf32>) -> () - "lmhlo.copy"(%0, %arg2) : (memref<1xf32>, memref<1xf32>) -> () - return -}