Implement RegionBranchOpInterface on air hierarchy ops (#1566)

* Implement RegionBranchOpInterface on air hierarchy ops

air.rank, air.launch, air.segment, and air.herd now implement
RegionBranchOpInterface so that upstream MLIR analyses (alias
analysis, dataflow, dead-code elimination, ...) can walk the
parent-operand → body-block-arg mapping that IsolatedFromAbove
already constrains explicitly.

The mapping is the same one HierarchyInterface already exposes via
getKernelOperands / getKernelArguments / getTiedKernelOperand, but
re-stated through the upstream interface so generic infrastructure
finds it. With this in place, mlir::LocalAliasAnalysis can now
trace through hierarchy block args, which removes the need for the
custom hierarchy walk in CanonicalizeAsyncOpDeps's getRoot.

Three caveats addressed:
- AIRDependency.cpp dispatched generic RegionBranchOpInterface ops
  through scf.if/affine.if-only handling. Skip air hierarchy ops
  there — they have their own async-token machinery.
- HerdOp::getKernelArguments hardcoded drop_front(4), assuming 2D.
  The new auto-verifier on RegionBranchOpInterface exercises this
  path on the existing 0-dim-herd test, exposing the latent bug.
  Replaced with drop_front(getNumDims() * 2) to match the other
  three ops.
- The body→parent successor yields no values; the async token
  result is a launch side-effect, not a yielded value, so
  getEntrySuccessorOperands and getSuccessorInputs return empty
  for that direction.

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

* Advertise hierarchy invocation bounds via getRegionInvocationBounds

air.launch / air.segment / air.herd / air.rank now implement the
RegionBranchOpInterface::getRegionInvocationBounds method. The body
executes once per (x, y, ...) coordinate of the iteration space; if
every size operand folds to a non-negative IntegerAttr we report a
tight {n, n} bound, otherwise Unknown.

This matches scf.forall's convention (each lattice point invokes the
body once, in parallel) and is needed so that upstream passes keyed on
RegionBranchOpInterface (ControlFlowSink, the inliner's repetition
heuristics, future LICM-style work) do not assume hierarchy bodies are
non-iterative just because there is no body->body edge in
getSuccessorRegions.

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

* Use upstream LocalAliasAnalysis in CanonicalizeAsyncOpDeps

Two changes that compose:

1. air.execute now implements MemoryEffectOpInterface. The body's
   effects are mirrored to the outer op (so RAW/WAR/WAW analyses do
   not assume the execute is side-effect free), and Allocate effects
   on yielded values are retargeted to the corresponding outer op
   result. With this in place upstream alias analyses can reason
   that two distinct `air.execute -> memref` ops produce non-aliasing
   storage, which previously needed the bespoke SSA-identity getRoot
   walk to recognise.

2. CanonicalizeAsyncOpDeps' mayAlias predicate now consults
   mlir::LocalAliasAnalysis first and falls back to the existing
   getRoot SSA-identity check only when alias() returns MayAlias.
   The fallback is still needed because LocalAliasAnalysis is
   intentionally conservative for two distinct entry-block func args
   (a caller could in principle pass the same buffer twice, and MLIR
   has no `noalias` annotation on func args by default), while AIR's
   domain assumption is that distinct top-level memref SSA values
   address distinct storage.

Net effect: strictly >= the previous precision. Cases LocalAliasAnalysis
can decide pick up the upstream analysis (subviews-of-distinct-allocs,
hierarchy-arg traversal, future passes that improve LocalAliasAnalysis);
cases it cannot decide fall through to the AIR-domain heuristic.

This is the followup hinted at in the previous commit's PR description.
The interface unlock would also let new passes use upstream alias
analysis directly, without re-implementing getRoot.

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

* Address Copilot review

- computeHierarchyBounds: pre-multiply overflow guard so a single
  oversized size operand cannot wrap uint64_t and silently report a
  small invocation bound.
