Collapse same-flow circuit S2MM memcpys onto one memtile DMA channel

[erwei-xilinx]

Summary

MemTileDMAAllocator::simpleDmaChannelAlloc (AIRToAIESchedulingUtils.cpp) collapses same-logical-flow S2MM memcpy ops (same air.channel declaration + same constant bundle indices) onto one physical DMA channel — but originally only for packet flows. For circuit flows, distinct ops of a single logical flow were round-robined onto separate physical channels.

This wastes the memtile's limited DMA channels: when several ops belong to one flow (e.g. an L2 buffer re-filled at one endpoint across loop iterations / rounds), each lands on its own S2MM channel. A memtile that also carries a wide broadcast can then exhaust its S2MM channels even though those ops share a single flow.

What this does

Extends the same-flow S2MM collapse to circuit flows. Since both directions of S2MM now use the identical same-logical-flow rule, the packet and circuit paths are folded into one branch:

• MM2S (source) side collapses promiscuously onto a packet-flow channel (broadcast fan-out / pkt_id multiplexing rely on it).
• S2MM (receiver) side collapses only flows proven identical (foundSameLogicalFlowInTile), for both packet and circuit flows. Distinct sources keep distinct channels since matching bundle indices are required.

This is behavior-preserving for packet flows: same-logical-flow implies same channel declaration, so a packet op can only ever match a packet allocation. The unification makes the memtile allocator smaller rather than larger.

The repeated "does this allocation host a dedicated channel" scan (previously written out three times) is extracted into allocation_info_t::containsDedicatedChannel(), reused by the unified memtile block and the shim packet-reuse walk.

getRepeatCounts / BD generation already emit the collapsed ops as sequential BD tasks on the shared channel, so no downstream change is needed.

Tradeoff

Collapsing same-flow S2MM trades channel count for serialization: same-logical-flow ops writing disjoint buffer regions that could previously overlap on two channels now share one channel as sequential BDs. This is the right call for the motivating case (the design does not fit otherwise) and mirrors the existing packet-flow behavior, but is applied unconditionally.

Tests

New circuit-flow lit tests mirroring the existing packet pair:

• memtile_circuit_same_flow_collapse.mlir — same flow → one S2MM channel (two BDs).

• memtile_circuit_two_sources_one_dst.mlir — distinct sources → distinct channels (guards against over-collapse).

• ninja check-air-mlir — full Conversion + Transform suites pass (0 failures)

• Verified on a multi-round temporal flash-attention deployment where the K/V-broadcast memtile previously overflowed its S2MM channels; with this change it fits and routes.

[Copilot]

Pull request overview

Extends MemTile DMA channel allocation so circuit (non-packet) S2MM memcpys that belong to the same logical flow (same air.channel declaration and same constant bundle indices) reuse a single physical S2MM DMA channel, reducing memtile DMA channel pressure in multi-round/refill scenarios.

Changes:

• Add a circuit-flow S2MM reuse path in MemTileDMAAllocator::simpleDmaChannelAlloc mirroring the existing packet-flow S2MM “same logical flow” collapse rule.
• Preserve the existing “never collapse dedicated DMA channels” behavior by checking both the incoming memcpy and any existing memcpys already assigned to the channel.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[erwei-xilinx]

Addressed in the follow-up commit: renamed isSamePacketFlowEndpoint → isSameLogicalFlowEndpoint and foundSamePacketFlowInTile → foundSameLogicalFlowInTile (and updated the header comment) to reflect that the predicate is flow-type-agnostic (same channel decl + same constant indices). Pure rename, no behavioral change.