[cute] Support explicit rank-1 column broadcasts

helion/_compiler/cute/aux_tensor.py

@@ -94,9 +94,9 @@ class Tcgen05AuxTensorDescriptor:
       the host tensor's FX node — the concrete shape / dtype / device
       the producer-body codegen needs at MMA-codegen time (the TMA
       atom and SMEM ring are sized from these).
-    - ``broadcast_axis``: ``None`` for the exact-shape rank-2 form
-      (``residual[tile_m, tile_n]``); ``1`` for the trailing-axis
-      rowvec broadcast form (``bias[tile_n]``). Matches the field on
+    - ``broadcast_axis``: ``None`` for exact-shape loads, ``1`` for a
+      trailing-axis row vector, and ``2`` for an explicit rank-1 column
+      vector. Matches the field on
       :class:`_AuxiliaryTensorStep` so the productive body uses the
       same axis convention as the existing per-thread splice.
     - ``store_value_node``: the FX node of the store value whose

helion/_compiler/cute/cute_epilogue.py

@@ -22,11 +22,11 @@
 - Auxiliary-tensor binary ops: same shape as above, but one or more
   steps are ``add/sub/mul/div`` with the chain carrier as one
   operand and a ``helion.language.load(aux_tensor, [...])``
-  call as the other. Two aux load shapes are accepted: the
-  exact-shape rank-2 form (``residual[tile_m, tile_n]``) and the
-  rank-1 trailing-axis (rowvec) broadcast form (``bias[tile_n]``).
+  call as the other. Accepted shapes include exact-shape rank-2 loads,
+  rank-1 trailing-axis row vectors, and explicit rank-1 column vectors
+  spelled ``bias[tile_m].unsqueeze(-1)``.
   See :class:`_AuxiliaryTensorStep` for the canonical contract.
-  Forms outside these two — 3-D underlying tensors with a static
+  Other forms — 3-D underlying tensors with a static
   collapse, mismatched indices, leading-axis rank-1
   (``bias[tile_m]``), kwargs — are rejected to the loud-failure
   backstop.
@@ -120,7 +120,8 @@ class _AuxiliaryTensorStep:
     ``broadcast_axis`` is ``None`` when the aux tensor matches the
     carrier rank exactly (``residual[tile_m, tile_n]``), ``1`` for
     a trailing-axis (rowvec) broadcast aux load (``bias[tile_n]``
-    with shape ``(N,)``), or ``0`` for a leading-axis (row /
+    with shape ``(N,)``), ``2`` for an explicit column vector
+    (``bias[tile_m].unsqueeze(-1)``), or ``0`` for a leading-axis (row /
     M-axis) broadcast aux load spelled explicitly as a ``(1, N)``
     tensor indexed ``bias[tile_m, tile_n]``. A *bare* rank-1
     operand on the RHS (``acc + bias[tile_m]``) is still rejected
@@ -132,11 +133,10 @@ class _AuxiliaryTensorStep:
     M axis themselves, so row 0 broadcasts over every output row.
     The splice site
     (``memory_ops._codegen_cute_store_tcgen05_tile``) owns the
