[spirv] Handle Status arg in Texture Gather methods and Refactor (#880)

* [spirv] Add Status arg in Texture Gather methods
* [spirv] Refactor handling of residency code struct

tools/clang/include/clang/SPIRV/ModuleBuilder.h

@@ -201,12 +201,15 @@ public:
   /// If compareVal is given a non-zero value, OpImageDrefGather or
   /// OpImageSparseDrefGather will be generated; otherwise, OpImageGather or
   /// OpImageSparseGather will be generated.
+  /// If residencyCodeId is not zero, the sparse version of the instructions will
+  /// be used, and the SPIR-V instruction for storing the resulting residency
+  /// code will also be emitted.
   uint32_t createImageGather(uint32_t texelType, uint32_t imageType,
                              uint32_t image, uint32_t sampler,
                              uint32_t coordinate, uint32_t component,
                              uint32_t compareVal, uint32_t constOffset,
                              uint32_t varOffset, uint32_t constOffsets,
-                             uint32_t sample, bool isSparse);
+                             uint32_t sample, uint32_t residencyCodeId);
 
   /// \brief Creates a select operation with the given values for true and false
   /// cases and returns the <result-id> for the result.
@@ -367,6 +370,11 @@ public:
   uint32_t getSampledImageType(uint32_t imageType);
   uint32_t getByteAddressBufferType(bool isRW);
 
+  /// \brief Returns a struct type with 2 members. The first member is an
+  /// unsigned integer type which can hold the 'Residency Code'. The second
+  /// member will be of the given type.
+  uint32_t getSparseResidencyStructType(uint32_t type);
+
   // === Constant ===
   uint32_t getConstantBool(bool value);
   uint32_t getConstantInt32(int32_t value);

tools/clang/lib/SPIRV/ModuleBuilder.cpp

