code/models/common/tests/modules/test_lazy_buffer.py

# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""Tests for LazyBuffer — lazy device buffer allocation for mutable state tensors."""

import pytest
import torch

import ttnn
from models.common.modules.lazy_buffer import LazyBuffer, resolve_lazy_buffer

# ==============================================================================
# Unit tests (no device)
# ==============================================================================


class TestLazyBufferUnit:
    def test_defaults(self):
        """Default dtype=int32, layout=TILE. device/mesh_mapper/memory_config=None."""
        source = torch.zeros(4, 8, dtype=torch.int32)
        buf = LazyBuffer(source=source)
        assert buf.dtype == ttnn.int32
        assert buf.layout == ttnn.TILE_LAYOUT
        assert buf.device is None
        assert buf.mesh_mapper is None
        assert buf.memory_config is None
        assert torch.equal(buf.source, source)

    def test_custom_fields(self):
        source = torch.ones(32, 1, dtype=torch.float32)
        buf = LazyBuffer(
            source=source,
            dtype=ttnn.bfloat16,
            layout=ttnn.ROW_MAJOR_LAYOUT,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
        )
        assert buf.dtype == ttnn.bfloat16
        assert buf.layout == ttnn.ROW_MAJOR_LAYOUT
        assert buf.memory_config == ttnn.DRAM_MEMORY_CONFIG

    def test_is_resolved_false_without_device(self):
        buf = LazyBuffer(source=torch.zeros(4, 8))
        assert not buf.is_resolved()

    def test_is_resolved_false_without_dtype(self):
        buf = LazyBuffer(source=torch.zeros(4, 8), dtype=None, device="fake")
        assert not buf.is_resolved()

    def test_is_resolved_false_without_layout(self):
        buf = LazyBuffer(source=torch.zeros(4, 8), layout=None, device="fake")
        assert not buf.is_resolved()

    def test_is_resolved_true_when_all_set(self):
        buf = LazyBuffer(source=torch.zeros(4, 8), device="fake", dtype=ttnn.int32, layout=ttnn.TILE_LAYOUT)
        assert buf.is_resolved()

    def test_raises_without_device(self):
        buf = LazyBuffer(source=torch.zeros(4, 8))
        with pytest.raises(ValueError, match="device must be set"):
            buf.get_device_buffer()

    def test_raises_without_layout(self):
        buf = LazyBuffer(source=torch.zeros(4, 8), layout=None, device="fake")
        with pytest.raises(ValueError, match="layout must be set"):
            buf.get_device_buffer()

    def test_update_before_materialize_replaces_source(self):
        """update() before get_device_buffer() just swaps the source tensor."""
        source = torch.zeros(4, 8)
        buf = LazyBuffer(source=source)
        new_source = torch.ones(4, 8)
        buf.update(new_source)
        assert torch.equal(buf.source, new_source)
        assert buf._value is None  # not yet materialized


