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

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);

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<air::ChannelInterface>(a.getOperation());
   auto chanB = dyn_cast_if_present<air::ChannelInterface>(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<air::MemcpyInterface>(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<air::MemcpyInterface>(o);
+    if (existingMc && memcpyIsDedicatedChannel(existingMc))
       return true;
   }
   return false;
@@ -1608,20 +1617,19 @@ FailureOr<air::allocation_info_t> 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<air::MemcpyInterface>(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<air::MemcpyInterface>(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;

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
+}

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
+}