diff --git a/tests/st/a2a3/host_build_graph/bgemm/README.md b/tests/st/a2a3/host_build_graph/bgemm/README.md index bb42b61e1e..ed5b5ff8bd 100644 --- a/tests/st/a2a3/host_build_graph/bgemm/README.md +++ b/tests/st/a2a3/host_build_graph/bgemm/README.md @@ -5,9 +5,12 @@ Tiled matrix multiplication example demonstrating Cube (AIC) and Vector (AIV) co ## Computation ```text -C = A @ B +C = C_init + A @ B ``` +`C` is an INOUT accumulator: every output tile is read-modify-written by +`tile_add`, starting from the host-supplied `C_init` staged to the device. + Tiled computation with 4x4x4 grid: - Tile size: 64 x 64 diff --git a/tests/st/a2a3/host_build_graph/bgemm/kernels/orchestration/bgemm_orch.cpp b/tests/st/a2a3/host_build_graph/bgemm/kernels/orchestration/bgemm_orch.cpp index a72e7ca081..0cb8114c93 100644 --- a/tests/st/a2a3/host_build_graph/bgemm/kernels/orchestration/bgemm_orch.cpp +++ b/tests/st/a2a3/host_build_graph/bgemm/kernels/orchestration/bgemm_orch.cpp @@ -11,7 +11,8 @@ /** * BGEMM orchestration — submit_task / TensorMap form * - * Tiled C = A @ B, tile 64x64, grid 4x4x4. Per output tile (m,n), for each k: + * Tiled C = C_init + A @ B, tile 64x64, grid 4x4x4. Per output tile (m,n), + * for each k: * P_k = A[m,k] @ B[k,n] (gemm_tile, AIC) — fresh runtime-allocated buffer * C[m,n] += P_k (tile_add, AIV, INOUT C tile) * @@ -20,9 +21,12 @@ * - add_{k-1} -> add_k via the C[m,n] tile (add_inout overlap). * Allocating a fresh P per k (instead of reusing one P buffer per tile) removes * the write-after-read hazard the explicit-edge version needed an extra edge - * for. C arrives zero-initialized (pure OUTPUT), so the accumulation is exact. + * for. * - * Arg layout: [A (IN), B (IN), C (OUT)] — flattened 1-D tile-first tensors. + * Arg layout: [A (IN), B (IN), C (INOUT)] — flattened 1-D tile-first tensors. + * C must be declared INOUT, not OUT: the k=0 tile_add reads C before any task + * has written it, and the runtime stages only IN/INOUT tensors to the device + * (a tensor declared OUT is handed to the kernel uninitialized). */ #include diff --git a/tests/st/a2a3/host_build_graph/bgemm/test_bgemm.py b/tests/st/a2a3/host_build_graph/bgemm/test_bgemm.py index a5c74e04c9..4a3fb0c90f 100644 --- a/tests/st/a2a3/host_build_graph/bgemm/test_bgemm.py +++ b/tests/st/a2a3/host_build_graph/bgemm/test_bgemm.py @@ -9,11 +9,10 @@ # ----------------------------------------------------------------------------------------------------------- """BGEMM — host_build_graph runtime with tiled matrix multiplication. -Computation: C = A @ B (4x4x4 grid, 64x64 tiles). +Computation: C = C_init + A @ B (4x4x4 grid, 64x64 tiles). Tests AIC (Cube) + AIV (Vector) cooperation with tile-first memory layout. """ -import pytest import torch from simpler.task_interface import ArgDirection as D @@ -26,15 +25,12 @@ GRID_K = 4 GRID_N = 4 BATCH = 1 +C_BASE = 0.25 -# The golden comparison reads memory this run never wrote: max_diff lands -# between 1e25 and 1e36 on inputs of magnitude 1e-2, roughly one full-suite run -# in two or three, and never when the case runs alone. Tracked in #1483. -@pytest.mark.skip(reason="flaky: reads uninitialised device memory in full-suite runs (#1483)") @scene_test(level=2, runtime="host_build_graph") class TestBgemmHostBuildGraph(SceneTestCase): - """BGEMM: tiled C = A @ B with AIC gemm + AIV tile add.""" + """BGEMM: tiled C = C_init + A @ B with AIC gemm + AIV tile add.""" RTOL = 1e-3 ATOL = 1e-3 @@ -43,7 +39,7 @@ class TestBgemmHostBuildGraph(SceneTestCase): "orchestration": { "source": "kernels/orchestration/bgemm_orch.cpp", "function_name": "aicpu_orchestration_entry", - "signature": [D.IN, D.IN, D.OUT], + "signature": [D.IN, D.IN, D.INOUT], }, "incores": [ { @@ -73,7 +69,10 @@ class TestBgemmHostBuildGraph(SceneTestCase): def generate_args(self, params): A = torch.randn(BATCH, GRID_M, GRID_K, TILE_M, TILE_K, dtype=torch.float32) * 0.01 B = torch.randn(BATCH, GRID_K, GRID_N, TILE_K, TILE_N, dtype=torch.float32) * 0.01 - C = torch.zeros(BATCH, GRID_M, GRID_N, TILE_M, TILE_N, dtype=torch.float32) + # C is an INOUT accumulator: the k=0 tile_add reads it before anything + # writes it. A non-zero base makes the host->device staging of C + # observable — a zeroed or unstaged device buffer fails the compare. + C = torch.full((BATCH, GRID_M, GRID_N, TILE_M, TILE_N), C_BASE, dtype=torch.float32) return TaskArgsBuilder( Tensor("A", A.flatten()), @@ -86,7 +85,6 @@ def compute_golden(self, args, params): B = args.B.reshape(BATCH, GRID_K, GRID_N, TILE_K, TILE_N) C = args.C.reshape(BATCH, GRID_M, GRID_N, TILE_M, TILE_N) - C[:] = 0.0 for batch in range(BATCH): for m_idx in range(GRID_M): for n_idx in range(GRID_N):