From 2cecb311a770082ee69aa04cdff95e985640368b Mon Sep 17 00:00:00 2001 From: erweiw Date: Tue, 28 Jul 2026 20:50:19 -0700 Subject: [PATCH 1/4] [AIR] Add air.shim_feed_no_pace opt-out for fire-and-free shim feeds A launch marked air.preserve_shim_dma_order currently propagates that marker to ALL of its air.channel.put/get ops, which makes AIRRtToNpu pace every shim MM2S feed with bounded double-buffered (await-on-drain) backpressure. That pacing is required for feeds coupled by a shared broadcast/multicast consumer that advances its destinations in lockstep, but it is wrong for feeds that are mutually independent: the per-task completion await serializes otherwise-concurrent channels and deadlocks when a single feed's BD exceeds the downstream ring depth. Add a per-op opt-out unit attr, air.shim_feed_no_pace. When present on a channel.put/get inside a preserve_shim_dma_order launch, the op is excluded from the preserve-marker propagation (AIROptimizeShimDMABDs) and therefore lowers to a fire-and-free MM2S feed (start, free BD, no per-task await), while still benefiting from the launch's no-fold guarantee. This lets a few large independent BDs stream concurrently, backpressured only by their downstream ring locks. copyChannelSteeringAttrs carries the marker across op rewrites so it survives to the shim-DMA-BD pass. The change is opt-in: behavior is unchanged for any op that does not carry the new attribute. --- mlir/include/air/Dialect/AIR/AIRDialect.h | 10 ++++++++++ mlir/lib/Dialect/AIR/IR/AIRDialect.cpp | 4 ++++ mlir/lib/Transform/AIRDependencyScheduleOpt.cpp | 15 ++++++++++++--- 3 files changed, 26 insertions(+), 3 deletions(-) diff --git a/mlir/include/air/Dialect/AIR/AIRDialect.h b/mlir/include/air/Dialect/AIR/AIRDialect.h index 2d7eff7b2..0532e77a3 100644 --- a/mlir/include/air/Dialect/AIR/AIRDialect.h +++ b/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"; diff --git a/mlir/lib/Dialect/AIR/IR/AIRDialect.cpp b/mlir/lib/Dialect/AIR/IR/AIRDialect.cpp index 66d093c26..b9f87ccfd 100644 --- a/mlir/lib/Dialect/AIR/IR/AIRDialect.cpp +++ b/mlir/lib/Dialect/AIR/IR/AIRDialect.cpp @@ -138,6 +138,10 @@ void air::copyChannelSteeringAttrs(Operation *src, Operation *dst) { 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(attrs::RefeedCount)) diff --git a/mlir/lib/Transform/AIRDependencyScheduleOpt.cpp b/mlir/lib/Transform/AIRDependencyScheduleOpt.cpp index 2ac2fa6d6..72ae5380f 100644 --- a/mlir/lib/Transform/AIRDependencyScheduleOpt.cpp +++ b/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(op)) - op->setAttr(air::attrs::PreserveShimDmaOrder, - mlir::UnitAttr::get(op->getContext())); + if (!isa(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; } From 44502f148a88e2b90bf2aa8692ab0b0bf3f2430f Mon Sep 17 00:00:00 2001 From: erweiw Date: Tue, 28 Jul 2026 21:11:35 -0700 Subject: [PATCH 2/4] [AIR] test: air.shim_feed_no_pace opt-out in air-opt-shim-dma-bds Covers the preserve-marker propagation and the per-op opt-out: an untagged feed in a preserve_shim_dma_order launch inherits the marker, while a feed tagged air.shim_feed_no_pace does not (it lowers fire-and-free) and keeps its opt-out attribute. --- .../opt_shim_dma_bds_no_pace_optout.mlir | 55 +++++++++++++++++++ 1 file changed, 55 insertions(+) create mode 100644 mlir/test/Transform/AIRDependencyScheduleOpt/opt_shim_dma_bds_no_pace_optout.mlir diff --git a/mlir/test/Transform/AIRDependencyScheduleOpt/opt_shim_dma_bds_no_pace_optout.mlir b/mlir/test/Transform/AIRDependencyScheduleOpt/opt_shim_dma_bds_no_pace_optout.mlir new file mode 100644 index 000000000..8ee10746d --- /dev/null +++ b/mlir/test/Transform/AIRDependencyScheduleOpt/opt_shim_dma_bds_no_pace_optout.mlir @@ -0,0 +1,55 @@ +//===- opt_shim_dma_bds_no_pace_optout.mlir --------------------*- MLIR -*-===// +// +// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +// +//===----------------------------------------------------------------------===// + +// RUN: air-opt %s -air-opt-shim-dma-bds="device=npu2" -split-input-file | FileCheck %s + +// air-opt-shim-dma-bds propagates a launch's air.preserve_shim_dma_order marker +// onto its air.channel.put/get ops (so airrt-to-npu paces them with bounded +// double-buffered backpressure), EXCEPT for ops that opt out via +// air.shim_feed_no_pace -- those lower fire-and-free instead. This covers both +// the default propagation and the per-op opt-out. + +// (1) An untagged feed in a preserve launch INHERITS the preserve marker. +// CHECK-LABEL: func.func @propagates_preserve +// CHECK: air.channel.put +// CHECK-SAME: air.preserve_shim_dma_order +air.channel @feed_paced [1] +func.func @propagates_preserve(%arg0: memref<512xbf16>) { + %c1 = arith.constant 1 : index + %0 = air.launch async (%i) in (%n=%c1) args(%a=%arg0) : memref<512xbf16> + attributes {air.preserve_shim_dma_order} { + %c0 = arith.constant 0 : index + %c64 = arith.constant 64 : index + %c1_0 = arith.constant 1 : index + %1 = air.channel.put async @feed_paced[] + (%a[%c0] [%c64] [%c1_0]) {metadata = @m0} : (memref<512xbf16>) + } + return +} + +// ----- + +// (2) A feed tagged air.shim_feed_no_pace is NOT given the preserve marker (it +// lowers fire-and-free), while still keeping its own opt-out attribute. +// CHECK-LABEL: func.func @respects_no_pace_optout +// CHECK: air.channel.put +// CHECK-SAME: air.shim_feed_no_pace +// CHECK-NOT: air.preserve_shim_dma_order +air.channel @feed_free [1] +func.func @respects_no_pace_optout(%arg0: memref<512xbf16>) { + %c1 = arith.constant 1 : index + %0 = air.launch async (%i) in (%n=%c1) args(%a=%arg0) : memref<512xbf16> + attributes {air.preserve_shim_dma_order} { + %c0 = arith.constant 0 : index + %c64 = arith.constant 64 : index + %c1_0 = arith.constant 1 : index + %1 = air.channel.put async @feed_free[] + (%a[%c0] [%c64] [%c1_0]) + {metadata = @m1, air.shim_feed_no_pace} : (memref<512xbf16>) + } + return +} From 9a5d00f02615cff79a7cb6de66a1d7c7d7a1b1c7 Mon Sep 17 00:00:00 2001 From: erweiw Date: Tue, 28 Jul 2026 21:39:58 -0700 Subject: [PATCH 3/4] [AIR] clang-format: reflow shim_feed_no_pace doc comment --- mlir/include/air/Dialect/AIR/AIRDialect.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/include/air/Dialect/AIR/AIRDialect.h b/mlir/include/air/Dialect/AIR/AIRDialect.h index 0532e77a3..1ec59f7f1 100644 --- a/mlir/include/air/Dialect/AIR/AIRDialect.h +++ b/mlir/include/air/Dialect/AIR/AIRDialect.h @@ -53,8 +53,8 @@ constexpr StringLiteral CoalescedShimFeed = "air.coalesced_shim_feed"; // 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. +// 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"; From 3624ffcb4d46ba5fa2acab90655d8198fe539974 Mon Sep 17 00:00:00 2001 From: erweiw Date: Wed, 29 Jul 2026 10:40:38 -0700 Subject: [PATCH 4/4] [AIR] test: cover shim_feed_no_pace end-to-end lowering and attr survival Two gaps in the opt-out coverage: - AIRRtToNpu lowering: a mixed preserve launch where one feed keeps the preserve marker and a sibling opts out. Asserts the kept feed is paced (issue_token + depth-2 completion-token awaits) while the opted-out feed lowers fire-and-free (no issue_token, no await, dma_free_task after start) -- the actual downstream consequence the propagation test could not show. - copyChannelSteeringAttrs round-trip: air.shim_feed_no_pace must survive the per-channel BD fold that rebuilds the op (two scf.for loops collapse into one wrap/stride put) in a non-preserve launch, so the marker reaches the preserve-marker propagation instead of being silently dropped on rebuild. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../shim_feed_no_pace_lowering.mlir | 58 +++++++++++++++++++ .../opt_shim_dma_bds_no_pace_optout.mlir | 32 ++++++++++ 2 files changed, 90 insertions(+) create mode 100644 mlir/test/Conversion/AIRRtToNpu/shim_feed_no_pace_lowering.mlir diff --git a/mlir/test/Conversion/AIRRtToNpu/shim_feed_no_pace_lowering.mlir b/mlir/test/Conversion/AIRRtToNpu/shim_feed_no_pace_lowering.mlir new file mode 100644 index 000000000..50ef954f6 --- /dev/null +++ b/mlir/test/Conversion/AIRRtToNpu/shim_feed_no_pace_lowering.mlir @@ -0,0 +1,58 @@ +//===- shim_feed_no_pace_lowering.mlir -------------------------*- MLIR -*-===// +// +// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT +// +//===----------------------------------------------------------------------===// + +// RUN: air-opt -airrt-to-npu %s | FileCheck %s + +// End-to-end consequence of the air.shim_feed_no_pace opt-out: airrt-to-npu +// gates its double-buffered pacing on air.preserve_shim_dma_order. An opted-out +// feed reaches this pass WITHOUT that marker (air-opt-shim-dma-bds excluded it), +// so it lowers fire-and-free (issue_token unset, dma_free_task after start, no +// await), while a sibling that kept the marker on the same channel is paced +// (issue_token set, depth-2 completion-token awaits). + +// CHECK-LABEL: aie.runtime_sequence @mixed_feeds +// Paced feed: issue_token set, bounded (depth=2) awaits. +// CHECK: %[[P0:.*]] = aiex.dma_configure_task_for @paced +// CHECK: issue_token = true +// CHECK: aiex.dma_start_task(%[[P0]]) +// CHECK: %[[P1:.*]] = aiex.dma_configure_task_for @paced +// CHECK: aiex.dma_start_task(%[[P1]]) +// CHECK: %[[P2:.*]] = aiex.dma_configure_task_for @paced +// CHECK: aiex.dma_await_task(%[[P0]]) +// CHECK: aiex.dma_start_task(%[[P2]]) +// CHECK: aiex.dma_await_task(%[[P1]]) +// CHECK: aiex.dma_await_task(%[[P2]]) +// Fire-and-free feed (opted out upstream, no preserve marker): no issue_token, +// no await, freed immediately after start. +// CHECK: %[[F0:.*]] = aiex.dma_configure_task_for @free +// CHECK-NOT: issue_token +// CHECK-NOT: aiex.dma_await_task +// CHECK: aiex.dma_start_task(%[[F0]]) +// CHECK-NEXT: aiex.dma_free_task(%[[F0]]) +module { + aie.device(npu1) { + %shim_noc_tile_0_0 = aie.tile(0, 0) + %shim_noc_tile_1_0 = aie.tile(1, 0) + aie.shim_dma_allocation @paced(%shim_noc_tile_0_0, MM2S, 0) + aie.shim_dma_allocation @free(%shim_noc_tile_1_0, MM2S, 0) + } {sym_name = "mixed"} + airrt.module_metadata{} + func.func @mixed_feeds(%arg0: memref<64xi32>, %arg1: memref<64xi32>) { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c64_i64 = arith.constant 64 : i64 + %c2_i32 = arith.constant 2 : i32 + %c3_i32 = arith.constant 3 : i32 + %p = airrt.segment_load "mixed" : i64 + %0 = airrt.dma_memcpy_nd(%c2_i32, %c0_i64, %c0_i64, %arg0[%c0_i64, %c0_i64, %c0_i64, %c0_i64], [%c1_i64, %c1_i64, %c1_i64, %c64_i64], [%c0_i64, %c0_i64, %c0_i64, %c1_i64]) {metadata = @paced, air.preserve_shim_dma_order} : (i32, i64, i64, memref<64xi32>, [i64, i64, i64, i64], [i64, i64, i64, i64], [i64, i64, i64, i64]) : !airrt.event + %1 = airrt.dma_memcpy_nd(%c2_i32, %c0_i64, %c0_i64, %arg0[%c0_i64, %c0_i64, %c0_i64, %c0_i64], [%c1_i64, %c1_i64, %c1_i64, %c64_i64], [%c0_i64, %c0_i64, %c0_i64, %c1_i64]) {metadata = @paced, air.preserve_shim_dma_order} : (i32, i64, i64, memref<64xi32>, [i64, i64, i64, i64], [i64, i64, i64, i64], [i64, i64, i64, i64]) : !airrt.event + %2 = airrt.dma_memcpy_nd(%c2_i32, %c0_i64, %c0_i64, %arg0[%c0_i64, %c0_i64, %c0_i64, %c0_i64], [%c1_i64, %c1_i64, %c1_i64, %c64_i64], [%c0_i64, %c0_i64, %c0_i64, %c1_i64]) {metadata = @paced, air.preserve_shim_dma_order} : (i32, i64, i64, memref<64xi32>, [i64, i64, i64, i64], [i64, i64, i64, i64], [i64, i64, i64, i64]) : !airrt.event + %3 = airrt.dma_memcpy_nd(%c3_i32, %c0_i64, %c0_i64, %arg1[%c0_i64, %c0_i64, %c0_i64, %c0_i64], [%c1_i64, %c1_i64, %c1_i64, %c64_i64], [%c0_i64, %c0_i64, %c0_i64, %c1_i64]) {metadata = @free} : (i32, i64, i64, memref<64xi32>, [i64, i64, i64, i64], [i64, i64, i64, i64], [i64, i64, i64, i64]) : !airrt.event + airrt.wait_all %0, %1, %2, %3 + return + } +} diff --git a/mlir/test/Transform/AIRDependencyScheduleOpt/opt_shim_dma_bds_no_pace_optout.mlir b/mlir/test/Transform/AIRDependencyScheduleOpt/opt_shim_dma_bds_no_pace_optout.mlir index 8ee10746d..299302b23 100644 --- a/mlir/test/Transform/AIRDependencyScheduleOpt/opt_shim_dma_bds_no_pace_optout.mlir +++ b/mlir/test/Transform/AIRDependencyScheduleOpt/opt_shim_dma_bds_no_pace_optout.mlir @@ -53,3 +53,35 @@ func.func @respects_no_pace_optout(%arg0: memref<512xbf16>) { } return } + +// ----- + +// (3) copyChannelSteeringAttrs must carry air.shim_feed_no_pace across the +// per-channel BD fold that rebuilds the op (here two scf.for loops collapse into +// one wrap/stride put). If the marker were dropped on rebuild it would be gone +// before the preserve-marker propagation reads it, silently re-pacing the feed. +// Non-preserve launch so folding actually runs. +// CHECK-LABEL: func.func @survives_bd_fold +// CHECK: air.channel.put +// CHECK-SAME: air.shim_feed_no_pace +air.channel @feed_fold [1] +func.func @survives_bd_fold(%arg0: memref<512x512xbf16>) { + %c1 = arith.constant 1 : index + %0 = air.launch async (%i) in (%n=%c1) args(%a=%arg0) : memref<512x512xbf16> { + %c0 = arith.constant 0 : index + %c1_0 = arith.constant 1 : index + %c256 = arith.constant 256 : index + %c512 = arith.constant 512 : index + %1 = air.wait_all async + %2 = scf.for %arg6 = %c0 to %c512 step %c256 iter_args(%arg7 = %1) -> (!air.async.token) { + %3 = scf.for %arg8 = %c0 to %c512 step %c256 iter_args(%arg9 = %arg7) -> (!air.async.token) { + %4 = air.channel.put async [%arg9] @feed_fold[] + (%a[%arg6, %arg8] [%c256, %c256] [%c512, %c1_0]) + {metadata = @m2, air.shim_feed_no_pace} : (memref<512x512xbf16>) + scf.yield %4 : !air.async.token + } + scf.yield %3 : !air.async.token + } + } + return +}