[PATCH] D109376: [OpenMP] Check OpenMP assumptions on call-sites as well

[Joseph Huber via Phabricator via llvm-commits]

jhuber6 created this revision.
jhuber6 added a reviewer: jdoerfert.
Herald added subscribers: ormris, dexonsmith, guansong, hiraditya, yaxunl.
jhuber6 requested review of this revision.
Herald added subscribers: llvm-commits, sstefan1.
Herald added a project: LLVM.

This patch adds functionality to check assumption attributes on call
sites as well.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D109376

Files:
  llvm/include/llvm/IR/Assumptions.h
  llvm/lib/IR/Assumptions.cpp
  llvm/lib/Transforms/IPO/OpenMPOpt.cpp


Index: llvm/lib/Transforms/IPO/OpenMPOpt.cpp
===================================================================
--- llvm/lib/Transforms/IPO/OpenMPOpt.cpp
+++ llvm/lib/Transforms/IPO/OpenMPOpt.cpp
@@ -3705,13 +3705,25 @@
     Function *Callee = getAssociatedFunction();
 
     // Helper to lookup an assumption string.
-    auto HasAssumption = [](Function *Fn, StringRef AssumptionStr) {
-      return Fn && hasAssumption(*Fn, AssumptionStr);
+    auto HasAssumption = [](Value *V, StringRef AssumptionStr) {
+      if (V) {
+        if (Function *Fn = dyn_cast<Function>(V))
+          return hasAssumption(*Fn, AssumptionStr);
+        if (CallBase *CB = dyn_cast<CallBase>(V))
+          return hasAssumption(*CB, AssumptionStr);
+      }
+      return false;
     };
 
     // Check for SPMD-mode assumptions.
-    if (HasAssumption(Callee, "ompx_spmd_amenable"))
+    if (HasAssumption(&CB, "ompx_spmd_amenable")) {
       SPMDCompatibilityTracker.indicateOptimisticFixpoint();
+      indicateOptimisticFixpoint();
+    }
+
+    if (HasAssumption(Callee, "ompx_spmd_amenable")) {
+      SPMDCompatibilityTracker.indicateOptimisticFixpoint();
+    }
 
     // First weed out calls we do not care about, that is readonly/readnone
     // calls, intrinsics, and "no_openmp" calls. Neither of these can reach a
Index: llvm/lib/IR/Assumptions.cpp
===================================================================
--- llvm/lib/IR/Assumptions.cpp
+++ llvm/lib/IR/Assumptions.cpp
@@ -11,6 +11,7 @@
 #include "llvm/IR/Assumptions.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/InstrTypes.h"
 
 using namespace llvm;
 
@@ -29,6 +30,21 @@
   });
 }
 
+bool llvm::hasAssumption(CallBase &CB,
+                         const KnownAssumptionString &AssumptionStr) {
+  const Attribute &A = CB.getFnAttr(AssumptionAttrKey);
+  if (!A.isValid())
+    return false;
+  assert(A.isStringAttribute() && "Expected a string attribute!");
+
+  SmallVector<StringRef, 8> Strings;
+  A.getValueAsString().split(Strings, ",");
+
+  return llvm::any_of(Strings, [=](StringRef Assumption) {
+    return Assumption == AssumptionStr;
+  });
+}
+
 StringSet<> llvm::KnownAssumptionStrings({
     "omp_no_openmp",          // OpenMP 5.1
     "omp_no_openmp_routines", // OpenMP 5.1
Index: llvm/include/llvm/IR/Assumptions.h
===================================================================
--- llvm/include/llvm/IR/Assumptions.h
+++ llvm/include/llvm/IR/Assumptions.h
@@ -21,6 +21,7 @@
 namespace llvm {
 
 class Function;
+class CallBase;
 
 /// The key we use for assumption attributes.
 constexpr StringRef AssumptionAttrKey = "llvm.assume";
@@ -45,6 +46,9 @@
 /// Return true if \p F has the assumption \p AssumptionStr attached.
 bool hasAssumption(Function &F, const KnownAssumptionString &AssumptionStr);
 
+/// Return true if \p CB has the assumption \p AssumptionStr attached.
+bool hasAssumption(CallBase &CB, const KnownAssumptionString &AssumptionStr);
+
 } // namespace llvm
 
 #endif


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D109376.371117.patch
Type: text/x-patch
Size: 3016 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20210907/ed2bb352/attachment.bin>