Refactor: ask async-DMA provisioning for one flag, not a set of engines

python/bindings/task_interface.cpp

@@ -64,7 +64,6 @@
 #include "common/host_span_scope.h"
 #include "host_log.h"
 #include "data_type.h"
-#include "dma_workspace.h"
 #include "worker_chip_orch_comm.h"
 #include "worker_bind.h"
 #include "task_args_wire.h"
@@ -2908,13 +2907,9 @@ NB_MODULE(_task_interface, m) {
                const std::string &aicore_path, const std::string &dispatcher_path, int device_id,
                std::optional<CallConfig> prewarm_config, bool enable_sdma, const std::string &sim_context_path,
                const std::string &sdma_warmup_path) {
-                // Translate the Python bool into a DmaWorkspaceKind bitmask so the
-                // platform-agnostic ChipWorker stays free of the enum. Empty mask
-                // when disabled leaves the Worker with no async-DMA provisioning.
-                uint32_t dma_workspace_mask = enable_sdma ? (uint32_t{1} << DMA_WORKSPACE_SDMA) : 0;
                 self.init(
                     host_lib_path, aicpu_path, aicore_path, dispatcher_path, device_id,
-                    prewarm_config.has_value() ? &(*prewarm_config) : nullptr, dma_workspace_mask, sim_context_path,
+                    prewarm_config.has_value() ? &(*prewarm_config) : nullptr, enable_sdma, sim_context_path,
                     sdma_warmup_path
                 );
             },

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

