Add a tool to dump/replay compute pipeline creation from gfx. (#3275)

* Add a tool to dump/replay compute pipeline creation from gfx.

* Fix x86 build.

---------

Co-authored-by: Yong He <yhe@nvidia.com>

3 files changed, 182 insertions, 0 deletions

diff --git a/tools/gfx/vulkan/vk-device.cpp b/tools/gfx/vulkan/vk-device.cpp
index f8bc9c3aa..fe3680eda 100644
--- a/tools/gfx/vulkan/vk-device.cpp
+++ b/tools/gfx/vulkan/vk-device.cpp
@@ -15,9 +15,12 @@
 #include "vk-swap-chain.h"
 #include "vk-transient-heap.h"
 #include "vk-vertex-layout.h"
+#include "vk-pipeline-dump-layer.h"
 
 #include "vk-helper-functions.h"
 
+#include "source/core/slang-platform.h"
+
 #ifdef GFX_NV_AFTERMATH
 #   include "GFSDK_Aftermath.h"
 #   include "GFSDK_Aftermath_Defines.h"
@@ -32,8 +35,20 @@ using namespace Slang;
 namespace vk
 {
 
+static bool shouldDumpPipeline()
+{
+    StringBuilder dumpPipelineSettings;
+    PlatformUtil::getEnvironmentVariable(toSlice("SLANG_GFX_DUMP_PIPELINE"), dumpPipelineSettings);
+    return dumpPipelineSettings.produceString() == "1";
+}
+
 DeviceImpl::~DeviceImpl()
 {
+    if (shouldDumpPipeline())
+    {
+        writePipelineDump(toSlice("gfx-vk-pipeline-dump.bin"));
+    }
+
     // Check the device queue is valid else, we can't wait on it..
     if (m_deviceQueue.isValid())
     {
@@ -782,6 +797,10 @@ Result DeviceImpl::initVulkanInstanceAndDevice(
 
     SLANG_RETURN_ON_FAIL(m_api.initDeviceProcs(m_device));
 
+    if (shouldDumpPipeline())
+    {
+        installPipelineDumpLayer(m_api);
+    }
     return SLANG_OK;
 }
 
diff --git a/tools/gfx/vulkan/vk-pipeline-dump-layer.cpp b/tools/gfx/vulkan/vk-pipeline-dump-layer.cpp
new file mode 100644
index 000000000..959eee15d
--- /dev/null
+++ b/tools/gfx/vulkan/vk-pipeline-dump-layer.cpp
@@ -0,0 +1,152 @@
+#include "vk-pipeline-dump-layer.h"
+#include "core/slang-basic.h"
+#include "core/slang-stream.h"
+namespace gfx {
+    using namespace Slang;
+
+    struct PipelineDumpContext
+    {
+        Dictionary<VkPipelineLayout, Index> pipelineLayouts;
+        Dictionary<VkShaderModule, Index> shaderModules;
+        Dictionary<VkDescriptorSetLayout, Index> descriptorSets;
+        Dictionary<VkPipeline, Index> computePipelines;
+        
+        List<uint8_t> serializedBytes;
+
+        VulkanApi api;
+
+        template<typename T>
+        void writeRaw(T v)
+        {
+            auto startIndex = serializedBytes.getCount();
+            serializedBytes.growToCount(startIndex + sizeof(T));
+            memcpy(serializedBytes.getBuffer() + startIndex, &v, sizeof(T));
+        }
+
+        template<typename T>
+        void writeArray(uint32_t elementCount, const T* data)
+        {
+            writeRaw(elementCount);
+
+            auto startIndex = serializedBytes.getCount();
+            serializedBytes.growToCount(startIndex + sizeof(T) * elementCount);
+            memcpy(serializedBytes.getBuffer() + startIndex, data, sizeof(T) * elementCount);
+        }
+
+        void writeStr(const char* str)
+        {
+            auto len = (uint32_t)strlen(str) + 1;
+            writeRaw(len);
+
+            auto startIndex = serializedBytes.getCount();
+            serializedBytes.growToCount(startIndex + len);
+            memcpy(serializedBytes.getBuffer() + startIndex, str, len - 1);
+            serializedBytes[startIndex + len - 1] = 0;
+        }
+
+        void writePipelineLayout(VkPipelineLayout layout, const VkPipelineLayoutCreateInfo* createInfo)
+        {
+            auto startIndex = serializedBytes.getCount();
+            writeRaw(createInfo->sType);
+            writeRaw(createInfo->flags);
+            writeRaw(createInfo->setLayoutCount);
+            for (uint32_t i = 0; i < createInfo->setLayoutCount; i++)
+                writeRaw(descriptorSets.getValue(createInfo->pSetLayouts[i]));
+            writeArray(createInfo->pushConstantRangeCount, createInfo->pPushConstantRanges);
+            pipelineLayouts[layout] = startIndex;
+        }
+
+        void writeShaderModule(VkShaderModule module, const VkShaderModuleCreateInfo* createInfo)
+        {
+            auto startIndex = serializedBytes.getCount();
+            writeRaw(createInfo->sType);
+            writeRaw(createInfo->flags);
+            writeArray((uint32_t)(createInfo->codeSize/sizeof(uint32_t)), createInfo->pCode);
+            shaderModules[module] = startIndex;
+        }
+
+        void writeDescriptorSetLayout(VkDescriptorSetLayout layout, const VkDescriptorSetLayoutCreateInfo* createInfo)
+        {
+            auto startIndex = serializedBytes.getCount();
+            writeRaw(createInfo->sType);
+            writeRaw(createInfo->flags);
+            writeArray(createInfo->bindingCount, createInfo->pBindings);
+            descriptorSets[layout] = startIndex;
+        }
+
+        void writePipeline(VkPipeline pipeline, const VkComputePipelineCreateInfo* createInfo)
+        {
+            auto startIndex = serializedBytes.getCount();
+            writeRaw(createInfo->sType);
+            writeRaw(createInfo->flags);
+            writeRaw(createInfo->stage.sType);
+            writeRaw(createInfo->stage.flags);
+            writeRaw(createInfo->stage.stage);
+            writeRaw(shaderModules.getValue(createInfo->stage.module));
+            writeStr(createInfo->stage.pName);
+            writeRaw(pipelineLayouts.getValue(createInfo->layout));
+            computePipelines[pipeline] = startIndex;
+        }
+
+        void writeToFile(UnownedStringSlice path)
+        {
+            RefPtr<FileStream> fs = new FileStream();
+            fs->init(path, FileMode::Create);
+            uint32_t pipelineCount = (uint32_t)computePipelines.getCount();
+            fs->write(&pipelineCount, sizeof(uint32_t));
+            for (auto& pair : computePipelines)
+            {
+                fs->write(KeyValueDetail::getValue(&pair), sizeof(Index));
+            }
+            Index blobSize = serializedBytes.getCount();
+            fs->write(&blobSize, sizeof(blobSize));
+            fs->write(serializedBytes.getBuffer(), serializedBytes.getCount());
+            fs->close();
+        }
+    };
+
+    PipelineDumpContext dumpContext;
+
+    VkResult SLANG_MCALL createPipelineLayout(VkDevice device, const VkPipelineLayoutCreateInfo* createInfo, const VkAllocationCallbacks* callbacks, VkPipelineLayout* outLayout)
+    {
+        auto result = dumpContext.api.vkCreatePipelineLayout(device, createInfo, callbacks, outLayout);
+        dumpContext.writePipelineLayout(*outLayout, createInfo);
+        return result;
+    }
+
+    VkResult SLANG_MCALL createComputePipelines(VkDevice device, VkPipelineCache cache, uint32_t createInfoCount, const VkComputePipelineCreateInfo* createInfos, const VkAllocationCallbacks* callbacks, VkPipeline* outPipelines)
+    {
+        auto result = dumpContext.api.vkCreateComputePipelines(device, cache, createInfoCount, createInfos, callbacks, outPipelines);
+        for (uint32_t i = 0; i < createInfoCount; i++)
+            dumpContext.writePipeline(outPipelines[i], createInfos + i);
+        return result;
+    }
+
+    VkResult SLANG_MCALL createShaderModule(VkDevice device, const VkShaderModuleCreateInfo* createInfo, const VkAllocationCallbacks* callbacks, VkShaderModule* outShaderModule)
+    {
+        auto result = dumpContext.api.vkCreateShaderModule(device, createInfo, callbacks, outShaderModule);
+        dumpContext.writeShaderModule(*outShaderModule, createInfo);
+        return result;
+    }
+
+    VkResult SLANG_MCALL createDescriptorSetLayout(VkDevice device, const VkDescriptorSetLayoutCreateInfo* createInfo, const VkAllocationCallbacks* callbacks, VkDescriptorSetLayout* outDescSetLayout)
+    {
+        auto result = dumpContext.api.vkCreateDescriptorSetLayout(device, createInfo, callbacks, outDescSetLayout);
+        dumpContext.writeDescriptorSetLayout(*outDescSetLayout, createInfo);
+        return result;
+    }
+
+    void installPipelineDumpLayer(VulkanApi& api)
+    {
+        dumpContext.api = api;
+        api.vkCreatePipelineLayout = createPipelineLayout;
+        api.vkCreateComputePipelines = createComputePipelines;
+        api.vkCreateShaderModule = createShaderModule;
+        api.vkCreateDescriptorSetLayout = createDescriptorSetLayout;
+    }
+
+    void writePipelineDump(UnownedStringSlice path)
+    {
+        dumpContext.writeToFile(path);
+    }
+} // renderer_test
diff --git a/tools/gfx/vulkan/vk-pipeline-dump-layer.h b/tools/gfx/vulkan/vk-pipeline-dump-layer.h
new file mode 100644
index 000000000..c514f7f3e
--- /dev/null
+++ b/tools/gfx/vulkan/vk-pipeline-dump-layer.h
@@ -0,0 +1,11 @@
+// vk-api.cpp
+#include "vk-api.h"
+
+#include "core/slang-string.h"
+
+namespace gfx {
+
+void installPipelineDumpLayer(VulkanApi& api);
+void writePipelineDump(Slang::UnownedStringSlice path);
+
+} // renderer_test