[AIR] getRefeedCount: explicit per-emission override wins for any value

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

@@ -3216,7 +3216,7 @@
 // The following methods are required by TilingInterface.
 SmallVector<Range> air::ChannelPutOp::getIterationDomain(OpBuilder &builder) {
   air::ChannelInterface put =
       dyn_cast_if_present<air::ChannelInterface>(getOperation());
   return getIterationDomainFromChanIf(builder, put);
 }
 
@@ -3245,10 +3245,34 @@
   return success();
 }
 
+// Validate an air.refeed_count attribute (single-buffer count-free
+// re-broadcast). Legal on the channel declaration and as a per-emission
+// override on a channel put/get; getRefeedCount reads either carrier. The count
+// sizes 32-bit AIE semaphore locks, so it must be a positive integer in
+// [1, INT32_MAX] -- reject a malformed value here rather than silently
+// narrowing it downstream. No-op when the attribute is absent.
+static LogicalResult verifyRefeedCountAttr(Operation *op) {
+  auto rc = op->getAttr(air::attrs::RefeedCount);
+  if (!rc)
+    return success();
+  auto rcInt = dyn_cast<IntegerAttr>(rc);
+  if (!rcInt)
+    return op->emitOpError()
+           << "\"" << air::attrs::RefeedCount << "\" must be an integer "
+           << "attribute";
+  if (rcInt.getInt() < 1 ||
+      rcInt.getInt() > std::numeric_limits<int32_t>::max())
+    return op->emitOpError() << "\"" << air::attrs::RefeedCount << "\" ("
+                             << rcInt.getInt() << ") must be in [1, INT32_MAX]";
+  return success();
+}
+
 LogicalResult air::ChannelPutOp::verify() {
   if (failed(verifySizesStridesRank(getOperation(), getSrcSizes(),
                                     getSrcStrides(), "src")))
     return failure();
+  if (failed(verifyRefeedCountAttr(getOperation())))
+    return failure();
 
   // Channel bundle indices must not be temporal loop induction variables.
   for (auto [i, idx] : llvm::enumerate(getIndices())) {
@@ -3349,6 +3373,8 @@
   if (failed(verifySizesStridesRank(getOperation(), getDstSizes(),
                                     getDstStrides(), "dst")))
     return failure();
+  if (failed(verifyRefeedCountAttr(getOperation())))
+    return failure();
 
   // Channel bundle indices must not be temporal loop induction variables.
   for (auto [i, idx] : llvm::enumerate(getIndices())) {
@@ -3462,19 +3488,10 @@
 
   // air.refeed_count (single-buffer count-free re-broadcast) must be a positive
   // integer count of re-sends. The channel declaration is its authoritative
-  // carrier.
-  if (auto rc = (*this)->getAttr(attrs::RefeedCount)) {
-    auto rcInt = dyn_cast<IntegerAttr>(rc);
-    if (!rcInt)
-      return emitOpError() << "\"" << attrs::RefeedCount
-                           << "\" must be an integer attribute";
-    // The count sizes 32-bit AIE semaphore locks; reject values that would not
-    // fit so a malformed count is caught here rather than silently narrowed.
-    if (rcInt.getInt() < 1 ||
-        rcInt.getInt() > std::numeric_limits<int32_t>::max())
-      return emitOpError() << "\"" << attrs::RefeedCount << "\" ("
-                           << rcInt.getInt() << ") must be in [1, INT32_MAX]";
-  }
+  // carrier; a put/get may also carry a per-emission override (validated the
+  // same way in its own verifier).
+  if (failed(verifyRefeedCountAttr(getOperation())))
+    return failure();
 
   // packet_ids pins explicit switchbox routing ids; only meaningful for
   // packet-switched channels. Each id must be a 5-bit AIE pkt_id and unique.

mlir/lib/Util/Util.cpp

@@ -704,9 +704,14 @@ int64_t air::getRefeedCount(air::ChannelInterface op) {
   if (!op)
     return 1;
   // Per-emission override on the put/get takes precedence over the
-  // channel-level declaration.
-  if (int64_t n = getRefeedCount(op.getOperation()); n > 1)
-    return n;
+  // channel-level declaration. An EXPLICIT per-emission integer attribute wins
+  // for ANY value (including 1), so a shared channel's re-feed default can be
+  // locally disabled on a specific emission (e.g. a value-1 re-broadcast that
+  // must not inherit the channel's count). Require an IntegerAttr (not just
+  // hasAttr) so a malformed/wrong-typed attribute falls through to the channel
+  // declaration rather than silently disabling the re-feed.
+  if (op.getOperation()->getAttrOfType<IntegerAttr>(air::attrs::RefeedCount))
+    return getRefeedCount(op.getOperation());
   return getRefeedCount(getChannelDeclarationThroughSymbol(op).getOperation());
 }
 

mlir/test/Conversion/AIRToAIE/air_channel_producer_refeed.mlir

@@ -203,6 +203,51 @@ func.func @mt_refeed_independent(%arg0: memref<64xi32>, %arg1: memref<64xi32>) {
 
 // -----
 
+// Per-emission override: the channel declares air.refeed_count=5 but the
+// producer put sets air.refeed_count=1, which wins for ANY value (not only
+// >1). The override disables the channel's re-feed on this emission, so the
+// producing core release stays 1 and NO lock inits to / acquires / releases 5
+// (the channel default). A regression that only honored per-op overrides >1
+// would fall back to the channel's 5 and re-introduce the 5-credit locks. 5 is
+// unique to this file, so these CHECK-NOTs are unambiguous.
+
+// CHECK-LABEL: aie.device
+// CHECK-NOT: {init = 5 : i32}
+// CHECK-NOT: AcquireGreaterEqual, 5
+// CHECK-NOT: Release, 5
+
+#set_ovr = affine_set<()[s0, s1] : (s0 >= 0, -s0 + 1 >= 0, s1 == 0)>
+air.channel @channel_ovr [1, 1] {air.refeed_count = 5 : i32}
+func.func @refeed_override() {
+  %c1 = arith.constant 1 : index
+  %0 = air.launch async (%arg4, %arg5) in (%arg6=%c1, %arg7=%c1) {
+    %1 = air.segment async {
+      %c2 = arith.constant 2 : index
+      %c1_0 = arith.constant 1 : index
+      %2 = air.herd @herd_0 async tile (%arg8, %arg9) in (%arg10=%c1_0, %arg11=%c2) {
+        %c0 = arith.constant 0 : index
+        %async_token_6, %results_7 = air.execute -> (memref<32x32xbf16, 2>) {
+          %alloc = memref.alloc() : memref<32x32xbf16, 2>
+          air.execute_terminator %alloc : memref<32x32xbf16, 2>
+        }
+        %3 = affine.if #set_ovr()[%arg8, %arg9] -> !air.async.token {
+          %4 = air.channel.put async [%async_token_6]  @channel_ovr[] (%results_7[] [] []) {air.refeed_count = 1 : i32} : (memref<32x32xbf16, 2>)
+          affine.yield %4 : !air.async.token
+        } else {
+          %4 = air.channel.get async [%async_token_6]  @channel_ovr[] (%results_7[] [] []) : (memref<32x32xbf16, 2>)
+          affine.yield %4 : !air.async.token
+        }
+        %async_token_8 = air.execute [%3] {
+          memref.dealloc %results_7 : memref<32x32xbf16, 2>
+        }
+      }
+    }
+  }
+  return
+}
+
+// -----
+
 // Without air.refeed_count the producing core release stays 1 (no re-feed),
 // the buf-free acquire stays 1, and no lock inits to 4. (This section is last
 // so its CHECK-NOT window is not entered by the positive sections above.)

mlir/test/Dialect/AIR/air_channel_invalid.mlir

@@ -179,6 +179,60 @@ air.channel @refeed_ok [1, 1] {air.refeed_count = 4 : i32}
 
 // -----
 
+// A per-emission air.refeed_count override on a put is an authoritative carrier
+// (getRefeedCount honors it), so it is validated the same as the channel-level
+// count: zero is rejected.
+air.channel @refeed_op_zero [1, 1]
+func.func @channel_put_refeed_zero(%m: memref<64xi32>) {
+  // expected-error @+1 {{'air.channel.put' op "air.refeed_count" (0) must be in [1, INT32_MAX]}}
+  air.channel.put @refeed_op_zero[] (%m[] [] []) {air.refeed_count = 0 : i32} : (memref<64xi32>)
+  return
+}
+
+// -----
+
+// Per-emission override on a put that overflows the 32-bit lock range is
+// rejected (would silently deadlock the DMA self-loop otherwise).
+air.channel @refeed_op_huge [1, 1]
+func.func @channel_put_refeed_huge(%m: memref<64xi32>) {
+  // expected-error @+1 {{'air.channel.put' op "air.refeed_count" (4294967296) must be in [1, INT32_MAX]}}
+  air.channel.put @refeed_op_huge[] (%m[] [] []) {air.refeed_count = 4294967296 : i64} : (memref<64xi32>)
+  return
+}
+
+// -----
+
+// Non-integer per-emission override on a put is rejected (getRefeedCount would
+// otherwise fall through and silently ignore the override).
+air.channel @refeed_op_str [1, 1]
+func.func @channel_put_refeed_str(%m: memref<64xi32>) {
+  // expected-error @+1 {{'air.channel.put' op "air.refeed_count" must be an integer attribute}}
+  air.channel.put @refeed_op_str[] (%m[] [] []) {air.refeed_count = "4"} : (memref<64xi32>)
+  return
+}
+
+// -----
+
+// Per-emission override on a get is validated too: negative is rejected.
+air.channel @refeed_op_get_neg [1, 1]
+func.func @channel_get_refeed_neg(%m: memref<64xi32>) {
+  // expected-error @+1 {{'air.channel.get' op "air.refeed_count" (-2) must be in [1, INT32_MAX]}}
+  air.channel.get @refeed_op_get_neg[] (%m[] [] []) {air.refeed_count = -2 : i32} : (memref<64xi32>)
+  return
+}
+
+// -----
+
+// Valid per-emission override of 1 (disables a channel's re-feed default on
+// this emission - positive test, should pass).
+air.channel @refeed_op_ok [1, 1] {air.refeed_count = 4 : i32}
+func.func @channel_put_refeed_ovr_ok(%m: memref<64xi32>) {
+  air.channel.put @refeed_op_ok[] (%m[] [] []) {air.refeed_count = 1 : i32} : (memref<64xi32>)
+  return
+}
+
+// -----
+
 // Test: packet_ids on a non-packet channel is rejected.
 // expected-error @+1 {{'air.channel' op "packet_ids" is only valid on a "npu_dma_packet" channel}}
 air.channel @pids_wrong_type [] {channel_type = "npu_dma_stream", packet_ids = [1, 2]}