[AIR] Add air.shim_feed_no_pace opt-out for fire-and-free shim feeds

mlir/include/air/Dialect/AIR/AIRDialect.h

@@ -46,6 +46,16 @@ constexpr StringLiteral PreserveShimDmaOrder = "air.preserve_shim_dma_order";
 // before starting the next group), reproducing the drain schedule of the
 // un-coalesced feed.
 constexpr StringLiteral CoalescedShimFeed = "air.coalesced_shim_feed";
+// Per-op opt-out (unit attr) on an air.channel.put/get inside a
+// preserve_shim_dma_order launch: this feed is NOT lockstep-coupled to its
+// siblings (no shared broadcast/multicast consumer), so it must NOT inherit the
+// launch's preserve marker and its bounded double-buffered pacing. Keeping it
+// out of the pacing lets it lower to a fire-and-free MM2S feed (start, free BD,
+// no per-task completion await), so a few large independent BDs (e.g. separate
+// K/V readback streams) run concurrently and are backpressured only by their
+// downstream ring locks -- the pacing's await-on-drain would otherwise serialize
+// them and deadlock when a single BD exceeds the ring depth.
+constexpr StringLiteral ShimFeedNoPace = "air.shim_feed_no_pace";
 constexpr StringLiteral TileDmaChannel = "air.tile_dma_channel";
 constexpr StringLiteral MemtileDmaChannelMin = "air.memtile_dma_channel_min";
 constexpr StringLiteral DedicatedDmaChannel = "air.dedicated_dma_channel";

mlir/lib/Dialect/AIR/IR/AIRDialect.cpp

@@ -138,6 +138,10 @@
     dst->setAttr(attrs::AwaitAppends, aa);
   if (auto ab = src->getAttr(attrs::AppendBarrier))
     dst->setAttr(attrs::AppendBarrier, ab);
+  // Per-op preserve-pacing opt-out (fire-and-free shim feed), consumed by
+  // AIROptimizeShimDMABDs's preserve-marker propagation.
+  if (auto np = src->getAttr(attrs::ShimFeedNoPace))
+    dst->setAttr(attrs::ShimFeedNoPace, np);
   // Producer-side re-feed count (single-buffer count-free re-broadcast), read
   // by AIRToAIE's lock allocators.
   if (auto rc = src->getAttrOfType<IntegerAttr>(attrs::RefeedCount))
@@ -3203,7 +3207,7 @@
   // (they are discardable attrs, otherwise lost on rebuild and the AIRToAIE /
   // AIRRtToNpu consumers never fire), then replace -- so `op` is not read after
   // it is erased.
   auto newOp = rewriter.create<OpT>(
       op.getLoc(), op->getResultTypes(), op.getAsyncDependencies(),
       op.getChanName(), op.getIndices(), input_memref, offsets, sizes, strides,
       /*pad_before=*/padBefore, /*pad_after=*/padAfter);

mlir/lib/Transform/AIRDependencyScheduleOpt.cpp

@@ -7292,9 +7292,18 @@ class AIROptimizeShimDMABDs
         // double-buffered awaits (backpressure) for these lockstep-coupled
         // shim feeds.
         launch.walk([&](Operation *op) {
-          if (isa<air::ChannelPutOp, air::ChannelGetOp>(op))
-            op->setAttr(air::attrs::PreserveShimDmaOrder,
-                        mlir::UnitAttr::get(op->getContext()));
+          if (!isa<air::ChannelPutOp, air::ChannelGetOp>(op))
+            return;
+          // Per-op opt-out: a feed explicitly marked air.shim_feed_no_pace is
+          // NOT lockstep-coupled to its siblings, so it must stay out of the
+          // preserve marker's bounded double-buffered pacing (it lowers to a
+          // fire-and-free MM2S feed instead). It still benefits from the
+          // launch's no-fold guarantee (the early return below skips
+          // per-channel BD folding for the whole launch region).
+          if (op->hasAttr(air::attrs::ShimFeedNoPace))
+            return;
+          op->setAttr(air::attrs::PreserveShimDmaOrder,
+                      mlir::UnitAttr::get(op->getContext()));
         });
         return;
       }