Check for function calls and local as global variables

Author: PietroGhg

llvm/include/llvm/SYCLLowerIR/CheckNDRangeSYCLNativeCPU.h

@@ -6,10 +6,11 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// A transformation pass that:
-// * Handles the kernel calling convention and attributes.
-// * Materializes the spirv builtins so that they can be handled by the host
-//   runtime.
+// Checks if the kernel uses features from nd_item such as:
+// * local id
+// * local range
+// * local memory
+// * work group barrier
 //===----------------------------------------------------------------------===//
 
 #pragma once

llvm/include/llvm/SYCLLowerIR/UtilsSYCLNativeCPU.h

@@ -40,6 +40,8 @@ inline bool isSYCLNativeCPU(const Module &M) {
   return M.getModuleFlag("is-native-cpu") != nullptr;
 }
 
+constexpr unsigned SyclNativeCpuLocalAS = 3;
+
 } // namespace utils
 } // namespace sycl
 } // namespace llvm

llvm/lib/SYCLNativeCPUUtils/CheckNDRangeSYCLNativeCPU.cpp

@@ -14,16 +14,22 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/SYCLLowerIR/CheckNDRangeSYCLNativeCPU.h"
+#include "llvm/ADT/PriorityWorklist.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/IR/CallingConv.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/Function.h"
 #include "llvm/IR/InstrTypes.h"
+#include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Metadata.h"
+#include "llvm/SYCLLowerIR/UtilsSYCLNativeCPU.h"
+#include "llvm/Support/Casting.h"
 
 using namespace llvm;
 
-static std::array<const char *, 13> NdFunctions{
+static std::array<const char *, 13> NdBuiltins{
     "_Z23__spirv_WorkgroupSize_xv",     "_Z23__spirv_WorkgroupSize_yv",
     "_Z23__spirv_WorkgroupSize_zv",     "_Z23__spirv_NumWorkgroups_xv",
     "_Z23__spirv_NumWorkgroups_yv",     "_Z23__spirv_NumWorkgroups_zv",
@@ -42,6 +48,55 @@ static void addNDRangeMetadata(Function &F, bool Value) {
 PreservedAnalyses
 CheckNDRangeSYCLNativeCPUPass::run(Module &M, ModuleAnalysisManager &MAM) {
   bool ModuleChanged = false;
+  SmallPtrSet<Function *, 5> NdFuncs; // Functions that use NDRange features
+  SmallPtrSet<Function *, 5> Visited;
+  SmallPriorityWorklist<Function *, 5> WorkList;
+
+  // Add builtins to the set of functions that may use NDRange features
+  for (auto &FName : NdBuiltins) {
+    auto F = M.getFunction(FName);
+    if (F == nullptr)
+      continue;
+    WorkList.insert(F);
+    NdFuncs.insert(F);
+  }
+
+  // Add users of local AS global var to the set of functions that may use
+  // NDRange features
+  for (auto &GV : M.globals()) {
+    if (GV.getAddressSpace() != sycl::utils::SyclNativeCpuLocalAS)
+      continue;
+
+    for (auto U : GV.users()) {
+      if (auto I = dyn_cast<Instruction>(U)) {
+        auto F = I->getFunction();
+        if (F != nullptr && NdFuncs.insert(F).second) {
+          WorkList.insert(F);
+          NdFuncs.insert(F);
+        }
+      }
+    }
+  }
+
+  // Traverse the use chain to find Functions that may use NDRange features
+  // (or, recursively, Functions that call Functions that may use NDRange
+  // features)
+  while (!WorkList.empty()) {
+    auto F = WorkList.pop_back_val();
+
+    for (User *U : F->users()) {
+      if (auto CI = dyn_cast<CallInst>(U)) {
+        auto Caller = CI->getFunction();
+        if (!Caller)
+          continue;
+        if (!Visited.contains(Caller)) {
+          WorkList.insert(Caller);
+          NdFuncs.insert(Caller);
+        }
+      }
+    }
+    Visited.insert(F);
+  }
 
   for (auto &F : M) {
     if (F.getCallingConv() == llvm::CallingConv::SPIR_KERNEL) {
@@ -55,23 +110,11 @@ CheckNDRangeSYCLNativeCPUPass::run(Module &M, ModuleAnalysisManager &MAM) {
         }
       }
 
-      for (auto &BB : F) {
-        for (auto &I : BB) {
-          if (auto CI = dyn_cast<CallInst>(&I)) {
-            auto CalleeName = CI->getCalledFunction()->getName();
-            if (std::find(NdFunctions.begin(), NdFunctions.end(), CalleeName) !=
-                NdFunctions.end()) {
-              IsNDRange = true;
-              break;
-            }
-          }
-        }
-        if (IsNDRange) {
-          break;
-        }
-      }
+      // Check if the kernel calls one of the ND Range builtins
+      IsNDRange |= NdFuncs.contains(&F);
 
       addNDRangeMetadata(F, IsNDRange);
+      ModuleChanged = true;
     }
   }
   return ModuleChanged ? PreservedAnalyses::none() : PreservedAnalyses::all();

sycl/test/check_device_code/native_cpu/nd_range_attr.cpp

@@ -1,12 +1,21 @@
 // REQUIRES: native_cpu
 // RUN: %clangxx -fsycl-device-only -fsycl-targets=native_cpu %s -S -o - | FileCheck %s
+// RUN: %clangxx -fsycl-device-only -O0 -fsycl-targets=native_cpu %s -S -o - | FileCheck %s
 
 #include <sycl/sycl.hpp>
 using namespace sycl;
 
 class Test;
 class Test2;
 class Test3;
+class Test4;
+
+template <typename AccT>
+void use_local_acc(nd_item<1> it, AccT &acc,
+                   const local_accessor<int, 1> &local_acc) {
+  local_acc[it.get_local_id()[0]] = 1;
+  acc[it.get_local_id()[0]] = local_acc[0];
+}
 
 int main() {
   sycl::queue deviceQueue;
@@ -22,18 +31,25 @@ int main() {
     deviceQueue.submit([&](handler &h) {
       auto acc = buf.template get_access<access::mode::write>(h);
       local_accessor<int, 1> local_acc(1, h);
-      h.parallel_for<Test2>(r, [=](nd_item<1> it) {
-        local_acc[0] = 1;
-        acc[0] = local_acc[0];
-      });
+      h.parallel_for<Test2>(
+          r, [=](nd_item<1> it) { use_local_acc(it, acc, local_acc); });
     });
     // CHECK-DAG: @_ZTS5Test2({{.*}} !is_nd_range [[MDID:![0-9]*]]
   }
+
   deviceQueue.submit([&](handler &h) {
     h.parallel_for<Test3>(1, [=](item<1> it) { it.get_id(); });
   });
   // CHECK-DAG: @_ZTS5Test3({{.*}} !is_nd_range [[MDNOT:![0-9]*]]
 
+  buffer<int, 1> buf(&res, 1);
+  deviceQueue.submit([&](sycl::handler &cgh) {
+    auto acc = sycl::accessor(buf, cgh, sycl::write_only);
+    cgh.parallel_for_work_group<Test4>(
+        range<1>(1), range<1>(1),
+        [=](auto group) { acc[group.get_group_id()] = 42; });
+  });
+  // CHECK-DAG: @_ZTS5Test4({{.*}} !is_nd_range [[MDID:![0-9]*]]
 }
 
 //CHECK:[[MDID]] = !{i1 true}