[SYCL][ESIMD]Limit propagation of ESIMD attributes

llvm/include/llvm/SYCLLowerIR/ESIMD/ESIMDUtils.h

@@ -9,12 +9,18 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Demangle/ItaniumDemangle.h"
 #include "llvm/IR/Function.h"
 
 namespace llvm {
 namespace esimd {
 
 constexpr char ESIMD_MARKER_MD[] = "sycl_explicit_simd";
+// This is the prefixes of the names generated from
+// sycl/ext/oneapi/experimental/invoke_simd.hpp::__builtin_invoke_simd
+// overloads instantiations:
+constexpr char INVOKE_SIMD_PREF[] = "_Z33__regcall3____builtin_invoke_simd";
 
 // Tells whether given function is a ESIMD kernel.
 bool isESIMDKernel(const Function &F);
@@ -68,5 +74,40 @@ void collectUsesLookThroughCastsAndZeroGEPs(const Value *V,
 /// in it. Returns nullptr if failed to do so.
 Type *getVectorTyOrNull(StructType *STy);
 
+/// Tracks the use of a function pointer being stored in a memory.
+/// Returns false if the function pointer is used as an argument for invoke_simd
+/// function call, true otherwise.
+bool filterFunctionPointer(Value *address);
+
+// Simplest possible implementation of an allocator for the Itanium demangler
+class SimpleAllocator {
+protected:
+  SmallVector<void *, 128> Ptrs;
+
+public:
+  void reset() {
+    for (void *Ptr : Ptrs) {
+      // Destructors are not called, but that is OK for the
+      // itanium_demangle::Node subclasses
+      std::free(Ptr);
+    }
+    Ptrs.resize(0);
+  }
+
+  template <typename T, typename... Args> T *makeNode(Args &&...args) {
+    void *Ptr = std::calloc(1, sizeof(T));
+    Ptrs.push_back(Ptr);
+    return new (Ptr) T(std::forward<Args>(args)...);
+  }
+
+  void *allocateNodeArray(size_t sz) {
+    void *Ptr = std::calloc(sz, sizeof(itanium_demangle::Node *));
+    Ptrs.push_back(Ptr);
+    return Ptr;
+  }
+
+  ~SimpleAllocator() { reset(); }
+};
+
 } // namespace esimd
 } // namespace llvm

llvm/include/llvm/SYCLLowerIR/SYCLUtils.h

@@ -18,6 +18,7 @@ namespace llvm {
 namespace sycl {
 namespace utils {
 using CallGraphNodeAction = std::function<void(Function *)>;
+using CallGraphFunctionFilter = std::function<bool(Value *)>;
 
 // Traverses call graph starting from given function up the call chain applying
 // given action to each function met on the way. If \c ErrorOnNonCallUse
@@ -26,24 +27,33 @@ using CallGraphNodeAction = std::function<void(Function *)>;
 // call graph as if the use was a call.
 // Functions which are part of the visited set ('Visited' parameter) are not
 // traversed.
-void traverseCallgraphUp(llvm::Function *F, CallGraphNodeAction NodeF,
-                         SmallPtrSetImpl<Function *> &Visited,
-                         bool ErrorOnNonCallUse);
+void traverseCallgraphUp(
+    llvm::Function *F, CallGraphNodeAction NodeF,
+    SmallPtrSetImpl<Function *> &Visited, bool ErrorOnNonCallUse,
+    const CallGraphFunctionFilter &functionFilter = [](Value *) {
+      return true;
+    });
 
 template <class CallGraphNodeActionF>
-void traverseCallgraphUp(Function *F, CallGraphNodeActionF ActionF,
-                         SmallPtrSetImpl<Function *> &Visited,
-                         bool ErrorOnNonCallUse) {
+void traverseCallgraphUp(
+    Function *F, CallGraphNodeActionF ActionF,
+    SmallPtrSetImpl<Function *> &Visited, bool ErrorOnNonCallUse,
+    const CallGraphFunctionFilter &functionFilter = [](Value *) {
+      return true;
+    }) {
   traverseCallgraphUp(F, CallGraphNodeAction(ActionF), Visited,
-                      ErrorOnNonCallUse);
+                      ErrorOnNonCallUse, functionFilter);
 }
 
 template <class CallGraphNodeActionF>
-void traverseCallgraphUp(Function *F, CallGraphNodeActionF ActionF,
-                         bool ErrorOnNonCallUse = true) {
+void traverseCallgraphUp(
+    Function *F, CallGraphNodeActionF ActionF, bool ErrorOnNonCallUse = true,
+    const CallGraphFunctionFilter &functionFilter = [](Value *) {
+      return true;
+    }) {
   SmallPtrSet<Function *, 32> Visited;
   traverseCallgraphUp(F, CallGraphNodeAction(ActionF), Visited,
-                      ErrorOnNonCallUse);
+                      ErrorOnNonCallUse, functionFilter);
 }
 } // namespace utils
 } // namespace sycl

llvm/lib/SYCLLowerIR/ESIMD/ESIMDUtils.cpp

@@ -123,5 +123,62 @@ Type *getVectorTyOrNull(StructType *STy) {
   return Res;
 }
 
+bool isInvokeSimdBuiltinCall(const CallInst *CI) {
+  Function *F = CI->getCalledFunction();
+
+  if (F && F->getName().startswith(INVOKE_SIMD_PREF)) {
+    return true;
+  }
+  return false;
+}
+
+// Tracks the use of a function pointer being stored in a memory.
+// Returns false if the function pointer is used as an argument for invoke_simd
+// function call, true otherwise.
+bool filterFunctionPointer(Value *address) {
+  if (address == nullptr) {
+    return true;
+  }
+
+  SmallPtrSet<const Use *, 4> Uses;
+  collectUsesLookThroughCasts(address, Uses);
+
+  for (const Use *U : Uses) {
+    Value *V = U->getUser();
+
+    if (auto *StI = dyn_cast<StoreInst>(V)) {
+      if (U != &StI->getOperandUse(StoreInst::getPointerOperandIndex())) {
+        // this is double indirection - not supported
+        return false;
+      }
+      V = stripCasts(StI->getPointerOperand());
+      if (!isa<AllocaInst>(V)) {
+        return false; // unsupported case of data flow through non-local memory
+      }
+
+      if (auto *LI = dyn_cast<LoadInst>(V)) {
+        // A value loaded from another address is stored at this address -
+        // recurse into the other address
+        if (!filterFunctionPointer(LI->getPointerOperand())) {
+          return false;
+        }
+      }
+    } else if (const auto *CI = dyn_cast<CallInst>(V)) {
+      // if __builtin_invoke_simd uses the pointer, do not traverse the function
+      if (isInvokeSimdBuiltinCall(CI)) {
+        return false;
+      }
+    } else if (isa<LoadInst>(V)) {
+      if (!filterFunctionPointer(V)) {
+        return false;
+      }
+    } else {
+      return false;
+    }
+  }
+
+  return true;
+}
+
 } // namespace esimd
 } // namespace llvm

