From e6b4768b5d2054eee45163b29cb245752192d402 Mon Sep 17 00:00:00 2001 From: Harris Hua Date: Mon, 6 Jul 2026 03:59:00 +0000 Subject: [PATCH 1/2] fix: Correct phantom method call in TensorTensor.execute Commit c38d790 ("NKI Lib 2026-04-13") added TensorTensor.execute() with a loop over self._dst.get_num_tiles_without_virtual_batches(), a method not defined on TileStream or any other class. This resulted in every call to the primitive raising AttributeError, making TensorTensor unusable. The sibling primitives (TensorScalar, Reciprocal) all iterate with self._dst.get_num_tiles(), the method actually defined on TileStream, which returns the product of the tile grid. Fix by calling get_num_tiles(). Add integration tests covering all five ops (multiply/add/subtract/maximum/minimum) across single-tile and multi-tile (partial-tile) shapes, verified in simulation. --- .../primitives/blas/tensor_tensor.py | 2 +- .../experimental/primitives/__init__.py | 13 ++ .../experimental/primitives/blas/__init__.py | 13 ++ .../primitives/blas/test_tensor_tensor.py | 195 ++++++++++++++++++ 4 files changed, 222 insertions(+), 1 deletion(-) create mode 100644 test/integration/nkilib/experimental/primitives/__init__.py create mode 100644 test/integration/nkilib/experimental/primitives/blas/__init__.py create mode 100644 test/integration/nkilib/experimental/primitives/blas/test_tensor_tensor.py 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/__init__.py b/test/integration/nkilib/experimental/primitives/__init__.py new file mode 100644 index 0000000..ce63f8f --- /dev/null +++ b/test/integration/nkilib/experimental/primitives/__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/__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..3687c64 --- /dev/null +++ b/test/integration/nkilib/experimental/primitives/blas/test_tensor_tensor.py @@ -0,0 +1,195 @@ +# 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, pytest_test_metadata +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_test_metadata(name="TensorTensor") +@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, + ) From c41bf61a430e23ac625d54edd183f782622f2c82 Mon Sep 17 00:00:00 2001 From: Harris Hua Date: Fri, 24 Jul 2026 00:36:54 +0000 Subject: [PATCH 2/2] test: remove pytest_test_metadata decorator per review (keep out of pipeline) --- .../nkilib/experimental/primitives/blas/test_tensor_tensor.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/test/integration/nkilib/experimental/primitives/blas/test_tensor_tensor.py b/test/integration/nkilib/experimental/primitives/blas/test_tensor_tensor.py index 3687c64..01448d5 100644 --- a/test/integration/nkilib/experimental/primitives/blas/test_tensor_tensor.py +++ b/test/integration/nkilib/experimental/primitives/blas/test_tensor_tensor.py @@ -32,7 +32,7 @@ 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, pytest_test_metadata +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 @@ -131,7 +131,6 @@ def _output_tensors(kernel_input): # ============================================================================= -@pytest_test_metadata(name="TensorTensor") @pytest_marks(["tensor_tensor"]) class TestTensorTensorPrimitive: """Tests class for the tensor_tensor BLAS primitive kernel."""