[mlir] Make `LocalAliasAnalysis` extesible

[Hardcode84]

Summary

This is an alternative to https://reviews.llvm.org/D138761 . Instead of adding ad-hoc attributes to existing LocalAliasAnalysis, expose aliasImpl method so user can override it.

[bondhugula]

mlir/test/lib/Analysis/TestAliasAnalysis.cpp

162

getParentOp can return null.

isa_and_non_null.

168

Likewise.

221–224

Add in sorted order.

[Hardcode84]

Add null check

[Hardcode84]

mlir/test/lib/Analysis/TestAliasAnalysis.cpp

168

We can only get there if previous check succeeded so we will never get null.

[bondhugula]

LGTM. On a note unrelated to this PR but on this topic, besides aliasImpl, I feel even collectUnderlyingAddressValues should be exposed. Currently, the alias analysis framework can only answer "can two values alias", but not "provide me all the values that could potentially alias with a given Value" (at least locally).

[bondhugula]

mlir/test/lib/Analysis/TestAliasAnalysis.cpp

166

Please add doc comment to this one and the one above although they may appear trivial now.

191

Doc comment.

[rriddle]

Looks reasonable to me as a step forward.

[Hardcode84]

Add comments

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);
+
+protected:
+  /// Given the two values, return their aliasing behavior.
+  virtual AliasResult aliasImpl(Value lhs, Value rhs);
 };
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_ALIASANALYSIS_LOCALALIASANALYSIS_H_

mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp

getAllocEffectFor(Value value,
   // 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();
 }
 
 /// Given the two values, return their aliasing behavior.
-static AliasResult aliasImpl(Value lhs, Value rhs) {
+AliasResult LocalAliasAnalysis::aliasImpl(Value lhs, Value rhs) {
   if (lhs == rhs)
     return AliasResult::MustAlias;
   Operation *lhsAllocScope = nullptr, *rhsAllocScope = nullptr;
   std::optional<MemoryEffects::EffectInstance> lhsAlloc, rhsAlloc;
 
   // Handle the case where lhs is a constant.
   Attribute lhsAttr, rhsAttr;
   if (matchPattern(lhs, m_Constant(&lhsAttr))) {

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

+// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-alias-analysis-extending))' -split-input-file -allow-unregistered-dialect 2>&1 | FileCheck %s
+
+// 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
+}

mlir/test/lib/Analysis/TestAliasAnalysis.cpp

 //===- TestAliasAnalysis.cpp - Test alias analysis results ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains test passes for constructing and testing alias analysis
 // results.
 //
 //===----------------------------------------------------------------------===//
 
 #include "TestAliasAnalysis.h"
 #include "mlir/Analysis/AliasAnalysis.h"
+#include "mlir/Analysis/AliasAnalysis/LocalAliasAnalysis.h"
+#include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 /// Print a value that is used as an operand of an alias query.
 static void printAliasOperand(Operation *op) {
   llvm::errs() << op->getAttrOfType<StringAttr>("test.ptr").getValue();
 }
struct TestAliasAnalysisModRefPass
   void runOnOperation() override {
     AliasAnalysis &aliasAnalysis = getAnalysis<AliasAnalysis>();
     runAliasAnalysisOnOperation(getOperation(), aliasAnalysis);
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
+// Testing LocalAliasAnalysis extending
+//===----------------------------------------------------------------------===//
+
+/// Check if value is function argument.
+static bool isFuncArg(Value val) {
+  auto blockArg = val.dyn_cast<BlockArgument>();
+  if (!blockArg)
+    return false;
+
+  return mlir::isa_and_nonnull<FunctionOpInterface>(
+      blockArg.getOwner()->getParentOp());
+}
+
+/// Check if value has "restrict" attribute. Value must be a function argument.
+static bool isRestrict(Value val) {
+  auto blockArg = val.cast<BlockArgument>();
+  auto func =
+      mlir::cast<FunctionOpInterface>(blockArg.getOwner()->getParentOp());
+  return !!func.getArgAttr(blockArg.getArgNumber(),
+                           "local_alias_analysis.restrict");
+}
+
+namespace {
+/// LocalAliasAnalysis extended to support "restrict" attreibute.
+class LocalAliasAnalysisRestrict : public LocalAliasAnalysis {
+protected:
+  AliasResult aliasImpl(Value lhs, Value rhs) override {
+    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;
+
+    return LocalAliasAnalysis::aliasImpl(lhs, rhs);
+  }
+};
+
+/// This pass tests adding additional analysis impls to the AliasAnalysis.
+struct TestAliasAnalysisExtendingPass
+    : public test::TestAliasAnalysisBase,
+      PassWrapper<TestAliasAnalysisExtendingPass, OperationPass<>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAliasAnalysisExtendingPass)
+
+  StringRef getArgument() const final {
+    return "test-alias-analysis-extending";
+  }
+  StringRef getDescription() const final {
+    return "Test alias analysis extending.";
+  }
+  void runOnOperation() override {
+    AliasAnalysis aliasAnalysis(getOperation());
+    aliasAnalysis.addAnalysisImplementation(LocalAliasAnalysisRestrict());
+    runAliasAnalysisOnOperation(getOperation(), aliasAnalysis);
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
 // Pass Registration
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace test {
 void registerTestAliasAnalysisPass() {
-  PassRegistration<TestAliasAnalysisPass>();
+  PassRegistration<TestAliasAnalysisExtendingPass>();
   PassRegistration<TestAliasAnalysisModRefPass>();
+  PassRegistration<TestAliasAnalysisPass>();
 }
 } // namespace test
 } // namespace mlir