diff --git a/backends/test/harness/tester.py b/backends/test/harness/tester.py index 1b15a503dae..483579d7265 100644 --- a/backends/test/harness/tester.py +++ b/backends/test/harness/tester.py @@ -381,11 +381,16 @@ def _dump_golden_artifacts( logger = logging.getLogger(__name__) os.makedirs(artifact_dir, exist_ok=True) + def _to_physical_layout(t: torch.Tensor) -> torch.Tensor: + # The .bin has to match the dim_order the .pte declares; contiguous() + # alone would re-lay a channels_last tensor back to NCHW. + return t.detach().permute(t.dim_order()).contiguous() + for i, inp in enumerate(inputs): if isinstance(inp, torch.Tensor): suffix = "" if len(inputs) == 1 else f"_{i}" path = os.path.join(artifact_dir, f"{artifact_name}_input{suffix}.bin") - inp.detach().contiguous().numpy().tofile(path) + _to_physical_layout(inp).numpy().tofile(path) logger.info(f"Saved golden input to {path}") if isinstance(reference_output, torch.Tensor): diff --git a/backends/test/harness/tests/test_golden_artifacts.py b/backends/test/harness/tests/test_golden_artifacts.py new file mode 100644 index 00000000000..d869d0291d1 --- /dev/null +++ b/backends/test/harness/tests/test_golden_artifacts.py @@ -0,0 +1,61 @@ +import os +import tempfile +import unittest + +import numpy as np +import torch +from executorch.backends.test.harness.tester import Tester + + +class GoldenArtifactTests(unittest.TestCase): + def _dump(self, inputs, reference_output): + artifact_dir = tempfile.mkdtemp() + Tester._dump_golden_artifacts(artifact_dir, "m", inputs, reference_output) + return artifact_dir + + def _read(self, artifact_dir, name): + return np.fromfile(os.path.join(artifact_dir, name), dtype=np.float32) + + def test_channels_last_input_is_written_as_nhwc(self): + x = torch.arange(24, dtype=torch.float32).reshape(1, 3, 2, 4) + channels_last = x.to(memory_format=torch.channels_last) + + artifact_dir = self._dump((channels_last,), channels_last) + + self.assertTrue( + np.array_equal( + self._read(artifact_dir, "m_input.bin"), + channels_last.permute(0, 2, 3, 1).reshape(-1).numpy(), + ) + ) + + def test_contiguous_input_is_unchanged(self): + x = torch.arange(24, dtype=torch.float32).reshape(1, 3, 2, 4) + + artifact_dir = self._dump((x,), x) + + self.assertTrue( + np.array_equal( + self._read(artifact_dir, "m_input.bin"), x.reshape(-1).numpy() + ) + ) + + def test_output_of_a_view_is_materialized(self): + # reference_output is the eager result, so a model ending in permute + # returns a view. The program materializes that contiguously, so the + # golden output has to follow the values rather than the source layout. + x = torch.arange(6, dtype=torch.float32).reshape(2, 3) + permuted = x.permute(1, 0) + + artifact_dir = self._dump((x,), permuted) + + self.assertTrue( + np.array_equal( + self._read(artifact_dir, "m_expected_output.bin"), + torch.tensor([0.0, 3.0, 1.0, 4.0, 2.0, 5.0]).numpy(), + ) + ) + + +if __name__ == "__main__": + unittest.main()