Updates LLVM usage to match
[8396aeb07cdd](https://github.com/llvm/llvm-project/commit/8396aeb07cdd)

PiperOrigin-RevId: 366034463
This commit is contained in:
A. Unique TensorFlower 2021-03-31 08:00:52 -07:00 committed by TensorFlow MLIR Team
parent bbe0aa204c
commit af3bc47a8b
6 changed files with 64 additions and 6 deletions

View File

@ -15,9 +15,9 @@
load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") 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) LLVM_BAZEL_TAG = "llvm-project-{commit}".format(commit = LLVM_COMMIT)

View File

@ -1,2 +1,2 @@
afed50a14b34eb619624aed5c85f4f610f360650 8396aeb07cddd8ab9a6a154a4ab7ac56fc24bda5

View File

@ -212,6 +212,8 @@ def LHLO_ReduceOp: LHLO_Op<"reduce", [SameVariadicOperandSize]>, BASE_HLO_Reduce
); );
let regions = (region SizedRegion<1>:$body); let regions = (region SizedRegion<1>:$body);
let hasCanonicalizer = 1;
} }
def LHLO_ReduceWindowOp: LHLO_Op<"reduce_window", []>, BASE_HLO_ReduceWindowOp { def LHLO_ReduceWindowOp: LHLO_Op<"reduce_window", []>, BASE_HLO_ReduceWindowOp {

View File

@ -36,6 +36,7 @@ limitations under the License.
#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h" #include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/IR/Attributes.h" #include "mlir/IR/Attributes.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/IR/Builders.h" #include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Dialect.h" #include "mlir/IR/Dialect.h"
@ -225,6 +226,61 @@ static LogicalResult Verify(CustomCallOp op) {
return success(); 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<ReduceOp> {
using OpRewritePattern<ReduceOp>::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<lmhlo::CopyOp>(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<lmhlo::CopyOp>(op) || llvm::isa<memref::DeallocOp>(op) ||
llvm::isa<memref::AllocOp>(op))
continue;
rewriter.clone(op, bvm);
}
rewriter.eraseOp(reduce);
return success();
}
};
void ReduceOp::getCanonicalizationPatterns(OwningRewritePatternList& results,
MLIRContext* context) {
results.insert<RemoveCopyInReduceBody>(context);
}
} // namespace lmhlo } // namespace lmhlo
} // namespace mlir } // namespace mlir

View File

@ -1,7 +1,7 @@
// RUN: mlir-hlo-opt %s -chlo-legalize-to-hlo -hlo-legalize-to-lhlo \ // RUN: mlir-hlo-opt %s -chlo-legalize-to-hlo -hlo-legalize-to-lhlo \
// RUN: -std-bufferize -tensor-bufferize -finalizing-bufferize \ // RUN: -std-bufferize -tensor-bufferize -finalizing-bufferize \
// RUN: --canonicalize -buffer-hoisting -buffer-deallocation \ // RUN: -canonicalize -buffer-hoisting -buffer-deallocation \
// RUN: -copy-removal -canonicalize -cse -lhlo-legalize-to-linalg \ // RUN: -canonicalize -cse -lhlo-legalize-to-linalg \
// RUN: -lhlo-fuse-linalg -convert-linalg-to-loops -canonicalize -cse \ // RUN: -lhlo-fuse-linalg -convert-linalg-to-loops -canonicalize -cse \
// RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-std \ // RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-std \
// RUN: -convert-std-to-llvm \ // RUN: -convert-std-to-llvm \

View File

@ -1,6 +1,6 @@
// RUN: mlir-hlo-opt %s -chlo-legalize-to-hlo \ // RUN: mlir-hlo-opt %s -chlo-legalize-to-hlo \
// RUN: -hlo-legalize-to-lhlo -buffer-hoisting \ // 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: -lhlo-legalize-to-linalg -lhlo-fuse-linalg -convert-linalg-to-loops \
// RUN: -lower-affine -convert-scf-to-std -canonicalize -cse \ // RUN: -lower-affine -convert-scf-to-std -canonicalize -cse \
// RUN: -convert-std-to-llvm | mlir-cpu-runner -e main \ // RUN: -convert-std-to-llvm | mlir-cpu-runner -e main \