perf: Don't do fused-multiply-add on scalars

crates/polars-ops/src/series/ops/fused.rs

@@ -4,7 +4,7 @@ use polars_core::prelude::*;
 use polars_core::utils::align_chunks_ternary;
 use polars_core::with_match_physical_numeric_polars_type;
 
-// a + (b * c)
+// (a * b) + c
 fn fma_arr<T: NumericNative>(
     a: &PrimitiveArray<T>,
     b: &PrimitiveArray<T>,
@@ -22,7 +22,7 @@ fn fma_arr<T: NumericNative>(
         .iter()
         .zip(b.iter())
         .zip(c.iter())
-        .map(|((a, b), c)| *a + (*b * *c))
+        .map(|((a, b), c)| *a * *b + *c)
         .collect::<Vec<_>>();
     PrimitiveArray::from_data_default(out.into(), validity)
 }
@@ -51,10 +51,7 @@ pub fn fma_columns(a: &Column, b: &Column, c: &Column) -> Column {
             fma_ca(a, b, c).into_column()
         })
     } else {
-        (a.as_materialized_series()
-            + &(b.as_materialized_series() * c.as_materialized_series()).unwrap())
-            .unwrap()
-            .into()
+        (&(a * b).unwrap() + c).unwrap()
     }
 }
 
@@ -105,10 +102,7 @@ pub fn fsm_columns(a: &Column, b: &Column, c: &Column) -> Column {
             fsm_ca(a, b, c).into_column()
         })
     } else {
-        (a.as_materialized_series()
-            - &(b.as_materialized_series() * c.as_materialized_series()).unwrap())
-            .unwrap()
-            .into()
+        (a - &(b * c).unwrap()).unwrap()
     }
 }
 
@@ -158,9 +152,6 @@ pub fn fms_columns(a: &Column, b: &Column, c: &Column) -> Column {
             fms_ca(a, b, c).into_column()
         })
     } else {
-        (&(a.as_materialized_series() * b.as_materialized_series()).unwrap()
-            - c.as_materialized_series())
-        .unwrap()
-        .into()
+        (&(a * b).unwrap() - c).unwrap()
     }
 }

crates/polars-plan/src/plans/optimizer/fused.rs

