diff --git a/mlir/include/mlir/Analysis/AliasAnalysis/LocalAliasAnalysis.h b/mlir/include/mlir/Analysis/AliasAnalysis/LocalAliasAnalysis.h --- a/mlir/include/mlir/Analysis/AliasAnalysis/LocalAliasAnalysis.h +++ b/mlir/include/mlir/Analysis/AliasAnalysis/LocalAliasAnalysis.h @@ -28,6 +28,10 @@ /// Return the modify-reference behavior of `op` on `location`. ModRefResult getModRef(Operation *op, Value location); + + /// For function arguments marked with this attribute LocalAliasAnalysis will + /// assume no aliasing with any other function argument. + static StringRef getRestrictArgName(); }; } // namespace mlir diff --git a/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp b/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp --- a/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp +++ b/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp @@ -243,10 +243,33 @@ return success(); } +static bool isFuncArg(Value val) { + auto blockArg = val.dyn_cast(); + if (!blockArg) + return false; + + return mlir::isa(blockArg.getOwner()->getParentOp()); +} + +static bool isRestrict(Value val) { + auto blockArg = val.cast(); + auto func = + mlir::cast(blockArg.getOwner()->getParentOp()); + return !!func.getArgAttr(blockArg.getArgNumber(), + LocalAliasAnalysis::getRestrictArgName()); +} + /// Given the two values, return their aliasing behavior. static AliasResult aliasImpl(Value lhs, Value rhs) { if (lhs == rhs) return AliasResult::MustAlias; + + // Assume no aliasing if both values are function arguments and any of them + // have restrict attr. + if (isFuncArg(lhs) && isFuncArg(rhs)) + if (isRestrict(lhs) || isRestrict(rhs)) + return AliasResult::NoAlias; + Operation *lhsAllocScope = nullptr, *rhsAllocScope = nullptr; Optional lhsAlloc, rhsAlloc; @@ -394,3 +417,7 @@ } return result; } + +StringRef LocalAliasAnalysis::getRestrictArgName() { + return "local_alias_analysis.restrict"; +} diff --git a/mlir/test/Analysis/test-alias-analysis.mlir b/mlir/test/Analysis/test-alias-analysis.mlir --- a/mlir/test/Analysis/test-alias-analysis.mlir +++ b/mlir/test/Analysis/test-alias-analysis.mlir @@ -256,3 +256,19 @@ return } + +// ----- + +// CHECK-LABEL: Testing : "restrict" +// CHECK-DAG: func.region0#0 <-> func.region0#1: NoAlias + +// CHECK-DAG: view1#0 <-> view2#0: NoAlias +// CHECK-DAG: view1#0 <-> func.region0#0: MustAlias +// CHECK-DAG: view1#0 <-> func.region0#1: NoAlias +// CHECK-DAG: view2#0 <-> func.region0#0: NoAlias +// CHECK-DAG: view2#0 <-> func.region0#1: MustAlias +func.func @restrict(%arg: memref, %arg1: memref {local_alias_analysis.restrict}) attributes {test.ptr = "func"} { + %0 = memref.subview %arg[0][2][1] {test.ptr = "view1"} : memref to memref<2xf32> + %1 = memref.subview %arg1[0][2][1] {test.ptr = "view2"} : memref to memref<2xf32> + return +}