diff --git a/src/nkilib_src/nkilib/experimental/primitives/blas/broadcast.py b/src/nkilib_src/nkilib/experimental/primitives/blas/broadcast.py index d42b3de..06c0ad7 100644 --- a/src/nkilib_src/nkilib/experimental/primitives/blas/broadcast.py +++ b/src/nkilib_src/nkilib/experimental/primitives/blas/broadcast.py @@ -17,6 +17,8 @@ import nki.isa as nisa import nki.language as nl +from ....core.utils.kernel_assert import kernel_assert +from ....core.utils.tensor_view import TensorView from .. import tile_stream from ..iter_order import RowMajor from ..tile_stream import TileStream, get_logical_shape @@ -37,6 +39,13 @@ def __init__( self._name = f"Broadcast(dst={dst.get_name()}, src={src.get_name()})" def execute(self) -> None: + # Single-tile only: src is one partition row and dst is <= pmax rows. For a + # larger dst, broadcast once into a <= pmax buffer and reuse it across tiles. + kernel_assert( + self._src.get_num_tiles() == 1 and self._dst.get_num_tiles() == 1, + f"Broadcast '{self._name}': src and dst must each be a single tile", + ) + self._dst.reset_cur_tile() self._src.reset_cur_tile() @@ -45,9 +54,12 @@ def execute(self) -> None: src_tile = self._src.get_tile() dst_npar = dst_tile.shape[0] - shuffle_mask = [] - for _ in range(32): - shuffle_mask.append(self._src_partition) + # The src view below is sliced down to the single chosen partition row, so + # within that view the source row is at index 0. shuffle_mask[i] selects the + # source partition (within the 32-partition quadrant) that output row i copies + # from -- it must therefore be 0, not src_partition (the slice already applied + # that offset; putting src_partition here too would double-apply it). + shuffle_mask = [0] * 32 for i in range((dst_npar + 31) // 32): cur_npar = min(32, dst_npar - i * 32) nisa.nc_stream_shuffle( diff --git a/test/integration/nkilib/experimental/primitives/blas/__init__.py b/test/integration/nkilib/experimental/primitives/blas/__init__.py new file mode 100644 index 0000000..ce63f8f --- /dev/null +++ b/test/integration/nkilib/experimental/primitives/blas/__init__.py @@ -0,0 +1,13 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/test/integration/nkilib/experimental/primitives/blas/test_broadcast.py b/test/integration/nkilib/experimental/primitives/blas/test_broadcast.py new file mode 100644 index 0000000..146f9d9 --- /dev/null +++ b/test/integration/nkilib/experimental/primitives/blas/test_broadcast.py @@ -0,0 +1,262 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration tests for the broadcast BLAS primitive. + +Replicates a single source partition row to all P rows of the destination: +dst[p, :] = src[src_partition, :] for every p (a partition broadcast). The source +tile has SRC_ROWS rows; src_partition selects which one is broadcast. +""" + +import ml_dtypes +import nki +import nki.language as nl +import numpy as np +import pytest +import torch +from nki.language import tile_size + +from nkilib_src.nkilib.experimental.primitives import blas, dma, tile_stream +from nkilib_src.nkilib.experimental.primitives.iter_order import RowMajor +from test.utils.common_dataclasses import CompilerArgs, Platforms +from test.utils.coverage_parametrized_tests import FilterResult +from test.utils.pytest_test_metadata import pytest_marks +from test.utils.test_orchestrator import Orchestrator +from test.utils.unit_test_framework import UnitTestFramework, torch_ref_wrapper + +# ============================================================================= +# Test Kernels (thin @nki.jit wrappers around the primitive under test) +# ============================================================================= + + +def _make_broadcast_kernel(P: int, src_rows: int, src_partition: int): + """Build a wrapper that broadcasts src row `src_partition` to P destination rows. + P, src_rows, and src_partition are closed over because they set the tensor shapes: + src is (src_rows, F), dst is (P, F).""" + + @nki.jit + def kernel_broadcast(src: nl.ndarray) -> nl.ndarray: + _, f = src.shape # src is (src_rows, F) + y = nl.ndarray((P, f), dtype=src.dtype, buffer=nl.shared_hbm) + + src_sb = tile_stream.alloc_logical((src_rows, f), src_rows, src.dtype, "src") + dst_sb = tile_stream.alloc_logical((P, f), P, src.dtype, "dst") + + dma.load(src_sb, src) + blas.broadcast(dst_sb, src_sb, src_partition=src_partition) + dma.store(y, dst_sb) + return y + + return kernel_broadcast + + +def _make_broadcast_class_kernel(P: int, src_rows: int, src_partition: int): + """Build a class-API (blas.Broadcast) wrapper covering both the valid single-tile case + and the rejected multi-tile case with the same code. + + Broadcast is a single-tile op by design: it spreads one per-feature source row across + <= pmax destination rows. For a larger destination the correct pattern is to broadcast + once into a <= pmax buffer and reuse it across the tile loop, so multi-tile streams are + rejected rather than silently doing an unintended per-tile broadcast.""" + + @nki.jit + def kernel_broadcast_class(src: nl.ndarray) -> nl.ndarray: + _, f = src.shape + p_tile = tile_size.pmax + y = nl.ndarray((P, f), dtype=src.dtype, buffer=nl.shared_hbm) + + src_sb = tile_stream.alloc_logical((src_rows, f), src_rows, src.dtype, "src") + dst_sb = tile_stream.alloc_logical((P, f), p_tile, src.dtype, "dst") + + dma.load(src_sb, src) + blas.Broadcast( + dst=tile_stream.tile(dst_sb, (p_tile, f), iter_order=RowMajor(), logical_p=P), + src=tile_stream.tile(src_sb, (src_rows, f), iter_order=RowMajor()), + src_partition=src_partition, + ).execute() + dma.Store( + tile_stream.tile_hbm(y, (p_tile, f), iter_order=RowMajor()), + tile_stream.tile(dst_sb, (p_tile, f), iter_order=RowMajor(), logical_p=P), + ).execute() + return y + + return kernel_broadcast_class + + +# ============================================================================= +# Torch Reference +# ============================================================================= + + +def _make_broadcast_torch_ref(P: int, src_partition: int): + """Reference: replicate source row `src_partition` to P rows (exact copy).""" + + def broadcast_torch_ref(src: torch.Tensor) -> torch.Tensor: + return src[src_partition : src_partition + 1, :].expand(P, src.shape[1]).clone() + + return broadcast_torch_ref + + +# ============================================================================= +# Inputs +# ============================================================================= + + +def _generate_inputs(src_rows, F, dtype): + """Generate a source tile (src_rows, F). Broadcast is exact, so any values work.""" + np.random.seed(42) + src = np.random.randn(src_rows, F).astype(dtype) + return {"src": src.astype(dtype)} + + +def _output_tensors(P): + def output_tensors(kernel_input): + return {"out": np.zeros((P, kernel_input["src"].shape[1]), dtype=kernel_input["src"].dtype)} + + return output_tensors + + +def _filter_src_partition(P=None, F=None, dtype=None, SRC_ROWS=None, SP=None): + """The broadcast row SP is a 0-based index into the SRC_ROWS-row source tile, so only + SP < SRC_ROWS is meaningful. Drop the rest (REDUNDANT = skip, not a negative test).""" + if SRC_ROWS is not None and SP is not None and SP >= SRC_ROWS: + return FilterResult.REDUNDANT + return FilterResult.VALID + + +# fmt: off +# P = destination partition-row count (1..pmax); F = free dim; SRC_ROWS = rows in the +# source tile; SP = which source row (< SRC_ROWS) to broadcast. +FAST_PARAM_NAMES = "P, F, dtype, SRC_ROWS, SP" +FAST_TEST_PARAMS = [ + pytest.param(1, 128, np.float32, 1, 0, id="1_128_float32"), + pytest.param(128, 512, np.float32, 1, 0, id="128_512_float32"), + pytest.param(64, 256, ml_dtypes.bfloat16, 1, 0, id="64_256_bfloat16"), + pytest.param(32, 1, np.float32, 1, 0, id="32_1_float32"), # F=1 edge + pytest.param(128, 64, np.float32, 4, 2, id="128_64_float32_sp2"), # non-zero src_partition + pytest.param(64, 32, ml_dtypes.bfloat16, 8, 5, id="64_32_bfloat16_sp5"), # non-zero src_partition +] + +# Single-tile class-API cases, chosen at the edges. P=pmax is the single-tile boundary +# (one more row -> 2 tiles -> guard fires); F=1 min free dim; SP = last valid source row. +CLASS_SINGLE_TILE_PARAMS = [ + pytest.param(128, 1, np.float32, 128, 127, id="Ppmax_F1_srcfull_splast"), + pytest.param(1, 2048, ml_dtypes.bfloat16, 16, 15, id="P1_Fmax_bfloat16_splast"), +] +# fmt: on + + +# ============================================================================= +# Tests +# ============================================================================= + + +@pytest_marks(["broadcast"]) +class TestBroadcastPrimitive: + """Tests class for the broadcast BLAS primitive kernel.""" + + def _run(self, test_manager, platform_target, P, F, dtype, src_rows, src_partition): + # Broadcast is an exact copy of the source row, so tolerances are 0. + def input_generator(test_config): + return _generate_inputs(src_rows, F, dtype) + + framework = UnitTestFramework( + test_manager=test_manager, + kernel_entry=_make_broadcast_kernel(P, src_rows, src_partition), + torch_ref=torch_ref_wrapper(_make_broadcast_torch_ref(P, src_partition)), + kernel_input_generator=input_generator, + output_tensor_descriptor=_output_tensors(P), + ) + framework.run_test( + test_config=None, + compiler_args=CompilerArgs(platform_target=platform_target, logical_nc_config=1), + atol=0.0, + rtol=0.0, + ) + + @pytest.mark.fast + @pytest.mark.parametrize(FAST_PARAM_NAMES, FAST_TEST_PARAMS) + def test_broadcast_fast(self, test_manager: Orchestrator, platform_target: Platforms, P, F, dtype, SRC_ROWS, SP): + """Broadcast row SP to P rows: dst[p, :] = src[SP, :].""" + self._run(test_manager, platform_target, P, F, dtype, SRC_ROWS, SP) + + @pytest.mark.coverage_parametrize( + P=[1, 8, 32, 64, 96, 128], + F=[1, 64, 128, 512, 1024, 2048], + dtype=[np.float32, ml_dtypes.bfloat16], + SRC_ROWS=[1, 4, 8, 16], + SP=[0, 1, 5, 15], + filter=_filter_src_partition, + coverage="pairs", + enable_automatic_boundary_tests=False, + ) + def test_broadcast_sweep( + self, test_manager: Orchestrator, platform_target: Platforms, P, F, dtype, SRC_ROWS, SP, is_negative_test_case + ): + """Sweep with pairwise coverage over dst row count P, free dim F, dtype, source + tile row count SRC_ROWS, and broadcast row SP (filtered to SP < SRC_ROWS).""" + self._run(test_manager, platform_target, P, F, dtype, src_rows=SRC_ROWS, src_partition=SP) + + @pytest.mark.fast + @pytest.mark.parametrize(FAST_PARAM_NAMES, CLASS_SINGLE_TILE_PARAMS) + def test_broadcast_class_single_tile( + self, test_manager: Orchestrator, platform_target: Platforms, P, F, dtype, SRC_ROWS, SP + ): + """Single-tile broadcast via the class API (blas.Broadcast) at the edges. Confirms the + single-tile guard passes the valid case (each stream has one tile) -- notably at P=pmax, + the boundary just below where the guard starts rejecting -- and produces correct output.""" + is_bf16 = dtype == ml_dtypes.bfloat16 + + def input_generator(test_config): + return _generate_inputs(SRC_ROWS, F, dtype) + + framework = UnitTestFramework( + test_manager=test_manager, + kernel_entry=_make_broadcast_class_kernel(P, SRC_ROWS, SP), + torch_ref=torch_ref_wrapper(_make_broadcast_torch_ref(P, SP)), + kernel_input_generator=input_generator, + output_tensor_descriptor=_output_tensors(P), + ) + framework.run_test( + test_config=None, + compiler_args=CompilerArgs(platform_target=platform_target, logical_nc_config=1), + atol=0.0, + rtol=0.0, + ) + + @pytest.mark.fast + def test_broadcast_multitile_rejected(self, test_manager: Orchestrator, platform_target: Platforms): + """Broadcast is single-tile by design (one per-feature row -> <= pmax rows). Feeding + the class API multi-tile src/dst streams (P > pmax) must be rejected by the + single-tile guard rather than silently doing a per-tile broadcast. This pins that + contract; the correct pattern for a larger destination is broadcast-once-and-reuse.""" + P, F, dtype = 256, 8, np.float32 # P > pmax -> dst has 2 tiles -> guard fires + + def input_generator(test_config): + return _generate_inputs(1, F, dtype) + + framework = UnitTestFramework( + test_manager=test_manager, + kernel_entry=_make_broadcast_class_kernel(P, src_rows=1, src_partition=0), + torch_ref=torch_ref_wrapper(_make_broadcast_torch_ref(P, 0)), + kernel_input_generator=input_generator, + output_tensor_descriptor=_output_tensors(P), + ) + with pytest.raises(Exception, match="src and dst must each be a single tile"): + framework.run_test( + test_config=None, + compiler_args=CompilerArgs(platform_target=platform_target, logical_nc_config=1), + atol=0.0, + rtol=0.0, + )