Strip segment/herd bodies during shim DMA BD loop unrolling

Add loopUnrollFullLightweight() which strips air.segment and air.herd
bodies from cloned iterations during loop unrolling in the shim DMA BD
pass. Iteration 0 keeps full bodies (needed by air-to-std channel
matching), while iterations 1..N get empty shells with just the
terminator. This is safe because:
- BD folding only operates on L3 channel ops, never segment/herd bodies
- ChannelCounterpartCache walks the entire module and needs only the
  first get/put per channel name (iteration 0 provides this)
- removeDeadDeviceComputeOps in airrt-to-npu strips these bodies anyway

The NPU instruction binary is byte-for-byte identical to the original.

Profiling on flash attention (12 heads, 1024 LQ/LK, NPU1):
  Total aircc: 6,400 ms -> 1,496 ms (4.3x faster)

Verified on local NPU1 hardware:
  test/xrt/12_matmul_transform_1x4_bf16 - PASS
  test/xrt/17_gemm_8x16_transform_vec_4x4 - PASS
  test/xrt/25_batch_matmul_bf16 - PASS
  air.insts.bin: identical for all 3 tests

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: erwei-xilinx

mlir/include/air/Util/Dependency.h

@@ -79,6 +79,10 @@ scf::ForOp hoistTargetOpsToNewSCFFor(PatternRewriter &rewriter,
 LogicalResult loopUnrollFullWithAsyncTokenPreserved(
     scf::ForOp forOp,
     function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn = nullptr);
+// Lightweight variant that clones empty shells for SegmentOp/HerdOp on
+// iterations 1..N, avoiding O(N*body_size) IR explosion from deep-cloning
+// hierarchy bodies that shim DMA BD folding never touches.
+LogicalResult loopUnrollFullLightweight(scf::ForOp forOp);
 
 // Unrolls an `scf.for` loop by a given factor while preserving async token
 // dependencies.

mlir/lib/Transform/AIRDependencyScheduleOpt.cpp

@@ -6373,9 +6373,10 @@ class AIROptimizeShimDMABDs
     // Canonicalize IR to make loop bounds explicitly static.
     applyCanonicalizationPatterns(ctx, func.getBody());
 
-    // Unroll outer scf.for loop nest.
+    // Unroll outer scf.for loop nest. Use lightweight clone to avoid
+    // deep-copying segment/herd bodies that BD folding never touches.
     for (auto scfFor : forLoopsToUnroll) {
-      if (failed(air::loopUnrollFullWithAsyncTokenPreserved(scfFor)))
+      if (failed(air::loopUnrollFullLightweight(scfFor)))
         signalPassFailure();
     }
     // Canonicalize IR to make loop bounds explicitly static.

mlir/lib/Util/Dependency.cpp

@@ -1006,6 +1006,134 @@ LogicalResult loopUnrollFullWithAsyncTokenPreserved(
   return success();
 }
 
+// Lightweight unroll: clones empty shells for SegmentOp/HerdOp on iterations
+// 1..N-1, keeping the full body only for iteration 0. This avoids the
+// O(N * body_size) IR explosion from deep-cloning segment/herd bodies that
+// BD folding never touches. The downstream air-to-std channel matching only
+// needs one copy of the segment body (via ChannelCounterpartCache which
+// stores the first get/put per channel name).
+LogicalResult loopUnrollFullLightweight(scf::ForOp forOp) {
+  auto tripCount = air::getStaticScfForTripCountAsInt(forOp);
+  if (!tripCount) {
+    // Dynamic loop bound — fall back to standard unroll.
+    return loopUnrollFullWithAsyncTokenPreserved(forOp);
+  }
+
+  // Pre-processing: label yield-defining ops for async token fixup.
+  labelYieldDefiningOpsOfForLoop(forOp, "scf.for_result_id");
+  Block *parentBlock = forOp->getBlock();
+
+  // Trivial case: single iteration — just promote.
+  if (*tripCount == 1) {
+    IRRewriter rewriter(forOp.getContext());
+    (void)forOp.promoteIfSingleIteration(rewriter);
+    preserveAsyncDependenciesAfterUnroll(*parentBlock);
+    return success();
+  }
+
+  Block *body = forOp.getBody();
+  OpBuilder builder(forOp.getContext());
+  builder.setInsertionPoint(body->getTerminator());
+
+  // Identify the last op before the yield (the "end" of the body to clone).
+  Block::iterator srcBlockEnd = std::prev(body->end(), 2);
+
+  Value iv = forOp.getInductionVar();
+  Value lb = forOp.getLowerBound();
+  Value step = forOp.getStep();
+  auto loc = forOp.getLoc();
+
+  // Track yielded values for iter_arg remapping across iterations.
+  SmallVector<Value> lastYielded(
+      forOp.getBody()->getTerminator()->getOperands());
+
+  // Clone iterations 1..N-1 into the body (iteration 0 is the original).
+  for (unsigned i = 1; i < *tripCount; i++) {
+    IRMapping operandMap;
+
+    // Map iter_args to the previous iteration's yielded values.
+    for (auto [iterArg, yielded] :
+         llvm::zip(forOp.getRegionIterArgs(), lastYielded))
+      operandMap.map(iterArg, yielded);
+
+    // Map IV to lb + i * step.
+    if (!iv.use_empty()) {
+      Value iterConst = arith::ConstantIndexOp::create(builder, loc, i);
+      Value offset = arith::MulIOp::create(builder, loc, step, iterConst);
+      Value newIV = arith::AddIOp::create(builder, loc, lb, offset);
+      operandMap.map(iv, newIV);
+    }
+
+    // Clone each op in the body. After cloning, strip segment/herd bodies
+    // in the cloned copy (they are deep inside inner loops but BD folding
+    // never touches them — only L3 channel ops at the launch level matter).
+    for (auto it = body->begin(); it != std::next(srcBlockEnd); it++) {
+      Operation *clonedOp = builder.clone(*it, operandMap);
+
+      // Walk the cloned op and strip segment/herd bodies, but only those
+      // inside a dummyLaunch (from AIRLaunchToScfForPattern). Standalone
+      // herds at the function level must keep their bodies.
+      clonedOp->walk([&](Operation *nested) {
+        if (!isa<air::SegmentOp, air::HerdOp>(nested))
+          return;
+        auto parentLaunch = nested->getParentOfType<air::LaunchOp>();
+        if (!parentLaunch || !parentLaunch->hasAttr("dummyLaunch"))
+          return;
+        for (Region &region : nested->getRegions()) {
+          if (region.empty())
+            continue;
+          Block &blk = region.front();
+          // Erase all ops except the terminator.
+          while (blk.getOperations().size() > 1) {
+            Operation &op = *std::prev(blk.end(), 2);
+            op.dropAllUses();
+            op.erase();
+          }
+        }
+      });
+    }
+
+    // Update yielded values for the next iteration.
+    auto *yield = forOp.getBody()->getTerminator();
+    for (unsigned j = 0; j < lastYielded.size(); j++)
+      lastYielded[j] = operandMap.lookupOrDefault(yield->getOperand(j));
+  }
+
+  // Replace the loop's results with the last iteration's yielded values,
+  // then inline the body into the parent block and erase the loop.
+  {
+    IRRewriter rewriter(forOp.getContext());
+
+    // Replace IV uses in iteration 0 with lb (the first iteration's IV value).
+    if (!iv.use_empty())
+      rewriter.replaceAllUsesWith(iv, lb);
+
+    // Replace iter_arg uses in iteration 0 with init values.
+    for (auto [iterArg, initVal] :
+         llvm::zip(forOp.getRegionIterArgs(), forOp.getInitArgs()))
+      rewriter.replaceAllUsesWith(iterArg, initVal);
+
+    // Replace loop results with the last iteration's yielded values.
+    for (auto [loopResult, yielded] :
+         llvm::zip(forOp.getResults(), lastYielded))
+      rewriter.replaceAllUsesWith(loopResult, yielded);
+
+    // Move all body ops (except the yield) before the loop in the parent.
+    Block *parentBlock2 = forOp->getBlock();
+    auto &bodyOps = body->getOperations();
+    auto yieldIt = std::prev(bodyOps.end()); // the yield terminator
+    parentBlock2->getOperations().splice(Block::iterator(forOp), bodyOps,
+                                         bodyOps.begin(), yieldIt);
+
+    // Erase the now-empty loop.
+    rewriter.eraseOp(forOp);
+  }
+
+  // Post-processing: reconnect async token dependencies.
+  preserveAsyncDependenciesAfterUnroll(*parentBlock);
+  return success();
+}
+
 // Unrolls an `scf.for` loop by a given factor while preserving async token
 // dependencies.
 //