gpuav: Select pipeline to instrument post creation

layers/CMakeLists.txt

@@ -211,6 +211,7 @@ target_sources(vvl PRIVATE
     chassis/chassis_modification_state.h
     chassis/chassis_manual.cpp
     chassis/dispatch_object_manual.cpp
+    chassis/dispatch_object.h
     containers/range.h
     containers/range_map.h
     core_checks/cc_android.cpp

layers/VkLayer_khronos_validation.json.in

@@ -817,7 +817,7 @@
                                         {
                                             "key": "gpuav_select_instrumented_shaders",
                                             "label": "Enable instrumenting shaders selectively",
-                                            "description": "Select which shaders to instrument by passing a VkValidationFeaturesEXT struct with GPU-AV enabled in the VkShaderModuleCreateInfo pNext or using a regex matching a shader/pipeline debug name. Because this only validates the selected shaders, it will allow GPU-AV to run much faster.",
+                                            "description": "Select which shaders to instrument by passing a VkValidationFeaturesEXT struct with GPU-AV enabled in the VkShaderModuleCreateInfo pNext or using a regex matching a shader/pipeline debug name. Because this only validates the selected shaders/pipelines, it will allow GPU-AV to run much faster.",
                                             "url": "https://github.com/KhronosGroup/Vulkan-ValidationLayers/blob/main/docs/gpu_av_selective_shader.md",
                                             "type": "BOOL",
                                             "default": false,

layers/chassis/dispatch_object.h

@@ -231,6 +231,17 @@ class HandleWrapper : public Logger {
         }
     }
 
+    // Replaces the "driver handle" in the "wrapped handle" to "driver handle" mapping.
+    // Returns the old "driver handle" if found.
+    template <typename HandleType>
+    HandleType Replace(HandleType wrapped_handle, HandleType new_driver_handle) {
+        const HandleType old_driver_handle = Find(wrapped_handle);
+        const uint64_t wrapped_handle_id = CastToUint64(wrapped_handle);
+        assert(wrapped_handle_id != 0);  // can't be 0, otherwise unwrap will apply special rule for VK_NULL_HANDLE
+        unique_id_mapping.insert_or_assign(wrapped_handle_id, CastToUint64(new_driver_handle));
+        return old_driver_handle;
+    }
+
     void UnwrapPnextChainHandles(const void* pNext);
 
     static std::atomic<uint64_t> global_unique_id;

layers/gpuav/instrumentation/gpuav_shader_instrumentor.cpp

@@ -19,6 +19,7 @@
 #include <vulkan/vulkan_core.h>
 #include <cstdint>
 
+#include "containers/container_utils.h"
 #include "error_message/error_location.h"
 #include "generated/vk_extension_helper.h"
 #include "generated/dispatch_functions.h"
@@ -405,6 +406,120 @@ void GpuShaderInstrumentor::PreCallRecordCreatePipelineLayout(VkDevice device, c
     }
 }
 
