Lightweight herd cloning during shim DMA BD loop unrolling

[erwei-xilinx]

Summary

Avoid O(N × body_size) IR explosion during shim-level loop unrolling in air-opt-shim-dma-bds by skipping deep-clone of herd compute bodies.

When loopUnrollFullWithAsyncTokenPreserved detects a shim-level loop (containing air.SegmentOp or air.HerdOp), it uses a custom manual unroller that:

1. Fully clones segment bodies — L3 channel ops needed by BD folding are preserved
2. Creates lightweight herd copies — only channel ops, allocs, deallocs, wait_alls, their transitive operand-defining ops, and the terminator are cloned; heavy compute ops (vector, arith, linalg) are skipped entirely
3. Uses OperationState to create empty herd shells, avoiding custom builder API issues

Profiling (flash attention 12×4 launch, tiles=2,2)

Metric	Main	This PR
air-opt-shim-dma-bds	~50 ms	~43 ms

Current limitation

The lightweight unroller only fires when annotateFn is null (the non-tiled unroll path in loopUnrollFullWithAsyncTokenPreserved). The tiled path, which uses loopUnrollByFactor with an annotateFn callback, still performs full deep-clone unrolling. Extending lightweight cloning to the tiled path is left as a follow-up.

Test plan

• ninja check-air-mlir passes (365/374, only pre-existing failures)
• New LIGHTWEIGHT FileCheck test in opt_shim_dma_bds.mlir verifies channel ops preserved, compute ops stripped
• Flash attention compiles end-to-end on NPU2
• Profiling confirms no regression vs main

[Copilot]

Pull request overview

This PR optimizes the air-opt-shim-dma-bds pass to avoid IR blow-up by skipping runtime-loop tiling and full unrolling when shim-dma-tile-sizes is empty or all-1 (the default path), while preserving scf.for loops for later/lower-cost unrolling in downstream passes.

Changes:

• Add a fast path in AIROptimizeShimDMABDs to skip tiling/unrolling when tile sizes are empty or all-1, while still performing L3 DMA folding.
• Add a new MLIR regression test covering all-1, non-trivial tiling (2), and empty-tile-size behavior.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
mlir/lib/Transform/AIRDependencyScheduleOpt.cpp	Adds fast-path logic to skip tiling/unrolling for empty/all-1 tile sizes and attempts to preserve downstream barrier behavior.
mlir/test/Transform/AIRDependencyScheduleOpt/opt_shim_dma_bds_skip_unroll.mlir	New test validating skip-unroll behavior and non-trivial tiling behavior.

---

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

[erwei-xilinx]

Closing this PR. Hardware E2E testing on local NPU1 confirmed the change produces incorrect DMA BD configurations.

Root cause: Skipping BD folding (applyAIRL3DmaFoldingPatterns) is not safe — the BD folding performs essential loop-to-BD-dimension conversion that changes the channel ops' wraps/strides. Without it, the runtime loop gets unrolled in airrt-to-npu with per-iteration BDs that don't match the hardware's expected multi-dimensional BD format. This produces wrong air.insts.bin (18 KB vs 6.5 KB reference, different offset patterns) and wrong numerical output on NPU1.

What doesn't work: Simply skipping the tiling/unrolling/BD-folding for all-1 tile sizes. The inner BD folding patterns (AIRSpecializeChannelWrapAndStrideInScfFor + AIRUnrollScfForIntoBDChain) are coupled — the isolation pattern splits loops so each channel op gets its own loop, the specialization pattern folds loops into BD dimensions, and the unroll pattern handles cases that can't be folded. Skipping all of them produces fundamentally different (and incorrect) BD configurations.

What would work: A more surgical approach that runs BD folding on the channel ops but prevents AIRUnrollScfForIntoBDChain from unrolling the specific outer runtime loops (from launch conversion). This requires modifying the BD folding patterns to distinguish between "BD loops" (inner loops that should be folded into dimensions) and "runtime loops" (outer loops that should be preserved). This is a larger change that needs careful design.

The compilation speed optimization plan at workspace/compilation_speed_optimization_plan.md remains valid as a design document for the correct approach.

[erwei-xilinx]

Closing this PR. The compilation speed optimization (5x on MLIR passes for flash attention) is real and validated by profiling, but all implementation approaches explored so far hit correctness issues:

1. Skip-unroll approach (defer loop unrolling to airrt-to-npu): Fails because loopUnrollFull in airrt-to-npu doesn't handle airrt::EventType loop-carried values, and purgeSCFParContainingOnlyWaitAllOps unconditionally erases all scf.parallel ops including ones containing DMA ops.

2. Attribute-based approach (mark runtime loops, skip them in AIRUnrollScfForIntoBDChain): Produces different (incorrect) NPU instruction binaries because BD folding is effectively a no-op on multi-channel-op loop bodies, so the channel ops retain symbolic offsets that downstream passes can't resolve.

3. Strip segment/herd bodies before unrolling (Strategy A): Fails because air-to-std runs after the shim BD pass on the same module and needs segment bodies to match channel put/get pairs.

The correct approach requires either:

• Lightweight cloning in loopUnrollFullWithAsyncTokenPreserved that skips segment/herd body contents during the clone step (only clones the shell + L3 channel ops)
• Pipeline restructuring to run air-to-std channel matching before the shim BD pass's unrolling

The profiling analysis and detailed plan remain at workspace/compilation_speed_optimization_plan.md.

[erwei-xilinx]

Known limitation: tiled path not covered

The lightweight unroller currently only fires when annotateFn is null — i.e., the non-tiled unroll path inside loopUnrollFullWithAsyncTokenPreserved. The tiled path (triggered by non-trivial shim-dma-tile-sizes like 2,2 or 4,4) goes through loopUnrollByFactor with an annotateFn callback, which still performs full O(N × body_size) deep-clone unrolling.

For the default aircc invocation (tile sizes 1,1), the shim BD pass doesn't tile or unroll — it directly applies BD folding. The lightweight unroller fires when the launch conversion creates scf.for loops that contain segments/herds and need to be unrolled without tiling.

Extending lightweight cloning to the tiled loopUnrollByFactor path is left as a follow-up.