[spirv] support SPV_EXT_shader_image_int64 extension for 64-bit integer resource (#3787)

This commit adds `Int64ImageEXT` capability and
`SPV_EXT_shader_image_int64` extension to the SPIRV backend
when 64-bit integer types are used on an image resource.

Author: LukasBanana

tools/clang/include/clang/SPIRV/FeatureManager.h

@@ -51,6 +51,7 @@ enum class Extension {
   NV_ray_tracing,
   NV_mesh_shader,
   KHR_ray_query,
+  EXT_shader_image_int64,
   Unknown,
 };
 

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -171,7 +171,16 @@ void CapabilityVisitor::addCapabilityForType(const SpirvType *type,
     if (imageType->isArrayedImage() && imageType->isMSImage())
       addCapability(spv::Capability::ImageMSArray);
 
-    addCapabilityForType(imageType->getSampledType(), loc, sc);
+    if (const auto *sampledType = imageType->getSampledType()) {
+      addCapabilityForType(sampledType, loc, sc);
+      if (const auto *sampledIntType = dyn_cast<IntegerType>(sampledType)) {
+        if (sampledIntType->getBitwidth() == 64) {
+          addCapability(spv::Capability::Int64ImageEXT);
+          addExtension(Extension::EXT_shader_image_int64,
+                       "64-bit image types in resource", loc);
+        }
+      }
+    }
   }
   // Sampled image type
   else if (const auto *sampledImageType = dyn_cast<SampledImageType>(type)) {

tools/clang/lib/SPIRV/FeatureManager.cpp

@@ -142,6 +142,7 @@ Extension FeatureManager::getExtensionSymbol(llvm::StringRef name) {
       .Case("SPV_KHR_ray_query", Extension::KHR_ray_query)
       .Case("SPV_KHR_fragment_shading_rate",
             Extension::KHR_fragment_shading_rate)
+      .Case("SPV_EXT_shader_image_int64", Extension::EXT_shader_image_int64)
       .Default(Extension::Unknown);
 }
 
@@ -193,6 +194,8 @@ const char *FeatureManager::getExtensionName(Extension symbol) {
     return "SPV_KHR_ray_query";
   case Extension::KHR_fragment_shading_rate:
     return "SPV_KHR_fragment_shading_rate";
+  case Extension::EXT_shader_image_int64:
+    return "SPV_EXT_shader_image_int64";
   default:
     break;
   }

tools/clang/lib/SPIRV/LowerTypeVisitor.cpp

@@ -758,6 +758,12 @@ LowerTypeVisitor::translateSampledTypeToImageFormat(QualType sampledType,
         return elemCount == 1 ? spv::ImageFormat::R32f
                               : elemCount == 2 ? spv::ImageFormat::Rg32f
                                                : spv::ImageFormat::Rgba32f;
+      case BuiltinType::LongLong:
+        if (elemCount == 1)
+          return spv::ImageFormat::R64i;
+      case BuiltinType::ULongLong:
+        if (elemCount == 1)
+          return spv::ImageFormat::R64ui;
       default:
         // Other sampled types unimplemented or irrelevant.
         break;

tools/clang/test/CodeGenSPIRV/type.rwtexture.with.64bit.scalar.hlsl

@@ -0,0 +1,25 @@
+// Run: %dxc -T cs_6_0 -E main
+
+// CHECK: %type_2d_image = OpTypeImage %ulong 2D 2 0 0 2 R64ui
+// CHECK: %_ptr_UniformConstant_type_2d_image = OpTypePointer UniformConstant %type_2d_image
+
+// CHECK: %type_2d_image_0 = OpTypeImage %long 2D 2 0 0 2 R64i
+// CHECK: %_ptr_UniformConstant_type_2d_image_0 = OpTypePointer UniformConstant %type_2d_image_0
+
+// CHECK: %type_2d_image_1 = OpTypeImage %int 2D 2 0 0 2 Rgba32i
+// CHECK: %_ptr_UniformConstant_type_2d_image_1 = OpTypePointer UniformConstant %type_2d_image_1
+
+// CHECK: %tex_ui = OpVariable %_ptr_UniformConstant_type_2d_image UniformConstant
+RWTexture2D<uint64_t> tex_ui;
+
+// CHECK: %tex_i = OpVariable %_ptr_UniformConstant_type_2d_image_0 UniformConstant
+RWTexture2D<int64_t> tex_i;
+
+// CHECK: %texout = OpVariable %_ptr_UniformConstant_type_2d_image_1 UniformConstant
+RWTexture2D<min12int4> texout;
+
+[numthreads(8, 8, 1)]
+void main(uint2 id : SV_DispatchThreadID)
+{
+    texout[id] =  tex_i[id] + tex_ui[id];
+};

tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -72,6 +72,9 @@ TEST_F(FileTest, RWTextureTypes) { runFileTest("type.rwtexture.hlsl"); }
 TEST_F(FileTest, RWTextureTypesWithMinPrecisionScalarTypes) {
   runFileTest("type.rwtexture.with.min.precision.scalar.hlsl");
 }
+TEST_F(FileTest, RWTextureTypesWith64bitsScalarTypes) {
+  runFileTest("type.rwtexture.with.64bit.scalar.hlsl");
+}
 TEST_F(FileTest, BufferType) { runFileTest("type.buffer.hlsl"); }
 TEST_F(FileTest, BufferTypeStructError1) {
   runFileTest("type.buffer.struct.error1.hlsl", Expect::Failure);