Fix: provide SDMA workspace on simulator

examples/a2a3/tensormap_and_ringbuffer/prefetch_async_demo/test_prefetch_async_demo.py

@@ -7,21 +7,28 @@
 # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 # See LICENSE in the root of the software repository for the full text of the License.
 # -----------------------------------------------------------------------------------------------------------
-"""Single-card TPREFETCH_ASYNC smoke test for onboard a2a3.
+"""Single-card TPREFETCH_ASYNC smoke test.
 
 Exercises the runtime-injected SDMA workspace: a Worker created with
-``enable_sdma=True`` provisions the PTO-ISA async-SDMA workspace once at init and
-injects its address into every kernel's GlobalContext, so the kernel obtains it
-via ``get_dma_workspace(args, DMA_WORKSPACE_SDMA)`` -- no workspace is threaded
-as a user arg. A Worker without ``enable_sdma`` creates no SDMA streams and its
+``enable_sdma=True`` provisions the workspace once at init and injects its
+address into every kernel's GlobalContext, so the kernel obtains it via
+``get_dma_workspace(args, DMA_WORKSPACE_SDMA)`` -- no workspace is threaded as a
+user arg. A Worker without ``enable_sdma`` creates no SDMA streams and its
 kernels read a zero workspace address.
 The kernel prefetches ``in`` into L2, waits on the returned event, then copies
 ``in`` to ``out``.
 
 The prefetch is a pure cache hint that changes no value, so ``out == in``
 bit-exactly is the property under test -- together with the event wait actually
-completing rather than hanging. The device log (``[SDMA] Created 48 STARS
-streams OK``) confirms the real SDMA path ran rather than the skip branch.
+completing rather than hanging. Because the kernel returns without writing
+``out`` on a null workspace, that equality also proves the address it read was
+live.
+
+Onboard a2a3 runs the real engine; the device log (``[SDMA] Created 48 STARS
+streams OK``) confirms the real SDMA path ran rather than the skip branch. On
+a2a3sim the workspace is inert scratch and ``TPREFETCH_ASYNC`` is a no-op
+(pto-isa ``cpu/TPrefetchAsync.hpp``), so what the sim case covers is the
+injection path: a live address reaches the kernel.
 
 Unlike the SDMA completion demo this needs no comm domain: the workspace is a
 runtime-owned per-device resource, so a single device is enough.