llvm/lib/SYCLLowerIR/ESIMD/ESIMDVerifier.cpp

@@ -19,11 +19,13 @@
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
+#include "llvm/SYCLLowerIR/ESIMD/ESIMDUtils.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Regex.h"
 
 using namespace llvm;
+using namespace llvm::esimd;
 namespace id = itanium_demangle;
 
 #define DEBUG_TYPE "esimd-verifier"
@@ -64,36 +66,6 @@ static const char *LegalSYCLFunctionsInStatelessMode[] = {
 
 namespace {
 
-// Simplest possible implementation of an allocator for the Itanium demangler
-class SimpleAllocator {
-protected:
-  SmallVector<void *, 128> Ptrs;
-
-public:
-  void reset() {
-    for (void *Ptr : Ptrs) {
-      // Destructors are not called, but that is OK for the
-      // itanium_demangle::Node subclasses
-      std::free(Ptr);
-    }
-    Ptrs.resize(0);
-  }
-
-  template <typename T, typename... Args> T *makeNode(Args &&...args) {
-    void *Ptr = std::calloc(1, sizeof(T));
-    Ptrs.push_back(Ptr);
-    return new (Ptr) T(std::forward<Args>(args)...);
-  }
-
-  void *allocateNodeArray(size_t sz) {
-    void *Ptr = std::calloc(sz, sizeof(id::Node *));
-    Ptrs.push_back(Ptr);
-    return Ptr;
-  }
-
-  ~SimpleAllocator() { reset(); }
-};
-
 class ESIMDVerifierImpl {
   const Module &M;
   bool ForceStatelessMem;

llvm/lib/SYCLLowerIR/ESIMD/LowerESIMD.cpp

@@ -41,6 +41,7 @@
 
 using namespace llvm;
 namespace id = itanium_demangle;
+using namespace llvm::esimd;
 
 #define DEBUG_TYPE "lower-esimd"
 
@@ -678,36 +679,6 @@ static const ESIMDIntrinDesc &getIntrinDesc(StringRef SrcSpelling) {
   return It->second;
 }
 
-// Simplest possible implementation of an allocator for the Itanium demangler
-class SimpleAllocator {
-protected:
-  SmallVector<void *, 128> Ptrs;
-
-public:
-  void reset() {
-    for (void *Ptr : Ptrs) {
-      // Destructors are not called, but that is OK for the
-      // itanium_demangle::Node subclasses
-      std::free(Ptr);
-    }
-    Ptrs.resize(0);
-  }
-
-  template <typename T, typename... Args> T *makeNode(Args &&...args) {
-    void *Ptr = std::calloc(1, sizeof(T));
-    Ptrs.push_back(Ptr);
-    return new (Ptr) T(std::forward<Args>(args)...);
-  }
-
-  void *allocateNodeArray(size_t sz) {
-    void *Ptr = std::calloc(sz, sizeof(id::Node *));
-    Ptrs.push_back(Ptr);
-    return Ptr;
-  }
-
-  ~SimpleAllocator() { reset(); }
-};
-
 Type *parsePrimitiveTypeString(StringRef TyStr, LLVMContext &Ctx) {
   return llvm::StringSwitch<Type *>(TyStr)
       .Case("bool", IntegerType::getInt1Ty(Ctx))

llvm/lib/SYCLLowerIR/ESIMD/LowerESIMDKernelAttrs.cpp

@@ -8,16 +8,18 @@
 // Finds and adds  sycl_explicit_simd attributes to wrapper functions that wrap
 // ESIMD kernel functions
 
+#include "llvm/Demangle/Demangle.h"
+#include "llvm/Demangle/ItaniumDemangle.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Pass.h"
 #include "llvm/SYCLLowerIR/ESIMD/ESIMDUtils.h"
 #include "llvm/SYCLLowerIR/ESIMD/LowerESIMD.h"
 #include "llvm/SYCLLowerIR/SYCLUtils.h"
 
-#include "llvm/IR/Module.h"
-#include "llvm/Pass.h"
-
 #define DEBUG_TYPE "LowerESIMDKernelAttrs"
 
 using namespace llvm;
+using namespace llvm::esimd;
 
 namespace llvm {
 PreservedAnalyses
@@ -37,7 +39,7 @@ SYCLFixupESIMDKernelWrapperMDPass::run(Module &M, ModuleAnalysisManager &MAM) {
               Modified = true;
             }
           },
-          false);
+          false, filterFunctionPointer);
     }
   }
   return Modified ? PreservedAnalyses::none() : PreservedAnalyses::all();

