LocalAliasAnalysis

Author: Hardcode84

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
     ]> {

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

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));

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