@@ -1151,12 +1151,12 @@ int device_memory_info_ctx(DeviceContextHandle ctx, DeviceMemoryInfo *info) {
 }
 
 int simpler_provision_dma_workspace(
-    DeviceContextHandle ctx, uint32_t required_mask, const void *sdma_warmup_binary, uint64_t sdma_warmup_size
+    DeviceContextHandle ctx, int enable_sdma, const void *sdma_warmup_binary, uint64_t sdma_warmup_size
 ) {
     if (ctx == NULL) return PTO_RUNTIME_ERR_INTERNAL;
     try {
         return static_cast<DeviceRunnerBase *>(ctx)->provision_dma_workspace(
-            required_mask, sdma_warmup_binary, static_cast<size_t>(sdma_warmup_size)
+            enable_sdma != 0, sdma_warmup_binary, static_cast<size_t>(sdma_warmup_size)
         );
     } catch (...) {
         return PTO_RUNTIME_ERR_INTERNAL;

src/common/platform/onboard/host/device_runner_base.cpp

@@ -983,11 +983,35 @@ int DeviceRunnerBase::unregister_callable(int32_t callable_id) {
 bool DeviceRunnerBase::has_callable(int32_t callable_id) const { return callables_.count(callable_id) != 0; }
 
 int DeviceRunnerBase::provision_dma_workspace(
-    uint32_t required_mask, const void *sdma_warmup_binary, size_t sdma_warmup_size
+    bool enable_sdma, const void *sdma_warmup_binary, size_t sdma_warmup_size
 ) {
     const uint32_t supported = dma_workspace_supported_mask();
-    if ((required_mask & ~supported) != 0) {
-        LOG_ERROR("provision_dma_workspace: unsupported mask=0x%x (supported=0x%x)", required_mask, supported);
+    constexpr uint32_t kSdmaBit = uint32_t{1} << DMA_WORKSPACE_SDMA;
+    // Opting in on a device that cannot provide SDMA is a caller error, not a
+    // silent no-op: a Worker built for TPREFETCH_ASYNC must not reach its first
+    // run reading a zero workspace address.
+    if (enable_sdma && (supported & kSdmaBit) == 0) {
+        LOG_ERROR("provision_dma_workspace: SDMA requested where unsupported (supported=0x%x)", supported);
+        return PTO_RUNTIME_ERR_UNSUPPORTED;
+    }
+    // Everything this device supports, minus what the caller declined. SDMA is
+    // the only declinable engine: its workspace cannot be obtained without also
+    // creating 48 CP-process STARS streams, which halves this Worker's
+    // post-fault reset budget, so a Worker that did not ask for it must not end
+    // up holding them. Every other supported engine carries no such cost and is
+    // provisioned unconditionally.
+    const uint32_t required_mask = enable_sdma ? supported : (supported & ~kSdmaBit);
+    // Dormant while one engine is supported, since required_mask is a subset of
+    // supported. It arms itself on the day dma_workspace_supported_mask() widens,
+    // which is the day the single-handle contract breaks — dma_workspace_release()
+    // casts the opaque handle back to the one provider type it can be, so a second
+    // engine would be released as the type of the first. A rejection here beats
+    // that silent type confusion.
+    if ((required_mask & (required_mask - 1)) != 0) {
+        LOG_ERROR(
+            "provision_dma_workspace: mask=0x%x names %d engines; one handle owns one provider", required_mask,
+            __builtin_popcount(required_mask)
+        );
         return PTO_RUNTIME_ERR_UNSUPPORTED;
     }
     if (dma_workspace_handle_ != nullptr) {

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

@@ -396,13 +396,15 @@ class DeviceRunnerBase {
     bool has_callable(int32_t callable_id) const;
 
     /**
-     * Provision the async-DMA workspaces named in `required_mask` once at Worker
-     * init and latch their device addresses into the resident KernelArgs so every
-     * subsequent run carries them (AICPU injects them into GlobalContext via
-     * get_dma_workspace). Called only for a Worker created with SDMA enabled;
-     * `required_mask` bits outside dma_workspace_supported_mask() are rejected, so
-     * a platform/runtime without SDMA fails fast. The provider handle is released
-     * by finalize_common().
+     * Provision this device's async-DMA workspaces once at Worker init and latch
+     * their device addresses into the resident KernelArgs so every subsequent run
+     * carries them (AICPU injects them into GlobalContext via get_dma_workspace).
+     * Provisions every engine dma_workspace_supported_mask() names, except SDMA
+     * unless `enable_sdma` is set — SDMA is declinable because its workspace
+     * cannot be obtained without also creating 48 CP-process STARS streams, which
+     * halves this Worker's post-fault reset budget. `enable_sdma` on a
+     * platform/runtime without SDMA is rejected, so such a Worker fails fast. The
+     * provider handle is released by finalize_common().
      *
      * `sdma_warmup_binary` / `sdma_warmup_size`, when non-empty, are handed to
      * launch_sdma_warmup_kernel() once the workspace is live. An absent or
@@ -411,9 +413,8 @@ class DeviceRunnerBase {
      *
      * @return 0 on success, negative on unsupported/failed provisioning.
      */
-    int provision_dma_workspace(
-        uint32_t required_mask, const void *sdma_warmup_binary = nullptr, size_t sdma_warmup_size = 0
-    );
+    int
+    provision_dma_workspace(bool enable_sdma, const void *sdma_warmup_binary = nullptr, size_t sdma_warmup_size = 0);
 
     /**
      * Content-derived stable identity for a registered callable: the

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

@@ -956,15 +956,15 @@ int device_memory_info_ctx(DeviceContextHandle ctx, DeviceMemoryInfo *info) {
 }
 
 int simpler_provision_dma_workspace(
-    DeviceContextHandle ctx, uint32_t required_mask, const void *sdma_warmup_binary, uint64_t sdma_warmup_size
+    DeviceContextHandle ctx, int enable_sdma, const void *sdma_warmup_binary, uint64_t sdma_warmup_size
 ) {
-    // Simulation provides no async-DMA workspaces; a non-empty request fails
-    // fast so an SDMA-enabled Worker cannot come up on sim. With no workspace
-    // there is likewise nothing for the warmup ELF to warm.
+    // Simulation provides no async-DMA workspaces; opting into SDMA fails fast
+    // so such a Worker cannot come up on sim. With no workspace there is
+    // likewise nothing for the warmup ELF to warm.
     (void)ctx;
     (void)sdma_warmup_binary;
     (void)sdma_warmup_size;
-    return required_mask == 0 ? 0 : PTO_RUNTIME_ERR_UNSUPPORTED;
+    return enable_sdma == 0 ? 0 : PTO_RUNTIME_ERR_UNSUPPORTED;
 }
 
 }  // extern "C"

src/common/platform_comm/comm.h

@@ -107,6 +107,12 @@ uint32_t dma_workspace_channel_count(void);
  * dma_workspace_release(). Bits outside dma_workspace_supported_mask() are
  * rejected, so provisioning fails fast on a platform/runtime without SDMA.
  *
+ * At most one bit may be set. A single opaque handle cannot own two providers,
+ * and dma_workspace_release() recovers the concrete type by casting it, so a
+ * second engine would be released as the type of the first. Provisioning more
+ * than one engine per call needs a per-kind handle on both sides of this seam;
+ * until then callers get a rejection rather than that silent type confusion.
+ *
  * @return 0 on success, non-zero on invalid/unsupported input or provisioning
  *         failure.
  */

src/common/worker/chip_worker.cpp

@@ -164,7 +164,7 @@ ChipWorker::~ChipWorker() { finalize(); }
 
 void ChipWorker::init(
     const std::string &host_lib_path, const std::string &aicpu_path, const std::string &aicore_path,
-    const std::string &dispatcher_path, int device_id, const CallConfig *prewarm_config, uint32_t dma_workspace_mask,
+    const std::string &dispatcher_path, int device_id, const CallConfig *prewarm_config, bool enable_sdma,
     const std::string &sim_context_path, const std::string &sdma_warmup_path
 ) {
     if (finalized_) {
@@ -439,7 +439,7 @@ void ChipWorker::init(
     // addresses are latched into the resident KernelArgs so every run carries
     // them. On failure, roll the whole Worker back through finalize() so no
     // half-provisioned state leaks, then surface the error.
-    if (dma_workspace_mask != 0) {
+    if (enable_sdma) {
         // The warmup ELF rides provisioning because it needs the workspace it
         // warms. Absent on arches that build no such ELF and on tests driving
         // _ChipWorker.init directly; the platform then skips the warmup and the
@@ -458,13 +458,11 @@ void ChipWorker::init(
             }
         }
         const uint8_t *warmup_ptr = warmup_bytes.empty() ? nullptr : warmup_bytes.data();
-        int prov_rc =
-            simpler_provision_dma_workspace_fn_(device_ctx_, dma_workspace_mask, warmup_ptr, warmup_bytes.size());
+        int prov_rc = simpler_provision_dma_workspace_fn_(device_ctx_, 1, warmup_ptr, warmup_bytes.size());
         if (prov_rc != 0) {
             finalize();
             throw std::runtime_error(
-                "async-DMA workspace provisioning (mask=" + std::to_string(dma_workspace_mask) + ") failed with code " +
-                std::to_string(prov_rc)
+                "async-DMA (SDMA) workspace provisioning failed with code " + std::to_string(prov_rc)
             );
         }
     }

src/common/worker/chip_worker.h

@@ -72,22 +72,23 @@ class ChipWorker {
     /// runtime-arena for its ring sizing right after the device comes up (the
     /// sizing is fork-constant, delivered by COW into init). A no-op for
     /// runtimes without a prebuilt arena.
-    /// `dma_workspace_mask` (a bitmask of DmaWorkspaceKind bits, 0 = none)
-    /// provisions those async-DMA workspaces once at init, so kernels can use
-    /// get_dma_workspace. Empty by default; a Worker that does not opt in creates
-    /// no SDMA streams. Provisioning fails fast (init throws) on a
-    /// platform/runtime that does not support a requested engine. The mask stays
-    /// a raw integer here so this platform-agnostic worker needs no platform
-    /// headers; the binding derives it from the DmaWorkspaceKind enum.
+    /// `enable_sdma` opts this Worker into the async-DMA (SDMA) workspace, so
+    /// kernels can use get_dma_workspace. Off by default; a Worker that does not
+    /// opt in creates no SDMA streams and its kernels read a zero address for
+    /// that engine. It names one engine rather than a set because SDMA is the
+    /// only one a caller can decline — every other supported engine is
+    /// provisioned unconditionally, and SDMA is conditional only because its
+    /// workspace is inseparable from 48 CP-process STARS streams whose
+    /// post-fault release CANN does not bound. Provisioning fails fast (init
+    /// throws) on a platform/runtime without SDMA support.
     /// `sdma_warmup_path`, when non-empty, is the vector-only ELF that walks the
     /// SDMA control path once per channel during that provisioning, moving the
     /// cold-start cost off the first TPREFETCH_ASYNC. Optional: an empty path (or
     /// an arch that builds no such ELF) only costs that first-call latency.
     void init(
         const std::string &host_lib_path, const std::string &aicpu_path, const std::string &aicore_path,
         const std::string &dispatcher_path, int device_id, const CallConfig *prewarm_config = nullptr,
-        uint32_t dma_workspace_mask = 0, const std::string &sim_context_path = "",
-        const std::string &sdma_warmup_path = ""
+        bool enable_sdma = false, const std::string &sim_context_path = "", const std::string &sdma_warmup_path = ""
     );
 
     /// Tear down everything: device resources and runtime library.
@@ -274,7 +275,7 @@ class ChipWorker {
     using GetPipelineContractFn = const PipelineContract *(*)();
     using SimplerUnregisterCallableFn = int (*)(void *, int32_t);
     using GetAicpuDlopenCountFn = size_t (*)(void *);
-    using SimplerProvisionDmaWorkspaceFn = int (*)(void *, uint32_t, const void *, uint64_t);
+    using SimplerProvisionDmaWorkspaceFn = int (*)(void *, int, const void *, uint64_t);
     using FinalizeDeviceFn = int (*)(void *);
     using EnsureAclReadyFn = int (*)(void *, int);
     using CreateCommStreamFn = void *(*)(void *);

src/common/worker/runtime_c_api.h

@@ -496,12 +496,22 @@ size_t get_host_dlopen_count(DeviceContextHandle ctx);
 size_t get_run_stream_set_create_count(DeviceContextHandle ctx);
 
 /**
- * Provision the async-DMA workspaces named in `required_mask` (a bitmask of
- * DmaWorkspaceKind bits) once at Worker init, latching their device addresses
- * into the resident KernelArgs so every subsequent run carries them. Called only
- * for a Worker created with SDMA enabled. Bits unsupported by this
- * platform/runtime are rejected, so a Worker opting into SDMA on sim / a5 / hbg
- * fails fast. Returns 0 on success, negative on unsupported/failed provisioning.
+ * Provision this device's async-DMA workspaces once at Worker init, latching
+ * their device addresses into the resident KernelArgs so every subsequent run
+ * carries them. Every engine the platform supports is provisioned except SDMA,
+ * which is provisioned only when `enable_sdma` is non-zero.
+ *
+ * SDMA is the one engine a caller can decline, and the parameter is a flag
+ * rather than an engine set for that reason. Its workspace cannot be obtained
+ * without also creating 48 CP-process STARS streams, and a Worker holding those
+ * gets a single device-reset attempt after an AICore fault instead of three, so
+ * defaulting it on would put every Worker in that population. Selecting among
+ * provisioned engines is the kernel's job, through get_dma_workspace(args,
+ * kind); a kind this device did not provision reads back 0.
+ *
+ * `enable_sdma` on a platform/runtime without SDMA support is rejected, so a
+ * Worker opting in on sim / a5 / hbg fails fast. Returns 0 on success, negative
+ * on unsupported/failed provisioning.
  *
  * `sdma_warmup_binary` / `sdma_warmup_size` carry the vector-only ELF that walks
  * the SDMA control path once per channel against the workspace just provisioned,
@@ -511,7 +521,7 @@ size_t get_run_stream_set_create_count(DeviceContextHandle ctx);
  * or sync fails does fail provisioning, because that card is then poisoned.
  */
 int simpler_provision_dma_workspace(
-    DeviceContextHandle ctx, uint32_t required_mask, const void *sdma_warmup_binary, uint64_t sdma_warmup_size
+    DeviceContextHandle ctx, int enable_sdma, const void *sdma_warmup_binary, uint64_t sdma_warmup_size
 );
 
 #ifdef __cplusplus

tests/ut/cpp/hierarchical/test_pipeline_contract.cpp

@@ -28,9 +28,13 @@ namespace {
 // Appended (never inserted) since the last revision: `sdma_warmup_path`, the
 // vector-only ELF used to warm the SDMA control path during init-time workspace
 // provisioning.
+// Retyped in place, not moved: the async-DMA parameter is `bool enable_sdma`
+// rather than a DmaWorkspaceKind bitmask. SDMA is the only engine a caller can
+// decline, so a set was never the question being asked; every other supported
+// engine is provisioned unconditionally.
 using ExpectedChipWorkerInit = void (ChipWorker::*)(
-    const std::string &, const std::string &, const std::string &, const std::string &, int, const CallConfig *,
-    uint32_t, const std::string &, const std::string &
+    const std::string &, const std::string &, const std::string &, const std::string &, int, const CallConfig *, bool,
+    const std::string &, const std::string &
 );
 static_assert(
     std::is_same_v<decltype(&ChipWorker::init), ExpectedChipWorkerInit>,