[SYCL] Sub-group load/store for raw pointers

sycl/include/CL/__spirv/spirv_ops.hpp

@@ -169,6 +169,50 @@ __SPIRV_ATOMICS(__SPIRV_ATOMIC_UNSIGNED, unsigned long)
 __SPIRV_ATOMICS(__SPIRV_ATOMIC_UNSIGNED, unsigned long long)
 __SPIRV_ATOMICS(__SPIRV_ATOMIC_MINMAX, Min)
 __SPIRV_ATOMICS(__SPIRV_ATOMIC_MINMAX, Max)
+__SYCL_CONVERGENT__ extern SYCL_EXTERNAL __attribute__((opencl_generic)) void *
+__spirv_PtrCastToGeneric(const void *Ptr) noexcept;
+
+__SYCL_CONVERGENT__ extern SYCL_EXTERNAL __attribute__((opencl_global)) void *
+__spirv_GenericCastToPtrExplicit_ToGlobal(const void *Ptr,
+                                          __spv::StorageClass::Flag S) noexcept;
+
+__SYCL_CONVERGENT__ extern SYCL_EXTERNAL __attribute__((opencl_local)) void *
+__spirv_GenericCastToPtrExplicit_ToLocal(const void *Ptr,
+                                         __spv::StorageClass::Flag S) noexcept;
+
+__SYCL_CONVERGENT__ extern SYCL_EXTERNAL __attribute__((opencl_private)) void *
+__spirv_GenericCastToPtrExplicit_ToPrivate(
+    const void *Ptr, __spv::StorageClass::Flag S) noexcept;
+
+template <typename dataT>
+__SYCL_CONVERGENT__ extern SYCL_EXTERNAL __attribute__((opencl_generic)) dataT *
+__spirv_PtrCastToGeneric(const void *Ptr) noexcept {
+  return (__attribute__((opencl_generic)) dataT *)__spirv_PtrCastToGeneric(Ptr);
+}
+
+template <typename dataT>
+__SYCL_CONVERGENT__ extern SYCL_EXTERNAL __attribute__((opencl_global)) dataT *
+__spirv_GenericCastToPtrExplicit_ToGlobal(
+    const void *Ptr, __spv::StorageClass::Flag S) noexcept {
+  return (__attribute__((opencl_global))
+          dataT *)__spirv_GenericCastToPtrExplicit_ToGlobal(Ptr, S);
+}
+
+template <typename dataT>
+__SYCL_CONVERGENT__ extern SYCL_EXTERNAL __attribute__((opencl_local)) dataT *
+__spirv_GenericCastToPtrExplicit_ToLocal(const void *Ptr,
+                                         __spv::StorageClass::Flag S) noexcept {
+  return (__attribute__((opencl_local))
+          dataT *)__spirv_GenericCastToPtrExplicit_ToLocal(Ptr, S);
+}
+
+template <typename dataT>
+__SYCL_CONVERGENT__ extern SYCL_EXTERNAL __attribute__((opencl_private)) dataT *
+__spirv_GenericCastToPtrExplicit_ToPrivate(
+    const void *Ptr, __spv::StorageClass::Flag S) noexcept {
+  return (__attribute__((opencl_private))
+          dataT *)__spirv_GenericCastToPtrExplicit_ToPrivate(Ptr, S);
+}
 
 template <typename dataT>
 __SYCL_CONVERGENT__ extern SYCL_EXTERNAL dataT

sycl/include/CL/__spirv/spirv_types.hpp

@@ -36,6 +36,46 @@ struct Scope {
   Flag flag_value;
 };
 
