Fix AttributeError in TensorTensor.execute from undefined tile-count method

[RuifengHua]

Summary

TensorTensor.execute() loops over self._dst.get_num_tiles_without_virtual_batches(), but that method is not defined on TileStream (or any other class in the library). Any call to the primitive therefore raises AttributeError, making TensorTensor unusable:

AttributeError: 'TileStream' object has no attribute 'get_num_tiles_without_virtual_batches'

The sibling elementwise primitives (TensorScalar, Reciprocal) all iterate their tile loop with self._dst.get_num_tiles() — the method actually defined on TileStream, which returns the product of the tile grid. TensorTensor is the only one referencing the undefined name.

This PR fixes the call to get_num_tiles() and adds an integration test suite for the primitive (which previously had no direct test coverage).

Repro (before fix)

import nki, nki.language as nl
from nkilib_src.nkilib.experimental.primitives import blas, dma, tile_stream
from nkilib_src.nkilib.experimental.primitives.iter_order import RowMajor

@nki.jit
def k(x1, x2):
    p, f = x1.shape
    a = tile_stream.alloc_logical((p, f), p, x1.dtype, "a")
    b = tile_stream.alloc_logical((p, f), p, x2.dtype, "b")
    d = tile_stream.alloc_logical((p, f), p, x1.dtype, "d")
    dma.load(a, x1); dma.load(b, x2)
    blas.TensorTensor(
        dst=tile_stream.tile(d, (p, f), iter_order=RowMajor()),
        src1=tile_stream.tile(a, (p, f), iter_order=RowMajor()),
        src2=tile_stream.tile(b, (p, f), iter_order=RowMajor()),
        op=nl.add,
    ).execute()   # AttributeError before this fix
    return d

Test plan

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

• Fast tests pass — all five ops (multiply, add, subtract, maximum, minimum), single-tile and multi-tile shapes
• Pairwise sweep passes — op × partition P (incl. P > 128 with a partial last tile) × free dim F × dtype (float32, bfloat16)
• Confirmed the test suite fails against the pre-fix code (AttributeError) and passes after the fix
• Sabotage check: mutating the torch reference makes the tests fail, confirming they validate the output rather than merely running

Note: validated in simulation only (no perf claim — this is a correctness fix). Tolerances: 1e-4 (fp32) / 1e-2 (bf16).

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.