- ExecuteOp::getEffects: when an inner op does not implement
  MemoryEffectOpInterface, consult mlir::isMemoryEffectFree before
  falling back to the conservative Read+Write+Free. arith.constant /
  affine.apply / similar Pure-trait ops are common in execute bodies
  and were pessimising the declared effects without it.

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

* Trim verbose comments

Comments now name only the WHY when it isn't obvious from the code.
PR-description-style narration (history, alternatives, rationale chains)
removed from the source.

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

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: erwei-xilinx

mlir/include/air/Dialect/AIR/AIR.td

@@ -37,6 +37,7 @@ def air_RankOp : air_Op<"rank", [air_AsyncOpInterface,
                                   AttrSizedOperandSegments,
                                   IsolatedFromAbove,
                                   AffineScope,
+                                  DeclareOpInterfaceMethods<RegionBranchOpInterface, ["getEntrySuccessorOperands", "getSuccessorInputs", "getRegionInvocationBounds"]>,
                                   SingleBlockImplicitTerminator<"RankTerminatorOp">]>,
                          Arguments<(ins OptionalAttr<SymbolNameAttr>:$sym_name,
                                         Variadic<air_AsyncToken>:$async_dependencies,
@@ -142,6 +143,7 @@ def air_LaunchOp : air_Op<"launch", [air_AsyncOpInterface,
                                      AttrSizedOperandSegments,
                                      IsolatedFromAbove,
                                      AffineScope,
+                                     DeclareOpInterfaceMethods<RegionBranchOpInterface, ["getEntrySuccessorOperands", "getSuccessorInputs", "getRegionInvocationBounds"]>,
                                      SingleBlockImplicitTerminator<"LaunchTerminatorOp">]>,
                           Arguments<(ins OptionalAttr<SymbolNameAttr>:$sym_name,
                                          Variadic<air_AsyncToken>:$async_dependencies,
@@ -231,6 +233,7 @@ def air_SegmentOp : air_Op<"segment", [air_AsyncOpInterface,
                                            AttrSizedOperandSegments,
                                            IsolatedFromAbove,
                                            AffineScope,
+                                           DeclareOpInterfaceMethods<RegionBranchOpInterface, ["getEntrySuccessorOperands", "getSuccessorInputs", "getRegionInvocationBounds"]>,
                                            SingleBlockImplicitTerminator<"SegmentTerminatorOp">]>,
                         Arguments<(ins OptionalAttr<SymbolNameAttr>:$sym_name,
                                        Variadic<air_AsyncToken>:$async_dependencies,
@@ -351,6 +354,7 @@ def air_HerdOp : air_Op<"herd", [air_AsyncOpInterface,
                                  AttrSizedOperandSegments,
                                  IsolatedFromAbove,
                                  AffineScope,
+                                 DeclareOpInterfaceMethods<RegionBranchOpInterface, ["getEntrySuccessorOperands", "getSuccessorInputs", "getRegionInvocationBounds"]>,
                                  SingleBlockImplicitTerminator<"HerdTerminatorOp">]>,
                         Arguments<(ins OptionalAttr<SymbolNameAttr>:$sym_name,
                                        OptionalAttr<StrAttr>:$link_with,
@@ -805,7 +809,8 @@ def air_ChannelGetOp : air_Op<"channel.get", [air_AsyncOpInterface,
 // AIR asynchronous region for dynamic event dispatching.
 
 def air_ExecuteOp : air_Op<"execute", [SingleBlockImplicitTerminator<"ExecuteTerminatorOp">,
-                                       air_AsyncOpInterface]> {
+                                       air_AsyncOpInterface,
+                                       DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (
     ins Variadic<air_AsyncToken>:$async_dependencies
   );

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

@@ -8,6 +8,7 @@
 
 #include "air/Dialect/AIR/AIRDialect.h"
 
+#include "mlir/Analysis/AliasAnalysis/LocalAliasAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
@@ -20,12 +21,14 @@
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Iterators.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Transforms/RegionUtils.h"
 
 #include "llvm/ADT/TypeSwitch.h"
 
 #include <iostream>
+#include <limits>
 
 using namespace mlir;
 
@@ -233,24 +236,11 @@ static void printAsyncDependencies(OpAsmPrinter &printer, Operation *op,
 template <class OpT>
 static LogicalResult CanonicalizeAsyncOpDeps(OpT op,
                                              PatternRewriter &rewriter) {
-  // Walk view-like ops, air.hierarchy body args, and loop iter_args back
-  // to a root memref. The collectors below resolve every memref through
-  // this walk before inserting, so the RAW/WAR/WAW comparisons reduce to
-  // set-membership.
-  //
-  // Intentionally conservative — DO NOT TIGHTEN without a strong reason:
-  //   * Two disjoint views of the same root return the same root and
-  //     are reported as aliasing. Modeling subview offsets/sizes to
-  //     prove disjointness was the bug behind #1559: the original code
-  //     compared SSA identity (the strictest possible "alias" predicate)
-  //     and dropped real RAW edges between fill and channel.put.
-  //   * iter_args resolve to the loop init; a body that conditionally
-  //     yields a different memref will under-approximate roots seen
-  //     across iterations, but still over-approximates aliasing.
-  //
-  // Over-approximating aliasing is the safe direction here: this
-  // predicate gates dead-edge *removal*, so false positives only
-  // suppress an optimization, never invent a race.
+  // Fallback for the AIR-domain assumption that distinct top-level memref
+  // SSA values address distinct storage; upstream alias analysis cannot
+  // make that claim because MLIR has no noalias on func args. Aliasing
+  // here gates dead-edge removal, so over-approximation only suppresses
+  // an optimization (#1559).
   auto getRoot = [](Value v) -> Value {
     while (true) {
       if (auto view = v.getDefiningOp<ViewLikeOpInterface>()) {
@@ -275,6 +265,12 @@ static LogicalResult CanonicalizeAsyncOpDeps(OpT op,
       return v;
     }
   };
+  mlir::LocalAliasAnalysis aliasAnalysis;
+  auto mayAlias = [&aliasAnalysis, &getRoot](Value a, Value b) {
+    if (aliasAnalysis.alias(a, b).isNo())
+      return false;
+    return getRoot(a) == getRoot(b);
+  };
   auto getAllReadAccess = [](Operation *op) {
     SmallVector<Value> operands;
     if (auto linalgop = dyn_cast_if_present<linalg::LinalgOp>(op)) {
@@ -320,18 +316,15 @@ static LogicalResult CanonicalizeAsyncOpDeps(OpT op,
     }
     return operands;
   };
-  // Root memrefs accessed by `o` under `accessFn`. With `walkConsumers=true`
-  // (sink side), also unions in accesses from `o`'s transitive async-token
-  // consumers — required for sync primitives whose own accesses don't
-  // overlap a real barrier dep (issue #1559). Must be `false` on the source
-  // side, or the source would inherit its own sink's accesses.
+  // walkConsumers=true unions accesses from `o`'s transitive async-token
+  // consumers (required for sync primitives, #1559). Must be false on the
+  // source side or the source inherits its own sink's accesses.
   auto getAllMemrefsAccessedByOp =
-      [getRoot](Operation *o, bool walkConsumers,
-                llvm::function_ref<SmallVector<Value>(Operation *)> accessFn) {
+      [](Operation *o, bool walkConsumers,
+         llvm::function_ref<SmallVector<Value>(Operation *)> accessFn) {
         llvm::SetVector<Value> memrefs;
-        auto insertRoot = [&](Value v) { memrefs.insert(getRoot(v)); };
         for (Value v : accessFn(o))
-          insertRoot(v);
+          memrefs.insert(v);
         SmallVector<Region *> regions;
         for (auto &region : o->getRegions())
           regions.push_back(&region);
@@ -340,15 +333,15 @@ static LogicalResult CanonicalizeAsyncOpDeps(OpT op,
           air::walkAsyncTokenConsumers(o, consumers);
           for (auto user : consumers) {
             for (Value v : accessFn(user))
-              insertRoot(v);
+              memrefs.insert(v);
             for (auto &region : user->getRegions())
               regions.push_back(&region);
           }
         }
         for (auto region : regions) {
           visitUsedValuesDefinedAbove(*region, [&](OpOperand *use) {
             if (llvm::is_contained(accessFn(use->getOwner()), use->get()))
-              insertRoot(use->get());
+              memrefs.insert(use->get());
           });
         }
         return memrefs;
@@ -447,19 +440,20 @@ static LogicalResult CanonicalizeAsyncOpDeps(OpT op,
           llvm::range_size(llvm::concat<Value>(memrefsReadBySinkOp,
                                                memrefsWrittenBySinkOp)) != 0;
       if (sourceOpTouchesMemref && sinkOpTouchesMemref) {
-        // Sets already contain root memrefs (resolved by the collectors),
-        // so RAW/WAR/WAW checks reduce to plain set-membership.
+        auto anyAlias = [&](const llvm::SetVector<Value> &as,
+                            const llvm::SetVector<Value> &bs) {
+          for (Value a : as)
+            for (Value b : bs)
+              if (mayAlias(a, b))
+                return true;
+          return false;
+        };
         bool RAWNotFound =
-            llvm::none_of(memrefsWrittenBySourceOp, [&](Value v) {
-              return memrefsReadBySinkOp.contains(v);
-            });
-        bool WARNotFound = llvm::none_of(memrefsReadBySourceOp, [&](Value v) {
-          return memrefsWrittenBySinkOp.contains(v);
-        });
+            !anyAlias(memrefsWrittenBySourceOp, memrefsReadBySinkOp);
+        bool WARNotFound =
+            !anyAlias(memrefsReadBySourceOp, memrefsWrittenBySinkOp);
         bool WAWNotFound =
-            llvm::none_of(memrefsWrittenBySourceOp, [&](Value v) {
-              return memrefsWrittenBySinkOp.contains(v);
-            });
+            !anyAlias(memrefsWrittenBySourceOp, memrefsWrittenBySinkOp);
         bool noSharedResource = llvm::none_of(
             resourcesUsedBySourceOp, [&resourcesUsedBySinkOp](SymbolRefAttr r) {
               return llvm::is_contained(resourcesUsedBySinkOp, r);
@@ -845,6 +839,53 @@ unsigned air::LaunchOp::getNumDims() {
   return segment_sizes[1];
 }
 
+// Hierarchy body invokes once per iteration-space coordinate — product of
+// the size operands when all are constant, Unknown otherwise.
+static InvocationBounds computeHierarchyBounds(ArrayRef<Attribute> operands,
+                                               unsigned sizeStart,
+                                               unsigned numDims) {
+  constexpr uint64_t kMaxUnsigned = std::numeric_limits<unsigned>::max();
+  uint64_t product = 1;
+  for (unsigned i = 0; i < numDims; ++i) {
+    auto intAttr = dyn_cast_if_present<IntegerAttr>(operands[sizeStart + i]);
+    if (!intAttr)
+      return InvocationBounds::getUnknown();
+    int64_t v = intAttr.getInt();
+    if (v < 0)
+      return InvocationBounds::getUnknown();
+    uint64_t uv = static_cast<uint64_t>(v);
+    if (uv != 0 && product > kMaxUnsigned / uv)
+      return InvocationBounds::getUnknown();
+    product *= uv;
+  }
+  unsigned n = static_cast<unsigned>(product);
+  return InvocationBounds(n, n);
+}
+
+OperandRange air::LaunchOp::getEntrySuccessorOperands(RegionSuccessor succ) {
+  if (succ.getSuccessor() == &getBody())
+    return getKernelOperands();
+  auto end = (*this)->operand_end();
+  return OperandRange(end, end);
+}
+ValueRange air::LaunchOp::getSuccessorInputs(RegionSuccessor succ) {
+  if (succ.getSuccessor() == &getBody())
+    return ValueRange(getKernelArguments());
+  return {};
+}
+void air::LaunchOp::getSuccessorRegions(
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
+  if (point.isParent())
+    regions.push_back(RegionSuccessor(&getBody()));
+  else
+    regions.push_back(RegionSuccessor::parent());
+}
+void air::LaunchOp::getRegionInvocationBounds(
+    ArrayRef<Attribute> operands, SmallVectorImpl<InvocationBounds> &bounds) {
+  bounds.push_back(computeHierarchyBounds(
+      operands, getAsyncDependencies().size(), getNumDims()));
+}
+
 //
 // RankOp
 //
@@ -1223,6 +1264,30 @@ unsigned air::RankOp::getNumDims() {
   return segment_sizes[2];
 }
 
+OperandRange air::RankOp::getEntrySuccessorOperands(RegionSuccessor succ) {
+  if (succ.getSuccessor() == &getBody())
+    return getKernelOperands();
+  auto end = (*this)->operand_end();
+  return OperandRange(end, end);
+}
+ValueRange air::RankOp::getSuccessorInputs(RegionSuccessor succ) {
+  if (succ.getSuccessor() == &getBody())
+    return ValueRange(getKernelArguments());
+  return {};
+}
+void air::RankOp::getSuccessorRegions(
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
+  if (point.isParent())
+    regions.push_back(RegionSuccessor(&getBody()));
+  else
+    regions.push_back(RegionSuccessor::parent());
+}
+void air::RankOp::getRegionInvocationBounds(
+    ArrayRef<Attribute> operands, SmallVectorImpl<InvocationBounds> &bounds) {
+  unsigned sizeStart = getAsyncDependencies().size() + (getUniverse() ? 1 : 0);
+  bounds.push_back(computeHierarchyBounds(operands, sizeStart, getNumDims()));
+}
+
 LogicalResult air::RankOp::verify() {
   // RankOp may be nested inside air.launch (for multi-GPU parallelism),
   // but not inside air.segment, air.herd, or another air.rank.
@@ -1500,6 +1565,30 @@ unsigned air::SegmentOp::getNumDims() {
   return segment_sizes[1];
 }
 
+OperandRange air::SegmentOp::getEntrySuccessorOperands(RegionSuccessor succ) {
+  if (succ.getSuccessor() == &getBody())
+    return getKernelOperands();
+  auto end = (*this)->operand_end();
+  return OperandRange(end, end);
+}
+ValueRange air::SegmentOp::getSuccessorInputs(RegionSuccessor succ) {
+  if (succ.getSuccessor() == &getBody())
+    return ValueRange(getKernelArguments());
+  return {};
+}
+void air::SegmentOp::getSuccessorRegions(
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
+  if (point.isParent())
+    regions.push_back(RegionSuccessor(&getBody()));
+  else
+    regions.push_back(RegionSuccessor::parent());
+}
+void air::SegmentOp::getRegionInvocationBounds(
+    ArrayRef<Attribute> operands, SmallVectorImpl<InvocationBounds> &bounds) {
+  bounds.push_back(computeHierarchyBounds(
+      operands, getAsyncDependencies().size(), getNumDims()));
+}
+
 /// Utility function to verify that all memref.alloc operations within a region
 /// have a memory space at least as local as the specified minimum.
 /// For example, minMemorySpace=L2 means allocations must be L2 or L1;
@@ -1856,7 +1945,7 @@ Value air::HerdOp::getKernelOperand(unsigned i) {
 }
 
 ArrayRef<BlockArgument> air::HerdOp::getKernelArguments() {
-  return getBody().front().getArguments().drop_front(4);
+  return getBody().front().getArguments().drop_front(getNumDims() * 2);
 }
 
 BlockArgument air::HerdOp::getKernelArgument(unsigned i) {
@@ -1870,6 +1959,30 @@ unsigned air::HerdOp::getNumDims() {
   return segment_sizes[1];
 }
 
+OperandRange air::HerdOp::getEntrySuccessorOperands(RegionSuccessor succ) {
+  if (succ.getSuccessor() == &getBody())
+    return getKernelOperands();
+  auto end = (*this)->operand_end();
+  return OperandRange(end, end);
+}
+ValueRange air::HerdOp::getSuccessorInputs(RegionSuccessor succ) {
+  if (succ.getSuccessor() == &getBody())
+    return ValueRange(getKernelArguments());
+  return {};
+}
+void air::HerdOp::getSuccessorRegions(
+    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
+  if (point.isParent())
+    regions.push_back(RegionSuccessor(&getBody()));
+  else
+    regions.push_back(RegionSuccessor::parent());
+}
+void air::HerdOp::getRegionInvocationBounds(
+    ArrayRef<Attribute> operands, SmallVectorImpl<InvocationBounds> &bounds) {
+  bounds.push_back(computeHierarchyBounds(
+      operands, getAsyncDependencies().size(), getNumDims()));
+}
+
 uint64_t air::HerdOp::getNumCols() {
   auto cols = getSizeOperands()[0].getDefiningOp();
   return cast<arith::ConstantIndexOp>(cols).value();
@@ -1901,6 +2014,58 @@ LogicalResult air::ExecuteOp::verify() {
   return success();
 }
 
+// Mirror body effects up to the execute op. Inner Allocate effects on
+// yielded values get retargeted to the outer result so alias analysis
+// can prove distinct execute-yielded memrefs NoAlias. Inner ops with no
+// declared effects fall back to conservative R+W+Free unless upstream
+// can prove them effect-free.
+void air::ExecuteOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  if (getBody().empty())
+    return;
+  llvm::SmallDenseMap<Value, OpResult> yieldToResult;
+  if (auto execTerm = dyn_cast_if_present<air::ExecuteTerminatorOp>(
+          getBody().getTerminator())) {
+    unsigned numYields = execTerm->getNumOperands();
+    unsigned firstYieldResult = getNumResults() - numYields;
+    for (unsigned i = 0; i < numYields; ++i)
+      yieldToResult[execTerm->getOperand(i)] =
+          cast<OpResult>(getResult(firstYieldResult + i));
+  }
+  bool sawUnknownOp = false;
+  getBody().walk([&](Operation *inner) {
+    if (inner == getBody().getTerminator())
+      return;
+    auto effectOp = dyn_cast<MemoryEffectOpInterface>(inner);
+    if (!effectOp) {
+      if (mlir::isMemoryEffectFree(inner))
+        return;
+      sawUnknownOp = true;
+      return;
+    }
+    SmallVector<MemoryEffects::EffectInstance> innerEffects;
+    effectOp.getEffects(innerEffects);
+    for (auto &e : innerEffects) {
+      if (isa<MemoryEffects::Allocate>(e.getEffect())) {
+        Value v = e.getValue();
+        auto it = v ? yieldToResult.find(v) : yieldToResult.end();
+        if (it != yieldToResult.end())
+          effects.emplace_back(e.getEffect(), it->second, e.getResource());
+      } else {
+        effects.emplace_back(e.getEffect(), e.getResource());
+      }
+    }
+  });
+  if (sawUnknownOp) {
+    effects.emplace_back(MemoryEffects::Read::get(),
+                         SideEffects::DefaultResource::get());
+    effects.emplace_back(MemoryEffects::Write::get(),
+                         SideEffects::DefaultResource::get());
+    effects.emplace_back(MemoryEffects::Free::get(),
+                         SideEffects::DefaultResource::get());
+  }
+}
+
 static LogicalResult FoldExecute(air::ExecuteOp op, PatternRewriter &rewriter) {
 
   // if the terminator is the only thing in the ExecuteOp,