From d42cb02448bbf4e4c51fc6e49f11d03b0276d521 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Mon, 19 Dec 2022 22:26:07 +0100 Subject: [PATCH] [mlir] Make `LocalAliasAnalysis` extesible 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. Differential Revision: https://reviews.llvm.org/D140348 --- .../AliasAnalysis/LocalAliasAnalysis.h | 4 ++ .../AliasAnalysis/LocalAliasAnalysis.cpp | 2 +- .../test-alias-analysis-extending.mlir | 15 +++++ mlir/test/lib/Analysis/TestAliasAnalysis.cpp | 66 ++++++++++++++++++- 4 files changed, 85 insertions(+), 2 deletions(-) create mode 100644 mlir/test/Analysis/test-alias-analysis-extending.mlir diff --git a/mlir/include/mlir/Analysis/AliasAnalysis/LocalAliasAnalysis.h b/mlir/include/mlir/Analysis/AliasAnalysis/LocalAliasAnalysis.h index 9669a3c13b24..06012403cf29 100644 --- a/mlir/include/mlir/Analysis/AliasAnalysis/LocalAliasAnalysis.h +++ b/mlir/include/mlir/Analysis/AliasAnalysis/LocalAliasAnalysis.h @@ -28,6 +28,10 @@ public: /// 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 diff --git a/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp b/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp index 1b7dee919753..73ddd81aa12b 100644 --- a/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp +++ b/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp @@ -246,7 +246,7 @@ getAllocEffectFor(Value value, } /// 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; diff --git a/mlir/test/Analysis/test-alias-analysis-extending.mlir b/mlir/test/Analysis/test-alias-analysis-extending.mlir new file mode 100644 index 000000000000..20f2a7a3de17 --- /dev/null +++ b/mlir/test/Analysis/test-alias-analysis-extending.mlir @@ -0,0 +1,15 @@ +// 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, %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 +} diff --git a/mlir/test/lib/Analysis/TestAliasAnalysis.cpp b/mlir/test/lib/Analysis/TestAliasAnalysis.cpp index 04b2bc3906a8..b563be4a57d4 100644 --- a/mlir/test/lib/Analysis/TestAliasAnalysis.cpp +++ b/mlir/test/lib/Analysis/TestAliasAnalysis.cpp @@ -13,6 +13,8 @@ #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; @@ -148,6 +150,67 @@ struct TestAliasAnalysisModRefPass }; } // namespace +//===----------------------------------------------------------------------===// +// Testing LocalAliasAnalysis extending +//===----------------------------------------------------------------------===// + +/// Check if value is function argument. +static bool isFuncArg(Value val) { + auto blockArg = val.dyn_cast(); + if (!blockArg) + return false; + + return mlir::isa_and_nonnull( + blockArg.getOwner()->getParentOp()); +} + +/// Check if value has "restrict" attribute. Value must be a function argument. +static bool isRestrict(Value val) { + auto blockArg = val.cast(); + auto func = + mlir::cast(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> { + 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 //===----------------------------------------------------------------------===// @@ -155,8 +218,9 @@ struct TestAliasAnalysisModRefPass namespace mlir { namespace test { void registerTestAliasAnalysisPass() { - PassRegistration(); + PassRegistration(); PassRegistration(); + PassRegistration(); } } // namespace test } // namespace mlir