[ESIMD] Fix slm_gather/scatter. (#4771)

[ESIMD] Fix slm_gather/scatter.

- remove limitations on element size being 4 bytes only
- enable 1,8,16,32 elements as allowed by hardware
- share implementation with surface-based gather/scatter

Signed-off-by: Konstantin S Bobrovsky <[email protected]>

Author: kbobrovs

llvm/lib/SYCLLowerIR/LowerESIMD.cpp

@@ -367,6 +367,15 @@ class ESIMDIntrinDescTable {
         {"gather_scaled",
          {"gather.scaled", {ai1(0), t(3), t(4), aSI(1), a(2), a(3), u(-1)}}},
 
+        // arg0: i32 log2 num blocks, CONSTANT (0/1/2 for num blocks 1/2/4)
+        // arg1: i16 scale, CONSTANT
+        // arg2: i32 surface index
+        // arg3: i32 global offset in bytes
+        // arg4: vXi32 element offset in bytes (overloaded)
+        // arg5: vXi1 predicate (overloaded)
+        {"gather_masked_scaled2",
+         {"gather.masked.scaled2", {t(3), t(4), aSI(0), a(1), a(2), ai1(3)}}},
+
         // arg0: vXi1 predicate (overloaded)
         // arg1: i32 log2 num blocks, CONSTANT (0/1/2 for num blocks 1/2/4)
         // arg2: i16 scale, CONSTANT

sycl/include/sycl/ext/intel/experimental/esimd/detail/memory_intrin.hpp

@@ -47,8 +47,8 @@ class AccessorPrivateProxy {
 };
 
 template <int ElemsPerAddr,
-          typename = sycl::detail::enable_if_t<
-              (ElemsPerAddr == 1 || ElemsPerAddr == 2 || ElemsPerAddr == 4)>>
+          typename = std::enable_if_t<(ElemsPerAddr == 1 || ElemsPerAddr == 2 ||
+                                       ElemsPerAddr == 4)>>
 constexpr unsigned int ElemsPerAddrEncoding() {
   // encoding requires log2 of ElemsPerAddr
   if constexpr (ElemsPerAddr == 1)
@@ -342,7 +342,7 @@ __esimd_gather_scaled2(SurfIndAliasTy surf_ind, uint32_t global_offset,
     ;
 #else
 {
-  static_assert(N == 1 || N == 8 || N == 16);
+  static_assert(N == 1 || N == 8 || N == 16 || N == 32);
   static_assert(TySizeLog2 <= 2 && Scale == 0);
   static_assert(std::is_integral<Ty>::value || TySizeLog2 == 2);
   throw cl::sycl::feature_not_supported();
@@ -392,9 +392,6 @@ __esimd_scatter_scaled(__SEIEED::simd_mask_storage_t<N> pred,
 }
 #endif // __SYCL_DEVICE_ONLY__
 
-// TODO bring the parameter order of __esimd* intrinsics in accordance with the
-// correponsing BE intrinsicics parameter order.
-
 // flat_atomic: flat-address atomic
 template <__SEIEE::atomic_op Op, typename Ty, int N,
           __SEIEE::CacheHint L1H = __SEIEE::CacheHint::None,
@@ -473,7 +470,7 @@ __ESIMD_INTRIN void __esimd_fence(uint8_t cntl)
 
 // Scaled gather from a surface.
 template <typename Ty, int N, typename SurfIndAliasTy, int TySizeLog2,
-          int16_t SCALE = 0>
+          int16_t Scale = 0>
 __ESIMD_INTRIN __SEIEED::vector_type_t<Ty, N>
 __esimd_gather_scaled(__SEIEED::simd_mask_storage_t<N> pred,
                       SurfIndAliasTy surf_ind, uint32_t global_offset,
@@ -486,6 +483,40 @@ __esimd_gather_scaled(__SEIEED::simd_mask_storage_t<N> pred,
 }
 #endif // __SYCL_DEVICE_ONLY__
 
+/// Predicated (masked) scaled gather from a surface.
+///
+/// Template (compile-time constant) parameters:
+/// @tparam Ty - element type
+/// @tparam N  - the number of elements to read
+/// @tparam SurfIndAliasTy - "surface index alias" type - internal type in the
+///   accessor used to denote the surface
+/// @tparam TySizeLog2 - Log2 of the number of bytes written per element:
+///   0 - 1 byte, 1 - 2 bytes, 2 - 4 bytes
+/// @tparam Scale - offset scale; only 0 is supported for now
+///
+/// Formal parameters:
+/// @param surf_ind - the surface index, taken from the SYCL memory object
+/// @param global_offset - offset added to each individual element's offset to
+///   compute actual memory access offset for that element
+/// @param offsets - per-element offsets
+/// @param pred - per-element predicates; elements with zero corresponding
+///   predicates are not written
+/// @return - elements read ("gathered") from memory
+
+template <typename Ty, int N, typename SurfIndAliasTy, int TySizeLog2,
+          int16_t Scale = 0>
+__ESIMD_INTRIN __SEIEED::vector_type_t<Ty, N>
+__esimd_gather_masked_scaled2(SurfIndAliasTy surf_ind, uint32_t global_offset,
+                              __SEIEED::vector_type_t<uint32_t, N> offsets,
+                              __SEIEED::simd_mask_storage_t<N> pred)
+#ifdef __SYCL_DEVICE_ONLY__
+    ;
+#else
+{
+  throw cl::sycl::feature_not_supported();
+}
+#endif // __SYCL_DEVICE_ONLY__
+
 // Reads a block of data from given surface at given offset, offset must be
 // 16-byte-aligned.
 template <typename Ty, int N, typename SurfIndAliasTy, int32_t IsModified = 0>
@@ -705,7 +736,6 @@ __ESIMD_INTRIN void __esimd_media_st(TACC handle, unsigned x, unsigned y,
 }
 #endif // __SYCL_DEVICE_ONLY__
 
-#ifdef __SYCL_DEVICE_ONLY__
 /// \brief Converts given value to a surface index.
 /// The input must always be a result of
 ///   detail::AccessorPrivateProxy::getNativeImageObj(acc)
@@ -724,15 +754,17 @@ __ESIMD_INTRIN void __esimd_media_st(TACC handle, unsigned x, unsigned y,
 /// pointer, where we can do ptr to uint32_t conversion.
 /// This intrinsic can be called only from the device code, as
 /// accessor => memory handle translation for host is different.
-///
-/// @param SYCL accessor's native memory object extracted from it via
+/// @param acc the SYCL accessor.
 ///   getNativeImageObj.
-///
-/// Returns the surface index (binding table index) value 'sid' corresponds to.
-///
-template <typename SurfIndAliasTy>
-__ESIMD_INTRIN __SEIEE::SurfaceIndex
-__esimd_get_surface_index(SurfIndAliasTy sid);
+/// Returns the binding table index value.
+template <typename MemObjTy>
+__ESIMD_INTRIN __SEIEE::SurfaceIndex __esimd_get_surface_index(MemObjTy obj)
+#ifdef __SYCL_DEVICE_ONLY__
+    ;
+#else
+{
+  throw cl::sycl::feature_not_supported();
+}
 #endif // __SYCL_DEVICE_ONLY__
 
 /// \brief Raw sends load.

sycl/include/sycl/ext/intel/experimental/esimd/detail/simd_obj_impl.hpp

@@ -506,7 +506,7 @@ template <typename Ty, int N, class Derived, class SFINAE> class simd_obj_impl {
   /// elements in this object. Source memory location is represented via a
   /// global accessor and offset.
   /// @param acc accessor to copy from.
-  /// @param offset offset to copy from.
+  /// @param offset offset to copy from (in bytes).
   template <typename AccessorT>
   ESIMD_INLINE EnableIfAccessor<AccessorT, accessor_mode_cap::can_read,
                                 sycl::access::target::global_buffer, void>