diff --git a/mlir/include/air/Conversion/AIRToAIESchedulingUtils.h b/mlir/include/air/Conversion/AIRToAIESchedulingUtils.h index 9b425abc8..328af37b8 100644 --- a/mlir/include/air/Conversion/AIRToAIESchedulingUtils.h +++ b/mlir/include/air/Conversion/AIRToAIESchedulingUtils.h @@ -125,13 +125,17 @@ struct allocation_info_t { bool foundAlloc(AIE::TileLike tile, air::ChannelOp channel_op); bool foundAlloc(AIE::TileLike tile, AIE::DMAChannel channel); bool foundPacketFlowAllocInTile(AIE::TileLike tile); - // True if a packet-flow memcpy on `tile` belongs to the SAME logical flow as - // `memcpyOp` (same channel decl + same constant bundle indices). Used for - // S2MM-side collapse discrimination: distinct sources must not share a - // physical channel. A non-constant index cannot be proven equal and is - // treated as distinct (safe: forces a separate channel). - bool foundSamePacketFlowInTile(AIE::TileLike tile, - air::MemcpyInterface memcpyOp); + // True if a memcpy on `tile` belongs to the SAME logical flow as `memcpyOp` + // (same channel decl + same constant bundle indices), for either packet or + // circuit flows. Used for S2MM-side collapse discrimination: distinct sources + // must not share a physical channel. A non-constant index cannot be proven + // equal and is treated as distinct (safe: forces a separate channel). + bool foundSameLogicalFlowInTile(AIE::TileLike tile, + air::MemcpyInterface memcpyOp); + // True if any memcpy in this allocation is marked air.dedicated_dma_channel. + // A dedicated flow must own its physical channel, so an allocation that hosts + // one is never a collapse target. + bool containsDedicatedChannel(); bool foundAlloc(air::ChannelOp channel_op); bool foundAlloc(AIE::DMAChannel channel); diff --git a/mlir/lib/Conversion/AIRToAIESchedulingUtils.cpp b/mlir/lib/Conversion/AIRToAIESchedulingUtils.cpp index a0c30e79d..05333df4d 100644 --- a/mlir/lib/Conversion/AIRToAIESchedulingUtils.cpp +++ b/mlir/lib/Conversion/AIRToAIESchedulingUtils.cpp @@ -925,8 +925,8 @@ bool xilinx::air::allocation_info_t::foundPacketFlowAllocInTile( // Same-logical-flow test: same channel declaration and same constant bundle // indices. A non-constant index cannot be proven equal, so it is distinct. -static bool isSamePacketFlowEndpoint(air::MemcpyInterface a, - air::MemcpyInterface b) { +static bool isSameLogicalFlowEndpoint(air::MemcpyInterface a, + air::MemcpyInterface b) { auto chanA = dyn_cast_if_present(a.getOperation()); auto chanB = dyn_cast_if_present(b.getOperation()); if (!chanA || !chanB) @@ -947,13 +947,22 @@ static bool isSamePacketFlowEndpoint(air::MemcpyInterface a, return true; } -bool xilinx::air::allocation_info_t::foundSamePacketFlowInTile( +bool xilinx::air::allocation_info_t::foundSameLogicalFlowInTile( AIE::TileLike tile, air::MemcpyInterface memcpyOp) { if (!foundAlloc(tile)) return false; for (auto o : memcpyOps) { auto existingMc = dyn_cast_if_present(o); - if (existingMc && isSamePacketFlowEndpoint(existingMc, memcpyOp)) + if (existingMc && isSameLogicalFlowEndpoint(existingMc, memcpyOp)) + return true; + } + return false; +} + +bool xilinx::air::allocation_info_t::containsDedicatedChannel() { + for (auto o : memcpyOps) { + auto existingMc = dyn_cast_if_present(o); + if (existingMc && memcpyIsDedicatedChannel(existingMc)) return true; } return false; @@ -1608,20 +1617,19 @@ FailureOr air::ShimDMAAllocator::allocNewDmaChannel( continue; if (t.dma_channel.direction != dir) continue; - bool tDedicated = false; bool tPacket = false; for (auto o : t.memcpyOps) { auto mc = dyn_cast_if_present(o); if (!mc) continue; auto ct = air::getChannelType(mc); - if (succeeded(ct) && ct.value() == "npu_dma_packet") + if (succeeded(ct) && ct.value() == "npu_dma_packet") { tPacket = true; - if (memcpyIsDedicatedChannel(mc)) - tDedicated = true; + break; + } } // Never collapse onto a channel that hosts a dedicated flow. - if (tPacket && !tDedicated) { + if (tPacket && !t.containsDedicatedChannel()) { packetLT = lt; packetCh = (int)t.dma_channel.channel; return true; @@ -1934,25 +1942,24 @@ air::MemTileDMAAllocator::simpleDmaChannelAlloc(air::MemcpyInterface &memcpyOp, t.memcpyOps.push_back(memcpyOp.getOperation()); return t; } - // Reuse an existing packet-flow channel on this tile via time - // multiplexing. Never collapse a channel marked air.dedicated_dma_channel - // (either direction). MM2S (source) side collapses promiscuously - // (broadcast fan-out / pkt_id multiplexing rely on it); S2MM (receiver) - // side collapses only flows proven identical (same channel decl + indices), - // since distinct sources fanning into one L2 buffer each need their own - // physical S2MM channel. - if (isPacketFlowOp && t.foundPacketFlowAllocInTile(tile)) { - bool tDedicated = false; - for (auto o : t.memcpyOps) { - auto existingMc = dyn_cast_if_present(o); - if (existingMc && memcpyIsDedicatedChannel(existingMc)) { - tDedicated = true; - break; - } - } - bool canCollapse = !memcpyIsDedicatedChannel(memcpyOp) && !tDedicated; - if (canCollapse && !isMM2S.value()) - canCollapse = t.foundSamePacketFlowInTile(tile, memcpyOp); + // Reuse an existing DMA channel on this tile instead of allocating a new + // one. Never collapse a channel marked air.dedicated_dma_channel, nor + // collapse onto an allocation that already hosts one (either direction). + // - 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 (same + // channel decl + constant bundle indices), for BOTH packet and circuit + // flows: distinct sources fanning into one L2 buffer each need their + // own physical S2MM channel, but repeated invocations of one flow + // (e.g. an L2 buffer re-filled at one endpoint across loop iterations) + // time-multiplex a single channel as sequential BD tasks. Collapsing + // circuit flows too keeps a memtile that also carries a wide broadcast + // from exhausting its S2MM channels on same-flow refills. + if (!memcpyIsDedicatedChannel(memcpyOp) && !t.containsDedicatedChannel()) { + bool canCollapse = + isMM2S.value() + ? (isPacketFlowOp && t.foundPacketFlowAllocInTile(tile)) + : t.foundSameLogicalFlowInTile(tile, memcpyOp); if (canCollapse) { t.memcpyOps.push_back(memcpyOp.getOperation()); return t; diff --git a/mlir/test/Conversion/AIRToAIE/memtile_circuit_same_flow_collapse.mlir b/mlir/test/Conversion/AIRToAIE/memtile_circuit_same_flow_collapse.mlir new file mode 100644 index 000000000..93eb5b5cb --- /dev/null +++ b/mlir/test/Conversion/AIRToAIE/memtile_circuit_same_flow_collapse.mlir @@ -0,0 +1,74 @@ +//===- memtile_circuit_same_flow_collapse.mlir -----------------*- MLIR -*-===// +// +// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +// +//===----------------------------------------------------------------------===// + +// RUN: air-opt %s -air-to-aie="row-offset=3 col-offset=2 device=xcve2802" | FileCheck %s + +// Circuit-flow companion to memtile_packet_same_flow_collapse.mlir. Verifies the +// collapse-PRESERVING side of the S2MM discrimination in +// MemTileDMAAllocator::simpleDmaChannelAlloc for CIRCUIT (non-packet) flows: +// multiple invocations of the SAME logical flow (same channel decl + same +// indices, e.g. scf.for unroll or ping-pong) time-multiplex ONE physical S2MM +// channel as sequential BD tasks. Before the same-logical-flow collapse was +// extended to circuit flows, these round-robined onto separate physical +// channels, wasting the memtile's limited DMA channels. The two-distinct-source +// counterpart (memtile_circuit_two_sources_one_dst.mlir) confirms distinct +// sources still get distinct channels. +// +// Here one circuit channel (@w0) is used by two get invocations writing disjoint +// rows of one L2 buffer. Both share S2MM 0 (two BDs); no second S2MM channel is +// allocated. + +// CHECK: aie.memtile_dma +// CHECK: aie.dma_start(S2MM, 0 +// CHECK: aie.dma_bd(%{{.*}} : memref<2x8xbf16, 1> offset = 0 len = 8) +// CHECK: aie.dma_bd(%{{.*}} : memref<2x8xbf16, 1> offset = 8 len = 8) +// CHECK-NOT: aie.dma_start(S2MM, 1 + +air.channel @w0 [1, 1] +air.channel @r0 [1, 1] +func.func @memtile_circ_same_flow() { + %c1 = arith.constant 1 : index + air.launch (%a, %b) in (%c=%c1, %d=%c1) { + air.segment @seg { + %c1_0 = arith.constant 1 : index + %c0 = arith.constant 0 : index + %c8 = arith.constant 8 : index + %t, %l2 = air.execute -> (memref<2x8xbf16, 1>) { + %alloc = memref.alloc() : memref<2x8xbf16, 1> + air.execute_terminator %alloc : memref<2x8xbf16, 1> + } + // Two invocations of the SAME channel write disjoint rows of one buffer. + air.channel.get @w0[] (%l2[%c0, %c0] [%c1_0, %c8] [%c8, %c1_0]) : (memref<2x8xbf16, 1>) + air.channel.get @w0[] (%l2[%c1_0, %c0] [%c1_0, %c8] [%c8, %c1_0]) : (memref<2x8xbf16, 1>) + air.channel.put @r0[] (%l2[] [] []) : (memref<2x8xbf16, 1>) + %d_ = air.execute {memref.dealloc %l2 : memref<2x8xbf16, 1>} + air.herd @h0 tile (%tx0, %ty0) in (%sx0=%c1_0, %sy0=%c1_0) attributes {x_loc = 2 : i64, y_loc = 3 : i64} { + %tok0, %l1a = air.execute -> (memref<8xbf16, 2>) { + %aa = memref.alloc() : memref<8xbf16, 2> + air.execute_terminator %aa : memref<8xbf16, 2> + } + air.channel.put @w0[] (%l1a[] [] []) : (memref<8xbf16, 2>) + %tok1, %l1b = air.execute -> (memref<8xbf16, 2>) { + %bb = memref.alloc() : memref<8xbf16, 2> + air.execute_terminator %bb : memref<8xbf16, 2> + } + air.channel.put @w0[] (%l1b[] [] []) : (memref<8xbf16, 2>) + %d0 = air.execute {memref.dealloc %l1a : memref<8xbf16, 2>} + %d0b = air.execute {memref.dealloc %l1b : memref<8xbf16, 2>} + } + air.herd @hr tile (%txr, %tyr) in (%sxr=%c1_0, %syr=%c1_0) attributes {x_loc = 4 : i64, y_loc = 3 : i64} { + %tok, %l1 = air.execute -> (memref<16xbf16, 2>) { + %aa = memref.alloc() : memref<16xbf16, 2> + air.execute_terminator %aa : memref<16xbf16, 2> + } + air.channel.get @r0[] (%l1[] [] []) : (memref<16xbf16, 2>) + %dr = air.execute {memref.dealloc %l1 : memref<16xbf16, 2>} + } + } + } + return +} diff --git a/mlir/test/Conversion/AIRToAIE/memtile_circuit_two_sources_one_dst.mlir b/mlir/test/Conversion/AIRToAIE/memtile_circuit_two_sources_one_dst.mlir new file mode 100644 index 000000000..8fe6330df --- /dev/null +++ b/mlir/test/Conversion/AIRToAIE/memtile_circuit_two_sources_one_dst.mlir @@ -0,0 +1,74 @@ +//===- memtile_circuit_two_sources_one_dst.mlir ----------------*- MLIR -*-===// +// +// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +// +//===----------------------------------------------------------------------===// + +// RUN: air-opt %s -air-to-aie="row-offset=3 col-offset=2 device=xcve2802" | FileCheck %s + +// Circuit-flow companion to memtile_packet_two_sources_one_dst.mlir. Guards the +// NO-OVER-COLLAPSE side of the S2MM discrimination in +// MemTileDMAAllocator::simpleDmaChannelAlloc: extending the same-logical-flow +// collapse to circuit flows must NOT merge distinct sources. Each distinct +// (channel symbol, indices) endpoint is a different logical flow and must get +// its own physical S2MM channel; only multiple invocations of the SAME flow may +// share one (see memtile_circuit_same_flow_collapse.mlir). +// +// Here two distinct circuit channels (@w0, @w1) each write a disjoint sub-region +// of one L2 buffer at the receiver memtile. isSameLogicalFlowEndpoint requires +// matching channel decls, so the two land on TWO separate S2MM channels. + +// CHECK: aie.memtile_dma +// CHECK: aie.dma_start(S2MM, 0 +// CHECK: aie.dma_bd(%{{.*}} : memref<2x8xbf16, 1> offset = 0 len = 8) +// CHECK: aie.dma_start(S2MM, 1 +// CHECK: aie.dma_bd(%{{.*}} : memref<2x8xbf16, 1> offset = 8 len = 8) + +air.channel @w0 [1, 1] +air.channel @w1 [1, 1] +air.channel @r0 [1, 1] +func.func @memtile_circ_two_src() { + %c1 = arith.constant 1 : index + air.launch (%a, %b) in (%c=%c1, %d=%c1) { + air.segment @seg { + %c1_0 = arith.constant 1 : index + %c0 = arith.constant 0 : index + %c8 = arith.constant 8 : index + %t, %l2 = air.execute -> (memref<2x8xbf16, 1>) { + %alloc = memref.alloc() : memref<2x8xbf16, 1> + air.execute_terminator %alloc : memref<2x8xbf16, 1> + } + // Two distinct circuit channels write disjoint rows of the same L2 buffer. + air.channel.get @w0[] (%l2[%c0, %c0] [%c1_0, %c8] [%c8, %c1_0]) : (memref<2x8xbf16, 1>) + air.channel.get @w1[] (%l2[%c1_0, %c0] [%c1_0, %c8] [%c8, %c1_0]) : (memref<2x8xbf16, 1>) + air.channel.put @r0[] (%l2[] [] []) : (memref<2x8xbf16, 1>) + %d_ = air.execute {memref.dealloc %l2 : memref<2x8xbf16, 1>} + air.herd @h0 tile (%tx0, %ty0) in (%sx0=%c1_0, %sy0=%c1_0) attributes {x_loc = 2 : i64, y_loc = 3 : i64} { + %tok, %l1 = air.execute -> (memref<8xbf16, 2>) { + %aa = memref.alloc() : memref<8xbf16, 2> + air.execute_terminator %aa : memref<8xbf16, 2> + } + air.channel.put @w0[] (%l1[] [] []) : (memref<8xbf16, 2>) + %d0 = air.execute {memref.dealloc %l1 : memref<8xbf16, 2>} + } + air.herd @h1 tile (%tx1, %ty1) in (%sx1=%c1_0, %sy1=%c1_0) attributes {x_loc = 3 : i64, y_loc = 3 : i64} { + %tok, %l1 = air.execute -> (memref<8xbf16, 2>) { + %aa = memref.alloc() : memref<8xbf16, 2> + air.execute_terminator %aa : memref<8xbf16, 2> + } + air.channel.put @w1[] (%l1[] [] []) : (memref<8xbf16, 2>) + %d1 = air.execute {memref.dealloc %l1 : memref<8xbf16, 2>} + } + air.herd @hr tile (%txr, %tyr) in (%sxr=%c1_0, %syr=%c1_0) attributes {x_loc = 4 : i64, y_loc = 3 : i64} { + %tok, %l1 = air.execute -> (memref<16xbf16, 2>) { + %aa = memref.alloc() : memref<16xbf16, 2> + air.execute_terminator %aa : memref<16xbf16, 2> + } + air.channel.get @r0[] (%l1[] [] []) : (memref<16xbf16, 2>) + %dr = air.execute {memref.dealloc %l1 : memref<16xbf16, 2>} + } + } + } + return +}