-    canonical broadcast-view contract — for both broadcast axes it
-    builds a 2-D logical view with stride 0 on the broadcast
-    (M) axis and stride 1 on the data (N) axis so the existing
-    ``partition_C → flat_divide → partition_D`` pipeline can run
-    unchanged. Mirrors Quack's ``RowVecLoad`` epilogue
+    canonical broadcast-view contract. Row vectors use stride ``(0, 1)``;
+    compact column vectors use stride ``(1, 0)``. Both feed the existing
+    ``partition_C → flat_divide → partition_D`` pipeline unchanged. The
+    row-vector form mirrors Quack's ``RowVecLoad`` epilogue
     (``quack/quack/epi_ops.py``).
     """
 
@@ -370,65 +370,80 @@ def _is_helion_load_node(node: torch.fx.Node) -> bool:
     return node.op == "call_function" and node.target is helion_load
 
 
-def _canonical_aux_load_operand(node: torch.fx.Node) -> tuple[torch.fx.Node, str]:
-    """Return the underlying aux load plus its local-value template.
+@dataclasses.dataclass(frozen=True)
+class _AuxLoadOperand:
+    load_node: torch.fx.Node
+    value_template: str = "{aux}"
+    rank1_data_axis: int | None = None
+
+
+def _aux_load_operand(node: torch.fx.Node) -> _AuxLoadOperand | None:
+    """Normalize the accepted wrappers around one auxiliary load.
 
-    Canonical unwrapping currently accepts raw aux loads and fp32 casts of
-    aux loads; other dtype casts stay outside the fused aux pattern.
+    ``rank1_data_axis`` records the surviving axis after an explicit
+    rank-1 ``unsqueeze``. This distinguishes a real column-vector expression
+    from a bare rank-1 operand, whose broadcast direction follows PyTorch's
+    trailing-axis rule.
     """
     if _is_helion_load_node(node):
-        return node, "{aux}"
+        return _AuxLoadOperand(node)
+    if node.op != "call_function" or node.kwargs:
+        return None
+
     if (
-        node.op == "call_function"
-        and node.target is torch.ops.prims.convert_element_type.default
-        and not node.kwargs
+        node.target is torch.ops.prims.convert_element_type.default
         and len(node.args) == 2
         and node.args[1] is torch.float32
         and isinstance(node.args[0], torch.fx.Node)
     ):
-        inner = node.args[0]
-        if _is_helion_load_node(inner):
-            return inner, "({aux}).to(cutlass.Float32)"
-    return node, "{aux}"
-
+        inner = _aux_load_operand(node.args[0])
+        if inner is None:
+            return None
+        return dataclasses.replace(
+            inner, value_template=f"({inner.value_template}).to(cutlass.Float32)"
+        )
 
-def _aux_load_operand(node: torch.fx.Node) -> tuple[torch.fx.Node, str]:
-    """Return ``(load_node, aux_template)`` for accepted aux operands.
+    if (
+        node.target is torch.ops.aten.unsqueeze.default
+        and len(node.args) == 2
+        and isinstance(node.args[0], torch.fx.Node)
+        and isinstance(node.args[1], int)
+        and not isinstance(node.args[1], bool)
+    ):
+        inner = _aux_load_operand(node.args[0])
+        load_val = inner.load_node.meta.get("val") if inner is not None else None
+        dim = node.args[1] % 2
+        if (
+            inner is None
+            or inner.rank1_data_axis is not None
+            or not isinstance(load_val, torch.Tensor)
+            or load_val.ndim != 1
+            or node.args[1] not in (-2, -1, 0, 1)
+        ):
+            return None
+        return dataclasses.replace(inner, rank1_data_axis=1 - dim)
 
-    ``aux_template`` is a ``{aux}``-keyed expression applied to the loaded
-    per-thread aux local before the outer binary op. This covers wrapped
-    residual terms such as ``0.5 * residual[tile_m, tile_n].to(torch.float32)``
-    without broadening the general chain model to arbitrary non-scalar binary
-    reuse.
-    """
-    load_node, aux_template = _canonical_aux_load_operand(node)
-    if load_node is not node:
-        return load_node, aux_template
-    if node.op != "call_function" or node.target is not torch.ops.aten.mul.Tensor:
-        return node, "{aux}"
-    if node.kwargs or len(node.args) < 2:
-        return node, "{aux}"
-    lhs = node.args[0]
-    rhs = node.args[1]
+    if node.target is not torch.ops.aten.mul.Tensor or len(node.args) < 2:
+        return None
+    lhs, rhs = node.args[:2]
     if isinstance(lhs, torch.fx.Node):
         scalar = _extract_scalar(rhs)
