From ac2acd0dec8c931128f700a1482dc39bdadc6d98 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Thu, 24 Jun 2021 14:02:18 -0700 Subject: [PATCH] [Vulkan] Added debug saving of Vulkan shaders, environment variable documentation. 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. --- 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) {