From 18d5eecac72287d70c9e24ea2737d619b36eaa55 Mon Sep 17 00:00:00 2001 From: ycmath Date: Wed, 26 Aug 2026 14:24:37 +0900 Subject: [PATCH] Vectorize nested multiply-accumulate loops into @fmac* DSD operations Detect the two-variable nested MAC pattern Z[k] = Z[k] + A[affine(k, l)] * X[f(l)] in emit_for and lower it to a strided base DSD over A plus a per-l @increment_dsd_offset and the dtype-matched @fmac* builtin, instead of scalar loops. On samples/spatial/blas/gemv.sptl this gives 4.5-7.7x (issue #69). Per the discussion in #69: - Loop-header order is irrelevant: the variable indexing the accumulator takes the k role, the other is the reduction variable (ask 2). - Dtypes dispatch like FMADSDOp._as_csl: @fmach (f16), @fmachs (f16 multiply / f32 accumulate), @fmacs (f32) (ask 3). - No separate flag: the vectorizer is gated behind --disable-dsd like all other DSD detection (ask 5). Conservative by construction: falls back to scalar loops on aliasing of Z with A or X, non-affine indices, unsupported dtype combinations, non-unit steps, or a nonzero k start. Arbitrary nesting depth and generalization to other DSD ops (asks 1 and 4) are deferred to a follow-up PR pending the increment_dsd semantics question. Co-Authored-By: Claude Fable 5 --- spada/syntax/csl/statements.py | 248 +++++++++++++++++++++ tests/spatial_ir/test_mac_vectorization.py | 188 ++++++++++++++++ 2 files changed, 436 insertions(+) create mode 100644 tests/spatial_ir/test_mac_vectorization.py diff --git a/spada/syntax/csl/statements.py b/spada/syntax/csl/statements.py index 31761b9a..ccdf1e9d 100644 --- a/spada/syntax/csl/statements.py +++ b/spada/syntax/csl/statements.py @@ -250,17 +250,265 @@ def emit_assignment(statement: spir.AssignmentStatement, dsds: UniqueDSDDict, dt return dsd_ops.DSD_ASSIGNMENT_MAPPING[dsd_op]() +# Targeted vectorization of the nested multiply-accumulate loop +# for (k, l) or (l, k) in either order: Z[k] = Z[k] + A[k*Ck + l*Cl + C0] * X[f(l)] +# into the handwritten CSL idiom (strided base DSD + per-l @increment_dsd_offset + @fmac*). +# The accumulator index decides the roles: the loop variable indexing Z is k, the other +# is the reduction variable l. Dtypes dispatch like FMADSDOp._as_csl (@fmach / @fmachs / +# @fmacs). Conservative: fires only when Z is distinct from A and X; anything else falls +# through to scalar loops. Gated behind --disable-dsd like all other DSD detection. +# @increment_dsd_offset is available from SDK 1.x (used by Cerebras' own csl-examples +# v1.4.0 cholesky benchmark). + + +def _affine_of(expr: spir.Expression, varnames: set[str]) -> Optional[tuple[dict[str, int], int]]: + """ + Decomposes an index expression into an affine form over the given loop variables. + + :param expr: The index expression to decompose. + :param varnames: The loop-variable names that may carry coefficients. + :return: A ``(coefficients, constant)`` pair such that the expression equals + ``sum(coefficients[var] * var) + constant`` with integer coefficients, + or None if the expression is not affine in the given variables. + """ + e = expr.value if isinstance(expr, spir.Expression) else expr + if isinstance(e, spir.Identifier): + if e.name in varnames: + return {e.name: 1}, 0 + return None + if isinstance(e, (spir.ConstantLiteral, spir.Parameter)): + try: + v = e.eval() + except (ValueError, TypeError): + return None + return ({}, int(v)) if isinstance(v, int) else None + if isinstance(e, spir.UnaryOperator) and e.op == '-': + sub = _affine_of(e.value, varnames) + if sub is None: + return None + return {k: -c for k, c in sub[0].items()}, -sub[1] + if isinstance(e, spir.BinaryOperator): + left = _affine_of(e.left, varnames) + right = _affine_of(e.right, varnames) + if e.op == '+' and left and right: + coeffs = dict(left[0]) + for k, c in right[0].items(): + coeffs[k] = coeffs.get(k, 0) + c + return coeffs, left[1] + right[1] + if e.op == '-' and left and right: + coeffs = dict(left[0]) + for k, c in right[0].items(): + coeffs[k] = coeffs.get(k, 0) - c + return coeffs, left[1] - right[1] + if e.op == '*' and left and right: + if not left[0]: # left operand is a pure constant + s = left[1] + return {k: c * s for k, c in right[0].items()}, right[1] * s + if not right[0]: # right operand is a pure constant + s = right[1] + return {k: c * s for k, c in left[0].items()}, left[1] * s + return None + return None + + +def _ids_in(expr: spir.Expression) -> set[str]: + """ + Collects the names of all identifiers appearing in an expression. + + :param expr: The expression to walk. + :return: The set of identifier names. + """ + e = expr.value if isinstance(expr, spir.Expression) else expr + return {n.name for n in e.walk() if isinstance(n, spir.Identifier)} + + +def _single_index(node: spir.SpatialNode) -> Optional[spir.Expression]: + """ + Extracts the index of a one-dimensional single-element array access. + + :param node: The AST node to inspect. + :return: The index expression if ``node`` is an ArraySlice with exactly one + non-range index, otherwise None. + """ + if not isinstance(node, spir.ArraySlice) or len(node.indices) != 1: + return None + idx = node.indices[0] + if isinstance(idx, spir.RangeExpression): + return None + return idx + + +def _try_emit_vectorized_mac(statement: spir.ForStatement, dtypes: dict[spir.Identifier, spir.IRType], + header_code: StringIO) -> Optional[str]: + """ + Attempts to vectorize a two-variable nested multiply-accumulate loop into DSD operations. + + Recognizes ``Z[k] = Z[k] + A[affine(k, l)] * X[f(l)]`` (in either loop-header order, + and either operand order) and emits a strided base DSD over A plus a per-``l`` + ``@increment_dsd_offset`` and the dtype-matched ``@fmac*`` builtin. The loop variable + that indexes the destination with coefficient 1 takes the ``k`` role; the other is the + reduction variable ``l``. Falls back for anything it cannot prove safe: aliasing of Z + with A or X, non-affine indices, unsupported dtype combinations, non-unit steps, or a + nonzero ``k`` start. + + :param statement: The for-loop statement to inspect. + :param dtypes: The data types dictionary. + :param header_code: The header code to write the DSD declarations into. + :return: The generated CSL replacing the loop nest, or None to fall back to scalar code. + """ + if len(statement.variables) != 2 or len(statement.range_expression) != 2: + return None + if len(statement.body) != 1 or not isinstance(statement.body[0], spir.AssignmentStatement): + return None + names = [v.identifier.name for v in statement.variables] + + # Role assignment: the loop variable that alone indexes the destination Z with + # coefficient 1 is k; the other is the reduction variable l. The loop-header + # order (k, l) vs (l, k) is irrelevant. + assign = statement.body[0] + dst = assign.destination + dst_idx = _single_index(dst) + if dst_idx is None: + return None + dst_aff = _affine_of(dst_idx, set(names)) + if dst_aff is None or dst_aff[1] != 0 or len(dst_aff[0]) != 1: + return None + (k_var, k_coeff), = dst_aff[0].items() + if k_coeff != 1: + return None + k_pos = names.index(k_var) + l_pos = 1 - k_pos + l_var = names[l_pos] + + def _range_consts(r): + start = r.start.eval() if r.start is not None else 0 + stop = r.stop.eval() + step = r.step.eval() if r.step is not None else 1 + if not all(isinstance(v, int) for v in (start, stop, step)): + raise TypeError + return start, stop, step + try: + k0, kk, ks = _range_consts(statement.range_expression[k_pos]) + l0, ll, ls = _range_consts(statement.range_expression[l_pos]) + except (ValueError, TypeError, AttributeError): + return None + if (k0, ks, ls) != (0, 1, 1): + return None + + src = assign.source.value + if isinstance(src, spir.MultiplyAccumulateOperator): + acc, mul_b, mul_c = src.a.value, src.b.value, src.c.value + elif isinstance(src, spir.BinaryOperator) and src.op == '+': + acc = src.left.value + mul = src.right.value + if not (isinstance(mul, spir.BinaryOperator) and mul.op == '*'): + acc, mul = src.right.value, src.left.value + if not (isinstance(mul, spir.BinaryOperator) and mul.op == '*'): + return None + mul_b, mul_c = mul.left.value, mul.right.value + else: + return None + + # The accumulator term must be the same Z[k] as the destination. + if not (isinstance(acc, spir.ArraySlice) and acc.array.name == dst.array.name): + return None + acc_idx = _single_index(acc) + if acc_idx is None or _affine_of(acc_idx, {k_var}) != ({k_var: 1}, 0): + return None + + # Multiply operands: one is the matrix A accessed by an affine index in k and l, + # the other a scalar access X indexed by l only. + def classify(node): + if not isinstance(node, spir.ArraySlice): + return None + idx = _single_index(node) + if idx is None: + return None + aff = _affine_of(idx, {k_var, l_var}) + if aff is None: + return None + coeffs, const = aff + if coeffs.get(k_var) and coeffs[k_var] > 0: + return ('mat', node, coeffs.get(k_var), coeffs.get(l_var, 0), const) + if k_var not in coeffs: + return ('vec', node, idx) + return None + + cb, cc = classify(mul_b), classify(mul_c) + if cb and cb[0] == 'mat' and cc and cc[0] == 'vec': + mat, vec = cb, cc + elif cc and cc[0] == 'mat' and cb and cb[0] == 'vec': + mat, vec = cc, cb + else: + return None + _, mat_node, ck, cl, c0 = mat + _, vec_node, vec_idx = vec + if _ids_in(vec_idx) - {l_var} != set(): + return None + + # Aliasing guard: the DSD op reorders reads relative to the sequential scalar loop, + # so reading the written array through A or X must not be vectorized. + if mat_node.array.name == dst.array.name or vec_node.array.name == dst.array.name: + return None + + # Dtype dispatch — follows the mapping of FMADSDOp._as_csl exactly + # (accumulator = Z, DSD multiplicand = A, scalar = X). + def base_dtype(name_node): + t = dtypes.get(name_node.array) + return getattr(t, 'base_type', None) or getattr(t, 'element_type', None) + z_t, a_t, x_t = base_dtype(dst), base_dtype(mat_node), base_dtype(vec_node) + if z_t == a_t == spir.ScalarType.f16 and x_t == spir.ScalarType.f16: + fmac_builtin, elem_type = '@fmach', 'f16' + elif z_t == a_t == spir.ScalarType.f32 and x_t == spir.ScalarType.f16: + fmac_builtin, elem_type = '@fmachs', 'f32' # 16-bit multiplication, 32-bit addition + elif z_t == a_t == spir.ScalarType.f32 and x_t == spir.ScalarType.f32: + fmac_builtin, elem_type = '@fmacs', 'f32' + else: + return None + + # Deterministic per-rectangle numbering: each vectorized loop writes exactly one + # __vecmac_dst_ declaration into this rectangle's header. + uid = header_code.getvalue().count('__vecmac_dst_') + z_name = name_to_csl(dst.array) + a_name = name_to_csl(mat_node.array) + x_l = expr_to_csl(vec_idx) + dst_dsd = f'__vecmac_dst_{uid}' + src_dsd = f'__vecmac_src_{uid}' + base_expr = f'__index * {ck}' + (f' + {c0}' if c0 else '') + header_code.write( + f'const {dst_dsd} = @get_dsd(mem1d_dsd, ' + f'.{{ .tensor_access = |__index|{{{kk}}} -> {z_name}[__index] }});\n' + f'const {src_dsd} = @get_dsd(mem1d_dsd, ' + f'.{{ .tensor_access = |__index|{{{kk}}} -> {a_name}[{base_expr}] }});\n') + lname = name_to_csl(statement.variables[l_pos].identifier) + off = lname + (f' * {cl}' if cl != 1 else '') + return ( + f'// vectorized MAC: {z_name}[k] += {a_name}[k*{ck}+{lname}*{cl}+{c0}] * {name_to_csl(vec_node.array)}[{x_l}]\n' + f'for (@range(i16, {l0}, {ll}, 1)) |{lname}| {{\n' + f' const __vecmac_a_{uid} = @increment_dsd_offset({src_dsd}, {off}, {elem_type});\n' + f' {fmac_builtin}({dst_dsd}, {dst_dsd}, __vecmac_a_{uid}, {name_to_csl(vec_node.array)}[{x_l}]);\n' + f'}}\n') + + def emit_for(statement: spir.ForStatement, dsds: UniqueDSDDict, dtypes: dict[spir.Identifier, spir.IRType], header_code: StringIO) -> str: """ Generates a CSL for loop statement from a Spatial IR for loop statement. + Nested multiply-accumulate loops are first offered to the DSD vectorizer (see + ``_try_emit_vectorized_mac``); everything else lowers to scalar loops. + :param statement: The Spatial IR for loop statement to convert. :param dsds: The unique DSD dictionary. :param dtypes: The data types dictionary. :param header_code: The header code to include. :return: The generated CSL for loop statement. """ + if not dsd_ops.DISABLE_DSD: + vectorized = _try_emit_vectorized_mac(statement, dtypes, header_code) + if vectorized is not None: + return vectorized + ranges = statement.range_expression vars_ = statement.variables diff --git a/tests/spatial_ir/test_mac_vectorization.py b/tests/spatial_ir/test_mac_vectorization.py new file mode 100644 index 00000000..fc70d36f --- /dev/null +++ b/tests/spatial_ir/test_mac_vectorization.py @@ -0,0 +1,188 @@ +"""Tests for the targeted vectorization of nested multiply-accumulate loops. + +The lowering should turn + + for (k, l) in [0:K, 0:K]: + z[k] = z[k] + A[k*K + l] * x[l] + +(in either loop-header order) into a strided base DSD plus a per-column +``@increment_dsd_offset`` and the dtype-matched ``@fmac*`` builtin, and must +fall back to scalar loops for every shape it cannot prove safe (aliasing, +non-affine indices, unsupported dtype combinations, or ``--disable-dsd``). +""" +from typing import Optional + +from spada.lowering.spatial_ir_to_csl import lower_spatial_ir_to_csl +from spada.syntax.csl import dsd_ops +from spada.syntax.spatial_ir import parser, passes + + +def _kernel(body: str, decls: Optional[str] = None, + a_t: str = 'f32', x_t: str = 'f32', z_t: str = 'f32') -> str: + decls = decls if decls is not None else f''' + {a_t}[K*K] A_flat + {x_t}[K] x + {z_t}[K] z + ''' + return f''' + kernel @t( + stream<{a_t}, K*K>[2, 2] readonly inp_A, + stream<{x_t}, K>[2, 2] readonly inp_x, + stream<{z_t}, K>[2, 2] writeonly out + ) {{ + place i16 i, i16 j in [0:2, 0:2] {{ + {decls} + }} + phase {{ + compute i16 i, i16 j in [0:2, 0:2] {{ + await receive(A_flat, inp_A[i, j]) + await receive(x, inp_x[i, j]) + {body} + await send(z, out[i, j]) + }} + }} + }} + ''' + + +def _lower(code: str, **lower_kwargs): + kernel = parser.parse_string(code, 'test.sptl') + kernel = passes.concretize_parameters(kernel, K=4) + kernel = passes.constexpr_propagation(kernel) + return lower_spatial_ir_to_csl(kernel, **lower_kwargs) + + +def _all_code(csl_files) -> str: + return '\n'.join(f.code for f in csl_files) + + +MAC_BODY = ''' + for i16 k in [0:K] { + z[k] = 0.0 + } + for i16 k, i16 l in [0:K, 0:K] { + z[k] = z[k] + A_flat[k*K + l] * x[l] + } +''' + +# Same computation with the loop-header order swapped: l is the outer variable. +MAC_BODY_SWAPPED = ''' + for i16 k in [0:K] { + z[k] = 0.0 + } + for i16 l, i16 k in [0:K, 0:K] { + z[k] = z[k] + A_flat[k*K + l] * x[l] + } +''' + + +def test_mac_loop_is_vectorized(): + code = _all_code(_lower(_kernel(MAC_BODY))) + assert '@fmacs(' in code + assert '@increment_dsd_offset(' in code + + +def test_swapped_loop_order_is_vectorized(): + # Operand roles come from the accumulator index, not from the loop-header + # position, so `for (l, k)` vectorizes identically (issue #69). + code = _all_code(_lower(_kernel(MAC_BODY_SWAPPED))) + assert '@fmacs(' in code + assert '@increment_dsd_offset(' in code + + +def test_f16_mac_uses_fmach(): + # An all-f16 MAC dispatches to @fmach and offsets the DSD as f16. + code = _all_code(_lower(_kernel(MAC_BODY, a_t='f16', x_t='f16', z_t='f16'))) + assert '@fmach(' in code + assert '@fmacs(' not in code + assert ', f16);' in code + + +def test_mixed_precision_mac_uses_fmachs(): + # An f32 accumulate with an f16 scalar operand dispatches to @fmachs, + # mirroring FMADSDOp._as_csl. + code = _all_code(_lower(_kernel(MAC_BODY, a_t='f32', x_t='f16', z_t='f32'))) + assert '@fmachs(' in code + + +def test_unsupported_dtype_combo_falls_back(): + # f16 accumulator with f32 scalar operand has no @fmac* builtin: scalar loop. + code = _all_code(_lower(_kernel(MAC_BODY, a_t='f16', x_t='f32', z_t='f16'))) + assert '@increment_dsd_offset(' not in code + + +def test_scalar_fallback_when_dsd_disabled(): + # MAC vectorization is part of DSD detection and is disabled by + # --disable-dsd alone; there is no separate flag. + try: + code = _all_code(_lower(_kernel(MAC_BODY), disable_dsd=True)) + assert '@increment_dsd_offset(' not in code + assert '@fmacs(' not in code + finally: + dsd_ops.DISABLE_DSD = False # module flag is sticky, reset for other tests + + +def test_no_vectorization_when_accumulator_aliases_matrix(): + # z appears as the matrix operand: DSD reordering would change semantics. + body = ''' + for i16 k in [0:K] { + z[k] = 1.0 + } + for i16 k, i16 l in [0:K, 0:K] { + z[k] = z[k] + z[k*1 + l*0] * x[l] + } + ''' + code = _all_code(_lower(_kernel(body))) + assert '@increment_dsd_offset(' not in code + + +def test_no_vectorization_when_accumulator_aliases_matrix_swapped_order(): + body = ''' + for i16 k in [0:K] { + z[k] = 1.0 + } + for i16 l, i16 k in [0:K, 0:K] { + z[k] = z[k] + z[k*1 + l*0] * x[l] + } + ''' + code = _all_code(_lower(_kernel(body))) + assert '@increment_dsd_offset(' not in code + + +def test_no_vectorization_when_accumulator_aliases_vector(): + body = ''' + for i16 k in [0:K] { + z[k] = 1.0 + } + for i16 k, i16 l in [0:K, 0:K] { + z[k] = z[k] + A_flat[k*K + l] * z[l] + } + ''' + code = _all_code(_lower(_kernel(body))) + assert '@increment_dsd_offset(' not in code + + +def test_no_vectorization_for_nonaffine_index(): + body = ''' + for i16 k in [0:K] { + z[k] = 0.0 + } + for i16 k, i16 l in [0:K, 0:K] { + z[k] = z[k] + A_flat[k*l] * x[l] + } + ''' + code = _all_code(_lower(_kernel(body))) + assert '@increment_dsd_offset(' not in code + + +def test_no_vectorization_when_accumulator_differs(): + body = ''' + for i16 k in [0:K] { + z[k] = 0.0 + } + for i16 k, i16 l in [0:K, 0:K] { + z[k] = x[k] + A_flat[k*K + l] * x[l] + } + ''' + code = _all_code(_lower(_kernel(body))) + assert '@increment_dsd_offset(' not in code