[AIR] Add air.shim_feed_no_pace opt-out for fire-and-free shim feeds - #1754
Merged
Conversation
A launch marked air.preserve_shim_dma_order currently propagates that marker to ALL of its air.channel.put/get ops, which makes AIRRtToNpu pace every shim MM2S feed with bounded double-buffered (await-on-drain) backpressure. That pacing is required for feeds coupled by a shared broadcast/multicast consumer that advances its destinations in lockstep, but it is wrong for feeds that are mutually independent: the per-task completion await serializes otherwise-concurrent channels and deadlocks when a single feed's BD exceeds the downstream ring depth. Add a per-op opt-out unit attr, air.shim_feed_no_pace. When present on a channel.put/get inside a preserve_shim_dma_order launch, the op is excluded from the preserve-marker propagation (AIROptimizeShimDMABDs) and therefore lowers to a fire-and-free MM2S feed (start, free BD, no per-task await), while still benefiting from the launch's no-fold guarantee. This lets a few large independent BDs stream concurrently, backpressured only by their downstream ring locks. copyChannelSteeringAttrs carries the marker across op rewrites so it survives to the shim-DMA-BD pass. The change is opt-in: behavior is unchanged for any op that does not carry the new attribute.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR introduces a per-operation opt-out attribute (air.shim_feed_no_pace) to prevent specific air.channel.put/get ops from inheriting the pacing/backpressure behavior implied by an enclosing air.preserve_shim_dma_order launch, enabling “fire-and-free” shim MM2S feeds for mutually independent channels.
Changes:
- Adds
air.shim_feed_no_paceas a centralized AIR attribute constant. - Preserves the new marker across channel-op rewrites via
copyChannelSteeringAttrs. - Updates
AIROptimizeShimDMABDsto skip preserve-marker propagation onto opted-out channel ops.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| mlir/lib/Transform/AIRDependencyScheduleOpt.cpp | Skips preserve pacing propagation for air.channel.put/get ops tagged air.shim_feed_no_pace. |
| mlir/lib/Dialect/AIR/IR/AIRDialect.cpp | Copies air.shim_feed_no_pace across channel-op rewrites so it survives to later transforms. |
| mlir/include/air/Dialect/AIR/AIRDialect.h | Defines the new attrs::ShimFeedNoPace attribute name and documents intended semantics. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Covers the preserve-marker propagation and the per-op opt-out: an untagged feed in a preserve_shim_dma_order launch inherits the marker, while a feed tagged air.shim_feed_no_pace does not (it lowers fire-and-free) and keeps its opt-out attribute.
…ival Two gaps in the opt-out coverage: - AIRRtToNpu lowering: a mixed preserve launch where one feed keeps the preserve marker and a sibling opts out. Asserts the kept feed is paced (issue_token + depth-2 completion-token awaits) while the opted-out feed lowers fire-and-free (no issue_token, no await, dma_free_task after start) -- the actual downstream consequence the propagation test could not show. - copyChannelSteeringAttrs round-trip: air.shim_feed_no_pace must survive the per-channel BD fold that rebuilds the op (two scf.for loops collapse into one wrap/stride put) in a non-preserve launch, so the marker reaches the preserve-marker propagation instead of being silently dropped on rebuild. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
air.shim_feed_no_pacethat opts a singleair.channel.put/getout of the pacing implied by an enclosingair.preserve_shim_dma_orderlaunch.AIRRtToNpupace every shim MM2S feed with bounded double-buffered (await-on-drain) backpressure. That is correct for feeds coupled by a shared broadcast/multicast consumer, but wrong for mutually independent feeds: the per-task completion await serializes otherwise-concurrent channels and can deadlock when a single feed's BD exceeds the downstream ring depth.Details
AIRDialect.h: newattrs::ShimFeedNoPaceconstant.AIRDialect.cpp(copyChannelSteeringAttrs): carry the marker across op rewrites so it survives to the shim-DMA-BD pass (mirrors the existingAwaitAppends/AppendBarrierhandling).AIRDependencyScheduleOpt.cpp(AIROptimizeShimDMABDs): skip the preserve-marker propagation for ops carrying the opt-out.Compatibility
Opt-in: behavior is byte-identical for any op that does not carry the new attribute.
Test plan
ninja check-air-mlirpreserve_shim_dma_orderlaunch with two independent MM2S feeds: without the attr both are paced (start/await interleaved); with the attr on both they issue fire-and-free (all starts before awaits).