[CuTeDSL] Allow reassoc fastmath flag on vector reduction ops

python/CuTeDSL/_mlir_helpers/vector.py

@@ -808,6 +808,7 @@ def reduce(
         op: Literal["add", "mul", "min", "max"] = "add",
         *,
         dim: Optional[Union[int, list[int]]] = None,
+        fastmath: Optional[arith.FastMathFlags] = None,
         acc: Any = None,
         loc: Optional[ir.Location] = None,
         ip: Optional[ir.InsertionPoint] = None,
@@ -825,6 +826,12 @@ def reduce(
             vs unsigned integer).
         :param dim: Dimension(s) to reduce.  ``None`` reduces all dims to a
             scalar.  An int or list of ints reduces only those dims.
+        :param fastmath: Optional fast-math flags for the reduction. Pass
+            ``arith.FastMathFlags.reassoc`` to allow the compiler to lower
+            the reduction as a tree instead of a strict left-to-right chain.
+            Defaults to ``None``. Only supported for a full reduction of a
+            1-D vector to a scalar (``dim is None``), raises ValueError
+            for multi-dimensional reductions.
         :param acc: Optional accumulator.  For scalar reduction a scalar value;
             for multi-dim reduction a vector matching the result shape.
         :return: Scalar (when ``dim is None``) or :class:`Vector` (when
@@ -872,15 +879,19 @@ def reduce(
         kind = kind_fn(self)
         vec_ty = ir.VectorType(self.type)
         elem_ty = vec_ty.element_type
+        fmf_kwargs = {"fastmath": fastmath} if fastmath is not None else {}
 
         ndim = len(vec_ty.shape)
 
         if dim is None and ndim == 1:
             # 1-D full reduction to scalar — wrap in _dtype so type info is preserved
-            raw = vector.reduction(elem_ty, kind, self, acc=acc, loc=loc, ip=ip)
+            raw = vector.reduction(elem_ty, kind, self, acc=acc, **fmf_kwargs, loc=loc, ip=ip)
             return self._dtype(raw)
 
         # Multi-dimension reduction
+        if fastmath is not None:
+            raise ValueError("Fastmath flags are only supported for 1-D full reductions.")
+
         if dim is None:
             # Reduce all dims for N-D vector
             reduction_dims = list(range(ndim))

test/python/CuTeDSL/test_vector_reduction_fastmath.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# Use of this software is governed by the terms and conditions of the
+# NVIDIA End User License Agreement (EULA), available at:
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
+#
+# Any use, reproduction, disclosure, or distribution of this software
+# and related documentation outside the scope permitted by the EULA
+# is strictly prohibited.
+
+"""
+Unit test for the ``fastmath`` parameter on ``Vector.reduce``.
+"""
+
+import unittest
+
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import arith, func
+from cutlass._mlir_helpers.vector import Vector
+
+
+def _reduce_ir(fastmath=None):
+    """Return the IR for a ``Vector.reduce("add")`` over ``vector<16xf32>``."""
+    with ir.Context(), ir.Location.unknown():
+        module = ir.Module.create()
+        vec_ty = ir.VectorType.get([16], ir.F32Type.get())
+        with ir.InsertionPoint(module.body):
+            fn = func.FuncOp("test", ir.FunctionType.get([vec_ty], []))
+            with ir.InsertionPoint(fn.add_entry_block()):
+                Vector(fn.arguments[0]).reduce("add", fastmath=fastmath)
+                func.ReturnOp([])
+        module.operation.verify()
+        return str(module)
+
+
+class TestVectorReduceFastmath(unittest.TestCase):
+    def test_fastmath_flag(self):
+        self.assertIn("fastmath<reassoc>", _reduce_ir(arith.FastMathFlags.reassoc))
+        self.assertNotIn("reassoc", _reduce_ir())
+
+
+if __name__ == "__main__":
+    unittest.main()