[SYCL][DOC][bindless][cuda] gather_image extension doc and impl (#17322)

This PR presents a portable `gather_image` API for the bindless images
extension, with an initial cuda implementation for 2D images only.
Cubemap support is also possible but not yet added at this stage, since
- it is unclear whether this can map to the spirv cube case.
- l0 doesn't yet even support cubemaps anyway

For 2D images the mapping to `__spirv_SampledImageGather` is
straightforward and unambiguous. For complete details see the extension
documentation in this PR.

---------

Signed-off-by: JackAKirk <[email protected]>
Co-authored-by: Sean Stirling <[email protected]>
Co-authored-by: Peter Žužek <[email protected]>
Co-authored-by: Steffen Larsen <[email protected]>

Author: 4 people

libclc/libspirv/lib/ptx-nvidiacl/images/image.cl

@@ -2727,6 +2727,69 @@ _CLC_DEFINE_SAMPLEDIMAGE_BINDLESS_FETCH_BUILTIN(half4, 1, Dv4_DF16_, v4f16, i, i
 _CLC_DEFINE_SAMPLEDIMAGE_BINDLESS_FETCH_BUILTIN(half4, 2, Dv4_DF16_, v4f16, Dv2_i, int2 coord, coord.x, coord.y)
 _CLC_DEFINE_SAMPLEDIMAGE_BINDLESS_FETCH_BUILTIN(half4, 3, Dv4_DF16_, v4f16, Dv3_i, int4 coord, coord.x, coord.y, coord.z)
 
+float4
+__nvvm_tld4_r_2d_v4f32_f32(long, float,
+                           float) __asm("__clc_llvm_nvvm_tld4_r_2d_v4f32_f32");
+float4
+__nvvm_tld4_g_2d_v4f32_f32(long, float,
+                           float) __asm("__clc_llvm_nvvm_tld4_g_2d_v4f32_f32");
+float4
+__nvvm_tld4_b_2d_v4f32_f32(long, float,
+                           float) __asm("__clc_llvm_nvvm_tld4_b_2d_v4f32_f32");
+float4
+__nvvm_tld4_a_2d_v4f32_f32(long, float,
+                           float) __asm("__clc_llvm_nvvm_tld4_a_2d_v4f32_f32");
+
+int4 __nvvm_tld4_r_2d_v4s32_f32(long, float, float) __asm(
+    "__clc_llvm_nvvm_tld4_r_2d_v4s32_f32");
+int4 __nvvm_tld4_g_2d_v4s32_f32(long, float, float) __asm(
+    "__clc_llvm_nvvm_tld4_g_2d_v4s32_f32");
+int4 __nvvm_tld4_b_2d_v4s32_f32(long, float, float) __asm(
+    "__clc_llvm_nvvm_tld4_b_2d_v4s32_f32");
+int4 __nvvm_tld4_a_2d_v4s32_f32(long, float, float) __asm(
+    "__clc_llvm_nvvm_tld4_a_2d_v4s32_f32");
+
+uint4 __nvvm_tld4_r_2d_v4u32_f32(long, float, float) __asm(
+    "__clc_llvm_nvvm_tld4_r_2d_v4u32_f32");
+uint4 __nvvm_tld4_g_2d_v4u32_f32(long, float, float) __asm(
+    "__clc_llvm_nvvm_tld4_g_2d_v4u32_f32");
+uint4 __nvvm_tld4_b_2d_v4u32_f32(long, float, float) __asm(
+    "__clc_llvm_nvvm_tld4_b_2d_v4u32_f32");
+uint4 __nvvm_tld4_a_2d_v4u32_f32(long, float, float) __asm(
+    "__clc_llvm_nvvm_tld4_a_2d_v4u32_f32");
+
+#define _CLC_DEFINE_SAMPLEDIMAGE_BINDLESS_GATHER_BUILTIN(                      \
+    elem_t, elem_t_mangled, vec_size, type)                                    \
+  _CLC_DEF elem_t MANGLE_FUNC_IMG_HANDLE(                                      \
+      26, __spirv_SampledImageGather, I##elem_t_mangled,                       \
+      Dv2_f##ET_T0_T1_j)(ulong imageHandle, float2 coord, unsigned i) {        \
+    elem_t result;                                                             \
+    switch (i) {                                                               \
+    case 0:                                                                    \
+      result = __nvvm_tld4_r_2d_v4##type##_f32(imageHandle, coord.x, coord.y); \
+      break;                                                                   \
+    case 1:                                                                    \
+      result = __nvvm_tld4_g_2d_v4##type##_f32(imageHandle, coord.x, coord.y); \
+      break;                                                                   \
+    case 2:                                                                    \
+      result = __nvvm_tld4_b_2d_v4##type##_f32(imageHandle, coord.x, coord.y); \
+      break;                                                                   \
+    case 3:                                                                    \
+      result = __nvvm_tld4_a_2d_v4##type##_f32(imageHandle, coord.x, coord.y); \
+      break;                                                                   \
+    default:                                                                   \
+      __builtin_trap();                                                        \
+      __builtin_unreachable();                                                 \
+    }                                                                          \
+    return result;                                                             \
+  }
+
+_CLC_DEFINE_SAMPLEDIMAGE_BINDLESS_GATHER_BUILTIN(float4, Dv4_f, v4f32, f32)
+_CLC_DEFINE_SAMPLEDIMAGE_BINDLESS_GATHER_BUILTIN(int4, Dv4_i, v4i32, s32)
+_CLC_DEFINE_SAMPLEDIMAGE_BINDLESS_GATHER_BUILTIN(uint4, Dv4_j, v4j32, u32)
+
+#undef _CLC_DEFINE_SAMPLEDIMAGE_BINDLESS_GATHER_BUILTIN
+
 // <--- MIPMAP --->
 
 // Define functions to call intrinsic

libclc/libspirv/lib/ptx-nvidiacl/images/image_helpers.ll

@@ -345,6 +345,103 @@ entry:
   ret <4 x float> %1
 }
 