llvm/lib/SYCLLowerIR/LowerInvokeSimd.cpp

@@ -39,7 +39,6 @@ using namespace llvm::esimd;
 
 namespace {
 
-constexpr char ESIMD_MARKER_MD[] = "sycl_explicit_simd";
 constexpr char REQD_SUB_GROUP_SIZE_MD[] = "intel_reqd_sub_group_size";
 
 class SYCLLowerInvokeSimdLegacyPass : public ModulePass {
@@ -73,10 +72,6 @@ ModulePass *llvm::createSYCLLowerInvokeSimdPass() {
 
 namespace {
 // TODO support lambda and functor overloads
-// This is the prefixes of the names generated from
-// sycl/ext/oneapi/experimental/invoke_simd.hpp::__builtin_invoke_simd
-// overloads instantiations:
-constexpr char INVOKE_SIMD_PREF[] = "_Z33__regcall3____builtin_invoke_simd";
 
 using ValueSetImpl = SmallPtrSetImpl<Value *>;
 using ValueSet = SmallPtrSet<Value *, 4>;

llvm/lib/SYCLLowerIR/SYCLUtils.cpp

@@ -1,5 +1,4 @@
-//===------------ SYCLUtils.cpp - SYCL utility functions
-//------------------===//
+//===------------ SYCLUtils.cpp - SYCL utility functions ------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -10,13 +9,18 @@
 //===----------------------------------------------------------------------===//
 #include "llvm/SYCLLowerIR/SYCLUtils.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/SYCLLowerIR/ESIMD/ESIMDUtils.h"
 
 namespace llvm {
 namespace sycl {
 namespace utils {
+
+using namespace llvm::esimd;
+
 void traverseCallgraphUp(llvm::Function *F, CallGraphNodeAction ActionF,
                          SmallPtrSetImpl<Function *> &FunctionsVisited,
-                         bool ErrorOnNonCallUse) {
+                         bool ErrorOnNonCallUse,
+                         const CallGraphFunctionFilter &functionFilter) {
   SmallVector<Function *, 32> Worklist;
 
   if (FunctionsVisited.count(F) == 0)
@@ -43,6 +47,13 @@ void traverseCallgraphUp(llvm::Function *F, CallGraphNodeAction ActionF,
         } else {
           // ... non-call is OK - add using function to the worklist
           if (auto *I = dyn_cast<Instruction>(FCall)) {
+            if (auto *SI = dyn_cast<StoreInst>(I)) {
+              Value *addr = SI->getPointerOperand();
+              if (!functionFilter(addr)) {
+                continue;
+              }
+            }
+
             auto UseF = I->getFunction();
 
             if (FunctionsVisited.count(UseF) == 0) {