From ef7143e1fc1a12993f584ce37dbde38b75966b40 Mon Sep 17 00:00:00 2001 From: Lunderberg Date: Sat, 26 Jun 2021 01:10:45 -0700 Subject: [PATCH] [Vulkan] Added debug saving of Vulkan shaders, environment variable documentation. (#8333) Frequently, looking at the shaders generated by the Vulkan codegen is useful for debugging. While this can be done by checking the `mod.imported_modules[0].get_source()`, that requires the shader to first pass validation. Co-authored-by: Eric Lunderberg --- docs/dev/runtimes/vulkan.rst | 52 +++++++++++++++++++++++++++ src/runtime/vulkan/vulkan_device.cc | 22 ++++-------- src/runtime/vulkan/vulkan_instance.cc | 6 ++-- src/support/utils.h | 26 ++++++++++++++ src/target/spirv/build_vulkan.cc | 22 +++++++++++- 5 files changed, 108 insertions(+), 20 deletions(-) diff --git a/docs/dev/runtimes/vulkan.rst b/docs/dev/runtimes/vulkan.rst index ed0dbe33a305..36b1508f5448 100644 --- a/docs/dev/runtimes/vulkan.rst +++ b/docs/dev/runtimes/vulkan.rst @@ -205,3 +205,55 @@ not enabled in the ``Target``, an exception will be raised. * - ``supports_int64`` - - Int64 + + +Vulkan-Specific Environment Variables +------------------------------------- + +Both the SPIR-V code generation and the Vulkan runtime have +environment variables that can modify some of the runtime behavior. +These are intended for debugging purposes, both to more easily test +specific code paths, and to output more information as needed. All +boolean flags are true if the environment variable is set to a +non-zero integer. An unset variable, the integer zero, or an empty +string are all false boolean flags. + +.. _VK_KHR_push_descriptor: https://khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VK_KHR_push_descriptor.html + +.. _VK_KHR_descriptor_update_template: https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VK_KHR_descriptor_update_template.html + +.. _VK_KHR_dedicated_allocation: https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VK_KHR_dedicated_allocation.html + +.. _VkMemoryDedicatedRequirements: https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VkMemoryDedicatedRequirements.html + +.. _Vulkan validation layers: https://github.com/KhronosGroup/Vulkan-LoaderAndValidationLayers/blob/master/layers/README.md + +.. _spvValidate: https://github.com/KhronosGroup/SPIRV-Tools#validator + + +* ``TVM_VULKAN_DISABLE_PUSH_DESCRIPTOR`` - A boolean flag. If true, + TVM will explicitly allocate descriptors, and will not use the + `VK_KHR_push_descriptor`_ or `VK_KHR_descriptor_update_template`_ + extensions. If false, TVM will decide whether to use these + extensions based on their availability. + +* ``TVM_VULKAN_DISABLE_DEDICATED_ALLOCATION`` - A boolean flag. If + true, TVM will not mark memory allocations as being dedicated + allocations, and will not use the `VK_KHR_dedicated_allocation`_ + extension. If false, TVM will decide whether memory allocations + should be marked as dedicated based on the + `VkMemoryDedicatedRequirements`_ for that buffer. + +* ``TVM_VULKAN_ENABLE_VALIDATION_LAYERS`` - A boolean flag. If true, + TVM will enable `Vulkan validation layers`_ that the device + supports. If false, no validation layers are enabled. + +* ``TVM_VULKAN_DISABLE_SHADER_VALIDATION`` - A boolean flag. If true, + the SPIR-V shader validation done with `spvValidate`_ is skipped. + If false (default), all SPIR-V shaders generated by TVM are + validated with `spvValidate`_. + +* ``TVM_VULKAN_DEBUG_SHADER_SAVEPATH`` - A path to a directory. If + set to a non-empty string, the Vulkan codegen will save tir, binary + SPIR-V, and disassembled SPIR-V shaders to this directory, to be + used for debugging purposes. diff --git a/src/runtime/vulkan/vulkan_device.cc b/src/runtime/vulkan/vulkan_device.cc index 5e4be8209550..16987210b232 100644 --- a/src/runtime/vulkan/vulkan_device.cc +++ b/src/runtime/vulkan/vulkan_device.cc @@ -24,6 +24,7 @@ #include #include +#include "../../support/utils.h" #include "vulkan_common.h" #include "vulkan_device.h" #include "vulkan_device_api.h" @@ -121,24 +122,15 @@ VulkanDeviceProperties::VulkanDeviceProperties(const VulkanInstance& instance, // Support is available based on these extensions, but allow it to // be disabled based on an environment variable. supports_push_descriptor = device.HasExtension("VK_KHR_push_descriptor") && - device.HasExtension("VK_KHR_descriptor_update_template"); - { - const char* disable = std::getenv("TVM_VULKAN_DISABLE_PUSH_DESCRIPTOR"); - if (disable && *disable) { - supports_push_descriptor = false; - } - } + device.HasExtension("VK_KHR_descriptor_update_template") && + !support::BoolEnvironmentVar("TVM_VULKAN_DISABLE_PUSH_DESCRIPTOR"); // Support is available based on these extensions, but allow it to // be disabled based on an environment variable. - supports_dedicated_allocation = device.HasExtension("VK_KHR_get_memory_requirements2") && - device.HasExtension("VK_KHR_dedicated_allocation"); - { - const char* disable = std::getenv("TVM_VULKAN_DISABLE_DEDICATED_ALLOCATION"); - if (disable && *disable) { - supports_dedicated_allocation = false; - } - } + supports_dedicated_allocation = + device.HasExtension("VK_KHR_get_memory_requirements2") && + device.HasExtension("VK_KHR_dedicated_allocation") && + !support::BoolEnvironmentVar("TVM_VULKAN_DISABLE_DEDICATED_ALLOCATION"); // The check of VK_SHADER_STAGE_COMPUTE_BIT isn't technically // needed, since it will be set so long at least one queue has diff --git a/src/runtime/vulkan/vulkan_instance.cc b/src/runtime/vulkan/vulkan_instance.cc index 351319e0e898..b8295d2cd605 100644 --- a/src/runtime/vulkan/vulkan_instance.cc +++ b/src/runtime/vulkan/vulkan_instance.cc @@ -22,6 +22,7 @@ #include #include +#include "../../support/utils.h" #include "vulkan_common.h" namespace tvm { @@ -32,10 +33,7 @@ VulkanInstance::VulkanInstance() { const auto layers = []() { std::vector layers; - const char* validation_enabled_env = std::getenv("TVM_VULKAN_ENABLE_VALIDATION_LAYERS"); - bool validation_enabled = validation_enabled_env && *validation_enabled_env; - - if (validation_enabled) { + if (support::BoolEnvironmentVar("TVM_VULKAN_ENABLE_VALIDATION_LAYERS")) { uint32_t inst_layer_prop_count; VULKAN_CALL(vkEnumerateInstanceLayerProperties(&inst_layer_prop_count, nullptr)); std::vector inst_layer_prop(inst_layer_prop_count); diff --git a/src/support/utils.h b/src/support/utils.h index d807c5b8bb63..d8e3bf5f30ab 100644 --- a/src/support/utils.h +++ b/src/support/utils.h @@ -37,6 +37,7 @@ #include #include #include +#include #include #include #include @@ -199,6 +200,31 @@ inline uint64_t HashCombine(uint64_t key, const T& value) { return key ^ (uint64_t(value) + 0x9e3779b9 + (key << 6) + (key >> 2)); } +/*! + * \brief Return whether a boolean flag is set as an environment variable. + * + * Returns true if the environment variable is set to a non-zero + * integer, or to a non-empty string that is not an integer. + * + * Returns false if the environment variable is unset, if the + * environment variable is set to the integer zero, or if the + * environment variable is an empty string. + */ +inline bool BoolEnvironmentVar(const char* varname) { + const char* var = std::getenv(varname); + if (!var) { + return false; + } + + int x = 0; + std::istringstream is(var); + if (is >> x) { + return x; + } + + return *var; +} + } // namespace support } // namespace tvm #endif // TVM_SUPPORT_UTILS_H_ diff --git a/src/target/spirv/build_vulkan.cc b/src/target/spirv/build_vulkan.cc index c19b71d1540b..e922942e8acf 100644 --- a/src/target/spirv/build_vulkan.cc +++ b/src/target/spirv/build_vulkan.cc @@ -26,8 +26,12 @@ #include #include +#include +#include + #include "../../runtime/vulkan/vulkan_module.h" #include "../../runtime/vulkan/vulkan_shader.h" +#include "../../support/utils.h" #include "../build_common.h" #include "codegen_spirv.h" @@ -121,7 +125,23 @@ runtime::Module BuildSPIRV(IRModule mod, Target target, bool webgpu_restriction) VulkanShader shader = cg.BuildFunction(f, entry); - spirv_tools.ValidateShader(shader.data); + if (auto path = std::getenv("TVM_VULKAN_DEBUG_SHADER_SAVEPATH")) { + if (*path) { + std::stringstream ss; + ss << path << "/" << f_name << "_"; + std::string prefix = ss.str(); + + std::ofstream(prefix + "tir.txt") << f; + std::ofstream(prefix + "spv.txt") << spirv_tools.BinaryToText(shader.data); + std::ofstream(prefix + "spv.spv", std::ios::binary) + .write(reinterpret_cast(shader.data.data()), + sizeof(shader.data[0]) * shader.data.size()); + } + } + + if (!support::BoolEnvironmentVar("TVM_VULKAN_DISABLE_SHADER_VALIDATION")) { + spirv_tools.ValidateShader(shader.data); + } if (webgpu_restriction) { for (auto param : f->params) {