+void GpuShaderInstrumentor::PreCallRecordSetDebugUtilsObjectNameEXT(VkDevice device, const VkDebugUtilsObjectNameInfoEXT* pNameInfo,
+                                                                    const RecordObject& record_obj) {
+    if (!gpuav_settings.select_instrumented_shaders) {
+        return;
+    }
+
+    if (!gpuav_settings.IsSpirvModified()) {
+        return;
+    }
+    if (pNameInfo->objectType != VK_OBJECT_TYPE_PIPELINE || !pNameInfo->pObjectName) {
+        return;
+    }
+
+    if (!gpuav_settings.MatchesAnyShaderSelectionRegex(pNameInfo->pObjectName)) {
+        return;
+    }
+
+    if (disabled[handle_wrapping]) {
+        InternalError(LogObjectList(), record_obj.location,
+                      "For GPU-AV selective pipeline instrumentation post creation to work, handle wrapping needs to be enabled.");
+        return;
+    }
+
+    VkPipeline wrapped_pipeline = CastFromUint64<VkPipeline>(pNameInfo->objectHandle);
+    auto pipeline_state = Get<vvl::Pipeline>(wrapped_pipeline);
+    ASSERT_AND_RETURN(pipeline_state);
+
+    if (pipeline_state->instrumentation_data.was_instrumented) {
+        return;
+    }
+
+    if (!NeedPipelineCreationShaderInstrumentation(*pipeline_state, record_obj.location)) {
+        return;
+    }
+
+    // The pipeline was selected by name, not by individual shader name, so force all its shaders to be instrumented
+    for (const auto& stage_state : pipeline_state->stage_states) {
+        if (stage_state.module_state && stage_state.module_state->VkHandle() != VK_NULL_HANDLE) {
+            selected_instrumented_shaders.insert(stage_state.module_state->VkHandle());
+        }
+    }
+
+    VkPipeline instrumented_pipeline = VK_NULL_HANDLE;
+    // Can't instrument ray tracing pipeline post creation,
+    // As corresponding shader binding tables may have already been created.
+    if (pipeline_state->linking_shaders == 0 &&
+        IsValueIn(pipeline_state->pipeline_type, {VK_PIPELINE_BIND_POINT_GRAPHICS, VK_PIPELINE_BIND_POINT_COMPUTE})) {
+        std::vector<chassis::ShaderInstrumentationMetadata> shader_instrumentation_metadata;
+        if (pipeline_state->pipeline_type == VK_PIPELINE_BIND_POINT_GRAPHICS) {
+            vku::safe_VkGraphicsPipelineCreateInfo new_pipeline_ci(pipeline_state->GraphicsCreateInfo());
+            new_pipeline_ci.flags &= ~VK_PIPELINE_CREATE_FAIL_ON_PIPELINE_COMPILE_REQUIRED_BIT;
+            const bool success = PreCallRecordPipelineCreationShaderInstrumentation(
+                nullptr, *pipeline_state, new_pipeline_ci, uint32_t(pipeline_state->stage_states.size()), record_obj.location,
+                shader_instrumentation_metadata);
+            if (!success) {
+                InternalError(device, record_obj.location, "Failed to instrument graphics pipeline in SetDebugUtilsObjectNameEXT.");
+                return;
+            }
+
+            const VkResult result =
+                DispatchCreateGraphicsPipelines(device, VK_NULL_HANDLE, 1, new_pipeline_ci.ptr(), nullptr, &instrumented_pipeline);
+            if (result != VK_SUCCESS || instrumented_pipeline == VK_NULL_HANDLE) {
+                InternalError(device, record_obj.location,
+                              "Failed to create instrumented graphics pipeline in SetDebugUtilsObjectNameEXT.");
+                return;
+            }
+        } else if (pipeline_state->pipeline_type == VK_PIPELINE_BIND_POINT_COMPUTE) {
+            vku::safe_VkComputePipelineCreateInfo new_pipeline_ci(pipeline_state->ComputeCreateInfo());
+            new_pipeline_ci.flags &= ~VK_PIPELINE_CREATE_FAIL_ON_PIPELINE_COMPILE_REQUIRED_BIT;
+            const bool success = PreCallRecordPipelineCreationShaderInstrumentation(
+                nullptr, *pipeline_state, new_pipeline_ci, uint32_t(pipeline_state->stage_states.size()), record_obj.location,
+                shader_instrumentation_metadata);
+            if (!success) {
+                InternalError(device, record_obj.location, "Failed to instrument compute pipeline in SetDebugUtilsObjectNameEXT.");
+                return;
+            }
+
+            const VkResult result =
+                DispatchCreateComputePipelines(device, VK_NULL_HANDLE, 1, new_pipeline_ci.ptr(), nullptr, &instrumented_pipeline);
+            if (result != VK_SUCCESS || instrumented_pipeline == VK_NULL_HANDLE) {
+                InternalError(device, record_obj.location,
+                              "Failed to create instrumented compute pipeline in SetDebugUtilsObjectNameEXT.");
+                return;
+            }
+        }
+
+        PostCallRecordPipelineCreationShaderInstrumentation(*pipeline_state, uint32_t(pipeline_state->stage_states.size()),
+                                                            shader_instrumentation_metadata);
+    } else {
+        vku::safe_VkGraphicsPipelineCreateInfo new_pipeline_ci(pipeline_state->GraphicsCreateInfo());
+        const bool success =
+            PreCallRecordPipelineCreationShaderInstrumentationGPL(nullptr, *pipeline_state, new_pipeline_ci, record_obj.location);
+        if (!success) {
+            InternalError(device, record_obj.location,
+                          "Failed to instrument graphics pipeline library in SetDebugUtilsObjectNameEXT.");
+            return;
+        }
+
+        const VkResult result =
+            DispatchCreateGraphicsPipelines(device, VK_NULL_HANDLE, 1, new_pipeline_ci.ptr(), nullptr, &instrumented_pipeline);
+        if (result != VK_SUCCESS || instrumented_pipeline == VK_NULL_HANDLE) {
+            InternalError(device, record_obj.location,
+                          "Failed to create instrumented graphics pipeline in SetDebugUtilsObjectNameEXT.");
+            return;
+        }
+    }
+
+    auto layer_data = vvl::GetDispatchDevice(device);
+    ASSERT_AND_RETURN(layer_data);
+    const VkPipeline old_pipeline = layer_data->Replace(pipeline_state->VkHandle(), instrumented_pipeline);
+    gpuav::PipelineSubState& pipeline_sub_state = SubState(*pipeline_state);
+    pipeline_sub_state.AddHandleToDestroy(old_pipeline);
+}
+
 void GpuShaderInstrumentor::PostCallRecordCreateShaderModule(VkDevice device, const VkShaderModuleCreateInfo* pCreateInfo,
                                                              const VkAllocationCallbacks* pAllocator, VkShaderModule* pShaderModule,
                                                              const RecordObject& record_obj,
@@ -739,6 +854,10 @@ void GpuShaderInstrumentor::PreCallRecordCreateRayTracingPipelinesKHR(
             continue;
         }
 
+        if (!IsPipelineSelectedForInstrumentation(pCreateInfos[i].pNext, VK_NULL_HANDLE, create_info_loc)) {
+            continue;
+        }
+
         auto& shader_instrumentation_metadata = chassis_state.shader_instrumentations_metadata[i];
 
         // Ray tracing pipelines can be made of libraries, but contrary to GPL instrumentation is not postponed
@@ -1117,20 +1236,26 @@ void GpuShaderInstrumentor::BuildDescriptorSetLayoutInfo(const vvl::DescriptorSe
     }
 }
 
