|
| 1 | +<!--- SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. --> |
| 2 | +<!--- SPDX-License-Identifier: Apache-2.0 --> |
| 3 | + |
| 4 | +# CUDA Lang task scheduling |
| 5 | + |
| 6 | +This package separates a pure-Python scheduling model from CUDA Lang device |
| 7 | +lowering. Import `task_scheduling` as `ts`, define resources and named |
| 8 | +`@ts.consumer_work` or `@ts.producer_work` methods, and capture a schedule with |
| 9 | +`@ts.schedule`. The captured `Schedule` is immutable; `TaskManager` consumes it |
| 10 | +for host validation, and `Task.to_device()` lowers it for device execution. |
| 11 | + |
| 12 | +`Task.to_device()` freezes the validated schedule into CUDA Lang-compatible nodes. |
| 13 | +Directly decorated static work methods are used for device lowering. Explicit |
| 14 | +callback mappings remain available as low-level overrides and are keyed by |
| 15 | +readable captured labels (optionally `"resource.label"` or `(resource, label)` |
| 16 | +when a task needs disambiguation). Pipeline steps are generated from their |
| 17 | +`PipelineConfig` and `DevicePipelineBinding`; they do not require |
| 18 | +kernel-specific acquire/wait/commit/release callbacks. At runtime, an immutable |
| 19 | +`ExecutionContext` carries internal routed values and pipeline index/phase |
| 20 | +state while kernel resources remain available through a named `tasks_inputs` |
| 21 | +object. Domain loops, first/last/every guards, opaque guards, warp ranges, CTA |
| 22 | +pinning, and `setmaxnreg` are lowered from this frozen program. |
| 23 | + |
| 24 | +A work decorator placed above `@staticmethod` lets host task execution and |
| 25 | +device lowering share one implementation. Instance work methods remain usable |
| 26 | +for host schedule construction; low-level users must supply an explicit |
| 27 | +callback mapping to lower them because resource instances are not passed into |
| 28 | +CUDA Lang callbacks. |
| 29 | + |
| 30 | +`TaskManager.to_device()` infers static work callbacks across all host tasks, |
| 31 | +derives device pipeline bindings from each resource's validated |
| 32 | +`PipelineConfig`, preserves validated task order, assigns private |
| 33 | +pipeline-state slots and allocator-derived barrier offsets, and returns a |
| 34 | +frozen `DeviceTaskManager`. Explicit pipeline bindings remain available only |
| 35 | +through the low-level `Task.to_device()` interface. Tasks containing only |
| 36 | +generated scheduling or pipeline operations require no empty callback entry. |
| 37 | +Device code may use `run(inputs)` for a managed one-call lifecycle or manage |
| 38 | +resource setup explicitly: |
| 39 | + |
| 40 | +```python |
| 41 | +device_allocators = device_manager.setup_resources_and_tasks() |
| 42 | +device_manager.run(inputs, device_allocators) |
| 43 | +``` |
| 44 | + |
| 45 | +`Task` also supports a schedule-inferred form, `Task(warp_idx, num_warps, |
| 46 | +schedule=...)`. Consumer and producer resource roles are derived from captured |
| 47 | +stages. Dependency graphs are optional explicit validation metadata; omitting |
| 48 | +one leaves the graph empty. Passing `cta_warps` additionally fills uncovered |
| 49 | +warp ranges with internal padding tasks. SMEM, TMEM, and multi-pipeline barrier |
| 50 | +allocators are constructed from resource requirements when no explicit |
| 51 | +allocator is provided; explicit allocators remain authoritative for |
| 52 | +alias-sensitive layouts. |
| 53 | + |
| 54 | +`TaskManager.freeze()` returns a canonical immutable `ProgramIR` snapshot with |
| 55 | +stable resource, task, step, value, control-flow, and dependency records. Host |
| 56 | +tools can inspect this snapshot without depending on mutable resource identity, |
| 57 | +while existing device lowering continues to use the manager's proven routing |
| 58 | +and pipeline representation. |
| 59 | + |
| 60 | +`setup_resources_and_tasks()` materializes the dynamic SMEM and coalesced |
| 61 | +barrier arenas, initializes every pipeline's full/empty runs, and returns |
| 62 | +immutable `DeviceAllocators`. Its `smem_allocator.get(name, dtype, shape)` and |
| 63 | +`barrier_allocator.get_ptr(name)` methods expose host-laid-out named storage to |
| 64 | +kernel infrastructure such as explicit TMEM allocation and clustered |
| 65 | +deallocation. `run()` then creates each task's private routing context and |
| 66 | +executes the complete task set. Kernel authors neither construct per-task |
| 67 | +contexts nor place barrier arrays in a positional values tuple. The SMEM arena |
| 68 | +base remains available as `stage_info.context.smem_base`. |
| 69 | + |
| 70 | +With `verbose=True`, `TaskManager.print_verbose_report()` prints task register |
| 71 | +budgets, memory usage, captured schedules with routed values, |
| 72 | +exhaustive host safety summaries, flattened checker inputs, and a representative |
| 73 | +complete/deadlock/race timeline. Allocator reports include allocation names, |
| 74 | +sizes, alignments, offsets, alias groups, barrier bytes, and physical totals. |
| 75 | +When a schedule contains runtime domain bounds, |
| 76 | +`exhaustive_representative_domain=True` makes the host checker use a bounded |
| 77 | +domain derived from the loop starts, while device lowering resolves the actual |
| 78 | +trip count from kernel inputs. |
| 79 | + |
| 80 | +The manager also computes a warp-group-rounded initial register budget and |
| 81 | +threads it into device task lowering, so `setmaxnreg` direction and verbose |
| 82 | +register contributions use the same baseline. |
| 83 | + |
| 84 | +Automatic value routing hides storage positions from callback authors. Each |
| 85 | +callback receives its routed inputs after the `StageInfo` argument and returns |
| 86 | +the values declared by its captured work method. The frozen `DeviceStep` |
| 87 | +performs the corresponding value-stack operations, while |
| 88 | +`DeviceTask.make_context(tasks_inputs)` creates the immutable execution |
| 89 | +context. The tuple representation remains an internal compiler detail; no |
| 90 | +placeholder or named routing slot is exposed to resource authors. |
| 91 | + |
| 92 | +`outputs=N` creates `N` independent lexical `ScheduleValue` instances. A work |
| 93 | +call returns one value directly or a tuple that can be unpacked normally, and |
| 94 | +its device callback returns the same number of runtime values. A value produced |
| 95 | +inside a conditional or domain loop cannot escape that scope accidentally. A |
| 96 | +loop that needs state across iterations declares it explicitly with |
| 97 | +`domain_loop(..., carried={"state": initial})`, reads `loop.state`, assigns the |
| 98 | +next value back to `loop.state`, and uses `loop.state` after the block as the |
| 99 | +post-loop result. The device adapter preserves zero-trip pass-through and |
| 100 | +materializes the backedge without exposing routing positions to callbacks. |
| 101 | + |
| 102 | +`StageInfo` contains the current `stage_idx`, `phase`, selected full `barrier`, |
| 103 | +zero-based iteration count, loop offset/bounds, work label, and owning |
| 104 | +`ExecutionContext`; loop offset/bounds are `None` for peeled work |
| 105 | +outside a domain loop. This lets pipeline payload work use |
| 106 | +`stage_info.barrier` without knowing how barrier arrays are stored by the |
| 107 | +device manager. |
| 108 | + |
| 109 | +Work-call arguments are classified during schedule capture: `ScheduleValue` |
| 110 | +arguments are routed, while ordinary Python values are captured as compile-time |
| 111 | +operands. The frozen device step preserves the method's original parameter order |
| 112 | +when it combines both kinds. Work methods may therefore declare any number of |
| 113 | +static arguments and route any statically known number of inputs or results |
| 114 | +without annotating the argument kinds. Routed values may have different CUDA |
| 115 | +Lang types. Work methods declare their result count with `outputs=N`, and each |
| 116 | +call creates independent lexical SSA values. |
| 117 | + |
| 118 | +The public host API also includes identity-based `MemoryResource` metadata, |
| 119 | +pipeline configurations and groups, aligned SMEM/TMEM/barrier allocators, |
| 120 | +bounded exhaustive interleaving checks, and static pipeline auditing. These |
| 121 | +objects do not import CUTLASS or CuTe and can be constructed and validated on a |
| 122 | +machine without a CUDA device. |
| 123 | + |
| 124 | +`TmemAllocator` validates and reports the host-side TMEM layout, while kernels |
| 125 | +own the device lifecycle explicitly. A kernel allocates TMEM, synchronizes its |
| 126 | +participating warps, passes the pointer through its named task inputs, and |
| 127 | +performs CTA or clustered peer-CTA deallocation after task execution. |
| 128 | + |
| 129 | +## Device support boundary |
| 130 | + |
| 131 | +`AsyncAsync`, `TmaAsync`, `TmaUmma`, and `UmmaAsync` metadata have a generic |
| 132 | +CUDA Lang lowering path. It owns immutable producer/consumer index and phase |
| 133 | +state, initializes offset-addressed full/empty runs in one manager-owned |
| 134 | +mbarrier arena, lowers try/acquire/wait, and emits the pipeline-specific |
| 135 | +mbarrier or `tcgen05.commit` signals. The remaining UMMA, CLC, DPC, and DLC |
| 136 | +pipeline kinds remain available to host analysis but raise `NotImplementedError` |
| 137 | +when device pipeline creation is requested. |
| 138 | + |
| 139 | +Persistent `WorkTileLoop` execution requires a scheduler-specific work-queue |
| 140 | +adapter and is currently host-analysis-only. Domain-loop starts and ends may be |
| 141 | +runtime values declared with `dynamic_domain_bound(resolver)`. A schedule may |
| 142 | +also declare `stage_info` as its first parameter and use, for example, |
| 143 | +`stage_info.context.tasks_inputs.num_rows` directly as a dynamic loop bound. |
| 144 | +The resolver receives the manager-owned `tasks_inputs` object inside the |
| 145 | +kernel; host analysis continues to use an explicit representative-domain |
| 146 | +fallback. Device domain loops currently require a positive static step. Literal |
| 147 | +opaque guards need an explicit device-condition adapter, while routed |
| 148 | +`ScheduleValue` opaque guards are supported. |
| 149 | + |
| 150 | +The exhaustive checker is bounded by `max_states`. Its result reports |
| 151 | +`hit_state_limit=True` when the bound prevents a proof; callers must not treat |
| 152 | +that result as a successful exhaustive proof. |
| 153 | + |
| 154 | +## Current limitations |
| 155 | + |
| 156 | +The current checker proves acquire/wait availability and physical-range alias |
| 157 | +safety across expanded domain/work-tile loops and opaque assignments. It does |
| 158 | +not yet model PDL launch ordering or scheduler-specific CLC queue state. PDL |
| 159 | +and CLC metadata can be retained on host objects, but they are not currently |
| 160 | +part of the interleaving state, so those configurations require an external |
| 161 | +check. |
| 162 | + |
| 163 | +Pipeline groups, interleave strides, multicast signaling, and deferred barrier |
| 164 | +storage are validated as host metadata. Generic device materialization for |
| 165 | +those combinations is not implemented. The copy and GEMM tutorials exercise |
| 166 | +the shared TMA-to-async, TMA-to-UMMA, and UMMA-to-async lowering while their |
| 167 | +callbacks contain only pipeline payload work. Static pipeline auditing |
| 168 | +recognizes the corresponding CUDA Lang calls but does not infer effects hidden |
| 169 | +behind arbitrary user wrappers. |
0 commit comments