-        inner = lhs
+        inner_node = lhs
     elif isinstance(rhs, torch.fx.Node):
         scalar = _extract_scalar(lhs)
-        inner = rhs
+        inner_node = rhs
     else:
-        return node, "{aux}"
-    if scalar is None:
-        return node, "{aux}"
-    load_node, inner_template = _canonical_aux_load_operand(inner)
-    if load_node is inner and not _is_helion_load_node(load_node):
-        return node, "{aux}"
-    return load_node, f"(({inner_template}) * {scalar!r})"
+        return None
+    inner = _aux_load_operand(inner_node)
+    if scalar is None or inner is None:
+        return None
+    return dataclasses.replace(
+        inner, value_template=f"(({inner.value_template}) * {scalar!r})"
+    )
 
 
 def _is_aux_load_operand_node(node: torch.fx.Node) -> bool:
-    load_node, _ = _aux_load_operand(node)
-    return _is_helion_load_node(load_node)
+    return _aux_load_operand(node) is not None
 
 
 @dataclasses.dataclass(frozen=True)
@@ -749,32 +764,36 @@ def _classify_binary(
         # tensor load; the other is the chain carrier. If both look
         # like aux loads or neither does, bail — the chain has no
         # unique carrier.
-        lhs_load, lhs_aux_template = _aux_load_operand(lhs)
-        rhs_load, rhs_aux_template = _aux_load_operand(rhs)
-        lhs_kind = aux_tensor_load_kind(
-            lhs_load,
-            carrier_tile_shape=carrier_tile_shape,
-            carrier_tile_index_nodes=carrier_tile_index_nodes,
-            carrier_global_shape=carrier_global_shape,
-        )
-        rhs_kind = aux_tensor_load_kind(
-            rhs_load,
-            carrier_tile_shape=carrier_tile_shape,
-            carrier_tile_index_nodes=carrier_tile_index_nodes,
-            carrier_global_shape=carrier_global_shape,
-        )
+        lhs_operand = _aux_load_operand(lhs)
+        rhs_operand = _aux_load_operand(rhs)
+
+        def load_kind(operand: _AuxLoadOperand | None) -> tuple[str, int | None] | None:
+            if operand is None:
+                return None
+            return aux_tensor_load_kind(
+                operand.load_node,
+                carrier_tile_shape=carrier_tile_shape,
+                carrier_tile_index_nodes=carrier_tile_index_nodes,
+                carrier_global_shape=carrier_global_shape,
+                explicit_rank1_data_axis=operand.rank1_data_axis,
+            )
+
+        lhs_kind = load_kind(lhs_operand)
+        rhs_kind = load_kind(rhs_operand)
         aux_load: torch.fx.Node
         carrier: torch.fx.Node
         forward_form: bool
         if lhs_kind is not None and rhs_kind is None:
-            aux_load = lhs_load
-            aux_template = lhs_aux_template
+            assert lhs_operand is not None
+            aux_load = lhs_operand.load_node
+            aux_template = lhs_operand.value_template
             carrier = rhs
             forward_form = False  # carrier is the right operand
             aux_kind = lhs_kind
         elif rhs_kind is not None and lhs_kind is None:
-            aux_load = rhs_load
-            aux_template = rhs_aux_template
+            assert rhs_operand is not None
+            aux_load = rhs_operand.load_node
+            aux_template = rhs_operand.value_template
             carrier = lhs
             forward_form = True  # carrier is the left operand
             aux_kind = rhs_kind
@@ -968,9 +987,8 @@ def analyze_tcgen05_unary_epilogue_chain(
     (``add`` / ``sub`` / ``mul`` / ``div`` against a compile-time
     Python literal), and auxiliary-tensor binary in two forms:
     exact-shape (``residual[tile_m, tile_n]``, rank-2 aux matching
-    the carrier tile shape) and rank-1 trailing-axis (rowvec)
-    broadcast (``bias[tile_n]``, where the single load index
-    symbol matches the carrier's trailing tile-id symbol). Other
+    the carrier tile shape), rank-1 trailing-axis rowvec broadcast,
+    and explicit rank-1 column-vector broadcast via ``unsqueeze``. Other
     shapes — 3-D collapsed loads, indices that are not exactly the
     carrier trailing tile-id symbol, leading-axis rank-1
     (``bias[tile_m]``), kwargs — are rejected so the loud-failure

helion/_compiler/cute/cute_fx_walk.py

@@ -195,6 +195,7 @@ def aux_tensor_load_kind(
     carrier_tile_shape: tuple[object, ...] | None,
     carrier_tile_index_nodes: tuple[torch.fx.Node, ...] | None = None,
     carrier_global_shape: tuple[object, ...] | None = None,
+    explicit_rank1_data_axis: int | None = None,
 ) -> tuple[str, int | None] | None:
     """Classify ``node`` as a recognized auxiliary-tensor load kind, or
     return ``None`` when the load is not eligible for fusion.
@@ -217,9 +218,9 @@ def aux_tensor_load_kind(
       ``acc + bias[tile_m]`` is either a shape error (BM ≠ BN) or
       rowvec broadcast against the trailing axis (BM == BN). A
       user wanting M-axis (column-vector) broadcasting via a bare
-      rank-1 operand must write it explicitly (``bias[tile_m][:,
-      None]`` / ``.unsqueeze(-1)``); that is a separate, deferred
-      pattern. See :class:`Tcgen05UnaryEpilogueChain`
+      rank-1 operand must write it explicitly with ``.unsqueeze(-1)``;
+      the operand normalizer records that unambiguous direction. See
+      :class:`Tcgen05UnaryEpilogueChain`
       (``cute_epilogue.py``) for the splice-side broadcast-view
       contract.
     - ``("broadcast", 0)``: a rank-2 leading-axis broadcast aux
@@ -249,6 +250,11 @@ def aux_tensor_load_kind(
       residual has a non-zero trailing stride and falls through to the
       exact matcher.
 
+    An explicit rank-1 ``unsqueeze`` is represented by
+    ``explicit_rank1_data_axis``. Axis 0 is a column vector and reuses the
+    ``("broadcast", 2)`` codegen; axis 1 is the explicit spelling of the
+    ordinary row-vector case.
+
     A rank-3 exact residual is accepted when its leading index is the carrier's
     leading passthrough tile. The classifier returns ``None`` for everything
     else: 3-D tensors with a different/static leading index, broadcast variants
@@ -370,6 +376,18 @@ def aux_tensor_load_kind(
     if len(carrier_tile_shape) != 2:
         return None
 
+    if explicit_rank1_data_axis is not None:
+        return _matches_rank1_broadcast(
+            aux_shape=aux_shape,
+            aux_tensor_shape=aux_tensor_shape,
+            aux_tensor_val=aux_tensor_val,
+            index_list=index_list,
+            carrier_tile_shape=carrier_tile_shape,
+            carrier_tile_index_nodes=carrier_tile_index_nodes,
+            carrier_global_shape=carrier_global_shape,
+            data_axis=explicit_rank1_data_axis,
+        )
+
     colvec = _matches_colvec_broadcast(
         aux_shape=aux_shape,
         aux_tensor_shape=aux_tensor_shape,
@@ -434,10 +452,9 @@ def _matches_broadcast(
     aligns to the last dimension under PyTorch broadcasting rules,
     so ``acc + bias[tile_m]`` is not a column-vector broadcast — it
     is a shape error (BM ≠ BN) or trailing-axis broadcast (BM == BN).
-    Users wanting an explicit M-axis broadcast must write it
-    explicitly (``bias[tile_m][:, None]`` / ``.unsqueeze(-1)`` /
-    ``.expand(...)``); that is a separate pattern handler not yet
-    wired up.
+    Users wanting an M-axis broadcast must write it explicitly as
+    ``bias[tile_m].unsqueeze(-1)``; the operand normalizer supplies axis 0
+    to :func:`_matches_rank1_broadcast` for that form.
 
     The splice site emits ``cute.make_layout((m, n),
     stride=(0, 1))`` with stride 1 hard-coded on the data axis,
@@ -455,42 +472,54 @@ def _matches_broadcast(
     to be divisible by the tile extent passes the tile check yet
     reads OOB at runtime.
     """
-    if len(aux_tensor_shape) != 1 or len(aux_shape) != 1 or len(index_list) != 1:
-        return None
-    if carrier_tile_index_nodes is None:
-        # Broadcast classification is gated on knowing the carrier's
-        # tile-id symbols so we can pin which axis the load broadcasts
-        # along. Without that, the unpredicated splice could broadcast
-        # the wrong axis silently — bail to the loud-failure backstop.
-        return None
-    if len(carrier_tile_index_nodes) != 2:
+    return _matches_rank1_broadcast(
+        aux_shape=aux_shape,
+        aux_tensor_shape=aux_tensor_shape,
+        aux_tensor_val=aux_tensor_val,
+        index_list=index_list,
+        carrier_tile_shape=carrier_tile_shape,
+        carrier_tile_index_nodes=carrier_tile_index_nodes,
+        carrier_global_shape=carrier_global_shape,
+        data_axis=1,
+    )
+
+
+def _matches_rank1_broadcast(
+    *,
+    aux_shape: tuple[object, ...],
+    aux_tensor_shape: tuple[object, ...],
+    aux_tensor_val: torch.Tensor,
+    index_list: Sequence[object],
+    carrier_tile_shape: tuple[object, ...],
+    carrier_tile_index_nodes: tuple[torch.fx.Node, ...] | None,
+    carrier_global_shape: tuple[object, ...] | None,
+    data_axis: int,
+) -> tuple[str, int] | None:
+    """Match a contiguous rank-1 load against one carrier data axis."""
+    if (
+        data_axis not in (0, 1)
+        or len(aux_tensor_shape) != 1
+        or len(aux_shape) != 1
+        or len(index_list) != 1
+        or carrier_tile_index_nodes is None
+        or len(carrier_tile_index_nodes) != 2
+        or aux_tensor_val.stride() != (1,)
+    ):
         return None
     load_idx = index_list[0]
-    if not isinstance(load_idx, torch.fx.Node):
-        return None
-    if aux_tensor_val.stride() != (1,):
-        return None
-    aux_extent = aux_shape[0]
-    # Trailing-axis (rowvec) broadcast only — see docstring.
-    axis = 1
-    if load_idx is not carrier_tile_index_nodes[axis]:
+    if (
+        load_idx is not carrier_tile_index_nodes[data_axis]
+        or aux_shape[0] != carrier_tile_shape[data_axis]
+    ):
         return None
-    if aux_extent != carrier_tile_shape[axis]:
+    if carrier_global_shape is not None and (
+        len(carrier_global_shape) != 2
+        or aux_tensor_shape[0] != carrier_global_shape[data_axis]
+    ):
         return None
-    if carrier_global_shape is not None:
-        if len(carrier_global_shape) != 2:
-            return None
-        # The underlying rank-1 aux tensor's *total* length must
-        # match the output's global axis size on the broadcast
-        # axis. The earlier ``aux_extent == carrier_tile_shape``
-        # check matches the per-load *tile* extent against the
-        # carrier's tile shape; without the global-size check, an
-        # aux smaller than the output axis whose length happens to
-        # be divisible by the tile extent could pass classification
-        # yet read OOB at runtime.
-        if aux_tensor_shape[0] != carrier_global_shape[axis]:
-            return None
-    return ("broadcast", axis)
+    # The existing column-vector path is identified by 2; row vectors retain
+    # their historical axis-1 representation.
+    return ("broadcast", 2 if data_axis == 0 else 1)
 
 
 def _matches_leading_broadcast(