Set op-folding element limit via a config option

Instead of hard-coding the max number of elements in a folded operation,
set the limit via a config option. (The default value matches the old
constant's value of 65536.)

Author: mrguenther

stablehlo/transforms/optimization/Passes.h

@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef STABLEHLO_TRANSFORMS_OPTIMIZATION_PASSES_H
 #define STABLEHLO_TRANSFORMS_OPTIMIZATION_PASSES_H
 
+#include <cstdint>
 #include <memory>
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -42,6 +43,7 @@ void populateStablehloCanonicalizationPatterns(MLIRContext *context,
 void populateStablehloAggressiveFolderPatterns(RewritePatternSet *patterns,
                                                MLIRContext *context,
                                                bool foldFloat,
+                                               int64_t foldOpElementLimit,
                                                PatternBenefit benefit = 1);
 
 /// A subset of folding patterns for StableHLO that is necessary for shape

stablehlo/transforms/optimization/Passes.td

@@ -15,16 +15,23 @@ limitations under the License.
 
 include "mlir/Pass/PassBase.td"
 
+defvar stablehlo_aggressive_folder_options = [
+  Option<"foldFloat", "fold-float", "bool", /*default=*/"true",
+          "Allow for potentially lossy computations using float type.">,
+  Option<"foldOpElementLimit", "fold-op-element-limit", "int64_t",
+          /*default=*/"65536",
+          "Upper limit on how many elements may be folded by an op folder. "
+          "This limit doesn't apply in certain special cases such as when "
+          "adding 0, multiplying by 1, or operating on splats in some ways.">,
+];
+
 def StablehloAggressiveFolderPass
     : Pass<"stablehlo-aggressive-folder", "func::FuncOp"> {
   let summary = "Folds StableHLO operations";
   let dependentDialects = [
     "mlir::stablehlo::StablehloDialect",
   ];
-  let options = [
-    Option<"foldFloat", "fold-float", "bool", /*default=*/"true",
-           "Allow for potentially lossy computations using float type.">,
-  ];
+  let options = stablehlo_aggressive_folder_options;
 }
 
 def StablehloAggressiveSimplificationPass
@@ -138,4 +145,5 @@ def StablehloTargetIndependentOptimizationPass
   let dependentDialects = [
     "mlir::stablehlo::StablehloDialect",
   ];
+  let options = stablehlo_aggressive_folder_options;
 }

stablehlo/transforms/optimization/StablehloAggressiveFolder.cpp

@@ -65,11 +65,6 @@ namespace stablehlo {
 
 namespace {
 
-// This is an upper limit on how many elements can be folded by an op folder.
-// This limit doesn't apply to some special cases like adding a zero,
-// multiplying by one, doing many operations with splats.
-constexpr int64_t kFoldOpEltLimit = 65536;
-
 // DenseElementsAttr can be constructed from ArrayRef<APInt> but not from
 // ArrayRef<APSInt>. This helper bridges the gap.
 DenseIntElementsAttr getTensorAttr(ShapedType type, ArrayRef<APSInt> values) {
@@ -425,15 +420,19 @@ struct EvalCompareOpPattern : public OpRewritePattern<CompareOp> {
 
 struct FoldConcatenateOpPattern final
     : OpRewritePattern<mlir::stablehlo::ConcatenateOp> {
-  using OpRewritePattern::OpRewritePattern;
+  FoldConcatenateOpPattern(MLIRContext* context, int64_t foldOpElementLimit,
+                           PatternBenefit benefit = 1,
+                           ArrayRef<StringRef> generatedNames = {})
+      : OpRewritePattern(context, benefit, generatedNames),
+        foldOpElementLimit(foldOpElementLimit) {}
 
   LogicalResult matchAndRewrite(mlir::stablehlo::ConcatenateOp op,
                                 PatternRewriter& rewriter) const override {
     RankedTensorType type = op.getType();
     if (!type.hasStaticShape()) return failure();
 
     size_t numElems = type.getNumElements();
-    if (numElems > kFoldOpEltLimit) return failure();
+    if (numElems > foldOpElementLimit) return failure();
 
     // Fold concatenate when all inputs are constants.
     OperandRange inputs = op.getInputs();
@@ -463,6 +462,8 @@ struct FoldConcatenateOpPattern final
         op, DenseElementsAttr::get(op.getType(), newElems));
     return success();
   }
+
+  int64_t foldOpElementLimit;
 };
 
 struct EvalConcatenateOpPattern : public OpRewritePattern<ConcatenateOp> {
@@ -817,13 +818,18 @@ struct FoldSqrtOpPattern : public OpRewritePattern<mlir::stablehlo::SqrtOp> {
 };
 
 struct EvalIotaOpPattern : public OpRewritePattern<IotaOp> {
-  using OpRewritePattern::OpRewritePattern;
+  EvalIotaOpPattern(MLIRContext* context, int64_t foldOpElementLimit,
+                    PatternBenefit benefit,
+                    ArrayRef<StringRef> generatedNames = {})
+      : OpRewritePattern(context, benefit, generatedNames),
+        foldOpElementLimit(foldOpElementLimit) {}
+
   LogicalResult matchAndRewrite(IotaOp op,
                                 PatternRewriter& rewriter) const override {
     LLVM_DEBUG(llvm::dbgs() << "EvalIotaOpPattern folding: " << op << '\n');
     auto resultType = cast<RankedTensorType>(op.getType());
     size_t numElems = resultType.getNumElements();
-    if (numElems > kFoldOpEltLimit)
+    if (numElems > foldOpElementLimit)
       return rewriter.notifyMatchFailure(op, "too many elements to fold");
 
     auto elementType = resultType.getElementType();
@@ -864,6 +870,8 @@ struct EvalIotaOpPattern : public OpRewritePattern<IotaOp> {
         op, DenseIntElementsAttr::get(resultType, values));
     return success();
   }
+
+  int64_t foldOpElementLimit;
 };
 
 template <typename RangeType>
@@ -927,7 +935,8 @@ struct StablehloAggressiveFolderPass
 
   LogicalResult initialize(MLIRContext* context) override {
     RewritePatternSet patterns_(context);
-    populateStablehloAggressiveFolderPatterns(&patterns_, context, foldFloat);
+    populateStablehloAggressiveFolderPatterns(&patterns_, context, foldFloat,
+                                              foldOpElementLimit);
     patterns = std::move(patterns_);
 
     return success();
@@ -947,16 +956,18 @@ struct StablehloAggressiveFolderPass
 void populateStablehloAggressiveFolderPatterns(RewritePatternSet* patterns,
                                                MLIRContext* context,
                                                bool foldFloat,
+                                               int64_t foldOpElementLimit,
                                                PatternBenefit benefit) {
   populateStablehloShapeFolderPatterns(patterns, context, foldFloat, benefit);
-  patterns->add<EvalIotaOpPattern>(context, benefit);
+  patterns->add<EvalIotaOpPattern>(context, foldOpElementLimit, benefit);
   patterns->add<EvalTransposeOpPattern>(context, benefit);
 
   // TODO: Consolidate FoldOp patterns
   // One is used by Shape Refinement, the other is a generic folder.
   patterns->add<FoldAddOpPattern, FoldBroadcastInDimSplatPattern,
-                FoldConcatenateOpPattern, FoldMulOpPattern,
-                FoldSubtractOpPattern, FoldSqrtOpPattern>(context);
+                FoldMulOpPattern, FoldSubtractOpPattern, FoldSqrtOpPattern>(
+      context);
+  patterns->add<FoldConcatenateOpPattern>(context, foldOpElementLimit);
 }
 
 void populateStablehloShapeFolderPatterns(RewritePatternSet* patterns,

stablehlo/transforms/optimization/StablehloTargetIndependentOptimization.cpp

@@ -32,11 +32,6 @@ namespace stablehlo {
 #define GEN_PASS_DEF_STABLEHLOTARGETINDEPENDENTOPTIMIZATIONPASS
 #include "stablehlo/transforms/optimization/Passes.h.inc"
 
-// This is an upper limit on how many elements can be folded by an op folder.
-// This limit doesn't apply to some special cases like adding a zero,
-// multiplying by one, doing many operations with splats.
-constexpr int64_t kFoldOpEltLimit = 65536;
-
 struct StablehloTargetIndependentOptimizationPass
     : public impl::StablehloTargetIndependentOptimizationPassBase<
           StablehloTargetIndependentOptimizationPass> {
@@ -45,9 +40,9 @@ struct StablehloTargetIndependentOptimizationPass
 
   LogicalResult initialize(MLIRContext* context) override {
     RewritePatternSet patterns_(context);
-    bool foldFloat = true;
     populateStablehloCanonicalizationPatterns(context, &patterns_);
     populateStablehloAggressiveFolderPatterns(&patterns_, context, foldFloat,
+                                              foldOpElementLimit,
                                               /*benefit=*/2);
     patterns = std::move(patterns_);
 
@@ -58,7 +53,7 @@ struct StablehloTargetIndependentOptimizationPass
     GreedyRewriteConfig config;
     config.fold = true;
     config.cseConstants = true;
-    config.maxIterations = kFoldOpEltLimit;
+    config.maxIterations = foldOpElementLimit;
     config.useTopDownTraversal = false;
     if (failed(applyPatternsGreedily(getOperation(), patterns, config)))
       signalPassFailure();