diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td index 40b7d7e33d5c2..89bd0f103d9f3 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, + DistinctObjectsTrait, DeclareOpInterfaceMethods // ViewLikeOpInterface TODO: ViewLikeOpInterface only supports a single argument ]> { diff --git a/mlir/include/mlir/Interfaces/ViewLikeInterface.h b/mlir/include/mlir/Interfaces/ViewLikeInterface.h index db9c37fc3dc99..c1c2269051efd 100644 --- a/mlir/include/mlir/Interfaces/ViewLikeInterface.h +++ b/mlir/include/mlir/Interfaces/ViewLikeInterface.h @@ -230,6 +230,22 @@ LogicalResult verifyListOfOperandsOrIntegers(Operation *op, StringRef name, ArrayRef attr, ValueRange values); +namespace OpTrait { +/// This trai 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 trait 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. +template +class DistinctObjectsTrait + : public TraitBase {}; +} // namespace OpTrait + } // namespace mlir #endif // MLIR_INTERFACES_VIEWLIKEINTERFACE_H_ diff --git a/mlir/include/mlir/Interfaces/ViewLikeInterface.td b/mlir/include/mlir/Interfaces/ViewLikeInterface.td index ed213bfdae337..131c1a0d92b24 100644 --- a/mlir/include/mlir/Interfaces/ViewLikeInterface.td +++ b/mlir/include/mlir/Interfaces/ViewLikeInterface.td @@ -414,4 +414,16 @@ def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface }]; } +// This trai 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 trait 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. +def DistinctObjectsTrait : NativeOpTrait<"DistinctObjectsTrait">; + #endif // MLIR_INTERFACES_VIEWLIKEINTERFACE diff --git a/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp b/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp index 8062b474539fd..a84d10d5d609d 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()) + return op; + + return nullptr; +} + +static Value getDistinctObjectsOperand(Operation *op, Value value) { + unsigned argNumber = cast(value).getResultNumber(); + return op->getOperand(argNumber); +} + +static std::optional 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 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, %arg1: memref) attributes {test.ptr = "func"} { + %0, %1 = memref.distinct_objects %arg, %arg1 {test.ptr = "distinct"} : memref, memref + return +}