gpu: Fix when OpStore is used in first block

spencer-lunarg · Jul 11, 2024 · fa15b77 · fa15b77
1 parent 9cfdad2
commit fa15b77
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 1 deletion.
diff --git a/layers/gpu/spirv/instruction.h b/layers/gpu/spirv/instruction.h
@@ -64,6 +64,9 @@ struct Instruction {
 
     void ToBinary(std::vector<uint32_t>& out);
 
+    bool operator==(Instruction const& other) const { return words_ == other.words_; }
+    bool operator!=(Instruction const& other) const { return words_ != other.words_; }
+
     // Store minimal extra data
     uint32_t result_id_index_ = 0;
     uint32_t type_id_index_ = 0;

diff --git a/layers/gpu/spirv/pass.cpp b/layers/gpu/spirv/pass.cpp
@@ -364,8 +364,12 @@ void Pass::InjectFunctionCheck(BasicBlockIt block_it, InstructionIt* inst_it, co
 InstructionIt Pass::FindTargetInstruction(BasicBlock& block) const {
     const uint32_t target_id = target_instruction_->ResultId();
     for (auto inst_it = block.instructions_.begin(); inst_it != block.instructions_.end(); ++inst_it) {
+        // This has to re-loop the entire block to find the instruction, using the ResultID, we can quickly compare
         if ((*inst_it)->ResultId() == target_id) {
-            return inst_it;
+            // Things like OpStore will have a result id of zero, so need to do deep instruction comparison
+            if (*(*inst_it) == *target_instruction_) {
+                return inst_it;
+            }
         }
     }
     assert(false);

diff --git a/tests/unit/gpu_av_spirv_positive.cpp b/tests/unit/gpu_av_spirv_positive.cpp
@@ -353,3 +353,45 @@ TEST_F(PositiveGpuAVSpirv, VulkanMemoryModelDeviceScope) {
     m_default_queue->Submit(*m_commandBuffer);
     m_default_queue->Wait();
 }
+
+TEST_F(PositiveGpuAVSpirv, FindMultipleStores) {
+    TEST_DESCRIPTION("Catches bug when various OpStore are in top of a function");
+    AddDisabledFeature(vkt::Feature::robustBufferAccess);
+    RETURN_IF_SKIP(InitGpuAvFramework());
+    RETURN_IF_SKIP(InitState());
+
+    const char shader_source[] = R"glsl(
+      #version 450
+      layout(set = 0, binding = 0) buffer StorageBuffer { uint data[]; } Data;  // data[4]
+
+      int foo() {
+            return (gl_WorkGroupSize.x > 1) ? 1 : 0;
+      }
+
+      void main() {
+            int index = foo(); // first OpStore is not instrumented
+            Data.data[index] = 0xdeadca71;
+            Data.data[index + 1] = 0xdeadca71;
+      }
+    )glsl";
+
+    CreateComputePipelineHelper pipe(*this);
+    pipe.dsl_bindings_ = {{0, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_ALL, nullptr}};
+    pipe.cs_ = std::make_unique<VkShaderObj>(this, shader_source, VK_SHADER_STAGE_COMPUTE_BIT);
+    pipe.CreateComputePipeline();
+
+    vkt::Buffer write_buffer(*m_device, 64, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
+                             VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
+    pipe.descriptor_set_->WriteDescriptorBufferInfo(0, write_buffer.handle(), 0, VK_WHOLE_SIZE, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER);
+    pipe.descriptor_set_->UpdateDescriptorSets();
+
+    m_commandBuffer->begin();
+    vk::CmdBindPipeline(m_commandBuffer->handle(), VK_PIPELINE_BIND_POINT_COMPUTE, pipe.Handle());
+    vk::CmdBindDescriptorSets(m_commandBuffer->handle(), VK_PIPELINE_BIND_POINT_COMPUTE, pipe.pipeline_layout_.handle(), 0, 1,
+                              &pipe.descriptor_set_->set_, 0, nullptr);
+    vk::CmdDispatch(m_commandBuffer->handle(), 1, 1, 1);
+    m_commandBuffer->end();
+
+    m_default_queue->Submit(*m_commandBuffer);
+    m_default_queue->Wait();
+}