[WIP][BACKEND] Generalize the MemBar to consider cross-CTA ops

The semantics here are that it's the user's/compiler's responsability to
add the relevant synchronisation if they reuse the same shmem buffer,
but otherwise the compiler will do so.

Author: lezcano

include/triton/Analysis/Membar.h

@@ -2,6 +2,8 @@
 #define TRITON_ANALYSIS_MEMBAR_H
 
 #include "Allocation.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "triton/Analysis/Utility.h"
 
 #include <set>
 
@@ -15,7 +17,83 @@ class OpBuilder;
 using MembarFilterFn = std::function<bool(Operation *, Operation *)>;
 
 struct BlockInfo {
-  using IntervalMapT = std::map<Interval<size_t>, std::set<Operation *>>;
+  // UFDS to represent cross-CTA reads/writes
+  struct UFDS {
+    SmallVector<unsigned> parent;
+    SmallVector<unsigned> rank;
+    // Invariant: At the root of a class, minRep[i] is the smallest element in
+    // the class
+    SmallVector<unsigned> minRep;
+
+    UFDS() = default;
+    explicit UFDS(unsigned n) : rank(n, 0), minRep(n) {
+      assert(llvm::isPowerOf2_32(n) && n != 0);
+      parent = llvm::to_vector(llvm::seq(n));
+      for (unsigned i = 0; i < n; ++i)
+        minRep[i] = i;
+    }
+
+    unsigned find(unsigned x) const {
+      unsigned p = parent[x];
+      while (p != parent[p])
+        p = parent[p];
+      return p;
+    }
+
+    unsigned findMin(unsigned x) const { return minRep[find(x)]; }
+
+    void unite(unsigned x, unsigned y) {
+      x = find(x);
+      y = find(y);
+      if (x == y)
+        return;
+
+      if (rank[x] < rank[y])
+        std::swap(x, y);
+
+      parent[y] = x;
+      minRep[x] = std::min(minRep[x], minRep[y]);
+
+      if (rank[x] == rank[y])
+        ++rank[x];
+    }
+
+    UFDS join(const UFDS &other) const {
+      // Transitive closure of two UFDS
+      UFDS result = *this;
+      for (unsigned i = 0; i < size(); ++i)
+        result.unite(i, other.find(i));
+      return result;
+    }
+
+    SmallVector<unsigned> canonical() const {
+      SmallVector<unsigned> reps(size());
+      for (unsigned i = 0; i < size(); ++i)
+        reps[i] = findMin(i);
+      return reps;
+    }
+
+    bool isDistributed() const { return *this != UFDS(parent.size()); }
+
+    bool operator<(const UFDS &other) const {
+      return canonical() < other.canonical();
+    }
+    bool operator==(const UFDS &other) const {
+      return canonical() == other.canonical();
+    }
+    bool operator!=(const UFDS &other) const { return !(*this == other); }
+
+    void print(raw_ostream &os) const {
+      os << "UFDS(";
+      llvm::interleaveComma(canonical(), os, [&](unsigned x) { os << x; });
+      os << ")";
+    }
+
+    size_t size() const { return parent.size(); }
+  };
+
+  using IntervalMapT =
+      std::map<std::pair<Interval<size_t>, UFDS>, std::set<Operation *>>;
 
   IntervalMapT syncReadIntervals;
   IntervalMapT syncWriteIntervals;
@@ -24,48 +102,82 @@ struct BlockInfo {
 
   /// Unions two BlockInfo objects.
   BlockInfo &join(const BlockInfo &other) {
-    for (auto &interval : other.syncReadIntervals)
-      syncReadIntervals[interval.first].insert(interval.second.begin(),
-                                               interval.second.end());
-    for (auto &interval : other.syncWriteIntervals)
-      syncWriteIntervals[interval.first].insert(interval.second.begin(),
-                                                interval.second.end());
+    // We don't fold the intervals (we could tho)
+    for (auto &[key, ops] : other.syncReadIntervals)
+      syncReadIntervals[key].insert(ops.begin(), ops.end());
+    for (auto &[key, ops] : other.syncWriteIntervals)
+      syncWriteIntervals[key].insert(ops.begin(), ops.end());
     return *this;
   }
 
   void dump() {
     auto &err = llvm::errs();
+
+    auto printKey = [&](const std::pair<Interval<size_t>, UFDS> &key) {
+      const auto &[interval, ufds] = key;
+      err << "    [" << interval.start() << ", " << interval.end() << "] ";
+      if (ufds.isDistributed())
+        ufds.print(err);
+      else if (ufds.size() == 1)
+        err << " (CTA local)";
+    };
     err << "Block Interval:\n";
     err << "  Read Intervals:\n";
-    for (auto &[interval, ops] : syncReadIntervals) {
-      err << "    [" << interval.start() << ", " << interval.end() << "] ";
+    for (auto &[key, ops] : syncReadIntervals) {
+      printKey(key);
       for (auto &op : ops)
         err << op->getName() << " ";
       err << "\n";
     }
     err << "  Write Intervals:\n";
-    for (auto &[interval, ops] : syncWriteIntervals) {
-      err << "    [" << interval.start() << ", " << interval.end() << "] ";
+    for (auto &[key, ops] : syncWriteIntervals) {
+      printKey(key);
       for (auto &op : ops)
         err << op->getName() << " ";
       err << "\n";
     }
   }
 
   /// Returns true if intervals in two BlockInfo objects are intersected.
-  bool isIntersected(const BlockInfo &other, MembarFilterFn filter) const {
-    return /*RAW*/ isIntersected(syncWriteIntervals, other.syncReadIntervals,
-                                 filter) ||
-           /*WAR*/
-           isIntersected(syncReadIntervals, other.syncWriteIntervals, filter) ||
-           /*WAW*/
-           isIntersected(syncWriteIntervals, other.syncWriteIntervals, filter);
+  std::optional<UFDS> isIntersected(const BlockInfo &other,
+                                    MembarFilterFn filter) const {
+    auto raw =
+        isIntersected(syncWriteIntervals, other.syncReadIntervals, filter);
+    auto war =
+        isIntersected(syncReadIntervals, other.syncWriteIntervals, filter);
+    auto waw =
+        isIntersected(syncWriteIntervals, other.syncWriteIntervals, filter);
+    auto maybeJoin = [](const std::optional<UFDS> &lhs,
+                        const std::optional<UFDS> &rhs) -> std::optional<UFDS> {
+      if (!lhs.has_value())
+        return rhs;
+      if (!rhs.has_value())
+        return lhs;
+      return lhs.value().join(rhs.value());
+    };
+    return maybeJoin(raw, maybeJoin(war, waw));
   }
 
   /// Clears the intervals because a barrier is inserted.
-  void sync() {
-    syncReadIntervals.clear();
-    syncWriteIntervals.clear();
+  /// If `cluster` is true, the barrier synchronizes all CTAs in the cluster and
+  /// we can drop every pending dependency. Otherwise only CTA-local
+  /// dependencies are cleared; distributed ones remain until a cluster barrier
+  /// is observed.
+  void sync(bool cluster) {
+    if (cluster) {
+      syncReadIntervals.clear();
+      syncWriteIntervals.clear();
+      return;
+    } else {
+      auto eraseNotDistributed = [](auto &map) {
+        for (auto &[key, _] : llvm::make_early_inc_range(map)) {
+          if (!key.second.isDistributed())
+            map.erase(key);
+        }
+      };
+      eraseNotDistributed(syncReadIntervals);
+      eraseNotDistributed(syncWriteIntervals);
+    }
   }
 
   /// Compares two BlockInfo objects.
@@ -77,18 +189,79 @@ struct BlockInfo {
   bool operator!=(const BlockInfo &other) const { return !(*this == other); }
 
 private:
-  bool isIntersected(const IntervalMapT &lhsIntervalSet,
-                     const IntervalMapT &rhsIntervalSet,
-                     MembarFilterFn filter) const {
-    for (auto &lhs : lhsIntervalSet)
-      for (auto &rhs : rhsIntervalSet)
-        if (lhs.first.intersects(rhs.first))
-          for (auto lhsOp : lhs.second)
-            for (auto rhsOp : rhs.second)
-              if (!filter || !filter(lhsOp, rhsOp))
-                return true;
-
-    return false;
+  static bool haveSameAlloc(Operation *lhs, Operation *rhs) {
+    auto collectAllocOwners = [](Operation *op) -> std::optional<Operation *> {
+      SmallVector<Operation *> allocs;
+      if (auto mei = dyn_cast<MemoryEffectOpInterface>(op)) {
+        SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>> effs;
+        mei.getEffects(effs);
+        for (auto &eff : effs) {
+          if (eff.getResource() != triton::gpu::SharedMemory::get())
+            continue;
+          if (auto v = eff.getValue()) {
+            // Hacky way to skip barriers...
+            if (cast<triton::gpu::MemDescType>(v.getType()).getNumElements() ==
+                1)
+              continue;
+            auto alloc = findShmemAlloc(v);
+            assert(alloc && "Expected to find a shmem alloc");
+            allocs.push_back(alloc.getOperation());
+          }
+        }
+      }
+      assert(allocs.size() <= 1 && "Expected to find exactly one shmem alloc");
+      if (allocs.empty()) {
+        return std::nullopt;
+      } else {
+        return allocs[0];
+      }
+    };
+
+    auto lhsAllocs = collectAllocOwners(lhs);
+    auto rhsAllocs = collectAllocOwners(rhs);
+    return lhsAllocs.has_value() && rhsAllocs.has_value() &&
+           lhsAllocs.value() == rhsAllocs.value();
+  }
+
+  std::optional<UFDS> isIntersected(const IntervalMapT &lhsIntervalSet,
+                                    const IntervalMapT &rhsIntervalSet,
+                                    MembarFilterFn filter) const {
+    // They intersect whenever the intervals intersect. If they do, collect the
+    // union of CTA sets for any op pair that is not filtered out and does not
+    // share the exact same explicit shared value.
+    std::optional<UFDS> ret = std::nullopt;
+    for (const auto &[lhsKey, lhsOps] : lhsIntervalSet) {
+      const auto &[intervalLhs, ctasLhs] = lhsKey;
+      for (const auto &[rhsKey, rhsOps] : rhsIntervalSet) {
+        const auto &[intervalRhs, ctasRhs] = rhsKey;
+        if (!intervalLhs.intersects(intervalRhs))
+          continue;
+
+        auto joined = ctasLhs.join(ctasRhs);
+        bool needsBarrier = llvm::any_of(lhsOps, [&, rhsOpsPtr = &rhsOps](
+                                                     const auto &lhsOp) {
+          return llvm::any_of(*rhsOpsPtr, [&](const auto &rhsOp) {
+            // Skip if filtered or both ops touch the same explicit shared
+            // allocation (same local_alloc).
+            return !((filter && filter(lhsOp, rhsOp)) ||
+                     (joined.isDistributed() && haveSameAlloc(lhsOp, rhsOp)));
+          });
+        });
+        if (!needsBarrier)
+          continue;
+
+        if (!ret.has_value()) {
+          ret = joined;
+        } else {
+          ret = ret->join(joined);
+        }
+        // Single CTA case, we can early exit
+        if (ret->size() == 1) {
+          return ret;
+        }
+      }
+    }
+    return ret;
   }
 };
 
@@ -170,7 +343,8 @@ class MembarAnalysis : public MembarOrFenceAnalysis {
                       FuncBlockInfoMapT *funcBlockInfoMap,
                       OpBuilder *builder) override;
 
-  void insertBarrier(Operation *operation, OpBuilder *builder);
+  void insertBarrier(Operation *operation, OpBuilder *builder,
+                     const BlockInfo::UFDS &ctaClasses);
 };
 
 /// Postorder traversal on the callgraph to insert membar instructions

include/triton/Analysis/Utility.h

@@ -415,6 +415,11 @@ std::unique_ptr<DataFlowSolver> createDataFlowSolver();
 bool isCvtWarpSync(const triton::LinearLayout &srcLayout,
                    const triton::LinearLayout &dstLayout);
 
+/// Returns the defining ttg.local_alloc (shared-memory allocation) for a
+/// memdesc-derived value by walking through memdesc/convert layout ops and
+/// loop-carried block arguments. Returns null if none is found.
+triton::gpu::LocalAllocOp findShmemAlloc(Value operand);
+
 } // namespace mlir
 
 #endif // TRITON_ANALYSIS_UTILITY_H

include/triton/Dialect/TritonGPU/Transforms/Utility.h

@@ -221,11 +221,6 @@ Operation *convertDistributedOpEncoding(Attribute encoding, Operation *op);
 // Returns the original memory allocation for a memdesc value
 triton::gpu::LocalAllocOp findShmemAlloc(Value operand);
 
-// Returns MMAs inside a for loop that are multi-buffered for pipeline analysis
-SmallVector<Operation *>
-getMMAsWithMultiBufferredOperands(scf::ForOp forOp,
-                                  SmallVector<Operation *> &mmaOps);
-
 // Given a list of ops, find the naerest common dominator of all ops or return
 // null if one could not be found. The ops are allowed to be in different
 // regions. The result op is not necessarily one of the ops in the list.

include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td

@@ -293,7 +293,8 @@ def TTNG_AsyncTMACopyGlobalToLocalOp : TTNG_Op<"async_tma_copy_global_to_local">
     I1:$pred,
     DefaultValuedAttr<TT_CacheModifierAttr, "triton::CacheModifier::NONE">:$cache,
     DefaultValuedAttr<TT_EvictionPolicyAttr, "triton::EvictionPolicy::NORMAL">:$evict,
-    DefaultValuedAttr<BoolAttr, "false">:$isVolatile
+    DefaultValuedAttr<BoolAttr, "false">:$isVolatile,
+    DefaultValuedAttr<BoolAttr, "false">:$multicast
   );
 
   let assemblyFormat = [{