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

include/ur_api.h

@@ -99,6 +99,7 @@ Enums
     * ${X}_DEVICE_INFO_BINDLESS_UNIQUE_ADDRESSING_PER_DIM_EXP
     * ${X}_DEVICE_INFO_BINDLESS_SAMPLE_1D_USM_EXP
     * ${X}_DEVICE_INFO_BINDLESS_SAMPLE_2D_USM_EXP
+    * ${X}_DEVICE_INFO_BINDLESS_IMAGES_GATHER_EXP
 
 * ${x}_command_t
     * ${X}_COMMAND_EXTERNAL_SEMAPHORE_WAIT_EXP
@@ -148,6 +149,7 @@ Enums
 
 * ${x}_mem_type_t
     * ${X}_MEM_TYPE_IMAGE_CUBEMAP_EXP
+    * ${X}_MEM_TYPE_IMAGE_GATHER_EXP
 
 Types
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

include/ur_print.hpp

@@ -113,6 +113,9 @@ etors:
     - name: BINDLESS_SAMPLE_2D_USM_EXP
       value: "0x201C"
       desc: "[$x_bool_t] returns true if the device supports sampling USM backed 2D sampled image data."
+    - name: BINDLESS_IMAGES_GATHER_EXP
+      value: "0x201D"
+      desc: "[$x_bool_t] returns true if the device supports sampled image gather."
 --- #--------------------------------------------------------------------------
 type: enum
 extend: true
@@ -180,6 +183,15 @@ etors:
       desc: "Experimental cubemap image object"
 --- #--------------------------------------------------------------------------
 type: enum
+extend: True
+desc: "Memory types"
+name: $x_mem_type_t
+etors:
+    - name: IMAGE_GATHER_EXP
+      value: "0x2001"
+      desc: "Experimental gather image object"
+--- #--------------------------------------------------------------------------
+type: enum
 desc: "Sampler cubemap seamless filtering mode."
 class: $xBindlessImages
 name: $x_exp_sampler_cubemap_filter_mode_t

scripts/core/EXP-BINDLESS-IMAGES.rst

@@ -932,6 +932,10 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
     // CUDA does support sampling 1D USM sampled image data.
     return ReturnValue(static_cast<ur_bool_t>(true));
   }
+  case UR_DEVICE_INFO_BINDLESS_IMAGES_GATHER_EXP: {
+    // CUDA does support sampled image gather.
+    return ReturnValue(static_cast<ur_bool_t>(true));
+  }
   case UR_DEVICE_INFO_TIMESTAMP_RECORDING_SUPPORT_EXP: {
     // CUDA supports recording timestamp events.
     return ReturnValue(static_cast<ur_bool_t>(true));

scripts/core/exp-bindless-images.yml

@@ -370,6 +370,11 @@ UR_APIEXPORT ur_result_t UR_APICALL urBindlessImagesImageAllocateExp(
     array_desc.Depth = pImageDesc->arraySize; // Should be 6 ONLY
     array_desc.Flags |= CUDA_ARRAY3D_CUBEMAP;
     break;
+  case UR_MEM_TYPE_IMAGE_GATHER_EXP:
+    array_desc.Height = pImageDesc->height;
+    array_desc.Depth = pImageDesc->arraySize;
+    array_desc.Flags |= CUDA_ARRAY3D_TEXTURE_GATHER;
+    break;
   default:
     return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
   }

source/adapters/cuda/device.cpp

@@ -891,7 +891,10 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
     return ReturnValue(
         static_cast<ur_bool_t>(hDevice->supportsHardwareImages()));
   }
-
+  case UR_DEVICE_INFO_BINDLESS_IMAGES_GATHER_EXP: {
+    // HIP doesn't support sampled image gather.
+    return ReturnValue(static_cast<ur_bool_t>(false));
+  }
   case UR_DEVICE_INFO_KERNEL_SET_SPECIALIZATION_CONSTANTS: {
     return ReturnValue(ur_bool_t{false});
   }

source/adapters/cuda/image.cpp

@@ -1165,6 +1165,10 @@ ur_result_t urDeviceGetInfo(
     // L0 does not support sampling 2D USM sampled image data.
     return ReturnValue(false);
   }
+  case UR_DEVICE_INFO_BINDLESS_IMAGES_GATHER_EXP: {
+    // L0 doesn't support sampled image gather.
+    return ReturnValue(static_cast<ur_bool_t>(false));
+  }
   case UR_DEVICE_INFO_PROGRAM_SET_SPECIALIZATION_CONSTANTS:
     return ReturnValue(true);
   case UR_DEVICE_INFO_KERNEL_SET_SPECIALIZATION_CONSTANTS:

source/adapters/hip/device.cpp

@@ -1589,6 +1589,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
   case UR_DEVICE_INFO_BINDLESS_UNIQUE_ADDRESSING_PER_DIM_EXP:
   case UR_DEVICE_INFO_BINDLESS_SAMPLE_1D_USM_EXP:
   case UR_DEVICE_INFO_BINDLESS_SAMPLE_2D_USM_EXP:
+  case UR_DEVICE_INFO_BINDLESS_IMAGES_GATHER_EXP:
     return ReturnValue(false);
   case UR_DEVICE_INFO_IMAGE_PITCH_ALIGN_EXP:
   case UR_DEVICE_INFO_MAX_IMAGE_LINEAR_WIDTH_EXP: