[mlir] LocalAliasAnalysis: Allow adding restrict to function arguments

[rriddle]

I have somewhat of a concern that this is very adhoc. Having seen the pain that LLVM has gone through to support restrict, I'd like to see some thought out discussion before adding any kind of support.

[Hardcode84]

I have somewhat of a concern that this is very adhoc. Having seen the pain that LLVM has gone through to support restrict, I'd like to see some thought out discussion before adding any kind of support.

This is very adhoc, yes :). But properly modelling restrict in MLIR includes discussion of properly modelling pointers in general, which is probably a discussion for many months. Another approach instead of adding adhoc attributes can be an ability to extend existing LocalAliasAnalysis implementation (e.g. make aliasImpl virtual). But in this case we have an issue with default LocalAliasAnalysis is always added and we have to hope both implementations will agree. This will work for this specific case, but I don't know how it will work in general.

mlir/include/mlir/Analysis/AliasAnalysis/LocalAliasAnalysis.h

 /// checks to see if they alias.
 class LocalAliasAnalysis {
 public:
   /// Given two values, return their aliasing behavior.
   AliasResult alias(Value lhs, Value rhs);
 
   /// 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
 
 #endif // MLIR_ANALYSIS_ALIASANALYSIS_LOCALALIASANALYSIS_H_

mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp

getAllocEffectFor(Value value, Optional<MemoryEffects::EffectInstance> &effect,
   // TODO: Here we could look at the users to see if the resource is either
   // freed on all paths within the region, or is just not captured by anything.
   // For now assume allocation scope to the function scope (we don't care if
   // pointer escape outside function).
   allocScopeOp = op->getParentOfType<FunctionOpInterface>();
   return success();
 }
 
+static bool isFuncArg(Value val) {
+  auto blockArg = val.dyn_cast<BlockArgument>();
+  if (!blockArg)
+    return false;
+
+  return mlir::isa<FunctionOpInterface>(blockArg.getOwner()->getParentOp());
+}
+
+static bool isRestrict(Value val) {
+  auto blockArg = val.cast<BlockArgument>();
+  auto func =
+      mlir::cast<FunctionOpInterface>(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<MemoryEffects::EffectInstance> lhsAlloc, rhsAlloc;
 
   // Handle the case where lhs is a constant.
   Attribute lhsAttr, rhsAttr;
   if (matchPattern(lhs, m_Constant(&lhsAttr))) {
     // TODO: This is overly conservative. Two matching constants don't
     // necessarily map to the same address. For example, if the two values
if (isa<MemoryEffects::Read>(effect.getEffect())) {
       assert(isa<MemoryEffects::Write>(effect.getEffect()));
       result = result.merge(ModRefResult::getMod());
     }
     if (result.isModAndRef())
       break;
   }
   return result;
 }
+
+StringRef LocalAliasAnalysis::getRestrictArgName() {
+  return "local_alias_analysis.restrict";
+}

mlir/test/Analysis/test-alias-analysis.mlir

func.func @constants(%arg: memref<2xf32>) attributes {test.ptr = "func"} {
   %1 = memref.alloc() {test.ptr = "alloc_1"} : memref<8x64xf32>
 
   %c0 = arith.constant {test.ptr = "constant_1"} 0 : index
   %c0_2 = arith.constant {test.ptr = "constant_2"} 0 : index
   %c1 = arith.constant {test.ptr = "constant_3"} 1 : index
 
   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<?xf32>, %arg1: memref<?xf32> {local_alias_analysis.restrict}) attributes {test.ptr = "func"} {
+  %0 = memref.subview %arg[0][2][1] {test.ptr = "view1"} : memref<?xf32> to memref<2xf32>
+  %1 = memref.subview %arg1[0][2][1] {test.ptr = "view2"} : memref<?xf32> to memref<2xf32>
+  return
+}