Add wave.schedule for more fine grained control #333
Conversation
harsh-nod
force-pushed
the
manual_schedule
branch
3 times, most recently
from
b0308ec to
ee8995f
Compare
There was a problem hiding this comment.
Pull Request Overview
Adds a new wave.schedule construct to enable manual, fine-grained scheduling (e.g., explicit pipelining stages) of Wave kernels. Integrates schedule ops, tag propagation, and a manual scheduling mode into compilation, with new lit and e2e tests.
- Introduces WaveSchedule and custom schedule ops (get_node_by_tag, partition_by_address_space, pipeline, getitem) to trace and apply manual schedules.
- Adds SchedulingType.MANUAL and plumbs an optional schedule through wave_compile and kernel tracing; propagates tags and scheduling parameters to support stage assignment.
- Adds tests demonstrating manual scheduling and MLIR output validation; propagates op tags through promotion/optimization passes.
Reviewed Changes
Copilot reviewed 19 out of 19 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| wave_lang/kernel/wave/wave_schedule.py | Adds WaveSchedule to trace and execute manual schedules using ScheduleContext and ScheduleRegionGraph. |
| wave_lang/kernel/wave/wave.py | Adds manual scheduling path and run_manual_schedule; wires SchedulingType.MANUAL. |
| wave_lang/kernel/wave/templates/test_kernels.py | Adds example GEMM with explicit manual schedule using schedule ops; returns kernel, schedule, and compile options. |
| wave_lang/kernel/wave/scheduling/schedule_enums.py | Adds MANUAL enum for scheduling. |
| wave_lang/kernel/wave/scheduling/schedule.py | Extracts apply_pipelined_schedule and adds propagate_scheduling_parameters_to_iter_args. |
| wave_lang/kernel/wave/scheduling/loop_reconstruction.py | Removes node_map plumbing and cleans up logging. |
| wave_lang/kernel/wave/promotion.py | Preserves tags on promoted Read/Write ops. |
| wave_lang/kernel/wave/minimize_global_loads.py | Preserves tags when adding optimized Read/Write nodes. |
| wave_lang/kernel/wave/global_to_shared_gathers.py | Preserves tags for gather path. |
| wave_lang/kernel/wave/compile.py | Adds optional schedule parameter to wave_compile and threads it through. |
| wave_lang/kernel/wave/init.py | Exposes wave_schedule_ops at wave.* for schedule DSL usage. |
| wave_lang/kernel/ops/wave_schedule_ops.py | Implements schedule ops and PipelinedLoop context manager managing manual pipeline application. |
| wave_lang/kernel/ops/wave_ops.py | Adds tag support and unifies NestedRegionOp.handle for tag propagation and Iterate-special handling. |
| wave_lang/kernel/ops/base.py | Adds define_schedule_op helper for schedule op dispatching. |
| wave_lang/kernel/_support/tracing.py | Adds ScheduleTracer and ScheduleContext with custom-op dispatcher and proxy result resolution. |
| wave_lang/kernel/_support/regions.py | Adds ScheduleRegionGraph with ScheduleTracer creation. |
| tests/kernel/wave/wave_gemm_test.py | Adds e2e test that compiles a kernel with a manual schedule. |
| lit_tests/kernel/wave/wave_schedule.py | Lit test validates MLIR emitted for scheduled GEMM. |
| lit_tests/kernel/wave/trace_wave_schedule.py | Lit tests for tracing schedule graphs and ops. |
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
harsh-nod
force-pushed
the
manual_schedule
branch
2 times, most recently
from
612f472 to
f4d703e
Compare
|
@raikonenfnu - I added a page on read the docs explaining this with a gemm example. |
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 21 out of 21 changed files in this pull request and generated 4 comments.
Comments suppressed due to low confidence (1)
wave_lang/kernel/wave/scheduling/schedule.py:209
- Corrected spelling of 'atleast' to 'at least'.
# After scheduling has completed, we have enough information to decide
# whether to pipeline the loop. For pipelining to be possible, we need
# to have atleast N iterations of the loop where N > num_stages - 1 (because
# we will be peeling off num_stages iterations from the loop).
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
harsh-nod
force-pushed
the
manual_schedule
branch
from
f4d703e to
4d1c8a9
Compare
harsh-nod
force-pushed
the
manual_schedule
branch
from
4d1c8a9 to
183eb77
Compare
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 21 out of 21 changed files in this pull request and generated 5 comments.
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
raikonenfnu
left a comment
raikonenfnu
left a comment
There was a problem hiding this comment.
Hi Harsh, can we get some intuition to ScheduleRegionGraph and ScheduleContext in the PR description? :)
harsh-nod
force-pushed
the
manual_schedule
branch
from
183eb77 to
558ec75
Compare
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 21 out of 21 changed files in this pull request and generated 6 comments.
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
harsh-nod
force-pushed
the
manual_schedule
branch
from
8d620b3 to
f6a4eb9
Compare
harsh-nod
force-pushed
the
manual_schedule
branch
from
f6a4eb9 to
01fa2c5
Compare
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 21 out of 21 changed files in this pull request and generated 8 comments.
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
|
|
||
| @property | ||
| def tag(self) -> Optional[str]: | ||
| return getattr(self.fx_node, "tag", None) |
There was a problem hiding this comment.
For proxy-based ops created in handle, the tag is set on node.fx_node.node (the underlying fx.Node) but the tag property checks self.fx_node (the fx.Proxy), so tag will always appear None. Either set tag on both proxy and underlying node (e.g., node.fx_node.tag = tag) or have the property also fall back to self.fx_node.node.
| return getattr(self.fx_node, "tag", None) | |
| # Check tag on fx_node (could be fx.Proxy or fx.Node) | |
| tag = getattr(self.fx_node, "tag", None) | |
| if tag is not None: | |
| return tag | |
| # If fx_node is a Proxy, check its underlying node | |
| node = getattr(self.fx_node, "node", None) | |
| if node is not None: | |
| return getattr(node, "tag", None) | |
| return None |
| node.fx_node.node.tkw_op_name = cls.tkw_op_name | ||
| node.fx_node.node.location = capture_location(graph.location_capture_config) | ||
| if tag is not None: | ||
| node.fx_node.node.tag = tag |
There was a problem hiding this comment.
For proxy-based ops created in handle, the tag is set on node.fx_node.node (the underlying fx.Node) but the tag property checks self.fx_node (the fx.Proxy), so tag will always appear None. Either set tag on both proxy and underlying node (e.g., node.fx_node.tag = tag) or have the property also fall back to self.fx_node.node.
| node.fx_node.node.tag = tag | |
| node.fx_node.node.tag = tag | |
| node.fx_node.tag = tag |
| new_node = graph.node_copy(self.fx_node, arg_transform=arg_transform) | ||
| new_node.tkw_op = self | ||
| new_node.tkw_op_name = self.tkw_op_name | ||
| new_node.tag = self.tag |
There was a problem hiding this comment.
Because self.tag returns None for proxy-based ops (see earlier comment), tag propagation during copy will drop an existing tag set only on the underlying node. After fixing tag storage, ensure copy propagates correctly (may need to read from underlying node if proxy).
| new_node.tag = self.tag | |
| # Propagate tag from self or underlying node if present | |
| tag = self.tag if self.tag is not None else getattr(self.fx_node, "tag", None) | |
| new_node.tag = tag |
harsh-nod
force-pushed
the
manual_schedule
branch
from
01fa2c5 to
f9485c3
Compare
| custom.scheduling_parameters = { | ||
| "absolute_cycle": stage * self.initiation_interval + i, | ||
| "stage": stage, | ||
| "initiation_interval": self.initiation_interval, | ||
| "prefetch_stage": None, | ||
| } | ||
| custom.scheduling_parameters["cycle"] = ( |
There was a problem hiding this comment.
Per offline discussion, we are now setting the stage and then infer absolute cycle. This is slightly different API than before e.g in prefetch or fourstage pipeline which is we determine the cycle and stage gets inferred.
This PR adds a new language construct, the wave.schedule which gives users more control over the generated code. Specifically, this PR allows authors to control the pipelining and specify which operators or groups of operators belong to which stage of the pipeline. A lit test and e2e test are added. Signed-off-by: Harsh Menon <harsh@nod-labs.com>
Signed-off-by: Harsh Menon <harsh@nod-labs.com>
Now exposes clusters for operators and num stages and initiation interval are automatically inferred. Signed-off-by: Harsh Menon <harsh@nod-labs.com>
harsh-nod
force-pushed
the
manual_schedule
branch
from
f9485c3 to
b70db2e
Compare
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 21 out of 21 changed files in this pull request and generated 4 comments.
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
3 participants
This PR adds a new language construct, the wave.schedule which gives users more control over the generated code. Specifically, this PR allows authors to control the pipelining and specify which operators or groups of operators belong to which stage of the pipeline. A lit test and e2e test are added. The new schedule is traced to a ScheduleRegionGraph (distinct from RegionGraph which is only for the kernel) and a corresponding ScheduleContext for tracing the schedule graph.