Import reduce-constant-prop patterns from stablehlo aggressive folder

[SuryanshSS1011]

Summary of the PR

Closes #1084.

Lifts LowerBoolSplatConstantsIntoReduceOpRegion and FoldReduceOpToConstantInitializer from openxla/stablehlo's StablehloAggressiveFolder.cpp into EnzymeHLOOpt's default pipeline, per @wsmoses' suggestion on the issue.

Together with Enzyme's existing AndSimplify/OrSimplify patterns, the two new patterns let --enzyme-hlo-opt fold reduce-of-splat-constants with And/Or bodies into a single constant.

Example

%c_0 = stablehlo.constant dense<true> : tensor<1x140xi1>
%c_7 = stablehlo.constant dense<true> : tensor<i1>
%0 = stablehlo.reduce(%c_0 init: %c_7) applies stablehlo.and
     across dimensions = [0, 1]
     : (tensor<1x140xi1>, tensor<i1>) -> tensor<i1>
return %0 : tensor<i1>

now folds to:

%c = stablehlo.constant dense<true> : tensor<i1>
return %c : tensor<i1>

The fold proceeds in three steps:

1. LowerBoolSplatConstantsIntoReduceOpRegion rewrites the reduce body so the splat constants are materialized inside the region.
2. Enzyme's existing AndSimplify folds that to const_true.
3. FoldReduceOpToConstantInitializer then folds the reduce, whose body now returns a constant.

Scope

• Imports the two patterns directly relevant to the issue. Other FoldReduceOp* patterns from the upstream file (FoldReduceOpReducingZeroDims, FoldReduceOpWithRedundantResults) are out of scope.
• Restricted to And/Or body ops, matching upstream.

Test plan

• New lit test test/lit_tests/reduce_const_prop.mlir covering the issue's exact MLIR, a sibling or case, and a negative case (non-constant input).
• Existing reduce/const/and lit tests pass, verified 23 tests via bazel test:
  • addreduceslicefusion (2), and_const_prop, and_pad_pad, binop_const_lift_computation, binopcomplexconstsimplify, broadcastreduce, concatreduce (2, 3), constpadconcat_to_concat, constpropthroughbarrier, convert_to_splatted_constants, convertconst, elementwise_reduce_slice_fuse (2), foldgather, foldpad, fullreduce_nocrash, gatherconstprop, is_finite_const_prop, log_const_prop, math_const_prop.
• Build clean (bazel build -c opt :enzymexlamlir-opt), no new warnings introduced.

[wsmoses]

can we extend this to also support add/mul, including non constants? [and appropriate tests]

[SuryanshSS1011]

A clarification before I extend further: on closer look, this isn't already a pattern in upstream's StablehloAggressiveFolder. The existing LowerBoolSplatConstantsIntoReduceOpRegion only handles idempotent body ops (and/or, where f(x, x) = f(x, x, ..., x)), and the cascade only terminates correctly for those.

For add/mul the result depends on the reduce dimension count: reduce(splat<2>, init=0) {add} over 4 elements should fold to 8, not 2. So it needs a new pattern that computes the closed form (init + N*x for add, init * x^N for mul), with overflow and float-precision handling.

That's meaningfully bigger than what's in this PR. Happy to do it as a follow-up if you'd rather keep #2524 focused on the and/or fold lifted from upstream. Let me know.

[wsmoses]

make sure to also add this to tablegen [see the dev docs]

[SuryanshSS1011]

Thanks for the quick review @wsmoses! I will extend to add/mul, add tests, and wire up the tablegen registration. I might come back with a clarifying question on the "non-constants" scope once I've sketched it out, since the current pattern is structured around splat-constant inputs.