Test changing vulkan compute shader entry-point

Commit for original bug fix:
"Fix entry-point remapping when editing compute shaders on vulkan"
c81a522360

util/test/demos/vk/vk_shader_editing.cpp

@@ -59,6 +59,30 @@ void main()
 #endif
 }
 
+)EOSHADER";
+
+  const std::string comp = R"EOSHADER(
+
+struct PushData
+{
+  uint4 data;
+};
+
+[[vk::push_constant]]
+ConstantBuffer<PushData> push;
+
+StructuredBuffer<uint4> inbuf : register(b0);
+RWStructuredBuffer<uint4> outbuf : register(b1);
+
+[numthreads(1, 1, 1)]
+void hlsl_main ()
+{
+  outbuf[0].x += inbuf[0].x * push.data.x;
+  outbuf[0].y += inbuf[0].y * push.data.y;
+  outbuf[0].z += inbuf[0].z * push.data.z;
+  outbuf[0].w += inbuf[0].w * push.data.w;
+}
+
 )EOSHADER";
 
   int main()
@@ -138,6 +162,42 @@ void main()
 
     vb.upload(DefaultTri);
 
+    VkDescriptorSetLayout compSetLayout =
+        createDescriptorSetLayout(vkh::DescriptorSetLayoutCreateInfo({
+            {0, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_COMPUTE_BIT},
+            {1, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_COMPUTE_BIT},
+        }));
+
+    VkPipelineLayout compLayout = createPipelineLayout(vkh::PipelineLayoutCreateInfo(
+        {compSetLayout},
+        {
+            vkh::PushConstantRange(VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(Vec4i)),
+        }));
+
+    VkPipeline compPipe1 = createComputePipeline(vkh::ComputePipelineCreateInfo(
+        compLayout, CompileShaderModule(comp, ShaderLang::hlsl, ShaderStage::comp, "hlsl_main")));
+
+    AllocatedBuffer bufin(this, vkh::BufferCreateInfo(1024, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
+                                                                VK_BUFFER_USAGE_TRANSFER_DST_BIT),
+                          VmaAllocationCreateInfo({0, VMA_MEMORY_USAGE_CPU_TO_GPU}));
+
+    AllocatedBuffer bufout(this, vkh::BufferCreateInfo(1024, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
+                                                                 VK_BUFFER_USAGE_TRANSFER_DST_BIT),
+                           VmaAllocationCreateInfo({0, VMA_MEMORY_USAGE_CPU_TO_GPU}));
+
+    setName(bufin.buffer, "bufin");
+    setName(bufout.buffer, "bufout");
+
+    VkDescriptorSet compSet = allocateDescriptorSet(compSetLayout);
+
+    vkh::updateDescriptorSets(
+        device, {
+                    vkh::WriteDescriptorSet(compSet, 0, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
+                                            {vkh::DescriptorBufferInfo(bufin.buffer)}),
+                    vkh::WriteDescriptorSet(compSet, 1, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
+                                            {vkh::DescriptorBufferInfo(bufout.buffer)}),
+                });
+
     while(Running())
     {
       VkCommandBuffer cmd = GetCommandBuffer();
@@ -181,6 +241,28 @@ void main()
 
       FinishUsingBackbuffer(cmd, VK_ACCESS_TRANSFER_WRITE_BIT, VK_IMAGE_LAYOUT_GENERAL);
 
+      // Fill bufin with 111
+      vkCmdFillBuffer(cmd, bufin.buffer, 0, 1024, 111);
+      // Fill bufout with 222
+      vkCmdFillBuffer(cmd, bufout.buffer, 0, 1024, 222);
+      vkh::cmdPipelineBarrier(
+          cmd, {},
+          {
+              vkh::BufferMemoryBarrier(VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT,
+                                       bufin.buffer, 0, 1024),
+              vkh::BufferMemoryBarrier(VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_WRITE_BIT,
+                                       bufout.buffer, 0, 1024),
+          });
+
+      vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, compPipe1);
+      vkh::cmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, compLayout, 0, {compSet}, {});
+
+      Vec4i push = {5, 6, 7, 8};
+      vkCmdPushConstants(cmd, compLayout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(Vec4i), &push);
+      setMarker(cmd, "Pre-Dispatch");
+      vkCmdDispatch(cmd, 1, 1, 1);
+      setMarker(cmd, "Post-Dispatch");
+
       vkEndCommandBuffer(cmd);
 
       Submit(0, 1, {cmd});

util/test/tests/Vulkan/VK_Shader_Editing.py

@@ -1,5 +1,6 @@
 import copy
 import rdtest
+import struct
 import renderdoc as rd
 from typing import Tuple
 
@@ -162,3 +163,47 @@ class VK_Shader_Editing(rdtest.TestCase):
         self.controller.FreeTargetResource(offsetVS)
         self.controller.FreeTargetResource(FS1)
         self.controller.FreeTargetResource(FS2)
+
+        bufout = self.get_resource_by_name("bufout").resourceId
+
+        self.controller.SetFrameEvent(self.find_action("Pre-Dispatch").eventId, False)
+        pipe: rd.PipeState = self.controller.GetPipelineState()
+        csrefl: rd.ShaderReflection = pipe.GetShaderReflection(rd.ShaderStage.Compute)
+
+        uints = struct.unpack_from('=4L', self.controller.GetBufferData(bufout, 0, 0), 0)
+        if not rdtest.value_compare(uints, [222, 222, 222, 222]):
+            raise rdtest.TestFailureException(
+                'bufout data is incorrect before dispatch: {}'.format(uints))
+
+        eid = self.find_action("Post-Dispatch").eventId
+        self.controller.SetFrameEvent(eid, False)
+
+        uints = struct.unpack_from('=4L', self.controller.GetBufferData(bufout, 0, 0), 0)
+        if not rdtest.value_compare(uints, [777, 888, 999, 1110]):
+            raise rdtest.TestFailureException(
+                'bufout data is incorrect after dispatch: {}'.format(uints))
+
+        self.check(csrefl.debugInfo.encoding == rd.ShaderEncoding.HLSL)
+        self.check(csrefl.entryPoint == "hlsl_main")
+        rdtest.log.success("Values are as expected before compute shader edits")
+
+        raw_source: bytes = csrefl.rawBytes
+        # search-replace in the SPIR-V for 'hlsl_main' replace with a string of the same length
+        patched_entry_source = raw_source.replace(b'hlsl_main', b'main_hlsl')
+        newShader: Tuple[rd.ResourceId, str] = self.controller.BuildTargetShader('main_hlsl',
+                                                                                 csrefl.encoding, patched_entry_source,
+                                                                                 rd.ShaderCompileFlags(),
+                                                                                 rd.ShaderStage.Compute)
+        if len(newShader[1]) != 0:
+            raise rdtest.TestFailureException("Failed to compile edited compute shader: {}".format(newShader[1]))
+        nochangeCS = newShader[0]
+        self.controller.ReplaceResource(csrefl.resourceId, nochangeCS)
+        self.controller.SetFrameEvent(eid, False)
+        uints = struct.unpack_from('=4L', self.controller.GetBufferData(bufout, 0, 0), 0)
+        if not rdtest.value_compare(uints, [777, 888, 999, 1110]):
+            raise rdtest.TestFailureException(
+                'bufout data is incorrect after dispatch: {}'.format(uints))
+
+        rdtest.log.success("Values are as expected after compute shader entry point change")
+
+        self.controller.FreeTargetResource(nochangeCS)