[ESIMD] Fix invoke_simd calls case with pointer passed to it (#6696)

The helper function created during translation of invoke_simd must
accept a pointer to a function, not a reference to a pointer to a function.
That additional level of indirection is automatically resolved by compiler
for invoke_simd, but needs to be manually resolved/adjusted for the helper
function.

Signed-off-by: Vyacheslav N Klochkov <[email protected]>

Author: v-klochkov

llvm/lib/SYCLLowerIR/LowerInvokeSimd.cpp

@@ -250,14 +250,15 @@ bool collectUsesLookTrhoughMemAndCasts(Value *V,
   for (const Use *U : TmpVUses) {
     User *UU = U->getUser();
     assert(!isCast(UU));
+
     auto *St = dyn_cast<StoreInst>(UU);
 
     if (!St) {
       Uses.insert(U);
       continue;
     }
     // Current user is a store (of V) instruction, see if...
-    assert((V = St->getValueOperand()) &&
+    assert((V == St->getValueOperand()) &&
            "bad V param in collectUsesLookTrhoughMemAndCasts");
     Value *Addr = stripCasts(St->getPointerOperand());
 

sycl/include/sycl/ext/oneapi/experimental/invoke_simd.hpp

@@ -34,13 +34,16 @@
 /// to the SIMD target, as well as process SPMD arguments in the way described
 /// in the specification for `invoke_simd`.
 /// @tparam SpmdRet the return type. Can be `uniform<T>`.
-/// @tparam SimdCallee the type of the SIMD callee function (the "target"). Must
-///  be a function type (not lambda or functor).
-/// @tparam SpmdArgs The original SPMD arguments passed to the invoke_simd.
-template <bool IsFunc, class SpmdRet, class SimdCallee, class... SpmdArgs,
-          class = std::enable_if_t<!IsFunc>>
+/// @tparam HelperFunc the type of SIMD callee helper function. It is needed
+/// to convert the arguments of user's callee function and pass them to call
+/// of user's function.
+/// @tparam UserSimdFuncAndSpmdArgs is the pack that contains the user's SIMD
+/// target function and the original SPMD arguments passed to invoke_simd.
+template <bool IsFunc, class SpmdRet, class HelperFunc,
+          class... UserSimdFuncAndSpmdArgs, class = std::enable_if_t<!IsFunc>>
 SYCL_EXTERNAL __regcall SpmdRet
-__builtin_invoke_simd(SimdCallee target, const void *obj, SpmdArgs... args)
+__builtin_invoke_simd(HelperFunc helper, const void *obj,
+                      UserSimdFuncAndSpmdArgs... args)
 #ifdef __SYCL_DEVICE_ONLY__
     ;
 #else
@@ -51,10 +54,10 @@ __builtin_invoke_simd(SimdCallee target, const void *obj, SpmdArgs... args)
 }
 #endif // __SYCL_DEVICE_ONLY__
 
-template <bool IsFunc, class SpmdRet, class SimdCallee, class... SpmdArgs,
-          class = std::enable_if_t<IsFunc>>
-SYCL_EXTERNAL __regcall SpmdRet __builtin_invoke_simd(SimdCallee target,
-                                                      SpmdArgs... args)
+template <bool IsFunc, class SpmdRet, class HelperFunc,
+          class... UserSimdFuncAndSpmdArgs, class = std::enable_if_t<IsFunc>>
+SYCL_EXTERNAL __regcall SpmdRet
+__builtin_invoke_simd(HelperFunc helper, UserSimdFuncAndSpmdArgs... args)
 #ifdef __SYCL_DEVICE_ONLY__
     ;
 #else
@@ -231,19 +234,21 @@ static constexpr int get_sg_size() {
 // with captures. Note __regcall - this is needed for efficient argument
 // forwarding.
 template <int N, class Callable, class... T>
-SYCL_EXTERNAL __regcall detail::SimdRetType<N, Callable, T...>
-simd_obj_call_helper(const void *obj_ptr,
-                     typename detail::spmd2simd<T, N>::type... simd_args) {
+[[intel::device_indirectly_callable]] SYCL_EXTERNAL __regcall detail::
+    SimdRetType<N, Callable, T...>
+    simd_obj_call_helper(const void *obj_ptr,
+                         typename detail::spmd2simd<T, N>::type... simd_args) {
   auto f =
       *reinterpret_cast<const std::remove_reference_t<Callable> *>(obj_ptr);
   return f(simd_args...);
 }
 
 // This function is a wrapper around a call to a function.
 template <int N, class Callable, class... T>
-SYCL_EXTERNAL __regcall detail::SimdRetType<N, Callable, T...>
-simd_func_call_helper(Callable f,
-                      typename detail::spmd2simd<T, N>::type... simd_args) {
+[[intel::device_indirectly_callable]] SYCL_EXTERNAL __regcall detail::
+    SimdRetType<N, Callable, T...>
+    simd_func_call_helper(Callable f,
+                          typename detail::spmd2simd<T, N>::type... simd_args) {
   return f(simd_args...);
 }
 
@@ -288,6 +293,19 @@ static constexpr bool is_function_ptr_or_ref_v =
 #endif // __INVOKE_SIMD_USE_STD_IS_FUNCTION_WA
     ;
 
+template <typename Callable> struct remove_ref_from_func_ptr_ref_type {
+  using type = Callable;
+};
+
+template <typename Ret, typename... Args>
+struct remove_ref_from_func_ptr_ref_type<Ret(__regcall *&)(Args...)> {
+  using type = Ret(__regcall *)(Args...);
+};
+
+template <typename T>
+using remove_ref_from_func_ptr_ref_type_t =
+    typename remove_ref_from_func_ptr_ref_type<T>::type;
+
 } // namespace detail
 
 // --- The main API
@@ -308,7 +326,8 @@ static constexpr bool is_function_ptr_or_ref_v =
 /// @param args SPMD parameters to the invoked function, which undergo
 ///   transformation before actual passing to the simd function, as described in
 ///   the specification.
-// TODO works only for functions now, enable for other callables.
+// TODO works only for functions and pointers to functions now,
+// enable for lambda functions and functors.
 template <class Callable, class... T>
 __attribute__((always_inline)) auto invoke_simd(sycl::sub_group sg,
                                                 Callable &&f, T... args) {
@@ -321,9 +340,17 @@ __attribute__((always_inline)) auto invoke_simd(sycl::sub_group sg,
   constexpr bool is_function = detail::is_function_ptr_or_ref_v<Callable>;
 
   if constexpr (is_function) {
+    // The variables typed as pointer to a function become lvalue-reference
+    // when passed to invoke_simd() as universal pointers. That creates an
+    // additional indirection, which is resolved automatically by the compiler
+    // for the caller side of __builtin_invoke_simd, but which must be resolved
+    // manually during the creation of simd_func_call_helper.
+    // The class remove_ref_from_func_ptr_ref_type is used removes that
+    // unwanted indirection.
     return __builtin_invoke_simd<true /*function*/, RetSpmd>(
-        detail::simd_func_call_helper<N, Callable, T...>, f,
-        detail::unwrap_uniform<T>::impl(args)...);
+        detail::simd_func_call_helper<
+            N, detail::remove_ref_from_func_ptr_ref_type_t<Callable>, T...>,
+        f, detail::unwrap_uniform<T>::impl(args)...);
   } else {
     // TODO support functors and lambdas which are handled in this branch.
     // The limiting factor for now is that the LLVMIR data flow analysis