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.

[wyzaws]

Is there any reason this should only be single tile? There may be cases where we want to split up the broadcast over many free dimension tiles.

[RuifengHua]

Thanks for the review!

On the single-tile guard

I tested the multi-tile cases in simulation to check whether the restriction is
actually necessary. Both are broken in the current execute():

1 src tile → 2 dst tiles (mismatched counts): crashes with
TileStream 'src_tiled' tile grid exhausted. The loop calls self._src.get_tile()
every iteration, advancing the src stream in lockstep with dst — but a broadcast
has a single source tile, so it runs off the end on the second iteration.

2 src tiles → 2 dst tiles (equal counts): compiles and runs, but silently
produces wrong output. With src_partition=0 and a 256-row dst, the result comes
back as an src[0]/src[1] stripe instead of a uniform broadcast — e.g. out[31]
is src[1] rather than src[0]. Same code path: single-tile is correct, two-tile
is corrupted, so it's the multi-tile layout that breaks it, not the inputs.

I traced the cause: it's the destination-side quadrant slicing, not the source.
The sliced source view is identical in both cases (shape=(1, 8), offset=0). But
under a multi-tile alloc_logical container the partition stride changes — a
quadrant that sits at offset 256 in the single-tile layout sits at 512 in the
two-tile layout (and the second tile's quadrants start one row in, at offset 8).
So dst_tile.slice(i*32 : i*32+32) lands on the wrong addresses and the shuffle
writes get interleaved, producing the stripe.

So the guard is converting two real failure modes (a tile-exhaustion crash, and
silent corruption) into a clear message, rather than restricting something that
would otherwise work.

[wyzaws]

Let me double check on the format of unit tests, our test infrastructure automatically picks up tests to run in pipeline and I want to make sure this currently does not run in pipeline.

[wyzaws]

Hi @RuifengHua thanks for the contribution, this looks like it does fix a bug in the code. I've left a comment about the single tile guard you added and I'll get back to you on the test comment I left next week. Same testing concern applies to your other PRs, but no action needed from you on that front for now. Thanks!