+; <--- TEXTURE GATHER --->
+declare {float,float,float,float} @llvm.nvvm.tld4.unified.r.2d.v4f32.f32(i64, float, float)
+define <4 x float> @__clc_llvm_nvvm_tld4_r_2d_v4f32_f32(i64 %img, float %x, float %y) nounwind alwaysinline {
+entry:
+  %0 = tail call {float,float,float,float} @llvm.nvvm.tld4.unified.r.2d.v4f32.f32(i64 %img, float %x, float %y);
+  %1 = tail call <4 x float>@__clc_structf32_to_vector({float,float,float,float} %0)
+  ret <4 x float> %1
+}
+
+declare {float,float,float,float} @llvm.nvvm.tld4.unified.g.2d.v4f32.f32(i64, float, float)
+define <4 x float> @__clc_llvm_nvvm_tld4_g_2d_v4f32_f32(i64 %img, float %x, float %y) nounwind alwaysinline {
+entry:
+  %0 = tail call {float,float,float,float} @llvm.nvvm.tld4.unified.g.2d.v4f32.f32(i64 %img, float %x, float %y);
+  %1 = tail call <4 x float>@__clc_structf32_to_vector({float,float,float,float} %0)
+  ret <4 x float> %1
+}
+
+declare {float,float,float,float} @llvm.nvvm.tld4.unified.b.2d.v4f32.f32(i64, float, float)
+define <4 x float> @__clc_llvm_nvvm_tld4_b_2d_v4f32_f32(i64 %img, float %x, float %y) nounwind alwaysinline {
+entry:
+  %0 = tail call {float,float,float,float} @llvm.nvvm.tld4.unified.b.2d.v4f32.f32(i64 %img, float %x, float %y);
+  %1 = tail call <4 x float>@__clc_structf32_to_vector({float,float,float,float} %0)
+  ret <4 x float> %1
+}
+
+declare {float,float,float,float} @llvm.nvvm.tld4.unified.a.2d.v4f32.f32(i64, float, float)
+define <4 x float> @__clc_llvm_nvvm_tld4_a_2d_v4f32_f32(i64 %img, float %x, float %y) nounwind alwaysinline {
+entry:
+  %0 = tail call {float,float,float,float} @llvm.nvvm.tld4.unified.a.2d.v4f32.f32(i64 %img, float %x, float %y);
+  %1 = tail call <4 x float>@__clc_structf32_to_vector({float,float,float,float} %0)
+  ret <4 x float> %1
+}
+
+declare {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.r.2d.v4s32.f32(i64, float, float)
+define <4 x i32> @__clc_llvm_nvvm_tld4_r_2d_v4s32_f32(i64 %img, float %x, float %y) nounwind alwaysinline {
+entry:
+  %0 = tail call {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.r.2d.v4s32.f32(i64 %img, float %x, float %y);
+  %1 = tail call <4 x i32>@__clc_struct32_to_vector({i32,i32,i32,i32} %0)
+  ret <4 x i32> %1
+}
+
+declare {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.g.2d.v4s32.f32(i64, float, float)
+define <4 x i32> @__clc_llvm_nvvm_tld4_g_2d_v4s32_f32(i64 %img, float %x, float %y) nounwind alwaysinline {
+entry:
+  %0 = tail call {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.g.2d.v4s32.f32(i64 %img, float %x, float %y);
+  %1 = tail call <4 x i32>@__clc_struct32_to_vector({i32,i32,i32,i32} %0)
+  ret <4 x i32> %1
+}
+
+declare {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.b.2d.v4s32.f32(i64, float, float)
+define <4 x i32> @__clc_llvm_nvvm_tld4_b_2d_v4s32_f32(i64 %img, float %x, float %y) nounwind alwaysinline {
+entry:
+  %0 = tail call {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.b.2d.v4s32.f32(i64 %img, float %x, float %y);
+  %1 = tail call <4 x i32>@__clc_struct32_to_vector({i32,i32,i32,i32} %0)
+  ret <4 x i32> %1
+}
+
+declare {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.a.2d.v4s32.f32(i64, float, float)
+define <4 x i32> @__clc_llvm_nvvm_tld4_a_2d_v4s32_f32(i64 %img, float %x, float %y) nounwind alwaysinline {
+entry:
+  %0 = tail call {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.a.2d.v4s32.f32(i64 %img, float %x, float %y);
+  %1 = tail call <4 x i32>@__clc_struct32_to_vector({i32,i32,i32,i32} %0)
+  ret <4 x i32> %1
+}
+
+declare {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.r.2d.v4u32.f32(i64, float, float)
+define <4 x i32> @__clc_llvm_nvvm_tld4_r_2d_v4u32_f32(i64 %img, float %x, float %y) nounwind alwaysinline {
+entry:
+  %0 = tail call {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.r.2d.v4u32.f32(i64 %img, float %x, float %y);
+  %1 = tail call <4 x i32>@__clc_struct32_to_vector({i32,i32,i32,i32} %0)
+  ret <4 x i32> %1
+}
+
+declare {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.g.2d.v4u32.f32(i64, float, float)
+define <4 x i32> @__clc_llvm_nvvm_tld4_g_2d_v4u32_f32(i64 %img, float %x, float %y) nounwind alwaysinline {
+entry:
+  %0 = tail call {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.g.2d.v4u32.f32(i64 %img, float %x, float %y);
+  %1 = tail call <4 x i32>@__clc_struct32_to_vector({i32,i32,i32,i32} %0)
+  ret <4 x i32> %1
+}
+
+declare {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.b.2d.v4u32.f32(i64, float, float)
+define <4 x i32> @__clc_llvm_nvvm_tld4_b_2d_v4u32_f32(i64 %img, float %x, float %y) nounwind alwaysinline {
+entry:
+  %0 = tail call {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.b.2d.v4u32.f32(i64 %img, float %x, float %y);
+  %1 = tail call <4 x i32>@__clc_struct32_to_vector({i32,i32,i32,i32} %0)
+  ret <4 x i32> %1
+}
+
+declare {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.a.2d.v4u32.f32(i64, float, float)
+define <4 x i32> @__clc_llvm_nvvm_tld4_a_2d_v4u32_f32(i64 %img, float %x, float %y) nounwind alwaysinline {
+entry:
+  %0 = tail call {i32,i32,i32,i32} @llvm.nvvm.tld4.unified.a.2d.v4u32.f32(i64 %img, float %x, float %y);
+  %1 = tail call <4 x i32>@__clc_struct32_to_vector({i32,i32,i32,i32} %0)
+  ret <4 x i32> %1
+}
+
 ; <--- TEXTURE FETCHING (integer coordinates) --->
 declare {i32,i32,i32,i32} @llvm.nvvm.tex.unified.1d.v4s32.s32(i64, i32)
 define <4 x i32> @__clc_llvm_nvvm_tex_1d_v4i32_s32(i64 %img, i32 %x) nounwind alwaysinline {

llvm/include/llvm/SYCLLowerIR/DeviceConfigFile.td

@@ -87,6 +87,7 @@ def AspectExt_intel_fpga_task_sequence : Aspect<"ext_intel_fpga_task_sequence">;
 def AspectExt_oneapi_atomic16 : Aspect<"ext_oneapi_atomic16">;
 def AspectExt_oneapi_virtual_functions : Aspect<"ext_oneapi_virtual_functions">;
 def AspectExt_intel_spill_memory_size : Aspect<"ext_intel_spill_memory_size">;
+def AspectExt_oneapi_bindless_images_gather : Aspect<"ext_oneapi_bindless_images_gather">;
 // Deprecated aspects
 def AspectInt64_base_atomics : Aspect<"int64_base_atomics">;
 def AspectInt64_extended_atomics : Aspect<"int64_extended_atomics">;
@@ -144,6 +145,7 @@ def : TargetInfo<"__TestAspectList",
     AspectExt_oneapi_bindless_sampled_image_fetch_1d_usm, AspectExt_oneapi_bindless_sampled_image_fetch_1d,
     AspectExt_oneapi_bindless_sampled_image_fetch_2d_usm, AspectExt_oneapi_bindless_sampled_image_fetch_2d,
     AspectExt_oneapi_bindless_sampled_image_fetch_3d,
+    AspectExt_oneapi_bindless_images_gather,
     AspectExt_intel_esimd,
     AspectExt_oneapi_ballot_group, AspectExt_oneapi_fixed_size_group, AspectExt_oneapi_opportunistic_group,
     AspectExt_oneapi_tangle_group, AspectExt_intel_matrix, AspectExt_oneapi_is_composite, AspectExt_oneapi_is_component,

sycl/doc/extensions/experimental/sycl_ext_oneapi_bindless_images.asciidoc

@@ -30,7 +30,7 @@ permission by Khronos.
 
 == Dependencies
 
-This extension is written against the SYCL 2020 revision 6 specification.  All
+This extension is written against the SYCL 2020 revision 9 specification.  All
 references below to the "core SYCL specification" or to section numbers in the
 SYCL specification refer to that revision.
 
@@ -181,6 +181,7 @@ enum class image_type : /* unspecified */ {
   mipmap,
   array,
   cubemap,
+  gather,
 };
 
 struct image_descriptor {
@@ -1329,10 +1330,15 @@ DataT fetch_image(const unsampled_image_handle &ImageHandle,
 template <typename DataT, typename HintT = DataT, typename CoordT>
 DataT fetch_image(const sampled_image_handle &ImageHandle,
                   const CoordT &Coords);
+
 template <typename DataT, typename HintT = DataT, typename CoordT>
-DataT sample_image(const sampled_image_handle &ImageHandle, 
+DataT sample_image(const sampled_image_handle &ImageHandle,
                    const CoordT &Coords);
 
+template <typename DataT>
+DataT gather_image(const sampled_image_handle &ImageHandle,
+                   const sycl::float2 &Coords, const unsigned Component);
+
 template <typename DataT, typename CoordT>
 void write_image(unsampled_image_handle ImageHandle,
                  const CoordT &Coords, const DataT &Color);
@@ -1342,9 +1348,10 @@ void write_image(unsampled_image_handle ImageHandle,
 Inside a kernel, it's possible to retrieve data from an image via `fetch_image` 
 or `sample_image`, passing the appropriate image handle. The `fetch_image` API 
 is applicable to sampled and unsampled images, and the data will be fetched 
-exactly as is in device memory. The `sample_image` API is only applicable to 
-sampled images, the image data will be sampled according to the 
-`bindless_image_sampler` that was passed to the image upon construction.
+exactly as is in device memory. The `sample_image` and `gather_image` APIs are
+only applicable to sampled images. For both APIs the image data will be sampled
+according to the `bindless_image_sampler` that was passed to the image upon
+construction.
 
 When fetching from a sampled image handle, data exactly as is in memory, no 
 sampling operations will be performed, and the `bindless_image_sampler` passed 
@@ -1406,6 +1413,58 @@ Attempting to sample a standard sampled image with `sample_mipmap` or any other
 defined sampling function will result in undefined behaviour.
 ====
 
+==== `gather_image`
+
+We provide the following device aspect to retrieve support information for
+`gather_image`.
+
+The device aspect descriptor for this query is:
+
+[frame="none",options="header"]
+|======================
+|Device descriptor |Description
+|`aspect::ext_oneapi_bindless_images_gather` | Indicates if the device supports
+`gather_image`.
+|======================
+
+`DataT` is a vector type with four components. DataT can only be `sycl::float4`
+, `sycl::int4`, or `sycl::uint4`. `ImageHandle` must be a two dimensional image
+instantiated using an `image_descriptor` with non-zero `width` and `height`
+members. The `image_descriptor::depth` member must be zero. If these conditions
+are not satisfied then a `sycl::exception` is thrown with error code
+`sycl::errc::invalid`. The return values, `DataT`, consist of the four texture,
+`T`, coordinates that would be used for linear filtering of `ImageHandle` at
+the coordinate specified by `Coords`, and for channel component `Component`.
+
+These coordinates are returned as follows:
+
+- `DataT.x = T[i, j + 1]`
+- `DataT.y = T[i + 1, j + 1]`
+- `DataT.z = T[i + 1, j]`
+- `DataT.w = T[i, j]`
+
+where `i` and `j` are calculated depending on the
+`sycl::coordinate_normalization_mode` used to constuct `ImageHandle` via a
+`bindless_image_sampler`. For
+`sycl::coordinate_normalization_mode::unnormalized` (in units of the image
+discretization) they are calculated as:
+
+- `i = floor(Coords.x - 0.5)`
+- `j = floor(Coords.y - 0.5)`
+
+[NOTE]
+=====
+In the CUDA backend there are different restrictions on the sizes of an image
+that can be created for a given device when `gather_image` is being used.
+=====
+
+For guaranteed correctness and portability across backends, it is necessary to
+construct `ImageHandle` using an `image_descriptor` with member
+`image_type::gather`. Calling `gather_image` on an `ImageHandle` constructed
+with `image_descriptor::num_channels < 4` will result in undefined behavior.
+
+====
+
 === Recognized standard types [[recognized_standard_types]]
 
 For the purposes of this extension, the following are classified as recognized 
@@ -1418,6 +1477,8 @@ standard types.
 
 Any other types are classified as user-defined types.
 
+===
+
 ==== User-defined types
 
 Some examples of a user-defined types may be:
@@ -2576,4 +2637,5 @@ These features still need to be handled:
                  - Add support for USM to USM copies and sub-copies.
                  - Add support for host to host copies and sub-copies.
 |6.8|2025-03-13| - Add support for importing timeline semaphores.
-|======================
+|6.9|2025-03-18| - Add new `gather` image type and the accompanying
+                   `gather_image` function.

sycl/include/sycl/__spirv/spirv_ops.hpp

@@ -196,6 +196,10 @@ template <class RetT, typename ImageT, typename TempArgT>
 extern __DPCPP_SYCL_EXTERNAL RetT __spirv_SampledImageArrayFetch(ImageT,
                                                                  TempArgT, int);
 
+template <class RetT, typename ImageT, typename TempArgT>
+extern __DPCPP_SYCL_EXTERNAL RetT __spirv_SampledImageGather(ImageT, TempArgT,
+                                                             unsigned);
+
 template <class RetT, typename ImageT, typename TempArgT>
 extern __DPCPP_SYCL_EXTERNAL RetT __spirv_ImageArrayRead(ImageT, TempArgT, int);
 

sycl/include/sycl/detail/image_ocl_types.hpp

@@ -103,6 +103,22 @@ static RetType __invoke__SampledImageFetch(ImageT Img, CoordT Coords) {
           Img, TmpCoords));
 }
 
+template <typename RetType, typename ImageT, typename CoordT>
+static std::enable_if_t<std::is_same_v<RetType, sycl::vec<float, 4>> ||
+                            std::is_same_v<RetType, sycl::vec<int, 4>> ||
+                            std::is_same_v<RetType, sycl::vec<unsigned int, 4>>,
+                        RetType>
+__invoke__SampledImageGather(ImageT Img, CoordT Coords, unsigned Component) {
+
+  // Convert from sycl types to builtin types to get correct function mangling.
+  using TempRetT = sycl::detail::ConvertToOpenCLType_t<RetType>;
+  auto TmpCoords = sycl::detail::convertToOpenCLType(Coords);
+
+  return sycl::detail::convertFromOpenCLTypeFor<RetType>(
+      __spirv_SampledImageGather<TempRetT, ImageT, decltype(TmpCoords)>(
+          Img, TmpCoords, Component));
+}
+
 template <typename RetType, typename ImageT, typename CoordT>
 static RetType __invoke__ImageArrayFetch(ImageT Img, CoordT Coords,
                                          int ArrayLayer) {

sycl/include/sycl/ext/oneapi/bindless_images.hpp

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include <sycl/aspects.hpp>
 #include <sycl/context.hpp>                               // for context
 #include <sycl/detail/export.hpp>                         // for __SYCL_EXPORT
 #include <sycl/device.hpp>                                // for device
@@ -910,6 +911,31 @@ DataT fetch_image(const sampled_image_handle &imageHandle [[maybe_unused]],
 #endif
 }
 
+template <typename DataT>
+#ifdef __SYCL_DEVICE_ONLY__
+[[__sycl_detail__::__uses_aspects__(
+    sycl::aspect::ext_oneapi_bindless_images_gather)]]
+#endif
+std::enable_if_t<std::is_same_v<DataT, float4> || std::is_same_v<DataT, int4> ||
+                     std::is_same_v<DataT, uint4>,
+                 DataT> gather_image(const sampled_image_handle &imageHandle
+                                     [[maybe_unused]],
+                                     const float2 &coords [[maybe_unused]],
+                                     const unsigned i [[maybe_unused]]) {
+#if defined(__SYCL_DEVICE_ONLY__)
+#if defined(__NVPTX__)
+  return __invoke__SampledImageGather<DataT>(
+      CONVERT_HANDLE_TO_SAMPLED_IMAGE(imageHandle.raw_handle, float2::size()),
+      coords, i);
+#else
+  return {0, 0, 0, 0};
+#endif
+#else
+  throw exception{make_error_code(errc::feature_not_supported),
+                  "gather_image is not supported on the host"};
+#endif
+}
+
 /**
  *  @brief   Sample data from a sampled image using its handle
  *