[mlir][LocalAliasAnalysis] Check for memref.distinct_objects in LocalAliasAnalysis

[Hardcode84]

Continuation of #156913, I'm planning to use it to infer LLVM alias scope attributes.

• Introduce DistinctObjectsInterface so analysis won't need to depend on memref dialect directly.
• Check the interface in LocalAliasAnalysis::aliasImpl.

[llvmbot]

@llvm/pr-subscribers-mlir-memref

@llvm/pr-subscribers-mlir

Author: Ivan Butygin (Hardcode84)

Changes

Continuation of #156913, I'm planning to use it to infer LLVM alias scope attributes.

• Introduce DistinctObjectsInterface so analysis won't need to depend on memref dialect directly.
• Check the interface in LocalAliasAnalysis::aliasImpl.

---

Full diff: https://github.com/llvm/llvm-project/pull/161533.diff

4 Files Affected:

• (modified) mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td (+1)
• (modified) mlir/include/mlir/Interfaces/ViewLikeInterface.td (+40)
• (modified) mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp (+30)
• (modified) mlir/test/Analysis/test-alias-analysis.mlir (+16)

diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index d4d67bfb278d5..e79dbaa0b172c 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -184,6 +184,7 @@ def AssumeAlignmentOp : MemRef_Op<"assume_alignment", [
 
 def DistinctObjectsOp : MemRef_Op<"distinct_objects", [
       Pure,
+      DistinctObjectsInterface,
       DeclareOpInterfaceMethods<InferTypeOpInterface>
       // ViewLikeOpInterface TODO: ViewLikeOpInterface only supports a single argument
     ]> {
diff --git a/mlir/include/mlir/Interfaces/ViewLikeInterface.td b/mlir/include/mlir/Interfaces/ViewLikeInterface.td
index ed213bfdae337..2c3adb4a7fd6a 100644
--- a/mlir/include/mlir/Interfaces/ViewLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/ViewLikeInterface.td
@@ -414,4 +414,44 @@ def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface
   }];
 }
 
+def DistinctObjectsInterface : OpInterface<"DistinctObjectsInterface"> {
+  let description = [{
+    This intefaces indicates that pointer-like objects (such as memrefs) returned
+    from this operation will never alias with each other. This provides a
+    guarantee to optimization passes that accesses through different results
+    of this operation can be safely reordered, as they will never reference
+    overlapping memory locations.
+
+    Operations with this interface take multiple pointer-like operands
+    and return the same operands with additional non-aliasing guarantees.
+    If the access to the results of this operation aliases at runtime, the
+    behavior of such access is undefined.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{ Return input pointer-like objects. }],
+      /*retTy=*/"::mlir::ValueRange",
+      /*methodName=*/"getDistinctOperands",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return $_op->getOperands();
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{ Return result pointer-like objects. }],
+      /*retTy=*/"::mlir::ValueRange",
+      /*methodName=*/"getDistinctResults",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return $_op->getResults();
+      }]
+    >
+  ];
+}
+
+
 #endif // MLIR_INTERFACES_VIEWLIKEINTERFACE
diff --git a/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp b/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp
index 8062b474539fd..15686fc6d4031 100644
--- a/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp
+++ b/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp
@@ -258,6 +258,33 @@ getAllocEffectFor(Value value,
   return success();
 }
 
+static Value getDistinctObjectsOperand(DistinctObjectsInterface op,
+                                       Value value) {
+  unsigned argNumber = cast<OpResult>(value).getResultNumber();
+  return op.getDistinctOperands()[argNumber];
+}
+
+static std::optional<AliasResult> checkDistinctObjects(Value lhs, Value rhs) {
+  // We should already checked that lhs and rhs are different.
+  assert(lhs != rhs && "lhs and rhs must be different");
+
+  // Result and corresponding operand must alias.
+  auto lhsOp = lhs.getDefiningOp<DistinctObjectsInterface>();
+  if (lhsOp && getDistinctObjectsOperand(lhsOp, lhs) == rhs)
+    return AliasResult::MustAlias;
+
+  auto rhsOp = rhs.getDefiningOp<DistinctObjectsInterface>();
+  if (rhsOp && getDistinctObjectsOperand(rhsOp, rhs) == lhs)
+    return AliasResult::MustAlias;
+
+  // If two different values come from the same `DistinctObjects` operation,
+  // they don't alias.
+  if (lhsOp && lhsOp == rhsOp)
+    return AliasResult::NoAlias;
+
+  return std::nullopt;
+}
+
 /// Given the two values, return their aliasing behavior.
 AliasResult LocalAliasAnalysis::aliasImpl(Value lhs, Value rhs) {
   if (lhs == rhs)
@@ -289,6 +316,9 @@ AliasResult LocalAliasAnalysis::aliasImpl(Value lhs, Value rhs) {
                : AliasResult::MayAlias;
   }
 
+  if (std::optional<AliasResult> result = checkDistinctObjects(lhs, rhs))
+    return *result;
+
   // Otherwise, neither of the values are constant so check to see if either has
   // an allocation effect.
   bool lhsHasAlloc = succeeded(getAllocEffectFor(lhs, lhsAlloc, lhsAllocScope));
diff --git a/mlir/test/Analysis/test-alias-analysis.mlir b/mlir/test/Analysis/test-alias-analysis.mlir
index 8cbee61c78b45..d71adee05c7a3 100644
--- a/mlir/test/Analysis/test-alias-analysis.mlir
+++ b/mlir/test/Analysis/test-alias-analysis.mlir
@@ -256,3 +256,19 @@ func.func @constants(%arg: memref<2xf32>) attributes {test.ptr = "func"} {
 
   return
 }
+
+// -----
+
+// CHECK-LABEL: Testing : "distinct_objects"
+// CHECK-DAG: func.region0#0 <-> func.region0#1: MayAlias
+
+// CHECK-DAG: distinct#0 <-> distinct#1: NoAlias
+// CHECK-DAG: distinct#0 <-> func.region0#0: MustAlias
+// CHECK-DAG: distinct#1 <-> func.region0#0: MayAlias
+// CHECK-DAG: distinct#0 <-> func.region0#1: MayAlias
+// CHECK-DAG: distinct#1 <-> func.region0#1: MustAlias
+
+func.func @distinct_objects(%arg: memref<?xf32>, %arg1: memref<?xf32>) attributes {test.ptr = "func"} {
+  %0, %1 = memref.distinct_objects %arg, %arg1 {test.ptr = "distinct"} : memref<?xf32>, memref<?xf32>
+  return
+}

[krzysz00]

The code looks fine, just a comment about the interface phrasing

[krzysz00]

Can we phrase the interface such that it 's not limited to memref? It sounds like this entire mechanism - and, to my knowledge, the alias analysis being modified - isn't actually limited to memref and could work on things like ptr.ptr.

[Hardcode84]

Wait, I specifically tried to phrase it pointer-like objects so ptr.ptr or similar should qualify. I can remove (such as memref) if you think it will be better.