1 files changed, 36 insertions, 0 deletions
diff --git a/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp b/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp
index 8062b474..a84d10d 100644
--- a/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp
+++ b/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp
@@ -258,6 +258,39 @@ getAllocEffectFor(Value value,
   return success();
 }
 
+static Operation *isDistinctObjectsOp(Operation *op) {
+  if (op && op->hasTrait<OpTrait::DistinctObjectsTrait>())
+    return op;
+
+  return nullptr;
+}
+
+static Value getDistinctObjectsOperand(Operation *op, Value value) {
+  unsigned argNumber = cast<OpResult>(value).getResultNumber();
+  return op->getOperand(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 = isDistinctObjectsOp(lhs.getDefiningOp());
+  if (lhsOp && getDistinctObjectsOperand(lhsOp, lhs) == rhs)
+    return AliasResult::MustAlias;
+
+  auto rhsOp = isDistinctObjectsOp(rhs.getDefiningOp());
+  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 +322,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));