Integrate LLVM at llvm/llvm-project@2bc22ea02edd (#2875)

Author: GleasonK

BUILD.bazel

@@ -1107,6 +1107,24 @@ gentbl_cc_library(
     deps = ["@llvm-project//mlir:PassBaseTdFiles"],
 )
 
+cc_library(
+    name = "stablehlo_broadcast_lowering",
+    srcs = [
+        "stablehlo/transforms/StablehloBroadcastLowering.cpp",
+    ],
+    hdrs = [
+        "stablehlo/transforms/StablehloBroadcastLowering.h",
+    ],
+    strip_include_prefix = ".",
+    deps = [
+        ":stablehlo_ops",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
 cc_library(
     name = "stablehlo_pass_utils",
     srcs = [

WORKSPACE.bazel

@@ -17,9 +17,9 @@ workspace(name = "stablehlo")
 
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
-LLVM_COMMIT = "42a8ff877d47131ecb1280a1cc7e5e3c3bca6952"
+LLVM_COMMIT = "2bc22ea02edda5926f3e53f141def9bf212ac1db"
 
-LLVM_SHA256 = "f768c5c3b987f68318b8ab3dd4530e54988dfe7d6bfb9b7c9c96acf503367d50"
+LLVM_SHA256 = "4a034eda852b3c2d448d38e8661cbac45ae2233a29defeb55913fa5205cd29f7"
 
 http_archive(
     name = "llvm-raw",

build_tools/llvm_version.txt

@@ -1 +1 @@
-42a8ff877d47131ecb1280a1cc7e5e3c3bca6952
+2bc22ea02edda5926f3e53f141def9bf212ac1db

stablehlo/dialect/StablehloOps.cpp

@@ -3275,12 +3275,12 @@ Attribute StablehloDialect::parseAttribute(DialectAsmParser& parser,
 // Entry point for Attribute printing, TableGen generated code will handle the
 // dispatch to the individual classes.
 void StablehloDialect::printAttribute(Attribute attr,
-                                      DialectAsmPrinter& os) const {
+                                      DialectAsmPrinter& printer) const {
   if (auto type_extensions = dyn_cast<TypeExtensionsAttr>(attr)) {
-    hlo::printTypeExtensions(cast<hlo::BoundedAttrInterface>(attr), os);
+    hlo::printTypeExtensions(cast<hlo::BoundedAttrInterface>(attr), printer);
     return;
   }
-  LogicalResult result = generatedAttributePrinter(attr, os);
+  LogicalResult result = generatedAttributePrinter(attr, printer);
   (void)result;
   assert(succeeded(result));
 }

stablehlo/dialect/StablehloOps.h

@@ -93,13 +93,14 @@ class StablehloDialect : public Dialect {
   Type parseType(DialectAsmParser& parser) const override;
 
   // Prints a type registered to this dialect.
-  void printType(Type type, DialectAsmPrinter& os) const override;
+  void printType(Type type, DialectAsmPrinter& printer) const override;
 
   // Parses an attribute registered to this dialect.
   Attribute parseAttribute(DialectAsmParser& parser, Type type) const override;
 
   // Prints an attribute registered to this dialect.
-  void printAttribute(Attribute attr, DialectAsmPrinter& os) const override;
+  void printAttribute(Attribute attr,
+                      DialectAsmPrinter& printer) const override;
 
   // Get the set dialect version.
   std::optional<StablehloDialectVersion> getVersion() const;

stablehlo/integrations/python/mlir/dialects/InterpreterOps.td

@@ -17,6 +17,6 @@ limitations under the License.
 #ifndef STABLEHLO_INTEGRATIONS_PYTHON_INTERPRETER_OPS
 #define STABLEHLO_INTEGRATIONS_PYTHON_INTERPRETER_OPS
 
-include "third_party/stablehlo/stablehlo/reference/InterpreterOps.h"
+include "stablehlo/reference/InterpreterOps.h"
 
 #endif

stablehlo/tests/BUILD.bazel

@@ -102,6 +102,8 @@ cc_library(
     deps = [
         ":test_utils_inc_gen",
         "//:stablehlo_assembly_format",
+        "//:stablehlo_broadcast_lowering",
+        "//:stablehlo_ops",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",

stablehlo/tests/CMakeLists.txt

@@ -48,6 +48,7 @@ add_mlir_library(StablehloTestUtils
   MLIRSupport
   MLIRTransformUtils
   StablehloAssemblyFormat
+  StablehloBroadcastLowering
 )
 
 set(LLVM_TARGET_DEFINITIONS CheckOps.td)

stablehlo/tests/TestUtils.cpp

@@ -19,6 +19,7 @@ limitations under the License.
 #include <utility>
 
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Casting.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
@@ -28,26 +29,50 @@ limitations under the License.
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/TypeRange.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Rewrite/FrozenRewritePatternSet.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "stablehlo/dialect/StablehloOps.h"
+#include "stablehlo/transforms/StablehloBroadcastLowering.h"
 
 namespace mlir {
 namespace hlo {
 
 namespace {
 
+struct BroadcastValuesPattern : public RewritePattern {
+  explicit BroadcastValuesPattern(MLIRContext* context)
+      : RewritePattern("hlo_test_broadcast.numpy_broadcast", 1, context) {}
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& rewriter) const override {
+    // Process all operands
+    SmallVector<Value> operands = llvm::to_vector(op->getOperands());
+    auto broadcastedOperands =
+        stablehlo::numpyBroadcastIfNeeded(rewriter, operands);
+    if (failed(broadcastedOperands)) return failure();
+
+    // Replace with custom call to avoid pattern reapplication
+    auto customCall = stablehlo::CustomCallOp::create(
+        rewriter, op->getLoc(), op->getResultTypes(), *broadcastedOperands);
+    customCall.setCallTargetName("numpy_broadcasted");
+    customCall.setHasSideEffect(true);
+    rewriter.replaceOp(op, customCall);
+    return success();
+  }
+};
+
 struct InferReturnTypesPattern : public RewritePattern {
-  explicit InferReturnTypesPattern(MLIRContext *context)
+  explicit InferReturnTypesPattern(MLIRContext* context)
       : RewritePattern("hlo_test_infer.get_return_types", 1, context) {}
-  LogicalResult matchAndRewrite(Operation *op,
-                                PatternRewriter &rewriter) const override {
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& rewriter) const override {
     if (op->getNumOperands() != 1) return failure();
-    auto *definingOp = op->getOperand(0).getDefiningOp();
+    auto* definingOp = op->getOperand(0).getDefiningOp();
     auto definingOpInt =
         llvm::dyn_cast_or_null<InferTypeOpInterface>(definingOp);
     if (!definingOpInt) return failure();
@@ -62,8 +87,8 @@ struct InferReturnTypesPattern : public RewritePattern {
     OperationState state(op->getLoc(), "hlo_test_infer.return_types",
                          op->getOperands(), op->getResultTypes(),
                          op->getAttrs());
-    auto *newOp = rewriter.create(state);
-    for (const auto &it : llvm::enumerate(types))
+    auto* newOp = rewriter.create(state);
+    for (const auto& it : llvm::enumerate(types))
       newOp->setAttr((StringRef("types") + Twine(it.index())).str(),
                      TypeAttr::get(it.value()));
     rewriter.replaceOp(op, {newOp->getResults()});
@@ -72,10 +97,10 @@ struct InferReturnTypesPattern : public RewritePattern {
 };
 
 struct ReifyReturnTypeShapesPattern : public RewritePattern {
-  explicit ReifyReturnTypeShapesPattern(MLIRContext *context)
+  explicit ReifyReturnTypeShapesPattern(MLIRContext* context)
       : RewritePattern("hlo_test_infer.reify_return_type_shapes", 1, context) {}
-  LogicalResult matchAndRewrite(Operation *op,
-                                PatternRewriter &rewriter) const override {
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& rewriter) const override {
     if (op->getNumOperands() != 1) return failure();
     auto definingOp =
         op->getOperand(0).getDefiningOp<InferShapedTypeOpInterface>();
@@ -89,7 +114,7 @@ struct ReifyReturnTypeShapesPattern : public RewritePattern {
   }
 };
 
-LogicalResult checkSpeculatability(PatternRewriter &rewriter, Operation *op,
+LogicalResult checkSpeculatability(PatternRewriter& rewriter, Operation* op,
                                    mlir::Speculation::Speculatability spec) {
   if (op->getNumOperands() != 1) return failure();
   auto definingOp =
@@ -106,67 +131,86 @@ LogicalResult checkSpeculatability(PatternRewriter &rewriter, Operation *op,
 }
 
 struct IsSpeculatablePattern : public RewritePattern {
-  explicit IsSpeculatablePattern(MLIRContext *context)
+  explicit IsSpeculatablePattern(MLIRContext* context)
       : RewritePattern("hlo_test_speculatability.is_speculatable", 1, context) {
   }
-  LogicalResult matchAndRewrite(Operation *op,
-                                PatternRewriter &rewriter) const override {
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& rewriter) const override {
     return checkSpeculatability(rewriter, op, mlir::Speculation::Speculatable);
   }
 };
 
 struct IsRecursivelySpeculatablePattern : public RewritePattern {
-  explicit IsRecursivelySpeculatablePattern(MLIRContext *context)
+  explicit IsRecursivelySpeculatablePattern(MLIRContext* context)
       : RewritePattern("hlo_test_speculatability.is_recursively_speculatable",
                        1, context) {}
-  LogicalResult matchAndRewrite(Operation *op,
-                                PatternRewriter &rewriter) const override {
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& rewriter) const override {
     return checkSpeculatability(rewriter, op,
                                 mlir::Speculation::RecursivelySpeculatable);
   }
 };
 
 struct IsNotSpeculatablePattern : public RewritePattern {
-  explicit IsNotSpeculatablePattern(MLIRContext *context)
+  explicit IsNotSpeculatablePattern(MLIRContext* context)
       : RewritePattern("hlo_test_speculatability.is_not_speculatable", 1,
                        context) {}
-  LogicalResult matchAndRewrite(Operation *op,
-                                PatternRewriter &rewriter) const override {
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& rewriter) const override {
     return checkSpeculatability(rewriter, op,
                                 mlir::Speculation::NotSpeculatable);
   }
 };
 
+#define GEN_PASS_DEF_HLOTESTBROADCASTPASS
 #define GEN_PASS_DEF_HLOTESTINFERPASS
 #define GEN_PASS_DEF_HLOTESTSPECULATABILITYPASS
 #include "stablehlo/tests/TestUtils.h.inc"
 
+struct HloTestBroadcastPass
+    : public impl::HloTestBroadcastPassBase<HloTestBroadcastPass> {
+  LogicalResult initialize(MLIRContext* context) override {
+    RewritePatternSet patterns(context);
+    patterns.add<BroadcastValuesPattern>(context);
+    patterns_ = std::move(patterns);
+    return success();
+  }
+
+  void runOnOperation() override {
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns_))))
+      return signalPassFailure();
+  }
+
+ private:
+  FrozenRewritePatternSet patterns_;
+};
+
 struct HloTestInferPass : public impl::HloTestInferPassBase<HloTestInferPass> {
-  LogicalResult initialize(MLIRContext *context) override {
-    RewritePatternSet patterns_(context);
-    patterns_.add<InferReturnTypesPattern>(context);
-    patterns_.add<ReifyReturnTypeShapesPattern>(context);
-    patterns = std::move(patterns_);
+  LogicalResult initialize(MLIRContext* context) override {
+    RewritePatternSet patterns(context);
+    patterns.add<InferReturnTypesPattern>(context);
+    patterns.add<ReifyReturnTypeShapesPattern>(context);
+    patterns_ = std::move(patterns);
     return success();
   }
 
   void runOnOperation() override {
-    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns_))))
       return signalPassFailure();
   }
 
  private:
-  FrozenRewritePatternSet patterns;
+  FrozenRewritePatternSet patterns_;
 };
 
 struct HloTestSpeculatabilityPass
     : public impl::HloTestSpeculatabilityPassBase<HloTestSpeculatabilityPass> {
-  LogicalResult initialize(MLIRContext *context) override {
-    RewritePatternSet patterns_(context);
-    patterns_.add<IsSpeculatablePattern>(context);
-    patterns_.add<IsNotSpeculatablePattern>(context);
-    patterns_.add<IsRecursivelySpeculatablePattern>(context);
-    patterns = std::move(patterns_);
+  LogicalResult initialize(MLIRContext* context) override {
+    RewritePatternSet patterns(context);
+    patterns.add<IsSpeculatablePattern>(context);
+    patterns.add<IsNotSpeculatablePattern>(context);
+    patterns.add<IsRecursivelySpeculatablePattern>(context);
+    patterns_ = std::move(patterns);
     return success();
   }
 
@@ -175,11 +219,11 @@ struct HloTestSpeculatabilityPass
     config.setMaxIterations(1)
         .setUseTopDownTraversal(true)
         .setRegionSimplificationLevel(GreedySimplifyRegionLevel::Disabled);
-    (void)applyPatternsGreedily(getOperation(), std::move(patterns));
+    (void)applyPatternsGreedily(getOperation(), std::move(patterns_));
   }
 
  private:
-  FrozenRewritePatternSet patterns;
+  FrozenRewritePatternSet patterns_;
 };
 
 #define GEN_PASS_REGISTRATION

stablehlo/tests/TestUtils.td

@@ -16,6 +16,11 @@ limitations under the License.
 
 include "mlir/Pass/PassBase.td"
 
+def HloTestBroadcastPass : Pass<"hlo-test-broadcast", "func::FuncOp"> {
+  let summary = "Uses test ops to invoke BroadcastUtils methods.";
+  let dependentDialects = ["stablehlo::StablehloDialect"];
+}
+
 def HloTestInferPass : Pass<"hlo-test-infer", "func::FuncOp"> {
   let summary = "Uses test ops to invoke InferShapedTypeOpInterface methods.";
   let dependentDialects = ["shape::ShapeDialect"];