@@ -23,30 +23,23 @@ fn get_expr(input: &[Node], op: FusedOperator, expr_arena: &Arena<AExpr>) -> AEx
 }
 
 fn check_eligible(
-    left: &Node,
-    right: &Node,
+    nodes: &[Node],
     expr_arena: &Arena<AExpr>,
     schema: &Schema,
 ) -> PolarsResult<bool> {
-    let field_left = expr_arena
-        .get(*left)
-        .to_field(&ToFieldContext::new(expr_arena, schema))?;
-    let type_right = expr_arena
-        .get(*right)
-        .to_dtype(&ToFieldContext::new(expr_arena, schema))?;
-    let type_left = &field_left.dtype;
-    // Exclude literals for now as these will not benefit from fused operations downstream #9857
+    // Exclude scalars for now as these will not benefit from fused operations downstream #9857
     // This optimization would also interfere with the `col -> lit` type-coercion rules
     // And it might also interfere with constant folding which is a more suitable optimizations here
-    if type_left.is_primitive_numeric()
-        && type_right.is_primitive_numeric()
-        && !has_aexpr_literal(*left, expr_arena)
-        && !has_aexpr_literal(*right, expr_arena)
-    {
-        Ok(true)
-    } else {
-        Ok(false)
+    for node in nodes {
+        let field = expr_arena
+            .get(*node)
+            .to_field(&ToFieldContext::new(expr_arena, schema))?;
+        if !field.dtype.is_primitive_numeric() || is_scalar_ae(*node, expr_arena) {
+            return Ok(false);
+        }
     }
+
+    Ok(true)
 }
 
 impl OptimizationRule for FusedArithmetic {
@@ -69,80 +62,39 @@ impl OptimizationRule for FusedArithmetic {
         match expr {
             BinaryExpr {
                 left,
-                op: Operator::Plus,
+                op: outer_op @ (Operator::Plus | Operator::Minus),
                 right,
             } => {
-                // FUSED MULTIPLY ADD
-                // For fma the plus is always the out as the multiply takes prevalence
-                match expr_arena.get(*left) {
-                    // Argument order is a + b * c
-                    // so we must swap operands
-                    //
-                    // input
-                    // (a * b) + c
-                    // swapped as
-                    // c + (a * b)
-                    BinaryExpr {
-                        left: a,
-                        op: Operator::Multiply,
-                        right: b,
-                    } => Ok(check_eligible(left, right, expr_arena, schema)?.then(|| {
-                        let input = &[*right, *a, *b];
-                        get_expr(input, FusedOperator::MultiplyAdd, expr_arena)
-                    })),
-                    _ => match expr_arena.get(*right) {
-                        // input
-                        // (a + (b * c)
-                        // kept as input
-                        BinaryExpr {
-                            left: a,
-                            op: Operator::Multiply,
-                            right: b,
-                        } => Ok(check_eligible(left, right, expr_arena, schema)?.then(|| {
-                            let input = &[*left, *a, *b];
-                            get_expr(input, FusedOperator::MultiplyAdd, expr_arena)
-                        })),
-                        _ => Ok(None),
-                    },
-                }
-            },
+                let (a, b, other, mul_is_left) = if let BinaryExpr {
+                    left: a,
+                    op: Operator::Multiply,
+                    right: b,
+                } = expr_arena.get(*left)
+                {
+                    (*a, *b, *right, true)
+                } else if let BinaryExpr {
+                    left: a,
+                    op: Operator::Multiply,
+                    right: b,
+                } = expr_arena.get(*right)
+                {
+                    (*a, *b, *left, false)
+                } else {
+                    return Ok(None);
+                };
 
-            BinaryExpr {
-                left,
-                op: Operator::Minus,
-                right,
-            } => {
-                // FUSED SUB MULTIPLY
-                match expr_arena.get(*right) {
-                    // input
-                    // (a - (b * c)
-                    // kept as input
-                    BinaryExpr {
-                        left: a,
-                        op: Operator::Multiply,
-                        right: b,
-                    } => Ok(check_eligible(left, right, expr_arena, schema)?.then(|| {
-                        let input = &[*left, *a, *b];
-                        get_expr(input, FusedOperator::SubMultiply, expr_arena)
-                    })),
-                    _ => {
-                        // FUSED MULTIPLY SUB
-                        match expr_arena.get(*left) {
-                            // input
-                            // (a * b) - c
-                            // kept as input
-                            BinaryExpr {
-                                left: a,
-                                op: Operator::Multiply,
-                                right: b,
-                            } => Ok(check_eligible(left, right, expr_arena, schema)?.then(|| {
-                                let input = &[*a, *b, *right];
-                                get_expr(input, FusedOperator::MultiplySub, expr_arena)
-                            })),
-                            _ => Ok(None),
-                        }
-                    },
+                if !check_eligible(&[a, b], expr_arena, schema)? {
+                    return Ok(None);
                 }
+
+                let (input, fused_op) = match (outer_op, mul_is_left) {
+                    (Operator::Plus, _) => ([a, b, other], FusedOperator::MultiplyAdd),
+                    (Operator::Minus, true) => ([a, b, other], FusedOperator::MultiplySub),
+                    (Operator::Minus, false) => ([other, a, b], FusedOperator::SubMultiply),
+                    _ => unreachable!(),
+                };
+
+                Ok(Some(get_expr(&input, fused_op, expr_arena)))
             },
             _ => Ok(None),
         }

py-polars/tests/unit/operations/arithmetic/test_arithmetic.py

@@ -188,7 +188,7 @@ def test_fused_arithm() -> None:
     )
     # the extra aliases are because the fma does operation reordering
     assert (
-        """col("c").fma([col("a"), col("b")]).alias("a"), col("a").fma([col("b"), col("c")]).alias("2")"""
+        """col("a").fma([col("b"), col("c")]), col("b").fma([col("c"), col("a")]).alias("2")"""
         in q.explain()
     )
     assert q.collect().to_dict(as_series=False) == {