From af3bc47a8b52c03c2a86bb598253d81add343b4b Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 31 Mar 2021 08:00:52 -0700 Subject: [PATCH] Integrate LLVM at llvm/llvm-project@8396aeb07cdd Updates LLVM usage to match [8396aeb07cdd](https://github.com/llvm/llvm-project/commit/8396aeb07cdd) PiperOrigin-RevId: 366034463 --- WORKSPACE | 4 +- build_tools/llvm_version.txt | 2 +- include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td | 2 + lib/Dialect/mhlo/IR/lhlo_ops.cc | 56 ++++++++++++++++++++ tests/end2end/broadcast.mlir | 4 +- tests/end2end/reduce.mlir | 2 +- 6 files changed, 64 insertions(+), 6 deletions(-) diff --git a/WORKSPACE b/WORKSPACE index 2ecb9d3..6fad727 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -15,9 +15,9 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") -LLVM_COMMIT = "afed50a14b34eb619624aed5c85f4f610f360650" +LLVM_COMMIT = "8396aeb07cddd8ab9a6a154a4ab7ac56fc24bda5" -LLVM_SHA256 = "eccfb27acb4bdcd9177af54b0b673a19fcb92a0bfbc223a09e8cff1666441fe4" +LLVM_SHA256 = "af486a33012d65171c6cf5a676c8da3b48414f3f55b4021511c102f1042be531" LLVM_BAZEL_TAG = "llvm-project-{commit}".format(commit = LLVM_COMMIT) diff --git a/build_tools/llvm_version.txt b/build_tools/llvm_version.txt index 9b90556..0e00940 100644 --- a/build_tools/llvm_version.txt +++ b/build_tools/llvm_version.txt @@ -1,2 +1,2 @@ -afed50a14b34eb619624aed5c85f4f610f360650 +8396aeb07cddd8ab9a6a154a4ab7ac56fc24bda5 diff --git a/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td b/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td index db3aa43..baa1e0a 100644 --- a/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td +++ b/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.td @@ -212,6 +212,8 @@ def LHLO_ReduceOp: LHLO_Op<"reduce", [SameVariadicOperandSize]>, BASE_HLO_Reduce ); let regions = (region SizedRegion<1>:$body); + + let hasCanonicalizer = 1; } def LHLO_ReduceWindowOp: LHLO_Op<"reduce_window", []>, BASE_HLO_ReduceWindowOp { diff --git a/lib/Dialect/mhlo/IR/lhlo_ops.cc b/lib/Dialect/mhlo/IR/lhlo_ops.cc index 40ebb51..2c14eb8 100644 --- a/lib/Dialect/mhlo/IR/lhlo_ops.cc +++ b/lib/Dialect/mhlo/IR/lhlo_ops.cc @@ -36,6 +36,7 @@ limitations under the License. #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" #include "mlir/IR/Attributes.h" +#include "mlir/IR/BlockAndValueMapping.h" #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Dialect.h" @@ -225,6 +226,61 @@ static LogicalResult Verify(CustomCallOp op) { return success(); } +//===----------------------------------------------------------------------===// +// ReduceOp +//===----------------------------------------------------------------------===// + +// Removes `lmhlo.copy` inside ReduceOp body. +// +// TODO(b/183920887): Remove this pattern as soon as bufferization is fixed. +struct RemoveCopyInReduceBody : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(ReduceOp reduce, + PatternRewriter& rewriter) const override { + // Find the only `lmhlo.copy` in the body of `reduce`. + CopyOp the_only_copy; + for (auto& op : reduce.body().front()) { + if (auto copy = dyn_cast(op)) { + if (the_only_copy == nullptr) { + the_only_copy = copy; + } else { + the_only_copy = nullptr; + break; + } + } + } + if (!the_only_copy) return failure(); + + auto new_reduce = rewriter.cloneWithoutRegions(reduce); + Block* new_block = + rewriter.createBlock(&new_reduce.body(), new_reduce.body().end(), + reduce.body().front().getArgumentTypes()); + + mlir::BlockAndValueMapping bvm; + for (auto item : llvm::zip(reduce.body().front().getArguments(), + new_block->getArguments())) { + bvm.map(std::get<0>(item), std::get<1>(item)); + } + bvm.map(the_only_copy.operand(), bvm.lookup(the_only_copy.output())); + + rewriter.setInsertionPointToStart(new_block); + for (auto& op : reduce.body().front()) { + if (llvm::isa(op) || llvm::isa(op) || + llvm::isa(op)) + continue; + rewriter.clone(op, bvm); + } + rewriter.eraseOp(reduce); + return success(); + } +}; + +void ReduceOp::getCanonicalizationPatterns(OwningRewritePatternList& results, + MLIRContext* context) { + results.insert(context); +} + } // namespace lmhlo } // namespace mlir diff --git a/tests/end2end/broadcast.mlir b/tests/end2end/broadcast.mlir index 71bd028..1e0e30d 100644 --- a/tests/end2end/broadcast.mlir +++ b/tests/end2end/broadcast.mlir @@ -1,7 +1,7 @@ // RUN: mlir-hlo-opt %s -chlo-legalize-to-hlo -hlo-legalize-to-lhlo \ // RUN: -std-bufferize -tensor-bufferize -finalizing-bufferize \ -// RUN: --canonicalize -buffer-hoisting -buffer-deallocation \ -// RUN: -copy-removal -canonicalize -cse -lhlo-legalize-to-linalg \ +// RUN: -canonicalize -buffer-hoisting -buffer-deallocation \ +// RUN: -canonicalize -cse -lhlo-legalize-to-linalg \ // RUN: -lhlo-fuse-linalg -convert-linalg-to-loops -canonicalize -cse \ // RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-std \ // RUN: -convert-std-to-llvm \ diff --git a/tests/end2end/reduce.mlir b/tests/end2end/reduce.mlir index 52c59c4..69e47e7 100644 --- a/tests/end2end/reduce.mlir +++ b/tests/end2end/reduce.mlir @@ -1,6 +1,6 @@ // RUN: mlir-hlo-opt %s -chlo-legalize-to-hlo \ // RUN: -hlo-legalize-to-lhlo -buffer-hoisting \ -// RUN: -buffer-deallocation -copy-removal -canonicalize -cse \ +// RUN: -buffer-deallocation -canonicalize -cse \ // RUN: -lhlo-legalize-to-linalg -lhlo-fuse-linalg -convert-linalg-to-loops \ // RUN: -lower-affine -convert-scf-to-std -canonicalize -cse \ // RUN: -convert-std-to-llvm | mlir-cpu-runner -e main \