Upgrade LLVM-Project, StableHLO dependencies (#414)

This is a combination of the following commits (from older to most recent):

## [compiler/Dialect/Plan] Improve efficiency of TensorRT clustering

Runs TensorRT analysis conversion on the container op and uses that result
in the clustering analysis instead of checking conversion op-by-op, which
will no longer work well in after an upcoming LLVM upgrade.

This change also removes the "tensorrt major version" parameter for the
`ClusterKindAttrInterface::getClusteringOpts` method arguments since that
information can be passed directly as a parameter of the `TensorRTClusterKindAttr``

## Upgrade all dependencies

This upgrades dependencies to the following base commits:

```
LLVM_COMMIT_UPSTREAM = "6c64c8a6f3f77c30745c751d4163ff6bf2fc323b"
TORCH_MLIR_COMMIT = "30c519369ed7eabad0282d0f874500a9b41fcbbd"
STABLEHLO_COMMIT = "6e403b1aa6a71f5eaa09cc720e4ad42f692745e6"
```

## Update how cluster kinds and their options are populated

This change updates the 'plan-stablehlo-clustering' pass so that it no longer
populates default "cluster kinds" if the 'plan.cluster_kinds' attribute is
missing on the module. Instead, it is up to the top-level compilation entrypoints
or the frontend to populate these key metadata attributes.

This allows for significant simplification in how options are passed to
various passes and pipelines. For example, the 'stablehlo-clustering' and
'plan-segmentation-pipeline' both no longer need to accept pass options
specifically for the TensorRT clustering options. These options are dictated
purely by the TensorRT clustering attribute attached to the module. Only the
top level compilation entrypoint requires such options, which it can use to
populate the clustering backend attributes parameters.

GitOrigin-RevId: 66aa8a1e07b4c94195ed27bcd96f7fd2d00c6f3e

Author: christopherbate

mlir-tensorrt/CMakeLists.txt

@@ -124,12 +124,13 @@ if(PROJECT_IS_TOP_LEVEL)
   elseif(MLIR_TRT_LLVM_COMMIT)
     mtrt_llvm_project(
       NAME llvm_project
-      VERSION 0.0.20240812
+      VERSION 0.0.20241126
       URL "https://github.com/llvm/llvm-project/archive/${MLIR_TRT_LLVM_COMMIT}.zip"
       EXCLUDE_FROM_ALL TRUE
       SOURCE_SUBDIR "llvm"
-      PATCHES "${CMAKE_SOURCE_DIR}/build_tools/llvm-project.patch"
-
+      PATCHES
+        "${CMAKE_CURRENT_LIST_DIR}/build_tools/patches/mlir/000_fix_bufferization_tensor_encoding_memory_spaces.patch"
+        "${CMAKE_CURRENT_LIST_DIR}/build_tools/patches/mlir/001-mlir-Add-a-null-pointer-check-in-symbol-lookup-11516.patch"
       OPTIONS
         "LLVM_ENABLE_PROJECTS mlir"
         "MLIR_ENABLE_BINDINGS_PYTHON ${MLIR_TRT_ENABLE_PYTHON}"
@@ -180,12 +181,11 @@ include(HandleLLVMOptions)
 # Download Stablehlo if it isn't provided by a parent project.
 if(MLIR_TRT_ENABLE_HLO AND NOT TARGET StablehloOps)
   mtrt_add_stablehlo(
-    VERSION 1.6.4
-    GIT_TAG 1456dfa1e1a83aab0cc717714ba3695886f60302
+    VERSION 1.8.0
+    GIT_TAG 6e403b1aa6a71f5eaa09cc720e4ad42f692745e6
     GIT_REPOSITORY "https://github.com/openxla/stablehlo.git"
-    PATCHES 
-      "${MLIR_TENSORRT_ROOT_DIR}/build_tools/stablehlo.patch"
-      "${MLIR_TENSORRT_ROOT_DIR}/build_tools/stablehlo_aggressive_folder.patch"
+    PATCHES
+      "${CMAKE_CURRENT_LIST_DIR}/build_tools/patches/stablehlo/0001-transforms-Fix-simplification-patterns-for-stablehlo.patch"
     OPTIONS
       "STABLEHLO_ENABLE_BINDINGS_PYTHON ${MLIR_TRT_ENABLE_PYTHON}"
       "STABLEHLO_BUILD_EMBEDDED ON"

mlir-tensorrt/build_tools/cmake/LLVMCommit.cmake

@@ -1 +1 @@
-set(LLVM_COMMIT "c49770c60f26e449379447109f7d915bd8de0384")
+set(LLVM_COMMIT "6c64c8a6f3f77c30745c751d4163ff6bf2fc323b")

mlir-tensorrt/build_tools/llvm-project.patch renamed to mlir-tensorrt/build_tools/patches/mlir/000_fix_bufferization_tensor_encoding_memory_spaces.patch

@@ -72,10 +72,10 @@ index 1c70a4b8df92..c97a1fe819a3 100644
 
    let hasFolder = 1;
 diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
-index a610ddcc9899..18c4e5cedc8c 100644
+index afc193b5517d..b5c81a0e73e2 100644
 --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
 +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
-@@ -533,6 +533,10 @@ def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
+@@ -526,6 +526,10 @@ def OneShotBufferize : Pass<"one-shot-bufferize"> {
             /*default=*/"false",
             "The memory space of an memref types must always be inferred. If "
             "unset, a default memory space of 0 is used otherwise.">,
@@ -87,10 +87,10 @@ index a610ddcc9899..18c4e5cedc8c 100644
              /*default=*/"false",
             "Test only: Only run inplaceability analysis and annotate IR">,
 diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
-index d51d63f243ea..550ac7e83b9e 100644
+index 85604eef2f28..065739ea8e59 100644
 --- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
 +++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
-@@ -719,7 +719,7 @@ void bufferization::replaceOpWithBufferizedValues(RewriterBase &rewriter,
+@@ -718,7 +718,7 @@ void bufferization::replaceOpWithBufferizedValues(RewriterBase &rewriter,
        // loose all of its users and eventually DCE away.
        rewriter.setInsertionPointAfter(op);
        replacement = rewriter.create<bufferization::ToTensorOp>(
@@ -100,7 +100,7 @@ index d51d63f243ea..550ac7e83b9e 100644
      replacements.push_back(replacement);
    }
 diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
-index 04a8ff30ee94..d1d5d3b89b3e 100644
+index f1841b860ff8..774d1a2ec04b 100644
 --- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
 +++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
 @@ -23,6 +23,16 @@ using namespace mlir::bufferization;
@@ -121,7 +121,7 @@ index 04a8ff30ee94..d1d5d3b89b3e 100644
      OpBuilder &b, Value value, MemRefType destType,
      const BufferizationOptions &options) {
 diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
-index e422086c9fde..d4271c68cbcc 100644
+index 429695126a95..911c2e862e6f 100644
 --- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
 +++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
 @@ -67,10 +67,14 @@ BufferizeTypeConverter::BufferizeTypeConverter() {
@@ -153,12 +153,12 @@ index e422086c9fde..d4271c68cbcc 100644
 +        };
        }
        opt.printConflicts = printConflicts;
-       opt.testAnalysisOnly = testAnalysisOnly;
+       opt.bufferAlignment = bufferAlignment;
 diff --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
-index 87464ccb7172..c654dbba46e8 100644
+index c2b8614148bf..9797b73f534a 100644
 --- a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
 +++ b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
-@@ -479,10 +479,6 @@ struct FromElementsOpInterface
+@@ -480,10 +480,6 @@ struct FromElementsOpInterface
      auto fromElementsOp = cast<tensor::FromElementsOp>(op);
      auto tensorType = cast<RankedTensorType>(fromElementsOp.getType());
 
@@ -169,7 +169,7 @@ index 87464ccb7172..c654dbba46e8 100644
      // Allocate a buffer for the result.
      Location loc = op->getLoc();
      auto shape = tensorType.getShape();
-@@ -492,10 +488,12 @@ struct FromElementsOpInterface
+@@ -493,10 +489,12 @@ struct FromElementsOpInterface
          /*copy=*/false);
      if (failed(tensorAlloc))
        return failure();
@@ -223,7 +223,7 @@ index 5293977fe733..55e086ff0110 100644
  }
 
 diff --git a/mlir/test/Dialect/Bufferization/Transforms/finalizing-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/finalizing-bufferize.mlir
-index ab18ce05e355..148908536d6c 100644
+index ab18ce05e355..97a69e153e39 100644
 --- a/mlir/test/Dialect/Bufferization/Transforms/finalizing-bufferize.mlir
 +++ b/mlir/test/Dialect/Bufferization/Transforms/finalizing-bufferize.mlir
 @@ -4,8 +4,8 @@
@@ -287,32 +287,17 @@ index ab18ce05e355..148908536d6c 100644
    return %1 : memref<?xf32>
  }
 
-@@ -77,19 +77,20 @@ func.func @static_layout_to_no_layout_cast(%m: memref<?xf32, strided<[1], offset
- // TODO: to_memref with layout maps not supported yet. This should fold to a
- // memref.cast.
- func.func @no_layout_to_dyn_layout_cast(%m: memref<?xf32>) -> memref<?xf32, strided<[1], offset: ?>> {
--  %0 = bufferization.to_tensor %m : memref<?xf32>
-+  %0 = bufferization.to_tensor %m : memref<?xf32> -> tensor<?xf32>
-   // expected-error @+1 {{failed to legalize unresolved materialization from ('memref<?xf32>') to 'memref<?xf32, strided<[1], offset: ?>>' that remained live after conversion}}
--  %1 = bufferization.to_memref %0 : memref<?xf32, strided<[1], offset: ?>>
--  // expected-note @below{{see existing live user here}}
-+  %1 = bufferization.to_memref %0 : tensor<?xf32> -> memref<?xf32, strided<[1], offset: ?>>
-+  // expected-note @below {{see existing live user here:}}
-   return %1 : memref<?xf32, strided<[1], offset: ?>>
- }
- 
+@@ -87,9 +87,8 @@ func.func @no_layout_to_dyn_layout_cast(%m: memref<?xf32>) -> memref<?xf32, stri
  // -----
 
  func.func @illegal_unranked_to_rank(%m: memref<*xf32>) -> memref<?xf32> {
 -  // expected-note @+1 {{prior use here}}
 -  %0 = bufferization.to_tensor %m : memref<*xf32>
 -  // expected-error @+1 {{expects different type than prior uses: 'tensor<?xf32>' vs 'tensor<*xf32>'}}
 -  %1 = bufferization.to_memref %0 : memref<?xf32>
-+
-+  %0 = bufferization.to_tensor %m : memref<*xf32> -> tensor<*xf32>
++  %0 = bufferization.to_tensor %m : memref<*xf32> -> tensor<?xf32>
 +  // expected-error @+1 {{failed to legalize unresolved materialization from ('memref<*xf32>') to 'memref<?xf32>' that remained live after conversion}}
-+  %1 = bufferization.to_memref %0 : tensor<*xf32> -> memref<?xf32>
-+  // expected-note @below {{see existing live user here:}}
++  %1 = bufferization.to_memref %0 : tensor<?xf32> -> memref<?xf32>
    return %1 : memref<?xf32>
  }
 diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
@@ -800,10 +785,10 @@ index 4bc2ed140da9..af5e745cb3a9 100644
    return %pack_18 : tensor<1x1x8x4x4x8xi32>
  }
 diff --git a/mlir/test/Dialect/MemRef/normalize-memrefs.mlir b/mlir/test/Dialect/MemRef/normalize-memrefs.mlir
-index c7af033a22a2..1c80a9f6024d 100644
+index 11114bcf2b1a..76f6ee511093 100644
 --- a/mlir/test/Dialect/MemRef/normalize-memrefs.mlir
 +++ b/mlir/test/Dialect/MemRef/normalize-memrefs.mlir
-@@ -356,10 +356,10 @@ func.func @neg_map() -> memref<2x3xf32, #neg> {
+@@ -360,11 +360,11 @@ func.func @neg_map() -> memref<2x3xf32, #neg> {
  // CHECK-LABEL: func @memref_with_strided_offset
  func.func @memref_with_strided_offset(%arg0: tensor<128x512xf32>, %arg1: index, %arg2: index) -> tensor<16x512xf32> {
    %c0 = arith.constant 0 : index
@@ -816,6 +801,7 @@ index c7af033a22a2..1c80a9f6024d 100644
 +  %1 = bufferization.to_tensor %cast : memref<16x512xf32, strided<[?, ?], offset: ?>> -> tensor<16x512xf32>
    return %1 : tensor<16x512xf32>
  }
+ 
 diff --git a/mlir/test/Dialect/SCF/bufferize.mlir b/mlir/test/Dialect/SCF/bufferize.mlir
 index ff1612310255..c7ca1dcb031b 100644
 --- a/mlir/test/Dialect/SCF/bufferize.mlir
@@ -2072,19 +2058,6 @@ index 78e29979ca1a..6332b35ef6c0 100644
  // CHECK:           %[[VAL_12:.*]] = memref.load %[[VAL_11]][] : memref<f32>
  // CHECK:           %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
  // CHECK:           %[[VAL_14:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-diff --git a/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir b/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir
-index f819458e0385..a8398deb7a3b 100644
---- a/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir
-+++ b/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir
-@@ -85,7 +85,7 @@ func.func @sqsum(%arg0: tensor<?x?x?x?xi32, #COO>) -> tensor<i32> {
- // CHECK:           %[[VAL_3:.*]] = arith.constant 0 : index
- // CHECK:           %[[VAL_4:.*]] = arith.constant 0 : i32
- // CHECK:           %[[VAL_5:.*]] = arith.constant dense<0> : tensor<10xi32>
--// CHECK:           %[[VAL_6:.*]] = bufferization.to_memref %[[VAL_5]] : memref<10xi32>
-+// CHECK:           %[[VAL_6:.*]] = bufferization.to_memref %[[VAL_5]] : tensor<10xi32> -> memref<10xi32>
- // CHECK:           linalg.fill ins(%[[VAL_4]] : i32) outs(%[[VAL_6]] : memref<10xi32>)
- // CHECK:           %[[VAL_7:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 0 : index} : tensor<10xi32, #sparse{{.*}}> to memref<?xindex>
- // CHECK:           %[[VAL_8:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 0 : index} : tensor<10xi32, #sparse{{.*}}> to memref<?xindex>
 diff --git a/mlir/test/Dialect/SparseTensor/sparse_lower.mlir b/mlir/test/Dialect/SparseTensor/sparse_lower.mlir
 index c27df0078552..62ebc7ef3d96 100644
 --- a/mlir/test/Dialect/SparseTensor/sparse_lower.mlir
@@ -2739,10 +2712,10 @@ index 3a3c8af15e6e..13be67d0562e 100644
    // CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index
    // CHECK-DAG: %[[dim0:.*]] = memref.dim %[[m1]], %[[c0]]
 diff --git a/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir b/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir
-index e2169fe1404c..c0e6b415e1a6 100644
+index dc4306b8316a..bd71970cfb08 100644
 --- a/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir
 +++ b/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir
-@@ -387,7 +387,7 @@ func.func @tensor.reshape() -> tensor<2x2x5xf32> {
+@@ -402,7 +402,7 @@ func.func @tensor.reshape() -> tensor<2x2x5xf32> {
  // CHECK-SAME:    %[[INPUT:[a-zA-Z0-9]*]]: memref<2x2xf32, strided<[?, ?], offset: ?>, 3>,
  // CHECK-SAME:    %[[LAYOUT:[a-zA-Z0-9]*]]: memref<2xi32, strided<[?], offset: ?>>,
  func.func @reshape_with_non_identity_layout(%arg0: memref<2x2xf32, strided<[?, ?], offset: ?>, 3>, %arg1: tensor<2xi32>, %idx: index) -> f32 {

mlir-tensorrt/build_tools/patches/mlir/001-mlir-Add-a-null-pointer-check-in-symbol-lookup-11516.patch

@@ -0,0 +1,25 @@
+From efdb53b115f5ce7233482a344bfc0c26f9e47041 Mon Sep 17 00:00:00 2001
+From: Christopher Bate <[email protected]>
+Date: Fri, 15 Nov 2024 05:17:54 +0000
+Subject: [PATCH] [mlir] Add a null pointer check in symbol lookup #115165
+
+---
+ mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp | 2 ++
+ 1 file changed, 2 insertions(+)
+
+diff --git a/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
+index 3c190d4e9919..e805e21d878b 100644
+--- a/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
++++ b/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
+@@ -186,6 +186,8 @@ void DeadCodeAnalysis::initializeSymbolCallables(Operation *top) {
+       // If a callable symbol has a non-call use, then we can't be guaranteed to
+       // know all callsites.
+       Operation *symbol = symbolTable.lookupSymbolIn(top, use.getSymbolRef());
++      if (!symbol)
++        continue;
+       auto *state = getOrCreate<PredecessorState>(getProgramPointAfter(symbol));
+       propagateIfChanged(state, state->setHasUnknownPredecessors());
+     }
+-- 
+2.47.0
+

mlir-tensorrt/build_tools/patches/stablehlo/0001-transforms-Fix-simplification-patterns-for-stablehlo.patch

@@ -0,0 +1,120 @@
+From 00a3c5e6b9207bae81c6d401fa368fdfe270122b Mon Sep 17 00:00:00 2001
+From: Christopher Bate <[email protected]>
+Date: Fri, 22 Nov 2024 22:43:47 +0000
+Subject: [PATCH] [transforms] Fix simplification patterns for
+ `stablehlo.(and|or)`
+
+Fixes an issue in `stablehlo-aggressive-simplification` where `%1` in
+the below would get replaced by `%arg0`:
+
+```
+  %0 = stablehlo.constant dense<1> : tensor<2xi32>
+  %1 = stablehlo.and %0, %arg0 : tensor<2xi32>
+```
+
+The pattern was checking whether `%0` is equal to `0b1` and was
+only tested on bools. A similar bug existed for `stablehlo.and`. Fixed
+by just making sure the constant is integer with all bits set to 1.
+---
+ .../stablehlo_aggressive_simplification.mlir  | 38 +++++++++++++++++++
+ ...ablehloAggressiveSimplificationPatterns.td | 14 ++++++-
+ 2 files changed, 50 insertions(+), 2 deletions(-)
+
+diff --git a/stablehlo/tests/transforms/stablehlo_aggressive_simplification.mlir b/stablehlo/tests/transforms/stablehlo_aggressive_simplification.mlir
+index 809c0700..b2d05de3 100644
+--- a/stablehlo/tests/transforms/stablehlo_aggressive_simplification.mlir
++++ b/stablehlo/tests/transforms/stablehlo_aggressive_simplification.mlir
+@@ -63,6 +63,25 @@ func.func @and_one(%arg0: tensor<2xi1>) -> tensor<2xi1> {
+   return %1 : tensor<2xi1>
+ }
+ 
++// CHECK-LABEL: @and_i32_one
++func.func @and_i32_one(%arg0: tensor<2xi32>) -> tensor<2xi32> {
++  %0 = stablehlo.constant dense<1> : tensor<2xi32>
++  %1 = stablehlo.and %0, %arg0 : tensor<2xi32>
++  // CHECK: %[[AND:.+]] = stablehlo.and
++  // CHECK: return %[[AND]]
++  return %1 : tensor<2xi32>
++}
++
++// CHECK-LABEL: @and_i32_neg_one
++//  CHECK-SAME:  (%[[ARG0:.+]]: tensor<2xi32>)
++func.func @and_i32_neg_one(%arg0: tensor<2xi32>) -> tensor<2xi32> {
++  %0 = stablehlo.constant dense<-1> : tensor<2xi32>
++  %1 = stablehlo.and %0, %arg0 : tensor<2xi32>
++  // CHECK-NOT:  stablehlo.and
++  // CHECK: return %[[ARG0]]
++  return %1 : tensor<2xi32>
++}
++
+ // -----
+ 
+ /////////
+@@ -540,6 +559,25 @@ func.func @or_one(%arg0: tensor<2xi1>) -> tensor<2xi1> {
+   return %1 : tensor<2xi1>
+ }
+ 
++// CHECK-LABEL: @or_i32_one
++func.func @or_i32_one(%arg0: tensor<2xi32>) -> tensor<2xi32> {
++  %0 = stablehlo.constant dense<1> : tensor<2xi32>
++  %1 = stablehlo.or %0, %arg0 : tensor<2xi32>
++  // CHECK: %[[OR:.+]] = stablehlo.or
++  // CHECK: return %[[OR]]
++  return %1 : tensor<2xi32>
++}
++
++// CHECK-LABEL: @or_i32_neg_one
++func.func @or_i32_neg_one(%arg0: tensor<2xi32>) -> tensor<2xi32> {
++  %0 = stablehlo.constant dense<-1> : tensor<2xi32>
++  %1 = stablehlo.or %0, %arg0 : tensor<2xi32>
++  // CHECK-NOT: stablehlo.or
++  // CHECK: [[NEG_ONE:%.+]] = stablehlo.constant dense<-1> : tensor<2xi32>
++  // CHECK: return [[NEG_ONE]]
++  return %1 : tensor<2xi32>
++}
++
+ // -----
+ 
+ /////////
+diff --git a/stablehlo/transforms/StablehloAggressiveSimplificationPatterns.td b/stablehlo/transforms/StablehloAggressiveSimplificationPatterns.td
+index 31f1f475..cef9f303 100644
+--- a/stablehlo/transforms/StablehloAggressiveSimplificationPatterns.td
++++ b/stablehlo/transforms/StablehloAggressiveSimplificationPatterns.td
+@@ -41,6 +41,16 @@ def AnySplat : AttrConstraint<CPred<"$_self.isSplat()">, "is any splat">;
+ def AnyZero : AttrConstraint<
+    CPred<"::mlir::matchPattern($_self, m_AnyAttrOf(m_Zero(), m_AnyZeroFloat()))">, "is int or float zero">;
+ 
++def IntAllOnes : AttrConstraint<
++    CPred<[{
++    ::mlir::matchPattern($_self,
++        ::mlir::detail::constant_int_predicate_matcher{
++            [](const llvm::APInt &val) {
++                return val.isAllOnes();
++        }})
++    }]>,
++    "is integer with all bits set to 1">;
++
+ def IntZero : AttrConstraint<
+    CPred<"::mlir::matchPattern($_self, m_Zero())">, "is integer zero">;
+ 
+@@ -101,7 +111,7 @@ def : Pat<(StableHLO_AndOp $lhs, (StableHLO_ConstantOp:$zero IntZero:$value)),
+           (replaceWithValue $zero)>;
+ 
+ // Pattern: and(X, 1) -> X
+-def : Pat<(StableHLO_AndOp $lhs, (StableHLO_ConstantOp:$one IntOne:$value)),
++def : Pat<(StableHLO_AndOp $lhs, (StableHLO_ConstantOp:$one IntAllOnes:$value)),
+           (replaceWithValue $lhs)>;
+ 
+ ////////
+@@ -208,7 +218,7 @@ def : Pat<(StableHLO_MulOp $lhs, (StableHLO_ConstantOp IntOne:$value)),
+ def : CanonicalizeConstantToRhs<StableHLO_OrOp>;
+ 
+ // Pattern: or(X, 1) -> 1
+-def : Pat<(StableHLO_OrOp $lhs, (StableHLO_ConstantOp:$one IntOne:$value)),
++def : Pat<(StableHLO_OrOp $lhs, (StableHLO_ConstantOp:$one IntAllOnes:$value)),
+           (replaceWithValue $one)>;
+ 
+ // Pattern: or(X, 0) -> X
+-- 
+2.47.0
+