@@ -65,7 +72,7 @@ class TestPrefetchAsyncDemo(SceneTestCase):
     CASES = [
         {
             "name": "prefetch_copy",
-            "platforms": ["a2a3"],
+            "platforms": ["a2a3sim", "a2a3"],
             "config": {},
             "params": {},
         },

src/a2a3/platform/sim/host/device_runner.cpp

@@ -32,6 +32,7 @@
 #include "aicpu/platform_aicpu_affinity.h"
 #include "call_config.h"
 #include "callable_protocol.h"
+#include "common/dma_workspace.h"
 #include "common/host_log_binding.h"
 #include "common/memory_barrier.h"
 #include "common/platform_config.h"
@@ -204,6 +205,33 @@ int DeviceRunner::ensure_binaries_loaded() {
         if (!load_sym("set_platform_scope_stats_base", reinterpret_cast<void **>(&set_platform_scope_stats_base_func_)))
             return PTO_RUNTIME_ERR_INTERNAL;
 
+        // Per-device one-shot latch, mirroring the onboard InitArgs path that
+        // feeds set_dma_workspace_addr: the scheduler copies these addresses
+        // into every core's GlobalContext, where get_dma_workspace(args, kind)
+        // reads them. The block is inert scratch — CPU simulation runs
+        // TPREFETCH_ASYNC as a no-op (pto-isa cpu/TPrefetchAsync.hpp) and has no
+        // implementation of the async transfers that would read it — but it is
+        // a live address, so a kernel that rejects a null workspace as broken
+        // injection behaves the same here as it does onboard. Kinds other than
+        // SDMA keep the zero address they start with.
+        if (dma_workspace_requested_) {
+            using SetDmaWorkspaceAddrFunc = void (*)(int, unsigned long long);
+            SetDmaWorkspaceAddrFunc set_dma_workspace_addr_func = nullptr;
+            if (!load_sym("set_dma_workspace_addr", reinterpret_cast<void **>(&set_dma_workspace_addr_func)))
+                return PTO_RUNTIME_ERR_INTERNAL;
+            if (dma_workspace_block_ == nullptr) {
+                dma_workspace_block_ = mem_alloc_.alloc(kSimDmaWorkspaceBytes);
+                if (dma_workspace_block_ == nullptr) {
+                    LOG_ERROR("Failed to allocate the %zu-byte sim async-DMA workspace", kSimDmaWorkspaceBytes);
+                    return PTO_RUNTIME_ERR_INTERNAL;
+                }
+                std::memset(dma_workspace_block_, 0, kSimDmaWorkspaceBytes);
+            }
+            set_dma_workspace_addr_func(
+                DMA_WORKSPACE_SDMA, static_cast<unsigned long long>(reinterpret_cast<uintptr_t>(dma_workspace_block_))
+            );
+        }
+
         // The AICPU sim SO binds its private HostLogger before the compatibility
         // level setter can emit a clock anchor.
         using SetLogLevelFunc = void (*)(int);
@@ -843,6 +871,7 @@ int DeviceRunner::finalize() {
     }
 
     mem_alloc_.finalize();
+    dma_workspace_block_ = nullptr;
     clear_cpu_sim_shared_storage();
 
     device_id_ = -1;

src/a5/platform/sim/host/device_runner.cpp

@@ -33,6 +33,7 @@
 #include "aicpu/platform_aicpu_affinity.h"
 #include "call_config.h"
 #include "callable_protocol.h"
+#include "common/dma_workspace.h"
 #include "common/host_log_binding.h"
 #include "common/memory_barrier.h"
 #include "common/platform_config.h"
@@ -204,6 +205,33 @@ int DeviceRunner::ensure_binaries_loaded() {
         if (!load_sym("set_platform_scope_stats_base", reinterpret_cast<void **>(&set_platform_scope_stats_base_func_)))
             return PTO_RUNTIME_ERR_INTERNAL;
 
+        // Per-device one-shot latch, mirroring the onboard InitArgs path that
+        // feeds set_dma_workspace_addr: the scheduler copies these addresses
+        // into every core's GlobalContext, where get_dma_workspace(args, kind)
+        // reads them. The block is inert scratch — CPU simulation runs
+        // TPREFETCH_ASYNC as a no-op (pto-isa cpu/TPrefetchAsync.hpp) and has no
+        // implementation of the async transfers that would read it — but it is
+        // a live address, so a kernel that rejects a null workspace as broken
+        // injection behaves the same here as it does onboard. Kinds other than
+        // SDMA keep the zero address they start with.
+        if (dma_workspace_requested_) {
+            using SetDmaWorkspaceAddrFunc = void (*)(int, unsigned long long);
+            SetDmaWorkspaceAddrFunc set_dma_workspace_addr_func = nullptr;
+            if (!load_sym("set_dma_workspace_addr", reinterpret_cast<void **>(&set_dma_workspace_addr_func)))
+                return PTO_RUNTIME_ERR_INTERNAL;
+            if (dma_workspace_block_ == nullptr) {
+                dma_workspace_block_ = mem_alloc_.alloc(kSimDmaWorkspaceBytes);
+                if (dma_workspace_block_ == nullptr) {
+                    LOG_ERROR("Failed to allocate the %zu-byte sim async-DMA workspace", kSimDmaWorkspaceBytes);
+                    return PTO_RUNTIME_ERR_INTERNAL;
+                }
+                std::memset(dma_workspace_block_, 0, kSimDmaWorkspaceBytes);
+            }
+            set_dma_workspace_addr_func(
+                DMA_WORKSPACE_SDMA, static_cast<unsigned long long>(reinterpret_cast<uintptr_t>(dma_workspace_block_))
+            );
+        }
+
         // The AICPU sim SO binds its private HostLogger before the compatibility
         // level setter can emit a clock anchor.
         using SetLogLevelFunc = void (*)(int);
@@ -803,6 +831,7 @@ int DeviceRunner::finalize() {
     prebuilt_runtime_arena_cache_image_.clear();
 
     mem_alloc_.finalize();
+    dma_workspace_block_ = nullptr;
     clear_cpu_sim_shared_storage();
 
     if (device_wall_dev_ptr_ != nullptr) {

src/common/platform/sim/host/c_api_shared.cpp

@@ -385,17 +385,15 @@ int simpler_init(
     // and the dispatcher / preinstall load path on sim isn't taken anyway.
     (void)dispatcher_binary;
     (void)dispatcher_size;
-    // Simulation provides no async-DMA workspaces, so there is nothing for a
-    // warmup ELF to warm either.
+    // Simulation drives no SDMA control path, so the warmup ELF has nothing to
+    // walk.
     (void)sdma_warmup_binary;
     (void)sdma_warmup_size;
 
     if (ctx == NULL) return PTO_RUNTIME_ERR_INTERNAL;
-    // Opting into SDMA fails here rather than at the first kernel read, so such
-    // a Worker cannot come up on sim at all.
-    if (enable_sdma != 0) return PTO_RUNTIME_ERR_UNSUPPORTED;
 
     SimDeviceRunnerBase *runner = static_cast<SimDeviceRunnerBase *>(ctx);
+    runner->set_dma_workspace_request(enable_sdma != 0);
 
     int rc;
     try {

src/common/platform/sim/host/device_runner_base.h

@@ -251,6 +251,18 @@ class SimDeviceRunnerBase {
         aicpu_so_binary_ = std::move(aicpu_so_binary);
         aicore_kernel_binary_ = std::move(aicore_kernel_binary);
     }
+
+    /**
+     * Record whether this Worker asked for an async-DMA workspace.
+     *
+     * Simulation has no DMA engine, so the workspace it hands out is inert
+     * scratch: `TPREFETCH_ASYNC` is a no-op under CPU simulation (pto-isa
+     * `cpu/TPrefetchAsync.hpp`) and nothing else reads it. It is provisioned
+     * anyway so a kernel that treats a null workspace as broken injection sees
+     * a live address, which is the same contract onboard offers. A Worker that
+     * does not ask keeps a zero address.
+     */
+    void set_dma_workspace_request(bool enable_sdma) { dma_workspace_requested_ = enable_sdma; }
     int device_id() const { return device_id_; }
     uint64_t last_device_wall_ns() const { return device_wall_ns_; }
     // Per-phase AICPU wall (ns) from the most recent run; RunWall aliases
@@ -339,6 +351,15 @@ class SimDeviceRunnerBase {
     std::vector<uint8_t> aicpu_so_binary_;
     std::vector<uint8_t> aicore_kernel_binary_;
 
+    // Set by simpler_init; read when the AICPU SO is loaded to decide whether to
+    // publish a workspace address into it.
+    bool dma_workspace_requested_{false};
+    // Matches the 16 KB the a2a3 onboard SdmaWorkspaceManager provisions, so a
+    // kernel sizing its use of the workspace against the onboard budget stays
+    // in bounds here. Freed with every other tracked block by mem_alloc_.
+    static constexpr size_t kSimDmaWorkspaceBytes = 16 * 1024;
+    void *dma_workspace_block_{nullptr};
+
     MemoryAllocator mem_alloc_;
     std::array<void *, PTO_PIPELINE_MAX_DEPTH> retained_temp_addrs_{};
     std::array<size_t, PTO_PIPELINE_MAX_DEPTH> retained_temp_sizes_{};

tests/ut/py/test_worker/test_dma_workspace_sim.py

@@ -9,12 +9,20 @@
 # ruff: noqa: PLC0415
 """Sim-backend test for the async-DMA workspace request carried by ``simpler_init``.
 
-Simulation provides no async-DMA engine, so ``enable_sdma=True`` has no workspace
-to hand a kernel. The contract is that such a Worker fails to come up rather than
-reaching its first run reading a zero address it expected to be live — the
-rejection is the only thing standing between an unsupported request and a silent
-wrong answer, and it is easy to lose to a refactor that derives the provisioned
-set from what the platform supports (an empty set looks like success).
+Simulation has no async-DMA engine, but ``enable_sdma=True`` still hands the
+kernel a live workspace address: an inert block, published through the same
+``set_dma_workspace_addr`` path onboard feeds from ``InitArgs``. The contract a
+kernel relies on is that a non-zero address means the workspace is there to use —
+a zero address where a live one was expected is a silent wrong answer rather than
+an error, because a kernel that branches on null (see
+``examples/a2a3/tensormap_and_ringbuffer/prefetch_async_demo``) skips its work and
+fails a golden comparison instead of reporting anything. Provisioning inert
+scratch keeps that contract on sim, where nothing dereferences the block:
+``TPREFETCH_ASYNC`` is a no-op stub (pto-isa ``cpu/TPrefetchAsync.hpp``) and the
+async transfers that would read it have no CPU-sim implementation at all.
+
+A Worker that does not ask for SDMA keeps a zero address, which is what makes the
+address meaningful as a signal.
 """
 
 from __future__ import annotations
@@ -40,20 +48,24 @@ def _make_sim_worker(*, enable_sdma: bool):
     )
 
 
-class TestSimRejectsSdma:
-    def test_enable_sdma_fails_init(self):
+class TestSimProvisionsSdmaWorkspace:
+    def test_enable_sdma_init_succeeds(self):
+        # An SDMA-enabled Worker comes up on sim. It used to be rejected here
+        # with PTO_RUNTIME_ERR_UNSUPPORTED (-1001), which kept every kernel that
+        # merely hints a prefetch off the simulator.
         worker = _make_sim_worker(enable_sdma=True)
-        # -1001 is PTO_RUNTIME_ERR_UNSUPPORTED. Pinning the code, not just "some
-        # exception", is what distinguishes a rejected request from sim failing
-        # to start for an unrelated reason.
-        with pytest.raises(RuntimeError, match=r"simpler_init failed with code -1001"):
+        try:
             worker.init()
-        worker.close()
+        finally:
+            worker.close()
 
     def test_without_sdma_init_succeeds(self):
-        # Positive control: the same Worker comes up when it does not ask for an
-        # engine sim has no provider for, so the failure above is attributable to
-        # the request rather than to sim being unable to start at all.
+        # Control: coming up is not itself evidence the request was honoured, so
+        # pair it with the Worker that asks for nothing. The workspace address
+        # each one hands a kernel is covered end-to-end by the a2a3
+        # prefetch_async_demo scene test, which reads it on-device.
         worker = _make_sim_worker(enable_sdma=False)
-        worker.init()
-        worker.close()
+        try:
+            worker.init()
+        finally:
+            worker.close()