From b7d0fd800546056f01006d0a35cc46fb1170c908 Mon Sep 17 00:00:00 2001 From: Chao Wang <26245345+ChaoWao@users.noreply.github.com> Date: Sat, 25 Jul 2026 21:23:48 -0700 Subject: [PATCH] Fix: stage bgemm's C accumulator to the device (host_build_graph) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #1483 `tests/st/a2a3/host_build_graph/bgemm` declared its `C` argument `D.OUT`, but the task graph read-modify-writes it: `p_add.add_inout(c_view)` feeds `kernel_tile_add`, whose first `k` iteration `TLOAD`s `C` before any task has written it. The runtime stages only IN/INOUT tensors host->device (`runtime_maker.cpp`, "pure write-only OUTPUT buffers are never read by the kernel"), and since #1365 it no longer zero-fills OUT buffers either. So `C` accumulated onto whatever the device allocation happened to hold. On sim that allocation is a plain `malloc` block on a process-lifetime worker: fresh zeros when the case runs alone, a previous case's bytes in a full-suite run — hence a golden mismatch of 1e25-1e36 that reproduced only under the full suite. Same class as #1449. Declare `C` INOUT, matching the two sibling bgemm variants (`examples/a5/tensormap_and_ringbuffer/bgemm`, `examples/a2a3/tensormap_and_ringbuffer/benchmark_bgemm`), which already tag their accumulator that way. Seed `C` with a non-zero base (0.25) so the staging is load-bearing and the case is deterministic: golden is now `C_init + A @ B`, which a zeroed or unstaged device buffer cannot match. Without the signature fix this fails every run instead of one in two or three — verified, it reproduces the reported 1e35 magnitude on the first attempt. Drop the `pytest.mark.skip` that quarantined the case: the underlying defect is fixed, so the case runs in CI again. Verified: bgemm passes 3/3 in isolation on a2a3sim and onboard on a2a3; the full `pytest examples tests/st --platform a2a3sim` suite is green 3/3 runs. Co-Authored-By: Claude Opus 5 (1M context) --- tests/st/a2a3/host_build_graph/bgemm/README.md | 5 ++++- .../bgemm/kernels/orchestration/bgemm_orch.cpp | 10 +++++++--- .../a2a3/host_build_graph/bgemm/test_bgemm.py | 18 ++++++++---------- 3 files changed, 19 insertions(+), 14 deletions(-) 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):