Disable FMul associativity and factoring rules that cause egraph explosion

FMul associativity and FP factoring rules interact to cause exponential
egraph growth on shaders with many shared multiply chains (e.g. sky-shader
with 1,264 FMul nodes). FMul associativity alone creates 305K+ matches via
Catalan-number re-parenthesization across shared e-classes; combined with
factoring (304K matches), tuples explode from 2,875 to 916K in one
iteration, causing OOM/SIGKILL.

Constant chain merging (FMul(FMul(x, const_a), const_b) → FMul(x, a*b))
is handled by dedicated rules and does not need general associativity.

Author: LegNeato

rust/spirv-tools-opt/src/direct/emit.rs

@@ -1301,9 +1301,7 @@ fn emit_pattern(
             // Refine it to avoid signed/unsigned mismatches:
             // 1. Signed int ops (ConvertSToF, SLessThan, etc.) → signed int type
             // 2. Other cross-class ops → infer from Sym operand's actual type
-            let operand_type = if operand_class == TypeClass::Int
-                && result_class != operand_class
-            {
+            let operand_type = if operand_class == TypeClass::Int && result_class != operand_class {
                 match opcode {
                     Op::ConvertSToF
                     | Op::ConvertFToS
@@ -1315,9 +1313,9 @@ fn emit_pattern(
                     | Op::SRem
                     | Op::SMod
                     | Op::SConvert => ctx.signed_int32_type.unwrap_or(operand_type),
-                    _ => infer_operand_type_from_args(
-                        args, arity, operand_class, operand_type, ctx,
-                    ),
+                    _ => {
+                        infer_operand_type_from_args(args, arity, operand_class, operand_type, ctx)
+                    }
                 }
             } else if result_class != operand_class {
                 // Non-int cross-class: infer from Sym args for consistency

rust/spirv-tools-opt/src/direct/mod.rs

@@ -1037,8 +1037,7 @@ pub fn optimize_module_direct(module: &Module) -> Result<Module, EgglogOptError>
                         // SPIR-V types (e.g. two constants with the same bit
                         // pattern but different type IDs, or Vec4/Vec2 both as
                         // Expr). CopyObject requires operand type == result type.
-                        let type_matches = ctx.id_to_type.get(&id)
-                            == ctx.id_to_type.get(&alias_id);
+                        let type_matches = ctx.id_to_type.get(&id) == ctx.id_to_type.get(&alias_id);
                         if type_matches {
                             id_aliases.insert(id, alias_id);
                             used_ids.insert(alias_id);
@@ -1091,8 +1090,8 @@ pub fn optimize_module_direct(module: &Module) -> Result<Module, EgglogOptError>
                                 // (e.g. Vec4 and Vec2 both as Expr). Emitting a
                                 // CopyObject with mismatched types would cause
                                 // validation errors.
-                                let type_matches = ctx.id_to_type.get(&id)
-                                    == ctx.id_to_type.get(&final_id);
+                                let type_matches =
+                                    ctx.id_to_type.get(&id) == ctx.id_to_type.get(&final_id);
                                 if type_matches {
                                     id_aliases.insert(id, final_id);
                                     used_ids.insert(final_id);
@@ -2357,9 +2356,7 @@ fn find_spirv_type(module: &Module, opcode: Op, width: Option<u32>) -> Option<Wo
         .find(|inst| {
             inst.class.opcode == opcode
                 && match width {
-                    Some(w) => {
-                        inst.operands.first() == Some(&rspirv::dr::Operand::LiteralBit32(w))
-                    }
+                    Some(w) => inst.operands.first() == Some(&rspirv::dr::Operand::LiteralBit32(w)),
                     None => true,
                 }
         })

rust/spirv-tools-opt/src/egglog_opt/tests.rs

@@ -6587,10 +6587,7 @@ fn test_boolconst_float_comparison_reflexive() {
 
     // FOrdLt(x,x) IS correctly folded to BoolConst(0)
     let check2 = egraph.parse_and_run_program(None, "(check (= lt_root false_val))");
-    assert!(
-        check2.is_ok(),
-        "FOrdLt(x, x) should fold to BoolConst(0)"
-    );
+    assert!(check2.is_ok(), "FOrdLt(x, x) should fold to BoolConst(0)");
 }
 
 #[test]

rust/spirv-tools-opt/src/rules/arithmetic.egg

@@ -209,6 +209,11 @@
 ; We only want to FACTOR (collect terms), not DISTRIBUTE (expand terms).
 
 ; (x * a) + (x * b) => x * (a + b)  - factoring out common term
+; NOTE: These are safe because there is no general integer Mul associativity
+; rule (only constant-folding variants). Without associativity creating new
+; Mul groupings, the feedback loop that causes FP factoring explosion cannot
+; form. If a general (Mul (Mul a b) c) -> (Mul a (Mul b c)) rule is ever
+; added, these must be disabled too (see floating_point.egg factoring).
 (rule ((= e (Add (Mul x a) (Mul x b))))
       ((union e (Mul x (Add a b)))))
 (rule ((= e (Add (Mul a x) (Mul b x))))

rust/spirv-tools-opt/src/rules/floating_point.egg

@@ -201,18 +201,30 @@
 ;   2. Any FAdd(x, Y) matches FAdd(FAdd(x, 0), Y) and cascades endlessly
 ;
 ; Safe variant: Only apply when 'b' is an FAdd/FMul (not an FConst that could be 0 or 1).
+; NOTE: FAdd associativity is safe with current workloads (only 11 matches at iter 3
+; on sky-shader), but could exhibit Catalan blowup if a shader has deep FAdd chains
+; with shared sub-expressions. Monitor if new shaders cause FAdd explosion.
 (rule ((= e (FAdd (FAdd a b) c)) (= b (FAdd _ _)))
       ((union e (FAdd a (FAdd b c)))))
-(rule ((= e (FMul (FMul a b) c)) (= b (FMul _ _)))
-      ((union e (FMul a (FMul b c)))))
+; DISABLED: FMul associativity causes exponential blowup on its own.
+; With ~1,264 FMul nodes sharing e-classes (common in sky/lighting shaders),
+; re-parenthesization creates 305K+ matches and 455K FMul tuples in one
+; iteration via Catalan-number growth across shared sub-expressions.
+; Constant chain merging (FMul(FMul(x, const_a), const_b) → FMul(x, const_a*b))
+; is handled by dedicated rules below and does NOT need this.
+; (rule ((= e (FMul (FMul a b) c)) (= b (FMul _ _)))
+;       ((union e (FMul a (FMul b c)))))
 
 ; Factoring: a*c + b*c = (a+b)*c
-; One-directional to prevent expansion that causes e-graph explosion.
-; We recognize expanded patterns and factor them, but don't distribute.
-(rule ((= e (FAdd (FMul a c) (FMul b c))))
-      ((union e (FMul (FAdd a b) c))))
-(rule ((= e (FAdd (FMul c a) (FMul c b))))
-      ((union e (FMul c (FAdd a b)))))
+; DISABLED: Creates a feedback loop with FMul associativity above.
+; Factoring creates new FMul+FAdd nodes, associativity re-parenthesizes the
+; new FMul nodes, exposing more factoring opportunities → exponential blowup.
+; In the sky-shader, this explodes from 2,875 to 916,652 tuples in one iteration
+; (305K associativity matches + 304K factoring matches).
+; (rule ((= e (FAdd (FMul a c) (FMul b c))))
+;       ((union e (FMul (FAdd a b) c))))
+; (rule ((= e (FAdd (FMul c a) (FMul c b))))
+;       ((union e (FMul c (FAdd a b)))))
 
 ; x + x = 2*x - ONE-DIRECTIONAL to prevent explosion
 ; (FMul x 2 -> FAdd x x would expand forever)

rust/spirv-tools-opt/tests/opt_block_cli.rs

@@ -5443,20 +5443,18 @@ fn build_loop_with_selection_in_continue_block() -> Vec<u32> {
         std::iter::empty(),
     )
     .unwrap();
-    let val = b.load(int, None, counter, None, std::iter::empty()).unwrap();
-    let cond = b
-        .s_less_than(bool_ty, None, val, c10)
-        .expect("less than");
+    let val = b
+        .load(int, None, counter, None, std::iter::empty())
+        .unwrap();
+    let cond = b.s_less_than(bool_ty, None, val, c10).expect("less than");
     b.branch_conditional(cond, continue_label, merge_label, std::iter::empty())
         .unwrap();
 
     // Continue block: contains a selection construct
     b.begin_block(Some(continue_label)).unwrap();
     let c2 = b.constant_bit32(int, 2);
     let is_even_rem = b.s_mod(int, None, val, c2).unwrap();
-    let is_even = b
-        .i_equal(bool_ty, None, is_even_rem, c0)
-        .expect("is_even");
+    let is_even = b.i_equal(bool_ty, None, is_even_rem, c0).expect("is_even");
     b.selection_merge(sel_merge_label, SelectionControl::NONE)
         .unwrap();
     b.branch_conditional(is_even, sel_then_label, sel_else_label, std::iter::empty())
@@ -5536,8 +5534,7 @@ fn build_store_with_typed_value() -> Vec<u32> {
     let _sint = b.type_int(32, 1);
     let void = b.type_void();
     let func_ty = b.type_function(void, vec![uint]);
-    let ptr_uint =
-        b.type_pointer(None, rspirv::spirv::StorageClass::Function, uint);
+    let ptr_uint = b.type_pointer(None, rspirv::spirv::StorageClass::Function, uint);
 
     let func = b
         .begin_function(void, None, FunctionControl::NONE, func_ty)
@@ -5610,8 +5607,7 @@ fn build_float_to_signed_int_module() -> Vec<u32> {
     let float = b.type_float(32, None);
     let void = b.type_void();
     let func_ty = b.type_function(void, vec![sint]);
-    let ptr_sint =
-        b.type_pointer(None, rspirv::spirv::StorageClass::Function, sint);
+    let ptr_sint = b.type_pointer(None, rspirv::spirv::StorageClass::Function, sint);
 
     let func = b
         .begin_function(void, None, FunctionControl::NONE, func_ty)
@@ -5684,8 +5680,7 @@ fn build_matrix_multiply_chain() -> Vec<u32> {
     let mat4 = b.type_matrix(vec4, 4);
     let void = b.type_void();
     let func_ty = b.type_function(void, vec![mat4, mat4, mat4, vec4]);
-    let ptr_vec4 =
-        b.type_pointer(None, rspirv::spirv::StorageClass::Function, vec4);
+    let ptr_vec4 = b.type_pointer(None, rspirv::spirv::StorageClass::Function, vec4);
 
     let func = b
         .begin_function(void, None, FunctionControl::NONE, func_ty)