-bool GpuShaderInstrumentor::IsPipelineSelectedForInstrumentation(VkPipeline pipeline, const Location& loc) {
+bool GpuShaderInstrumentor::IsPipelineSelectedForInstrumentation(const void* pipeline_ci_pnext, VkPipeline pipeline,
+                                                                 const Location& loc) {
     if (!gpuav_settings.select_instrumented_shaders) {
         return true;
     }
 
     bool should_instrument_pipeline = false;
     {
         std::string pipeline_debug_name;
-        {
+        if (auto debug = vku::FindStructInPNextChain<VkDebugUtilsObjectNameInfoEXT>(pipeline_ci_pnext)) {
+            if (debug->pObjectName) {
+                pipeline_debug_name = debug->pObjectName;
+            }
+        } else if (pipeline != VK_NULL_HANDLE) {
             std::unique_lock<std::mutex> lock(debug_report->debug_output_mutex);
             pipeline_debug_name = debug_report->GetUtilsObjectNameNoLock(HandleToUint64(pipeline));
         }
-
-        should_instrument_pipeline = gpuav_settings.MatchesAnyShaderSelectionRegex(pipeline_debug_name);
+        if (!pipeline_debug_name.empty()) {
+            should_instrument_pipeline = gpuav_settings.MatchesAnyShaderSelectionRegex(pipeline_debug_name);
+        }
     }
     if (should_instrument_pipeline) {
         LogInfo("GPU-AV::Selective shader instrumentation", LogObjectList(), loc, "(%s) will be instrumented for validation.",
@@ -1237,6 +1362,9 @@ bool GpuShaderInstrumentor::PreCallRecordPipelineCreationShaderInstrumentation(
     // Can set this once for all shaders in the pipeline
     BuildDescriptorSetLayoutInfo(pipeline_state, interface.instrumentation_dsl);
 
+    bool is_pipeline_selected_for_instrumentation =
+        IsPipelineSelectedForInstrumentation(modified_pipeline_ci.pNext, VK_NULL_HANDLE, loc);
+
     for (uint32_t stage_state_i = 0; stage_state_i < stages_count; ++stage_state_i) {
         const auto& stage_state = pipeline_state.stage_states[stage_state_i];
         auto modified_module_state = std::const_pointer_cast<vvl::ShaderModule>(stage_state.module_state);
@@ -1259,8 +1387,9 @@ bool GpuShaderInstrumentor::PreCallRecordPipelineCreationShaderInstrumentation(
                 const_cast<vku::safe_VkShaderModuleCreateInfo*>(reinterpret_cast<const vku::safe_VkShaderModuleCreateInfo*>(
                     vku::FindStructInPNextChain<VkShaderModuleCreateInfo>(stage_ci.pNext)));
 
-            if (!IsShaderSelectedForInstrumentation(modified_shader_module_ci, modified_module_state->VkHandle(),
-                                                    loc.dot(vvl::Field::pStages, stage_state_i).dot(vvl::Field::module))) {
+            if (!(is_pipeline_selected_for_instrumentation ||
+                  IsShaderSelectedForInstrumentation(modified_shader_module_ci, modified_module_state->VkHandle(),
+                                                     loc.dot(vvl::Field::pStages, stage_state_i).dot(vvl::Field::module)))) {
                 continue;
             }
         }
@@ -1416,10 +1545,15 @@ bool GpuShaderInstrumentor::PreCallRecordPipelineCreationShaderInstrumentationGP
         // creation process no matter caching state.
         new_lib_ci.flags &= ~VK_PIPELINE_CREATE_FAIL_ON_PIPELINE_COMPILE_REQUIRED_BIT;
         bool is_library_instrumented = false;
-
         // If pipeline library is selected for instrumentation, force instrumentation of all its shaders
-        const bool force_pipeline_instrumentation =
-            IsPipelineSelectedForInstrumentation(modified_lib->VkHandle(), loc.dot(vvl::Field::pLibraries, modified_lib_i));
+        bool is_pipeline_selected_for_instrumentation = IsPipelineSelectedForInstrumentation(
+            modified_pipeline_ci.pNext, modified_lib->VkHandle(), loc.dot(vvl::Field::pLibraries, modified_lib_i));
+        // Case where pipeline is instrumented post creation. Otherwise, pipeline handle has yet to be created!
+        // currently only useful when this function is called in PreCallRecordSetDebugUtilsObjectNameEXT
+        if (linked_pipeline_state.VkHandle() != VK_NULL_HANDLE) {
+            is_pipeline_selected_for_instrumentation |=
+                IsPipelineSelectedForInstrumentation(modified_pipeline_ci.pNext, linked_pipeline_state.VkHandle(), loc);
+        }
 
         for (uint32_t stage_state_i = 0; stage_state_i < static_cast<uint32_t>(modified_lib->stage_states.size());
              ++stage_state_i) {
@@ -1457,9 +1591,9 @@ bool GpuShaderInstrumentor::PreCallRecordPipelineCreationShaderInstrumentationGP
                         vku::FindStructInPNextChain<VkShaderModuleCreateInfo>(modified_stage_ci->pNext)));
 
                 // TODO - this is in need of testing, when only selecting various library as well as selecting everything
-                if (!force_pipeline_instrumentation &&
-                    !IsShaderSelectedForInstrumentation(modified_shader_module_ci, modified_module_state->VkHandle(),
-                                                        loc.dot(vvl::Field::pStages, stage_state_i).dot(vvl::Field::module))) {
+                if (!(is_pipeline_selected_for_instrumentation ||
+                      IsShaderSelectedForInstrumentation(modified_shader_module_ci, modified_module_state->VkHandle(),
+                                                         loc.dot(vvl::Field::pStages, stage_state_i).dot(vvl::Field::module)))) {
                     continue;
                 }
 

layers/gpuav/instrumentation/gpuav_shader_instrumentor.h

@@ -103,6 +103,9 @@ class GpuShaderInstrumentor : public vvl::DeviceProxy {
                                            const VkAllocationCallbacks *pAllocator, VkPipelineLayout *pPipelineLayout,
                                            const RecordObject &record_obj, chassis::CreatePipelineLayout &chassis_state) override;
 
+    void PreCallRecordSetDebugUtilsObjectNameEXT(VkDevice device, const VkDebugUtilsObjectNameInfoEXT* pNameInfo,
+                                                 const RecordObject& record_obj) override;
+
     void PostCallRecordCreateShaderModule(VkDevice device, const VkShaderModuleCreateInfo *pCreateInfo,
                                           const VkAllocationCallbacks *pAllocator, VkShaderModule *pShaderModule,
                                           const RecordObject &record_obj, chassis::CreateShaderModule &chassis_state) override;
@@ -246,7 +249,7 @@ class GpuShaderInstrumentor : public vvl::DeviceProxy {
     bool set_null_descriptors_ = false;
 
   private:
-    bool IsPipelineSelectedForInstrumentation(VkPipeline pipeline, const Location &loc);
+    bool IsPipelineSelectedForInstrumentation(const void* pipeline_ci_pnext, VkPipeline pipeline, const Location& loc);
     bool IsShaderSelectedForInstrumentation(vku::safe_VkShaderModuleCreateInfo *modified_shader_module_ci,
                                             VkShaderModule modified_shader, const Location &loc);
     void AddDescriptorHeapMappings(VkBaseOutStructure *create_info);

layers/gpuav/resources/gpuav_state_trackers.cpp

@@ -671,9 +671,9 @@ ShaderObjectSubState::ShaderObjectSubState(vvl::ShaderObject& obj) : vvl::Shader
 PipelineSubState::PipelineSubState(Validator& gpuav, vvl::Pipeline& pipeline) : vvl::PipelineSubState(pipeline), gpuav_(gpuav) {}
 
 VkPipelineLayout PipelineSubState::GetPipelineLayoutUnion(const Location& loc, vvl::DescriptorMode mode) const {
-    std::unique_lock<std::mutex> recreated_layout_lock(recreated_layout_mutex);
-    if (recreated_layout != VK_NULL_HANDLE) {
-        return recreated_layout;
+    std::unique_lock<std::mutex> recreated_layout_lock(mutex_);
+    if (recreated_layout_ != VK_NULL_HANDLE) {
+        return recreated_layout_;
     }
 
     const std::shared_ptr<const vvl::PipelineLayout> pipeline_layout_state = base.PipelineLayoutState();
@@ -719,23 +719,35 @@ VkPipelineLayout PipelineSubState::GetPipelineLayoutUnion(const Location& loc, v
         pipeline_layout_ci.pPushConstantRanges = pipeline_layout_state->push_constant_ranges_layout->data();
     }
 
-    const VkResult result = DispatchCreatePipelineLayout(gpuav_.device, &pipeline_layout_ci, nullptr, &recreated_layout);
+    const VkResult result = DispatchCreatePipelineLayout(gpuav_.device, &pipeline_layout_ci, nullptr, &recreated_layout_);
     (void)result;
     assert(result == VK_SUCCESS);
 
     for (size_t i : recreated_desc_set_layouts_indices) {
         DispatchDestroyDescriptorSetLayout(gpuav_.device, set_layout_handles[i], nullptr);
     }
 
-    return recreated_layout;
+    return recreated_layout_;
 }
 
 void PipelineSubState::Destroy() {
-    std::unique_lock<std::mutex> recreated_layout_lock(recreated_layout_mutex);
-    if (recreated_layout != VK_NULL_HANDLE) {
-        DispatchDestroyPipelineLayout(gpuav_.device, recreated_layout, nullptr);
-        recreated_layout = VK_NULL_HANDLE;
+    std::unique_lock<std::mutex> recreated_layout_lock(mutex_);
+    if (recreated_layout_ != VK_NULL_HANDLE) {
+        DispatchDestroyPipelineLayout(gpuav_.device, recreated_layout_, nullptr);
+        recreated_layout_ = VK_NULL_HANDLE;
     }
+    if (uninstrumented_pipeline != VK_NULL_HANDLE) {
+        // vkDestroyPipeline expects an unwrapped handle,
+        // so cannot use DispatchDestroyPipeline as it will try to unwrap supplied pipeline handle
+        auto layer_data = vvl::GetDispatchDevice(gpuav_.device);
+        layer_data->device_dispatch_table.DestroyPipeline(gpuav_.device, uninstrumented_pipeline, nullptr);
+    }
+}
+
+void PipelineSubState::AddHandleToDestroy(VkPipeline pipeline) {
+    std::unique_lock<std::mutex> lock(mutex_);
+    assert(uninstrumented_pipeline == VK_NULL_HANDLE);
+    uninstrumented_pipeline = pipeline;
 }
 
 }  // namespace gpuav