[Mlir-commits] [mlir] [mlir][LocalAliasAnalysis] Check for `memref.distinct_objects` in `LocalAliasAnalysis` (PR #161533)

[Ivan Butygin]

https://github.com/Hardcode84 updated https://github.com/llvm/llvm-project/pull/161533

>From 33b5b6692c680310ee5ab972908a03c79fa2c277 Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Wed, 24 Sep 2025 12:34:15 +0200
Subject: [PATCH] LocalAliasAnalysis

---
 .../mlir/Dialect/MemRef/IR/MemRefOps.td       |  1 +
 .../mlir/Interfaces/ViewLikeInterface.td      | 40 +++++++++++++++++++
 .../AliasAnalysis/LocalAliasAnalysis.cpp      | 30 ++++++++++++++
 mlir/test/Analysis/test-alias-analysis.mlir   | 16 ++++++++
 4 files changed, 87 insertions(+)

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
+}