[Mlir-commits] [mlir] [milr][gpu] Make barrier elimination address-space aware (PR #178101)

[Jakub Kuderski]

================
@@ -78,14 +78,132 @@ static void addAllValuelessEffects(
   effects.emplace_back(MemoryEffects::Effect::get<MemoryEffects::Free>());
 }
 
+/// Looks through known "view-like" ops to find the base memref.
+static Value getBase(Value v) {
+  while (true) {
+    Operation *definingOp = v.getDefiningOp();
+    if (!definingOp)
+      break;
+
+    bool shouldContinue = TypeSwitch<Operation *, bool>(v.getDefiningOp())
+                              .Case([&](ViewLikeOpInterface op) {
+                                v = op.getViewSource();
+                                return true;
+                              })
+                              .Case([&](memref::TransposeOp op) {
+                                v = op.getIn();
+                                return true;
+                              })
+                              .Default(false);
+    if (!shouldContinue)
+      break;
+  }
+  return v;
+}
+
+/// Returns `true` if accesses to the given memory space could potentially be
+/// fenced by a barrier synchoronizing on the given `fencedAddressSpaces`. If
+/// the set of address spaces is not given, it is equal to all possible address
+/// spaces. Memory spaces that are not `#gpu.address_space` are deemed to
+/// overlap with all GPU address spaces.
+static bool isAddressSpacePotentiallyFenced(
+    Attribute memorySpace,
+    std::optional<ArrayRef<gpu::AddressSpaceAttr>> fencedAddressSpaces) {
+  if (!fencedAddressSpaces)
+    return true;
+
+  auto gpuMemSpace = dyn_cast_if_present<gpu::AddressSpaceAttr>(memorySpace);
+  if (!gpuMemSpace)
+    return true;
+
+  // Check if this GPU address space is in the fenced set.
+  return llvm::is_contained(*fencedAddressSpaces, gpuMemSpace);
+}
+
+/// Succeeds if the effect operates on a memref whose memory  space
+/// could be one of the given fenced address spaces. This will both look at the
+/// address space of the effect's operand and of the view-like operations that
+/// define that memref, so as to inspect any memory-space casts or similar
+/// operations (like amdgpu buffer casts) that may provide more information.
+/// This assumes that directly-conflicting casts (that is, for example, casting
+/// a memref in global memory to make it one in workspace memory) can't happen.
+static LogicalResult effectMightAffectAddressSpaces(
+    const MemoryEffects::EffectInstance &effect,
+    std::optional<ArrayRef<gpu::AddressSpaceAttr>> fencedAddressSpaces) {
+  if (!fencedAddressSpaces)
+    return success();
+
+  Value value = effect.getValue();
+  if (!value)
+    return success();
+
+  auto mightMatch = [&](Value v) {
+    auto memrefType = dyn_cast<BaseMemRefType>(v.getType());
+    if (!memrefType)
+      return true;
+    return isAddressSpacePotentiallyFenced(memrefType.getMemorySpace(),
+                                           fencedAddressSpaces);
+  };
+
+  if (!mightMatch(value))
+    return failure();
+
+  Value base = value;
+  while (auto viewLike = base.getDefiningOp<ViewLikeOpInterface>()) {
+    base = viewLike.getViewSource();
+    // We assume that we won't see directly incompatible casts, like global =>
+    // flat/null => workspace.
+    if (!mightMatch(base))
+      return failure();
+  }
+
+  return success();
+}
+
+/// Returns `true` if `op` is a `BarrierOp` that fences any address spaces that
+/// could overlap with the given fenced address spaces.
+static bool isBarrierWithCommonFencedMemory(
+    Operation *op,
+    std::optional<ArrayRef<gpu::AddressSpaceAttr>> fencedAddressSpaces) {
+  auto barrier = dyn_cast<BarrierOp>(op);
+  if (!barrier)
+    return false;
+
+  std::optional<ArrayAttr> otherFencedSpaces = barrier.getAddressSpaces();
+  // Barriers with unspecified fencing fence everything.
+  if (!otherFencedSpaces)
+    return true;
+  // while barriers that fence nothing can't close off our search.
----------------
kuhar wrote:

```suggestion
  // While barriers that fence nothing can't close off our search.
```

https://github.com/llvm/llvm-project/pull/178101