diff --git a/src/nkilib_src/nkilib/experimental/primitives/blas/tensor_tensor.py b/src/nkilib_src/nkilib/experimental/primitives/blas/tensor_tensor.py index b4ecad3..4de5137 100644 --- a/src/nkilib_src/nkilib/experimental/primitives/blas/tensor_tensor.py +++ b/src/nkilib_src/nkilib/experimental/primitives/blas/tensor_tensor.py @@ -64,7 +64,7 @@ def execute_tile(self): ) def execute(self): - for _ in range(self._dst.get_num_tiles_without_virtual_batches()): + for _ in range(self._dst.get_num_tiles()): self.execute_tile() self._dst.reset_cur_tile() self._src1.reset_cur_tile() 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_tensor_tensor.py b/test/integration/nkilib/experimental/primitives/blas/test_tensor_tensor.py new file mode 100644 index 0000000..01448d5 --- /dev/null +++ b/test/integration/nkilib/experimental/primitives/blas/test_tensor_tensor.py @@ -0,0 +1,194 @@ +# 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 tensor_tensor BLAS primitive. + +Element-wise binary op between two tensors: dst = op(src1, src2), where op is one of +multiply / add / subtract / maximum / minimum. Class-only primitive (no compact fn), so +there is a single invocation surface (blas.TensorTensor). One tiled kernel covers both +single-tile (P <= pmax) and multi-tile (P > pmax) inputs: the class API loops over +dst.get_num_tiles(), which is 1 for small P and > 1 (with a partial last tile) for large P. +""" + +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.pytest_test_metadata import pytest_marks +from test.utils.test_orchestrator import Orchestrator +from test.utils.unit_test_framework import UnitTestFramework, torch_ref_wrapper + +# Ops, one per operation named in the TensorTensor docstring. Each maps the kernel-side +# nl op to its torch reference (both operands are full tensors of the same shape). The refs +# use named (x1, x2) params so they match the kernel's input keys directly +OPS = { + "multiply": {"nl_op": nl.multiply, "ref": lambda x1, x2: torch.multiply(x1, x2)}, + "add": {"nl_op": nl.add, "ref": lambda x1, x2: torch.add(x1, x2)}, + "subtract": {"nl_op": nl.subtract, "ref": lambda x1, x2: torch.subtract(x1, x2)}, + "maximum": {"nl_op": nl.maximum, "ref": lambda x1, x2: torch.maximum(x1, x2)}, + "minimum": {"nl_op": nl.minimum, "ref": lambda x1, x2: torch.minimum(x1, x2)}, +} + + +# ============================================================================= +# Test Kernels (thin @nki.jit wrappers around the primitive under test) +# ============================================================================= + + +def _make_tensor_tensor_kernel(op_name): + """Build a TensorTensor wrapper: dst = op(src1, src2). `op_name` (a key of OPS) selects + the nl op, which is closed over -- so the kernel keeps a clean (x1, x2) tensor signature. + """ + op = OPS[op_name]["nl_op"] + + @nki.jit + def kernel_tensor_tensor(x1: nl.ndarray, x2: nl.ndarray) -> nl.ndarray: + p, f = x1.shape + p_tile = tile_size.pmax + y = nl.ndarray((p, f), dtype=x1.dtype, buffer=nl.shared_hbm) + + src1_sb = tile_stream.alloc_logical((p, f), p_tile, x1.dtype, "src1") + src2_sb = tile_stream.alloc_logical((p, f), p_tile, x2.dtype, "src2") + dst_sb = tile_stream.alloc_logical((p, f), p_tile, x1.dtype, "dst") + + dma.Load( + tile_stream.tile(src1_sb, (p_tile, f), iter_order=RowMajor(), logical_p=p), + tile_stream.tile_hbm(x1, (p_tile, f), iter_order=RowMajor()), + ).execute() + dma.Load( + tile_stream.tile(src2_sb, (p_tile, f), iter_order=RowMajor(), logical_p=p), + tile_stream.tile_hbm(x2, (p_tile, f), iter_order=RowMajor()), + ).execute() + blas.TensorTensor( + dst=tile_stream.tile(dst_sb, (p_tile, f), iter_order=RowMajor(), logical_p=p), + src1=tile_stream.tile(src1_sb, (p_tile, f), iter_order=RowMajor(), logical_p=p), + src2=tile_stream.tile(src2_sb, (p_tile, f), iter_order=RowMajor(), logical_p=p), + op=op, + ).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_tensor_tensor + + +# ============================================================================= +# Inputs +# ============================================================================= + + +def _generate_inputs(P, F, dtype): + """Generate two independent source tiles (P, F). Distinct seeds so src1 != src2, + which makes maximum/minimum select from both operands (not trivially one side).""" + x1 = np.random.RandomState(42).randn(P, F).astype(dtype) + x2 = np.random.RandomState(7).randn(P, F).astype(dtype) + return {"x1": x1.astype(dtype), "x2": x2.astype(dtype)} + + +def _output_tensors(kernel_input): + return {"out": np.zeros_like(kernel_input["x1"])} + + +# fmt: off +# op is orthogonal to shape (each op is the same nisa.tensor_tensor through the same tile +# loop), so the fast tier pairs each op 1:1 with a distinct shape -- every op AND every +# shape runs once, without the redundant op x shape cross-product. Shapes span single-tile +# (P <= 128) and multi-tile (P > 128, incl. a partial last tile). The sweep below does the +# thorough pairwise op x P x F x dtype coverage. +FAST_PARAM_NAMES = "P, F, dtype, op_name" +FAST_TEST_PARAMS = [ + pytest.param(1, 128, np.float32, "multiply", id="1_128_float32_multiply"), # single tile + pytest.param(128, 512, np.float32, "add", id="128_512_float32_add"), # single tile, full P + pytest.param(64, 256, ml_dtypes.bfloat16, "subtract", id="64_256_bfloat16_subtract"), # single tile, bf16 + pytest.param(256, 64, np.float32, "maximum", id="256_64_float32_maximum"), # 2 full tiles + pytest.param(300, 128, ml_dtypes.bfloat16, "minimum", id="300_128_bfloat16_minimum_partial"), # 128 + 128 + 44 +] +# fmt: on + + +# ============================================================================= +# Tests +# ============================================================================= + + +@pytest_marks(["tensor_tensor"]) +class TestTensorTensorPrimitive: + """Tests class for the tensor_tensor BLAS primitive kernel.""" + + def _run(self, test_manager, platform_target, kernel_entry, torch_ref, P, F, dtype): + is_bf16 = dtype == ml_dtypes.bfloat16 + + def input_generator(test_config): + return _generate_inputs(P, F, dtype) + + framework = UnitTestFramework( + test_manager=test_manager, + kernel_entry=kernel_entry, + torch_ref=torch_ref_wrapper(torch_ref), + kernel_input_generator=input_generator, + output_tensor_descriptor=_output_tensors, + ) + framework.run_test( + test_config=None, + compiler_args=CompilerArgs(platform_target=platform_target, logical_nc_config=1), + atol=1e-2 if is_bf16 else 1e-4, + rtol=1e-2 if is_bf16 else 1e-4, + ) + + @pytest.mark.fast + @pytest.mark.parametrize(FAST_PARAM_NAMES, FAST_TEST_PARAMS) + def test_tensor_tensor_fast(self, test_manager: Orchestrator, platform_target: Platforms, P, F, dtype, op_name): + """TensorTensor over every op, across single- and multi-tile shapes: dst = op(src1, src2).""" + self._run( + test_manager, + platform_target, + _make_tensor_tensor_kernel(op_name), + OPS[op_name]["ref"], + P, + F, + dtype, + ) + + # One tiled kernel, so the sweep spans single-tile (P <= 128) and multi-tile + # (P > 128, incl. partial last tiles) in a single pairwise sweep. + @pytest.mark.coverage_parametrize( + op_name=list(OPS.keys()), + P=[1, 32, 128, 200, 300], + F=[1, 64, 128, 512, 1024, 2048], + dtype=[np.float32, ml_dtypes.bfloat16], + coverage="pairs", + enable_automatic_boundary_tests=False, + ) + def test_tensor_tensor_sweep( + self, test_manager: Orchestrator, platform_target: Platforms, op_name, P, F, dtype, is_negative_test_case + ): + """Pairwise sweep over op, partition P (single- and multi-tile), free dim F, dtype.""" + self._run( + test_manager, + platform_target, + _make_tensor_tensor_kernel(op_name), + OPS[op_name]["ref"], + P, + F, + dtype, + )