opt: add StorageImageWriteWithoutFormat to trimm pass (#5860)

* opt: add StorageImageWriteWithoutFormat to trimm pass

---------

Signed-off-by: Nathan Gauër <brioche@google.com>

source/opt/trim_capabilities_pass.cpp

@@ -49,6 +49,7 @@ constexpr uint32_t kOpTypeImageMSIndex = kOpTypeImageArrayedIndex + 1;
 constexpr uint32_t kOpTypeImageSampledIndex = kOpTypeImageMSIndex + 1;
 constexpr uint32_t kOpTypeImageFormatIndex = kOpTypeImageSampledIndex + 1;
 constexpr uint32_t kOpImageReadImageIndex = 0;
+constexpr uint32_t kOpImageWriteImageIndex = 0;
 constexpr uint32_t kOpImageSparseReadImageIndex = 0;
 constexpr uint32_t kOpExtInstSetInIndex = 0;
 constexpr uint32_t kOpExtInstInstructionInIndex = 1;
@@ -338,6 +339,8 @@ Handler_OpImageRead_StorageImageReadWithoutFormat(
   const uint32_t dim = type->GetSingleWordInOperand(kOpTypeImageDimIndex);
   const uint32_t format = type->GetSingleWordInOperand(kOpTypeImageFormatIndex);
 
+  // If the Image Format is Unknown and Dim is SubpassData,
+  // StorageImageReadWithoutFormat is required.
   const bool is_unknown = spv::ImageFormat(format) == spv::ImageFormat::Unknown;
   const bool requires_capability_for_unknown =
       spv::Dim(dim) != spv::Dim::SubpassData;
@@ -346,6 +349,26 @@ Handler_OpImageRead_StorageImageReadWithoutFormat(
              : std::nullopt;
 }
 
+static std::optional<spv::Capability>
+Handler_OpImageWrite_StorageImageWriteWithoutFormat(
+    const Instruction* instruction) {
+  assert(instruction->opcode() == spv::Op::OpImageWrite &&
+         "This handler only support OpImageWrite opcodes.");
+  const auto* def_use_mgr = instruction->context()->get_def_use_mgr();
+
+  const uint32_t image_index =
+      instruction->GetSingleWordInOperand(kOpImageWriteImageIndex);
+  const uint32_t type_index = def_use_mgr->GetDef(image_index)->type_id();
+
+  // If the Image Format is Unknown, StorageImageWriteWithoutFormat is required.
+  const Instruction* type = def_use_mgr->GetDef(type_index);
+  const uint32_t format = type->GetSingleWordInOperand(kOpTypeImageFormatIndex);
+  const bool is_unknown = spv::ImageFormat(format) == spv::ImageFormat::Unknown;
+  return is_unknown
+             ? std::optional(spv::Capability::StorageImageWriteWithoutFormat)
+             : std::nullopt;
+}
+
 static std::optional<spv::Capability>
 Handler_OpImageSparseRead_StorageImageReadWithoutFormat(
     const Instruction* instruction) {
@@ -365,9 +388,10 @@ Handler_OpImageSparseRead_StorageImageReadWithoutFormat(
 }
 
 // Opcode of interest to determine capabilities requirements.
-constexpr std::array<std::pair<spv::Op, OpcodeHandler>, 12> kOpcodeHandlers{{
+constexpr std::array<std::pair<spv::Op, OpcodeHandler>, 13> kOpcodeHandlers{{
     // clang-format off
     {spv::Op::OpImageRead,         Handler_OpImageRead_StorageImageReadWithoutFormat},
+    {spv::Op::OpImageWrite,        Handler_OpImageWrite_StorageImageWriteWithoutFormat},
     {spv::Op::OpImageSparseRead,   Handler_OpImageSparseRead_StorageImageReadWithoutFormat},
     {spv::Op::OpTypeFloat,         Handler_OpTypeFloat_Float16 },
     {spv::Op::OpTypeFloat,         Handler_OpTypeFloat_Float64 },

source/opt/trim_capabilities_pass.h

@@ -100,6 +100,7 @@ class TrimCapabilitiesPass : public Pass {
       spv::Capability::Shader,
       spv::Capability::ShaderClockKHR,
       spv::Capability::StorageImageReadWithoutFormat,
+      spv::Capability::StorageImageWriteWithoutFormat,
       spv::Capability::StorageInputOutput16,
       spv::Capability::StoragePushConstant16,
       spv::Capability::StorageUniform16,

test/opt/trim_capabilities_pass_test.cpp

@@ -2366,6 +2366,94 @@ TEST_F(TrimCapabilitiesPassTest,
   EXPECT_EQ(std::get<1>(result), Pass::Status::SuccessWithChange);
 }
 
+TEST_F(TrimCapabilitiesPassTest,
+       StorageImageWriteWithoutFormat_RemainsWhenRequiredWithWrite) {
+  const std::string kTest = R"(
+               OpCapability StorageImageWriteWithoutFormat
+; CHECK:       OpCapability StorageImageWriteWithoutFormat
+               OpCapability Shader
+               OpCapability StorageImageExtendedFormats
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %main "main" %id %img
+               OpExecutionMode %main LocalSize 8 8 8
+               OpSource HLSL 670
+               OpName %type_image "type.3d.image"
+               OpName %img "img"
+               OpName %main "main"
+               OpDecorate %id BuiltIn GlobalInvocationId
+               OpDecorate %img DescriptorSet 0
+               OpDecorate %img Binding 0
+      %float = OpTypeFloat 32
+    %float_4 = OpConstant %float 4
+    %float_5 = OpConstant %float 5
+    %v2float = OpTypeVector %float 2
+          %9 = OpConstantComposite %v2float %float_4 %float_5
+ %type_image = OpTypeImage %float 3D 2 0 0 2 Unknown
+    %ptr_img = OpTypePointer UniformConstant %type_image
+       %uint = OpTypeInt 32 0
+     %v3uint = OpTypeVector %uint 3
+  %ptr_input = OpTypePointer Input %v3uint
+       %void = OpTypeVoid
+         %15 = OpTypeFunction %void
+        %img = OpVariable %ptr_img UniformConstant
+         %id = OpVariable %ptr_input Input
+       %main = OpFunction %void None %15
+         %16 = OpLabel
+         %17 = OpLoad %v3uint %id
+         %18 = OpLoad %type_image %img
+               OpImageWrite %18 %17 %9 None
+               OpReturn
+               OpFunctionEnd
+  )";
+  const auto result =
+      SinglePassRunAndMatch<TrimCapabilitiesPass>(kTest, /* skip_nop= */ false);
+  EXPECT_EQ(std::get<1>(result), Pass::Status::SuccessWithoutChange);
+}
+
+TEST_F(TrimCapabilitiesPassTest,
+       StorageImageWriteWithoutFormat_RemovedWithWriteOnKnownFormat) {
+  const std::string kTest = R"(
+               OpCapability StorageImageWriteWithoutFormat
+; CHECK-NOT:   OpCapability StorageImageWriteWithoutFormat
+               OpCapability Shader
+               OpCapability StorageImageExtendedFormats
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %main "main" %id %img
+               OpExecutionMode %main LocalSize 8 8 8
+               OpSource HLSL 670
+               OpName %type_image "type.3d.image"
+               OpName %img "img"
+               OpName %main "main"
+               OpDecorate %id BuiltIn GlobalInvocationId
+               OpDecorate %img DescriptorSet 0
+               OpDecorate %img Binding 0
+      %float = OpTypeFloat 32
+    %float_4 = OpConstant %float 4
+    %float_5 = OpConstant %float 5
+    %v2float = OpTypeVector %float 2
+          %9 = OpConstantComposite %v2float %float_4 %float_5
+ %type_image = OpTypeImage %float 3D 2 0 0 2 Rg32f
+    %ptr_img = OpTypePointer UniformConstant %type_image
+       %uint = OpTypeInt 32 0
+     %v3uint = OpTypeVector %uint 3
+  %ptr_input = OpTypePointer Input %v3uint
+       %void = OpTypeVoid
+         %15 = OpTypeFunction %void
+        %img = OpVariable %ptr_img UniformConstant
+         %id = OpVariable %ptr_input Input
+       %main = OpFunction %void None %15
+         %16 = OpLabel
+         %17 = OpLoad %v3uint %id
+         %18 = OpLoad %type_image %img
+               OpImageWrite %18 %17 %9 None
+               OpReturn
+               OpFunctionEnd
+  )";
+  const auto result =
+      SinglePassRunAndMatch<TrimCapabilitiesPass>(kTest, /* skip_nop= */ false);
+  EXPECT_EQ(std::get<1>(result), Pass::Status::SuccessWithChange);
+}
+
 TEST_F(TrimCapabilitiesPassTest, PhysicalStorageBuffer_RemovedWhenUnused) {
   const std::string kTest = R"(
                OpCapability PhysicalStorageBufferAddresses