+struct StorageClass {
+  enum Flag : uint32_t {
+    UniformConstant = 0,
+    Input = 1,
+    sUniform = 2,
+    Output = 3,
+    Workgroup = 4,
+    CrossWorkgroup = 5,
+    Private = 6,
+    Function = 7,
+    Generic = 8,
+    PushConstant = 9,
+    AtomicCounter = 10,
+    Image = 11,
+    StorageBuffer = 12,
+    CallableDataKHR = 5328,
+    CallableDataNV = 5328,
+    IncomingCallableDataKHR = 5329,
+    IncomingCallableDataNV = 5329,
+    RayPayloadKHR = 5338,
+    RayPayloadNV = 5338,
+    HitAttributeKHR = 5339,
+    HitAttributeNV = 5339,
+    IncomingRayPayloadKHR = 5342,
+    IncomingRayPayloadNV = 5342,
+    RecordBufferKHR = 5343,
+    ShaderRecordBufferNV = 5343,
+    PhysicalStorageBuffer = 5349,
+    PhysicalStorageBufferEXT = 5349,
+    CodeSectionINTEL = 5605,
+    DeviceOnlyINTEL = 5936,
+    HostOnlyINTEL = 5937,
+    Max = 0x7fffffff,
+    CapabilityUSMStorageClassesINTEL = 5935,
+  };
+  constexpr StorageClass(Flag flag) : flag_value(flag) {}
+  constexpr operator uint32_t() const { return flag_value; }
+  Flag flag_value;
+};
+
 struct MemorySemanticsMask {
 
   enum Flag : uint32_t {

sycl/include/CL/sycl/ONEAPI/sub_group.hpp

@@ -45,6 +45,11 @@ using AcceptableForLocalLoadStore =
     bool_constant<!std::is_same<void, SelectBlockT<T>>::value &&
                   Space == access::address_space::local_space>;
 
+template <typename T, access::address_space Space>
+using AcceptableForPrivateLoadStore =
+    bool_constant<!std::is_same<void, SelectBlockT<T>>::value &&
+                  Space == access::address_space::private_space>;
+
 #ifdef __SYCL_DEVICE_ONLY__
 template <typename T, access::address_space Space>
 T load(const multi_ptr<T, Space> src) {
@@ -224,6 +229,50 @@ struct sub_group {
 
   /* --- sub_group load/stores --- */
   /* these can map to SIMD or block read/write hardware where available */
+#ifdef __SYCL_DEVICE_ONLY__
+  // Method for decorated pointer
+  template <typename T>
+  detail::enable_if_t<
+      !std::is_same<typename detail::remove_AS<T>::type, T>::value, T>
+  load(T *src) const {
+    return load(sycl::multi_ptr<typename detail::remove_AS<T>::type,
+                                sycl::detail::deduce_AS<T>::value>(
+        (typename detail::remove_AS<T>::type *)src));
+  }
+
+#ifndef SYCL_USE_DECORATED_REF
+  // Method for raw pointer
+  template <typename T>
+  detail::enable_if_t<
+      std::is_same<typename detail::remove_AS<T>::type, T>::value, T>
+  load(T *src) const {
+
+    auto p = __spirv_GenericCastToPtrExplicit_ToPrivate<T>(
+        src, __spv::StorageClass::Function);
+    if (p)
+      return load(p);
+
+    auto l = __spirv_GenericCastToPtrExplicit_ToLocal<T>(
+        src, __spv::StorageClass::Workgroup);
+    if (l)
+      return load(l);
+
+    auto g = __spirv_GenericCastToPtrExplicit_ToGlobal<T>(
+        src, __spv::StorageClass::CrossWorkgroup);
+    if (g)
+      return load(g);
+
+    // Fallback for other address spaces to be mapped to global
+    return load(__spirv_PtrCastToGeneric<T>(src));
+  }
+#endif // SYCL_USE_DECORATED_REF
+#else  //__SYCL_DEVICE_ONLY__
+  template <typename T> T load(T *src) const {
+    (void)src;
+    throw runtime_error("Sub-groups are not supported on host device.",
+                        PI_INVALID_DEVICE);
+  }
+#endif //__SYCL_DEVICE_ONLY__
 
   template <typename T, access::address_space Space>
   sycl::detail::enable_if_t<
@@ -255,6 +304,20 @@ struct sub_group {
 #endif
   }
 
+  template <typename T, access::address_space Space>
+  sycl::detail::enable_if_t<
+      sycl::detail::sub_group::AcceptableForPrivateLoadStore<T, Space>::value,
+      T>
+  load(const multi_ptr<T, Space> src) const {
+#ifdef __SYCL_DEVICE_ONLY__
+    return src.get()[get_local_id()[0]];
+#else
+    (void)src;
+    throw runtime_error("Sub-groups are not supported on host device.",
+                        PI_INVALID_DEVICE);
+#endif
+  }
+
   template <int N, typename T, access::address_space Space>
   sycl::detail::enable_if_t<
       sycl::detail::sub_group::AcceptableForGlobalLoadStore<T, Space>::value &&
@@ -315,6 +378,59 @@ struct sub_group {
 #endif
   }
 
+#ifdef __SYCL_DEVICE_ONLY__
+  // Method for decorated pointer
+  template <typename T>
+  detail::enable_if_t<
+      !std::is_same<typename detail::remove_AS<T>::type, T>::value>
+  store(T *dst, const typename detail::remove_AS<T>::type &x) const {
+    store(sycl::multi_ptr<typename detail::remove_AS<T>::type,
+                          sycl::detail::deduce_AS<T>::value>(
+              (typename detail::remove_AS<T>::type *)dst),
+          x);
+  }
+
+#ifndef SYCL_USE_DECORATED_REF
+  // Method for raw pointer
+  template <typename T>
+  detail::enable_if_t<
+      std::is_same<typename detail::remove_AS<T>::type, T>::value>
+  store(T *dst, const typename detail::remove_AS<T>::type &x) const {
+
+    auto p = __spirv_GenericCastToPtrExplicit_ToPrivate<T>(
+        dst, __spv::StorageClass::Function);
+    if (p) {
+      store(p, x);
+      return;
+    }
+
+    auto l = __spirv_GenericCastToPtrExplicit_ToLocal<T>(
+        dst, __spv::StorageClass::Workgroup);
+    if (l) {
+      store(l, x);
+      return;
+    }
+
+    auto g = __spirv_GenericCastToPtrExplicit_ToGlobal<T>(
+        dst, __spv::StorageClass::CrossWorkgroup);
+    if (g) {
+      store(g, x);
+      return;
+    }
+
+    // Fallback for other address spaces to be mapped to global
+    store(__spirv_PtrCastToGeneric<T>(dst), x);
+  }
+#endif // SYCL_USE_DECORATED_REF
+#else  //__SYCL_DEVICE_ONLY__
+  template <typename T> void store(T *dst, const T &x) const {
+    (void)dst;
+    (void)x;
+    throw runtime_error("Sub-groups are not supported on host device.",
+                        PI_INVALID_DEVICE);
+  }
+#endif //__SYCL_DEVICE_ONLY__
+
   template <typename T, access::address_space Space>
   sycl::detail::enable_if_t<
       sycl::detail::sub_group::AcceptableForGlobalLoadStore<T, Space>::value>
@@ -347,6 +463,20 @@ struct sub_group {
 #endif
   }
 
+  template <typename T, access::address_space Space>
+  sycl::detail::enable_if_t<
+      sycl::detail::sub_group::AcceptableForPrivateLoadStore<T, Space>::value>
+  store(multi_ptr<T, Space> dst, const T &x) const {
+#ifdef __SYCL_DEVICE_ONLY__
+    dst.get()[get_local_id()[0]] = x;
+#else
+    (void)dst;
+    (void)x;
+    throw runtime_error("Sub-groups are not supported on host device.",
+                        PI_INVALID_DEVICE);
+#endif
+  }
+
   template <int N, typename T, access::address_space Space>
   sycl::detail::enable_if_t<
       sycl::detail::sub_group::AcceptableForGlobalLoadStore<T, Space>::value &&

sycl/include/CL/sycl/access/access.hpp

@@ -115,13 +115,15 @@ constexpr bool modeWritesNewData(access::mode m) {
 #define __OPENCL_LOCAL_AS__ __attribute__((opencl_local))
 #define __OPENCL_CONSTANT_AS__ __attribute__((opencl_constant))
 #define __OPENCL_PRIVATE_AS__ __attribute__((opencl_private))
+#define __OPENCL_GENERIC_AS__ __attribute__((opencl_generic))
 #else
 #define __OPENCL_GLOBAL_AS__
 #define __OPENCL_GLOBAL_DEVICE_AS__
 #define __OPENCL_GLOBAL_HOST_AS__
 #define __OPENCL_LOCAL_AS__
 #define __OPENCL_CONSTANT_AS__
 #define __OPENCL_PRIVATE_AS__
+#define __OPENCL_GENERIC_AS__
 #endif
 
 template <access::target accessTarget> struct TargetToAS {
@@ -187,17 +189,15 @@ template <typename ElementType>
 struct DecoratedType<ElementType, access::address_space::local_space> {
   using type = __OPENCL_LOCAL_AS__ ElementType;
 };
+template <class T> struct remove_AS { typedef T type; };
 
-template <class T>
-struct remove_AS {
-  typedef T type;
+template <class T> struct deduce_AS {
+  static const access::address_space value =
+      access::address_space::global_space;
 };
 
 #ifdef __SYCL_DEVICE_ONLY__
-template <class T>
-struct remove_AS<__OPENCL_GLOBAL_AS__ T> {
-  typedef T type;
-};
+template <class T> struct remove_AS<__OPENCL_GLOBAL_AS__ T> { typedef T type; };
 
 #ifdef __ENABLE_USM_ADDR_SPACE__
 template <class T> struct remove_AS<__OPENCL_GLOBAL_DEVICE_AS__ T> {
@@ -207,22 +207,50 @@ template <class T> struct remove_AS<__OPENCL_GLOBAL_DEVICE_AS__ T> {
 template <class T> struct remove_AS<__OPENCL_GLOBAL_HOST_AS__ T> {
   typedef T type;
 };
+
+template <class T> struct deduce_AS<__OPENCL_GLOBAL_DEVICE_AS__ T> {
+  static const access::address_space value =
+      access::address_space::global_device_space;
+};
+
+template <class T> struct deduce_AS<__OPENCL_GLOBAL_HOST_AS__ T> {
+  static const access::address_space value =
+      access::address_space::global_host_space;
+};
 #endif // __ENABLE_USM_ADDR_SPACE__
 
-template <class T>
-struct remove_AS<__OPENCL_PRIVATE_AS__ T> {
+template <class T> struct remove_AS<__OPENCL_PRIVATE_AS__ T> {
   typedef T type;
 };
 
-template <class T>
-struct remove_AS<__OPENCL_LOCAL_AS__ T> {
+template <class T> struct remove_AS<__OPENCL_LOCAL_AS__ T> { typedef T type; };
+
+template <class T> struct remove_AS<__OPENCL_CONSTANT_AS__ T> {
   typedef T type;
 };
 
-template <class T>
-struct remove_AS<__OPENCL_CONSTANT_AS__ T> {
+template <class T> struct remove_AS<__OPENCL_GENERIC_AS__ T> {
   typedef T type;
 };
+
+template <class T> struct deduce_AS<__OPENCL_PRIVATE_AS__ T> {
+  static const access::address_space value =
+      access::address_space::private_space;
+};
+
+template <class T> struct deduce_AS<__OPENCL_LOCAL_AS__ T> {
+  static const access::address_space value = access::address_space::local_space;
+};
+
+template <class T> struct deduce_AS<__OPENCL_CONSTANT_AS__ T> {
+  static const access::address_space value =
+      access::address_space::constant_space;
+};
+
+template <class T> struct deduce_AS<__OPENCL_GENERIC_AS__ T> {
+  static const access::address_space value =
+      access::address_space::global_space;
+};
 #endif
 
 #undef __OPENCL_GLOBAL_AS__
@@ -231,8 +259,8 @@ struct remove_AS<__OPENCL_CONSTANT_AS__ T> {
 #undef __OPENCL_LOCAL_AS__
 #undef __OPENCL_CONSTANT_AS__
 #undef __OPENCL_PRIVATE_AS__
-
+#undef __OPENCL_GENERIC_AS__
 } // namespace detail
 
-}  // namespace sycl
-}  // __SYCL_INLINE_NAMESPACE(cl)
+} // namespace sycl
+} // __SYCL_INLINE_NAMESPACE(cl)