@@ -438,11 +438,14 @@ uint32_t ModuleBuilder::createImageGather(
     uint32_t texelType, uint32_t imageType, uint32_t image, uint32_t sampler,
     uint32_t coordinate, uint32_t component, uint32_t compareVal,
     uint32_t constOffset, uint32_t varOffset, uint32_t constOffsets,
-    uint32_t sample, bool isSparse) {
+    uint32_t sample, uint32_t residencyCodeId) {
   assert(insertPoint && "null insert point");
 
-  if (isSparse)
+  uint32_t sparseRetType = 0;
+  if (residencyCodeId) {
     requireCapability(spv::Capability::SparseResidency);
+    sparseRetType = getSparseResidencyStructType(texelType);
+  }
 
   // An OpSampledImage is required to do the image sampling.
   const uint32_t sampledImgId = theContext.takeNextId();
@@ -456,12 +459,12 @@ uint32_t ModuleBuilder::createImageGather(
       llvm::Optional<spv::ImageOperandsMask>(composeImageOperandsMask(
           /*bias*/ 0, /*lod*/ 0, std::make_pair(0, 0), constOffset, varOffset,
           constOffsets, sample, &params));
-  const uint32_t texelId = theContext.takeNextId();
+  uint32_t texelId = theContext.takeNextId();
 
   if (compareVal) {
-    if (isSparse) {
+    if (residencyCodeId) {
       // Note: OpImageSparseDrefGather does not take the component parameter.
-      instBuilder.opImageSparseDrefGather(texelType, texelId, sampledImgId,
+      instBuilder.opImageSparseDrefGather(sparseRetType, texelId, sampledImgId,
                                           coordinate, compareVal, mask);
     } else {
       // Note: OpImageDrefGather does not take the component parameter.
@@ -469,8 +472,8 @@ uint32_t ModuleBuilder::createImageGather(
                                     coordinate, compareVal, mask);
     }
   } else {
-    if (isSparse) {
-      instBuilder.opImageSparseGather(texelType, texelId, sampledImgId,
+    if (residencyCodeId) {
+      instBuilder.opImageSparseGather(sparseRetType, texelId, sampledImgId,
                                       coordinate, component, mask);
     } else {
       instBuilder.opImageGather(texelType, texelId, sampledImgId, coordinate,
@@ -483,6 +486,14 @@ uint32_t ModuleBuilder::createImageGather(
   instBuilder.x();
   insertPoint->appendInstruction(std::move(constructSite));
 
+  if (residencyCodeId) {
+    // Write the Residency Code
+    const auto status = createCompositeExtract(getUint32Type(), texelId, {0});
+    createStore(residencyCodeId, status);
+    // Extract the real result from the struct
+    texelId = createCompositeExtract(texelType, texelId, {1});
+  }
+
   return texelId;
 }
 
@@ -820,6 +831,12 @@ ModuleBuilder::getStructType(llvm::ArrayRef<uint32_t> fieldTypes,
   return typeId;
 }
 
+uint32_t ModuleBuilder::getSparseResidencyStructType(uint32_t type) {
+  const auto uintType = getUint32Type();
+  return getStructType({uintType, type}, "SparseResidencyStruct",
+                       {"Residency.Code", "Result.Type"});
+}
+
 uint32_t ModuleBuilder::getArrayType(uint32_t elemType, uint32_t count,
                                      Type::DecorationSet decorations) {
   const Type *type = Type::getArray(theContext, elemType, count, decorations);

tools/clang/lib/SPIRV/SPIRVEmitter.cpp

@@ -2296,31 +2296,11 @@ uint32_t SPIRVEmitter::processTextureGatherRGBACmpRGBA(
         offsetType, {offset0, offset1, offset2, offset3});
   }
 
-  if (!hasStatusArg)
-    return theBuilder.createImageGather(
-        retTypeId, imageTypeId, image, sampler, coordinate,
-        theBuilder.getConstantInt32(component), compareVal, constOffset,
-        varOffset, constOffsets, /*sampleNumber*/ 0, /*isSparse*/ false);
-
-  // If the Status parameter is present, OpImageSparseGather should be used.
-  // The result type of this SPIR-V instruction is a struct in which the first
-  // member is an integer that holds the Residency Code.
-  const auto uintType = theBuilder.getUint32Type();
-  const auto sparseRetType =
-      theBuilder.getStructType({uintType, retTypeId}, "SparseResidencyStruct",
-                               {"Residency.Code", "Result.Type"});
-
-  // Perform ImageSparseGather
-  const auto sparseGather = theBuilder.createImageGather(
-      sparseRetType, imageTypeId, image, sampler, coordinate,
+  const auto status = hasStatusArg ? doExpr(expr->getArg(numArgs - 1)) : 0;
+  return theBuilder.createImageGather(
+      retTypeId, imageTypeId, image, sampler, coordinate,
       theBuilder.getConstantInt32(component), compareVal, constOffset,
-      varOffset, constOffsets, /*sampleNumber*/ 0, /*isSparse*/ true);
-  // Write the Residency Code
-  const auto status =
-      theBuilder.createCompositeExtract(uintType, sparseGather, {0});
-  theBuilder.createStore(doExpr(expr->getArg(numArgs - 1)), status);
-  // Return the results
-  return theBuilder.createCompositeExtract(retTypeId, sparseGather, {1});
+      varOffset, constOffsets, /*sampleNumber*/ 0, status);
 }
 
 uint32_t SPIRVEmitter::processTextureGatherCmp(const CXXMemberCallExpr *expr) {
@@ -2362,32 +2342,12 @@ uint32_t SPIRVEmitter::processTextureGatherCmp(const CXXMemberCallExpr *expr) {
 
   const auto retType = typeTranslator.translateType(callee->getReturnType());
   const auto imageType = typeTranslator.translateType(imageExpr->getType());
+  const auto status = hasStatusArg ? doExpr(expr->getArg(numArgs - 1)) : 0;
 
-  if(!hasStatusArg)
-    return theBuilder.createImageGather(
-        retType, imageType, image, sampler, coordinate,
-        /*component*/ 0, comparator, constOffset, varOffset, /*constOffsets*/ 0,
-        /*sampleNumber*/ 0, /*isSparse*/ false);
-
-  // If the Status parameter is present, OpImageSparseGather should be used.
-  // The result type of this SPIR-V instruction is a struct in which the first
-  // member is an integer that holds the Residency Code.
-  const auto uintType = theBuilder.getUint32Type();
-  const auto sparseRetType =
-      theBuilder.getStructType({uintType, retType}, "SparseResidencyStruct",
-                               {"Residency.Code", "Result.Type"});
-
-  // Perform ImageSparseGather
-  const auto sparseGather = theBuilder.createImageGather(
-      sparseRetType, imageType, image, sampler, coordinate,
+  return theBuilder.createImageGather(
+      retType, imageType, image, sampler, coordinate,
       /*component*/ 0, comparator, constOffset, varOffset, /*constOffsets*/ 0,
-      /*sampleNumber*/ 0, /*isSparse*/ true);
-  // Write the Residency Code
-  const auto status =
-      theBuilder.createCompositeExtract(uintType, sparseGather, {0});
-  theBuilder.createStore(doExpr(expr->getArg(numArgs - 1)), status);
-  // Return the results
-  return theBuilder.createCompositeExtract(retType, sparseGather, {1});
+      /*sampleNumber*/ 0, status);
 }
 
 SpirvEvalInfo SPIRVEmitter::processBufferTextureLoad(const Expr *object,
@@ -2825,35 +2785,52 @@ uint32_t SPIRVEmitter::processTextureSampleGather(const CXXMemberCallExpr *expr,
   //                           float Location
   //                           [, int Offset]);
   //
+  // For Texture2D/Texture2DArray:
+  // <Template Type>4 Object.Gather(sampler_state S,
+  //                                float2|3|4 Location,
+  //                                int2 Offset
+  //                                [, uint Status]);
+  //
+  // For TextureCube/TextureCubeArray:
   // <Template Type>4 Object.Gather(sampler_state S,
   //                                float2|3|4 Location
-  //                                [, int2 Offset]);
+  //                                [, uint Status]);
+  //
+  // Other Texture types do not have a Gather method.
+
+  const auto numArgs = expr->getNumArgs();
+  const bool hasStatusArg =
+      expr->getArg(numArgs - 1)->getType()->isUnsignedIntegerType();
+  const auto status = hasStatusArg ? doExpr(expr->getArg(numArgs - 1)) : 0;
+
+  // Subtract 1 for status (if it exists), and 2 for sampler_state and location.
+  const bool hasOffsetArg = numArgs - hasStatusArg - 2 > 0;
 
   const auto *imageExpr = expr->getImplicitObjectArgument();
   const uint32_t imageType = typeTranslator.translateType(imageExpr->getType());
-
   const uint32_t image = loadIfGLValue(imageExpr);
   const uint32_t sampler = doExpr(expr->getArg(0));
   const uint32_t coordinate = doExpr(expr->getArg(1));
-  // .Sample()/.Gather() has a third optional paramter for offset.
+  // .Sample()/.Gather() may have a third optional paramter for offset.
   uint32_t constOffset = 0, varOffset = 0;
-  handleOptionalOffsetInMethodCall(expr, 2, &constOffset, &varOffset);
+  if(hasOffsetArg)
+    handleOffsetInMethodCall(expr, 2, &constOffset, &varOffset);
 
   const auto retType =
       typeTranslator.translateType(expr->getDirectCallee()->getReturnType());
 
   if (isSample) {
+    // TODO: Handle sparse cases for this method.
     return theBuilder.createImageSample(
         retType, imageType, image, sampler, coordinate, /*compareVal*/ 0,
         /*bias*/ 0, /*lod*/ 0, std::make_pair(0, 0), constOffset, varOffset,
         /*constOffsets*/ 0, /*sampleNumber*/ 0);
   } else {
-    // TODO: update this function to handle sparse cases.
     return theBuilder.createImageGather(
         retType, imageType, image, sampler, coordinate,
         // .Gather() doc says we return four components of red data.
         theBuilder.getConstantInt32(0), /*compareVal*/ 0, constOffset,
-        varOffset, /*constOffsets*/ 0, /*sampleNumber*/ 0, /*isSparse*/ false);
+        varOffset, /*constOffsets*/ 0, /*sampleNumber*/ 0, status);
   }
 }
 

tools/clang/test/CodeGenSPIRV/texture.array.gather.hlsl

@@ -9,6 +9,10 @@ TextureCubeArray <float>  t8 : register(t8);
 // .Gather() does not support Texture1DArray.
 
 // CHECK: OpCapability ImageGatherExtended
+// CHECK: OpCapability SparseResidency
+
+// CHECK: %SparseResidencyStruct = OpTypeStruct %uint %v4int
+// CHECK: %SparseResidencyStruct_0 = OpTypeStruct %uint %v4float
 
 // CHECK: [[v4fc:%\d+]] = OpConstantComposite %v4float %float_0_1 %float_0_2 %float_0_3 %float_0_4
 
@@ -42,5 +46,29 @@ float4 main(float3 location: A, int2 offset: B) : SV_Target {
 // CHECK-NEXT:            {{%\d+}} = OpImageGather %v4float [[sampledImg]] [[v4fc]] %int_0
     float4 val8 = t8.Gather(gSampler, float4(0.1, 0.2, 0.3, 0.4));
 
+    uint status;
+// CHECK:                [[t6:%\d+]] = OpLoad %type_2d_image_array_0 %t6
+// CHECK-NEXT:     [[gSampler:%\d+]] = OpLoad %type_sampler %gSampler
+// CHECK-NEXT:          [[loc:%\d+]] = OpLoad %v3float %location
+// CHECK-NEXT:       [[offset:%\d+]] = OpLoad %v2int %offset
+// CHECK-NEXT:   [[sampledImg:%\d+]] = OpSampledImage %type_sampled_image_1 [[t6]] [[gSampler]]
+// CHECK-NEXT: [[structResult:%\d+]] = OpImageSparseGather %SparseResidencyStruct [[sampledImg]] [[loc]] %int_0 Offset [[offset]]
+// CHECK-NEXT:       [[status:%\d+]] = OpCompositeExtract %uint [[structResult]] 0
+// CHECK-NEXT:                         OpStore %status [[status]]
+// CHECK-NEXT:       [[result:%\d+]] = OpCompositeExtract %v4int [[structResult]] 1
+// CHECK-NEXT:                         OpStore %val9 [[result]]
+    int4 val9 = t6.Gather(gSampler, location, offset, status);
+
+// CHECK:                [[t8:%\d+]] = OpLoad %type_cube_image_array_0 %t8
+// CHECK-NEXT:     [[gSampler:%\d+]] = OpLoad %type_sampler %gSampler
+// CHECK-NEXT:   [[sampledImg:%\d+]] = OpSampledImage %type_sampled_image_2 [[t8]] [[gSampler]]
+// CHECK-NEXT: [[structResult:%\d+]] = OpImageSparseGather %SparseResidencyStruct_0 [[sampledImg]] [[v4fc]] %int_0 None
+// CHECK-NEXT:       [[status:%\d+]] = OpCompositeExtract %uint [[structResult]] 0
+// CHECK-NEXT:                         OpStore %status [[status]]
+// CHECK-NEXT:       [[result:%\d+]] = OpCompositeExtract %v4float [[structResult]] 1
+// CHECK-NEXT:                         OpStore %val10 [[result]]
+    float4 val10 = t8.Gather(gSampler, float4(0.1, 0.2, 0.3, 0.4), status);
+
     return 1.0;
 }
+

tools/clang/test/CodeGenSPIRV/texture.gather.hlsl

@@ -8,6 +8,11 @@ Texture2D   <int2>   t6 : register(t6);
 TextureCube <float>  t8 : register(t8);
 // .Gather() does not support Texture1D and Texture3D.
 
+// CHECK: OpCapability SparseResidency
+
+// CHECK: %SparseResidencyStruct = OpTypeStruct %uint %v4int
+// CHECK: %SparseResidencyStruct_0 = OpTypeStruct %uint %v4float
+
 // CHECK: [[v2ic:%\d+]] = OpConstantComposite %v2int %int_1 %int_2
 // CHECK: [[v3fc:%\d+]] = OpConstantComposite %v3float %float_0_1 %float_0_2 %float_0_3
 
@@ -39,5 +44,28 @@ float4 main(float2 location: A) : SV_Target {
 // CHECK-NEXT:            {{%\d+}} = OpImageGather %v4float [[sampledImg]] [[v3fc]] %int_0
     float4 val8 = t8.Gather(gSampler, float3(0.1, 0.2, 0.3));
 
+    uint status;
+
+// CHECK:                [[t6:%\d+]] = OpLoad %type_2d_image_0 %t6
+// CHECK-NEXT:     [[gSampler:%\d+]] = OpLoad %type_sampler %gSampler
+// CHECK-NEXT:          [[loc:%\d+]] = OpLoad %v2float %location
+// CHECK-NEXT:   [[sampledImg:%\d+]] = OpSampledImage %type_sampled_image_1 [[t6]] [[gSampler]]
+// CHECK-NEXT: [[structResult:%\d+]] = OpImageSparseGather %SparseResidencyStruct [[sampledImg]] [[loc]] %int_0 ConstOffset [[v2ic]]
+// CHECK-NEXT:       [[status:%\d+]] = OpCompositeExtract %uint [[structResult]] 0
+// CHECK-NEXT:                         OpStore %status [[status]]
+// CHECK-NEXT:       [[result:%\d+]] = OpCompositeExtract %v4int [[structResult]] 1
+// CHECK-NEXT:                         OpStore %val9 [[result]]
+    int4 val9 = t6.Gather(gSampler, location, int2(1, 2), status);
+
+// CHECK:                [[t8:%\d+]] = OpLoad %type_cube_image_0 %t8
+// CHECK-NEXT:     [[gSampler:%\d+]] = OpLoad %type_sampler %gSampler
+// CHECK-NEXT:   [[sampledImg:%\d+]] = OpSampledImage %type_sampled_image_2 [[t8]] [[gSampler]]
+// CHECK-NEXT: [[structResult:%\d+]] = OpImageSparseGather %SparseResidencyStruct_0 [[sampledImg]] [[v3fc]] %int_0 None
+// CHECK-NEXT:       [[status:%\d+]] = OpCompositeExtract %uint [[structResult]] 0
+// CHECK-NEXT:                         OpStore %status [[status]]
+// CHECK-NEXT:       [[result:%\d+]] = OpCompositeExtract %v4float [[structResult]] 1
+// CHECK-NEXT:                         OpStore %val10 [[result]]
+    float4 val10 = t8.Gather(gSampler, float3(0.1, 0.2, 0.3), status);
+
     return 1.0;
 }