Fix broadcast wrong output for src_partition != 0 and add single-tile guard

[RuifengHua]

Summary

Broadcast replicates one source partition row to every row of the destination; src_partition selects which row of the source tile is broadcast.

Broadcast.execute() slices the source down to the chosen row:

src_tile.slice(dim=0, start=self._src_partition, end=self._src_partition + 1)

so within that sliced view the target row is at index 0. It then built the shuffle mask as:

shuffle_mask = [self._src_partition] * 32

nc_stream_shuffle's shuffle_mask[i] selects which source partition (within the quadrant of the already-sliced view) output row i copies from. Because the slice has already applied the src_partition offset, the mask must be 0 — using src_partition double-applies the offset and produces wrong output for any src_partition != 0. (src_partition == 0 happens to work, which is why this stayed latent — no in-tree caller uses a non-zero value.)

This PR:

• fixes the mask to [0] * 32;
• adds a single-tile guard — broadcast spreads one per-feature row across <= pmax rows, so a multi-tile stream is rejected with an actionable message instead of silently doing an unintended per-tile broadcast (the correct pattern for a larger destination is to broadcast once into a <= pmax buffer and reuse it across the tile loop, as the in-tree callers do);
• adds an integration test suite for the primitive (previously untested).

Repro (before fix)

Broadcasting source row 2 to all rows should copy src[2] into every output row. Before the fix, a non-zero src_partition reads out of the sliced view and yields wrong output (nan here):

import nki, nki.language as nl
import numpy as np
from nkilib_src.nkilib.experimental.primitives import blas, dma, tile_stream

@nki.jit
def k(src):
    p, f = src.shape
    y = nl.ndarray((4, f), dtype=src.dtype, buffer=nl.shared_hbm)
    src_sb = tile_stream.alloc_logical((p, f), p, src.dtype, "src")
    dst_sb = tile_stream.alloc_logical((4, f), 4, src.dtype, "dst")
    dma.load(src_sb, src)
    blas.broadcast(dst_sb, src_sb, src_partition=2)  # broadcast row 2 to all 4 dst rows
    dma.store(y, dst_sb)
    return y

src = np.arange(8 * 4, dtype=np.float32).reshape(8, 4)  # 8 distinct rows
out = nki.simulate(k)(src)
# expected: every output row == src[2] == [8, 9, 10, 11]
# before fix: out is wrong (nan); after fix: every row == src[2]

Test plan

Added test/integration/nkilib/experimental/primitives/blas/test_broadcast.py, verified in simulation mode (--test-mode simulation):

• Fast + pairwise sweep pass — destination rows P, free dim F, dtype (float32, bfloat16), source tile row count SRC_ROWS, and broadcast row SP (filtered to SP < SRC_ROWS), including non-zero src_partition
• Confirmed the non-zero src_partition cases fail against the pre-fix mask and pass after the fix; src_partition == 0 passes both before and after
• Single-tile guard, class API (blas.Broadcast): a valid single-tile broadcast passes at the edges (notably P = pmax, the boundary just below where the guard rejects), and a multi-tile stream (P > pmax) is rejected with the guard message
• Confirmed the guard negative test fails when the guard is removed (it validates the guard, not incidental behavior)

Note: validated in simulation only (correctness fix, no performance claim). Broadcast is an exact copy, so tolerances are 0.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[aws-gjfreema]

Thanks for your contribution! I've reached out to our team to review your commit and we will get back to you shortly with any follow-up.