Fix identification of Vulkan images and buffers (#3253)

Fixes #3252

* Image and buffer queries did not account for optional level of
arrayness on the variable
  * new tests

source/opt/instruction.cpp

@@ -232,6 +232,14 @@ bool Instruction::IsVulkanStorageImage() const {
 
   Instruction* base_type =
       context()->get_def_use_mgr()->GetDef(GetSingleWordInOperand(1));
+
+  // Unpack the optional layer of arraying.
+  if (base_type->opcode() == SpvOpTypeArray ||
+      base_type->opcode() == SpvOpTypeRuntimeArray) {
+    base_type = context()->get_def_use_mgr()->GetDef(
+        base_type->GetSingleWordInOperand(0));
+  }
+
   if (base_type->opcode() != SpvOpTypeImage) {
     return false;
   }
@@ -258,6 +266,14 @@ bool Instruction::IsVulkanSampledImage() const {
 
   Instruction* base_type =
       context()->get_def_use_mgr()->GetDef(GetSingleWordInOperand(1));
+
+  // Unpack the optional layer of arraying.
+  if (base_type->opcode() == SpvOpTypeArray ||
+      base_type->opcode() == SpvOpTypeRuntimeArray) {
+    base_type = context()->get_def_use_mgr()->GetDef(
+        base_type->GetSingleWordInOperand(0));
+  }
+
   if (base_type->opcode() != SpvOpTypeImage) {
     return false;
   }
@@ -284,6 +300,14 @@ bool Instruction::IsVulkanStorageTexelBuffer() const {
 
   Instruction* base_type =
       context()->get_def_use_mgr()->GetDef(GetSingleWordInOperand(1));
+
+  // Unpack the optional layer of arraying.
+  if (base_type->opcode() == SpvOpTypeArray ||
+      base_type->opcode() == SpvOpTypeRuntimeArray) {
+    base_type = context()->get_def_use_mgr()->GetDef(
+        base_type->GetSingleWordInOperand(0));
+  }
+
   if (base_type->opcode() != SpvOpTypeImage) {
     return false;
   }
@@ -307,6 +331,13 @@ bool Instruction::IsVulkanStorageBuffer() const {
   Instruction* base_type =
       context()->get_def_use_mgr()->GetDef(GetSingleWordInOperand(1));
 
+  // Unpack the optional layer of arraying.
+  if (base_type->opcode() == SpvOpTypeArray ||
+      base_type->opcode() == SpvOpTypeRuntimeArray) {
+    base_type = context()->get_def_use_mgr()->GetDef(
+        base_type->GetSingleWordInOperand(0));
+  }
+
   if (base_type->opcode() != SpvOpTypeStruct) {
     return false;
   }
@@ -340,6 +371,14 @@ bool Instruction::IsVulkanUniformBuffer() const {
 
   Instruction* base_type =
       context()->get_def_use_mgr()->GetDef(GetSingleWordInOperand(1));
+
+  // Unpack the optional layer of arraying.
+  if (base_type->opcode() == SpvOpTypeArray ||
+      base_type->opcode() == SpvOpTypeRuntimeArray) {
+    base_type = context()->get_def_use_mgr()->GetDef(
+        base_type->GetSingleWordInOperand(0));
+  }
+
   if (base_type->opcode() != SpvOpTypeStruct) {
     return false;
   }

test/opt/instruction_test.cpp

@@ -35,6 +35,7 @@ using DescriptorTypeTest = PassTest<::testing::Test>;
 using OpaqueTypeTest = PassTest<::testing::Test>;
 using GetBaseTest = PassTest<::testing::Test>;
 using ValidBasePointerTest = PassTest<::testing::Test>;
+using VulkanBufferTest = PassTest<::testing::Test>;
 
 TEST(InstructionTest, CreateTrivial) {
   Instruction empty;
@@ -1143,6 +1144,260 @@ OpFunctionEnd
   EXPECT_TRUE(null_inst->IsValidBasePointer());
 }
 
+TEST_F(VulkanBufferTest, VulkanStorageBuffer) {
+  const std::string text = R"(
+OpCapability Shader
+OpCapability RuntimeDescriptorArray
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %1 "main"
+OpExecutionMode %1 LocalSize 1 1 1
+OpDecorate %2 Block
+OpMemberDecorate %2 0 Offset 0
+OpDecorate %3 BufferBlock
+OpMemberDecorate %3 0 Offset 0
+%4 = OpTypeVoid
+%5 = OpTypeInt 32 0
+%2 = OpTypeStruct %5
+%3 = OpTypeStruct %5
+
+%6 = OpTypePointer StorageBuffer %2
+%7 = OpTypePointer Uniform %2
+%8 = OpTypePointer Uniform %3
+
+%9 = OpConstant %5 1
+%10 = OpTypeArray %2 %9
+%11 = OpTypeArray %3 %9
+%12 = OpTypePointer StorageBuffer %10
+%13 = OpTypePointer Uniform %10
+%14 = OpTypePointer Uniform %11
+
+%15 = OpTypeRuntimeArray %2
+%16 = OpTypeRuntimeArray %3
+%17 = OpTypePointer StorageBuffer %15
+%18 = OpTypePointer Uniform %15
+%19 = OpTypePointer Uniform %16
+
+%50 = OpTypeFunction %4
+%1 = OpFunction %4 None %50
+%51 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  std::unique_ptr<IRContext> context =
+      BuildModule(SPV_ENV_UNIVERSAL_1_3, nullptr, text);
+  EXPECT_NE(context, nullptr);
+
+  // Standard SSBO and UBO
+  Instruction* inst = context->get_def_use_mgr()->GetDef(6);
+  EXPECT_EQ(true, inst->IsVulkanStorageBuffer());
+  inst = context->get_def_use_mgr()->GetDef(7);
+  EXPECT_EQ(false, inst->IsVulkanStorageBuffer());
+  inst = context->get_def_use_mgr()->GetDef(8);
+  EXPECT_EQ(true, inst->IsVulkanStorageBuffer());
+
+  // Arrayed SSBO and UBO
+  inst = context->get_def_use_mgr()->GetDef(12);
+  EXPECT_EQ(true, inst->IsVulkanStorageBuffer());
+  inst = context->get_def_use_mgr()->GetDef(13);
+  EXPECT_EQ(false, inst->IsVulkanStorageBuffer());
+  inst = context->get_def_use_mgr()->GetDef(14);
+  EXPECT_EQ(true, inst->IsVulkanStorageBuffer());
+
+  // Runtime arrayed SSBO and UBO
+  inst = context->get_def_use_mgr()->GetDef(17);
+  EXPECT_EQ(true, inst->IsVulkanStorageBuffer());
+  inst = context->get_def_use_mgr()->GetDef(18);
+  EXPECT_EQ(false, inst->IsVulkanStorageBuffer());
+  inst = context->get_def_use_mgr()->GetDef(19);
+  EXPECT_EQ(true, inst->IsVulkanStorageBuffer());
+}
+
+TEST_F(VulkanBufferTest, VulkanUniformBuffer) {
+  const std::string text = R"(
+OpCapability Shader
+OpCapability RuntimeDescriptorArray
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %1 "main"
+OpExecutionMode %1 LocalSize 1 1 1
+OpDecorate %2 Block
+OpMemberDecorate %2 0 Offset 0
+OpDecorate %3 BufferBlock
+OpMemberDecorate %3 0 Offset 0
+%4 = OpTypeVoid
+%5 = OpTypeInt 32 0
+%2 = OpTypeStruct %5
+%3 = OpTypeStruct %5
+
+%6 = OpTypePointer StorageBuffer %2
+%7 = OpTypePointer Uniform %2
+%8 = OpTypePointer Uniform %3
+
+%9 = OpConstant %5 1
+%10 = OpTypeArray %2 %9
+%11 = OpTypeArray %3 %9
+%12 = OpTypePointer StorageBuffer %10
+%13 = OpTypePointer Uniform %10
+%14 = OpTypePointer Uniform %11
+
+%15 = OpTypeRuntimeArray %2
+%16 = OpTypeRuntimeArray %3
+%17 = OpTypePointer StorageBuffer %15
+%18 = OpTypePointer Uniform %15
+%19 = OpTypePointer Uniform %16
+
+%50 = OpTypeFunction %4
+%1 = OpFunction %4 None %50
+%51 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  std::unique_ptr<IRContext> context =
+      BuildModule(SPV_ENV_UNIVERSAL_1_3, nullptr, text);
+  EXPECT_NE(context, nullptr);
+
+  // Standard SSBO and UBO
+  Instruction* inst = context->get_def_use_mgr()->GetDef(6);
+  EXPECT_EQ(false, inst->IsVulkanUniformBuffer());
+  inst = context->get_def_use_mgr()->GetDef(7);
+  EXPECT_EQ(true, inst->IsVulkanUniformBuffer());
+  inst = context->get_def_use_mgr()->GetDef(8);
+  EXPECT_EQ(false, inst->IsVulkanUniformBuffer());
+
+  // Arrayed SSBO and UBO
+  inst = context->get_def_use_mgr()->GetDef(12);
+  EXPECT_EQ(false, inst->IsVulkanUniformBuffer());
+  inst = context->get_def_use_mgr()->GetDef(13);
+  EXPECT_EQ(true, inst->IsVulkanUniformBuffer());
+  inst = context->get_def_use_mgr()->GetDef(14);
+  EXPECT_EQ(false, inst->IsVulkanUniformBuffer());
+
+  // Runtime arrayed SSBO and UBO
+  inst = context->get_def_use_mgr()->GetDef(17);
+  EXPECT_EQ(false, inst->IsVulkanUniformBuffer());
+  inst = context->get_def_use_mgr()->GetDef(18);
+  EXPECT_EQ(true, inst->IsVulkanUniformBuffer());
+  inst = context->get_def_use_mgr()->GetDef(19);
+  EXPECT_EQ(false, inst->IsVulkanUniformBuffer());
+}
+
+TEST_F(VulkanBufferTest, ImageQueries) {
+  const std::string text = R"(
+OpCapability Shader
+OpCapability ImageBuffer
+OpCapability RuntimeDescriptorArray
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %1 "main"
+OpExecutionMode %1 LocalSize 1 1 1
+%2 = OpTypeVoid
+%3 = OpTypeFloat 32
+
+%4 = OpTypeImage %3 Buffer 0 0 0 1 Rgba32f
+%5 = OpTypeImage %3 Buffer 0 0 0 2 Rgba32f
+%6 = OpTypeImage %3 2D 0 0 0 1 Rgba32f
+%7 = OpTypeImage %3 2D 0 0 0 2 Rgba32f
+
+%8 = OpTypePointer UniformConstant %4
+%9 = OpTypePointer UniformConstant %5
+%10 = OpTypePointer UniformConstant %6
+%11 = OpTypePointer UniformConstant %7
+
+%12 = OpTypeInt 32 0
+%13 = OpConstant %12 1
+%14 = OpTypeArray %4 %13
+%15 = OpTypeArray %5 %13
+%16 = OpTypeArray %6 %13
+%17 = OpTypeArray %7 %13
+%18 = OpTypePointer UniformConstant %14
+%19 = OpTypePointer UniformConstant %15
+%20 = OpTypePointer UniformConstant %16
+%21 = OpTypePointer UniformConstant %17
+
+%22 = OpTypeRuntimeArray %4
+%23 = OpTypeRuntimeArray %5
+%24 = OpTypeRuntimeArray %6
+%25 = OpTypeRuntimeArray %7
+%26 = OpTypePointer UniformConstant %22
+%27 = OpTypePointer UniformConstant %23
+%28 = OpTypePointer UniformConstant %24
+%29 = OpTypePointer UniformConstant %25
+
+%50 = OpTypeFunction %4
+%1 = OpFunction %4 None %50
+%51 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  std::unique_ptr<IRContext> context =
+      BuildModule(SPV_ENV_UNIVERSAL_1_3, nullptr, text);
+  EXPECT_NE(context, nullptr);
+
+  // Bare pointers
+  Instruction* inst = context->get_def_use_mgr()->GetDef(8);
+  EXPECT_EQ(false, inst->IsVulkanStorageImage());
+  EXPECT_EQ(false, inst->IsVulkanSampledImage());
+  EXPECT_EQ(false, inst->IsVulkanStorageTexelBuffer());
+
+  inst = context->get_def_use_mgr()->GetDef(9);
+  EXPECT_EQ(false, inst->IsVulkanStorageImage());
+  EXPECT_EQ(false, inst->IsVulkanSampledImage());
+  EXPECT_EQ(true, inst->IsVulkanStorageTexelBuffer());
+
+  inst = context->get_def_use_mgr()->GetDef(10);
+  EXPECT_EQ(false, inst->IsVulkanStorageImage());
+  EXPECT_EQ(true, inst->IsVulkanSampledImage());
+  EXPECT_EQ(false, inst->IsVulkanStorageTexelBuffer());
+
+  inst = context->get_def_use_mgr()->GetDef(11);
+  EXPECT_EQ(true, inst->IsVulkanStorageImage());
+  EXPECT_EQ(false, inst->IsVulkanSampledImage());
+  EXPECT_EQ(false, inst->IsVulkanStorageTexelBuffer());
+
+  // Array pointers
+  inst = context->get_def_use_mgr()->GetDef(18);
+  EXPECT_EQ(false, inst->IsVulkanStorageImage());
+  EXPECT_EQ(false, inst->IsVulkanSampledImage());
+  EXPECT_EQ(false, inst->IsVulkanStorageTexelBuffer());
+
+  inst = context->get_def_use_mgr()->GetDef(19);
+  EXPECT_EQ(false, inst->IsVulkanStorageImage());
+  EXPECT_EQ(false, inst->IsVulkanSampledImage());
+  EXPECT_EQ(true, inst->IsVulkanStorageTexelBuffer());
+
+  inst = context->get_def_use_mgr()->GetDef(20);
+  EXPECT_EQ(false, inst->IsVulkanStorageImage());
+  EXPECT_EQ(true, inst->IsVulkanSampledImage());
+  EXPECT_EQ(false, inst->IsVulkanStorageTexelBuffer());
+
+  inst = context->get_def_use_mgr()->GetDef(21);
+  EXPECT_EQ(true, inst->IsVulkanStorageImage());
+  EXPECT_EQ(false, inst->IsVulkanSampledImage());
+  EXPECT_EQ(false, inst->IsVulkanStorageTexelBuffer());
+
+  // Runtime array pointers
+  inst = context->get_def_use_mgr()->GetDef(26);
+  EXPECT_EQ(false, inst->IsVulkanStorageImage());
+  EXPECT_EQ(false, inst->IsVulkanSampledImage());
+  EXPECT_EQ(false, inst->IsVulkanStorageTexelBuffer());
+
+  inst = context->get_def_use_mgr()->GetDef(27);
+  EXPECT_EQ(false, inst->IsVulkanStorageImage());
+  EXPECT_EQ(false, inst->IsVulkanSampledImage());
+  EXPECT_EQ(true, inst->IsVulkanStorageTexelBuffer());
+
+  inst = context->get_def_use_mgr()->GetDef(28);
+  EXPECT_EQ(false, inst->IsVulkanStorageImage());
+  EXPECT_EQ(true, inst->IsVulkanSampledImage());
+  EXPECT_EQ(false, inst->IsVulkanStorageTexelBuffer());
+
+  inst = context->get_def_use_mgr()->GetDef(29);
+  EXPECT_EQ(true, inst->IsVulkanStorageImage());
+  EXPECT_EQ(false, inst->IsVulkanSampledImage());
+  EXPECT_EQ(false, inst->IsVulkanStorageTexelBuffer());
+}
+
 }  // namespace
 }  // namespace opt
 }  // namespace spvtools

test/opt/local_redundancy_elimination_test.cpp

@@ -154,6 +154,50 @@ TEST_F(LocalRedundancyEliminationTest, KeepInstructionsInDifferentBlocks) {
   SinglePassRunAndMatch<LocalRedundancyEliminationPass>(text, false);
 }
 
+TEST_F(LocalRedundancyEliminationTest, StorageBufferIdentification) {
+  const std::string text = R"(
+; CHECK: [[gep:%\w+]] = OpAccessChain
+; CHECK: [[ld:%\w+]] = OpLoad {{%\w+}} [[gep]]
+; CHECK: [[add:%\w+]] = OpIAdd {{%\w+}} [[ld]]
+; CHECK: OpStore [[gep]] [[add]]
+; CHECK: [[ld:%\w+]] = OpLoad {{%\w+}} [[gep]]
+; CHECK: [[add:%\w+]] = OpIAdd {{%\w+}} [[ld]]
+; CHECK: OpStore [[gep]] [[add]]
+
+OpCapability Shader
+OpCapability Linkage
+OpMemoryModel Logical GLSL450
+OpDecorate %block BufferBlock
+OpMemberDecorate %block 0 Offset 0
+%void = OpTypeVoid
+%int = OpTypeInt 32 0
+%int_0 = OpConstant %int 0
+%int_1 = OpConstant %int 1
+%block = OpTypeStruct %int
+%array = OpTypeArray %block %int_1
+%ptr_ssbo_array = OpTypePointer Uniform %array
+%ptr_ssbo_int = OpTypePointer Uniform %int
+%var = OpVariable %ptr_ssbo_array Uniform
+%void_fn = OpTypeFunction %void
+%fn = OpFunction %void None %void_fn
+%entry = OpLabel
+%gep1 = OpAccessChain %ptr_ssbo_int %var %int_0 %int_0
+%ld1 = OpLoad %int %gep1
+%add1 = OpIAdd %int %ld1 %int_1
+%gep2 = OpAccessChain %ptr_ssbo_int %var %int_0 %int_0
+OpStore %gep2 %add1
+%gep3 = OpAccessChain %ptr_ssbo_int %var %int_0 %int_0
+%ld3 = OpLoad %int %gep3
+%add3 = OpIAdd %int %ld3 %int_1
+%gep4 = OpAccessChain %ptr_ssbo_int %var %int_0 %int_0
+OpStore %gep4 %add3
+OpReturn
+OpFunctionEnd
+)";
+
+  SinglePassRunAndMatch<LocalRedundancyEliminationPass>(text, true);
+}
+
 }  // namespace
 }  // namespace opt
 }  // namespace spvtools