Integrate LLVM at llvm/llvm-project@8396aeb07c

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

PiperOrigin-RevId: 366034463

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)
 

build_tools/llvm_version.txt

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

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 {

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<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 mlir
 

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: -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: -convert-linalg-to-llvm -lower-affine -convert-scf-to-std \
 // RUN: -convert-std-to-llvm \

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 \