class TestResolveLazyBuffer:
    def test_fills_none_fields(self):
        buf = LazyBuffer(source=torch.zeros(4, 8), dtype=ttnn.bfloat16, layout=None, memory_config=None)
        resolved = resolve_lazy_buffer(buf, layout=ttnn.ROW_MAJOR_LAYOUT, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        assert resolved.layout == ttnn.ROW_MAJOR_LAYOUT
        assert resolved.memory_config == ttnn.DRAM_MEMORY_CONFIG

    def test_does_not_overwrite_non_none(self):
        buf = LazyBuffer(source=torch.zeros(4, 8), dtype=ttnn.bfloat16)
        resolved = resolve_lazy_buffer(buf, dtype=ttnn.int32)
        assert resolved.dtype == ttnn.bfloat16  # preserved, not overwritten

    def test_preserves_source(self):
        source = torch.randn(4, 8)
        buf = LazyBuffer(source=source)
        resolved = resolve_lazy_buffer(buf, device="fake")
        assert torch.equal(resolved.source, source)

    def test_returns_new_instance(self):
        buf = LazyBuffer(source=torch.zeros(4, 8))
        resolved = resolve_lazy_buffer(buf, device="fake")
        assert buf is not resolved
        assert buf.device is None  # original unchanged
        assert resolved.device == "fake"


# ==============================================================================
# Device tests
# ==============================================================================


@pytest.mark.parametrize("ttnn_mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
class TestLazyBufferDevice:
    def test_get_device_buffer_returns_ttnn_tensor(self, ttnn_mesh_device):
        buf = LazyBuffer(
            source=torch.zeros(32, 64, dtype=torch.int32),
            dtype=ttnn.int32,
            layout=ttnn.TILE_LAYOUT,
            device=ttnn_mesh_device,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
        )
        result = buf.get_device_buffer()
        assert isinstance(result, ttnn.Tensor)

    def test_get_device_buffer_idempotent(self, ttnn_mesh_device):
        """Second call returns the exact same Python object (cached handle)."""
        buf = LazyBuffer(
            source=torch.zeros(32, 64, dtype=torch.int32),
            dtype=ttnn.int32,
            layout=ttnn.TILE_LAYOUT,
            device=ttnn_mesh_device,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
        )
        first = buf.get_device_buffer()
        second = buf.get_device_buffer()
        assert first is second

    def test_update_after_materialize_preserves_handle(self, ttnn_mesh_device):
        """update() after materialization keeps the same tensor handle (no reallocation)."""
        buf = LazyBuffer(
            source=torch.zeros(32, 1, dtype=torch.float32),
            dtype=ttnn.bfloat16,
            layout=ttnn.TILE_LAYOUT,
            device=ttnn_mesh_device,
        )
        handle_before = buf.get_device_buffer()
        buf.update(torch.ones(32, 1, dtype=torch.float32))
        handle_after = buf.get_device_buffer()
        assert handle_before is handle_after

    def test_update_after_materialize_changes_device_data(self, ttnn_mesh_device):
        """update() actually writes new data to device — readback should match new source."""
        buf = LazyBuffer(
            source=torch.zeros(32, 1, dtype=torch.float32),
            dtype=ttnn.bfloat16,
            layout=ttnn.TILE_LAYOUT,
            device=ttnn_mesh_device,
        )
        buf.get_device_buffer()  # materialize with zeros

        # Update with ones
        new_source = torch.ones(32, 1, dtype=torch.float32)
        buf.update(new_source)

        # Readback and verify
        readback = ttnn.to_torch(buf.get_device_buffer()).float()
        assert torch.allclose(
            readback[:32, :1], new_source, atol=0.01
        ), f"Readback mismatch: expected ones, got {readback[:4, :1].flatten()}"

    def test_update_before_materialize_uses_new_source(self, ttnn_mesh_device):
        """update() before get_device_buffer() means materialization uses the new source."""
        buf = LazyBuffer(
            source=torch.zeros(32, 1, dtype=torch.float32),
            dtype=ttnn.bfloat16,
            layout=ttnn.TILE_LAYOUT,
            device=ttnn_mesh_device,
        )

        # Update BEFORE materializing
        new_source = torch.full((32, 1), 42.0, dtype=torch.float32)
        buf.update(new_source)

        # Now materialize — should use new_source (42.0), not original (0.0)
        readback = ttnn.to_torch(buf.get_device_buffer()).float()
        assert readback[0, 0].item() == pytest.approx(42.0, abs=0.5), f"Expected ~42.0, got {readback[0, 0].item()}"

    def test_with_dram_memory_config(self, ttnn_mesh_device):
        buf = LazyBuffer(
            source=torch.randn(32, 64, dtype=torch.float32),
            dtype=ttnn.bfloat16,
            layout=ttnn.TILE_LAYOUT,
            device=ttnn_mesh_device,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
        )
        t = buf.get_device_buffer()
        assert isinstance(t, ttnn.Tensor)

    def test_row_major_layout(self, ttnn_mesh_device):
        buf = LazyBuffer(
            source=torch.zeros(32, 1, dtype=torch.int32),
            dtype=ttnn.int32,
            layout=ttnn.ROW_MAJOR_LAYOUT,
            device=ttnn_mesh_device,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
        )
        t = buf.get_device_buffer()
        assert isinstance(t, ttnn.Tensor)