From 81927e6673f71a16eb8062ceef4bffde049d9ce1 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 ed0dbe33a3052..36b1508f54486 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 5e4be82095500..16987210b2324 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 351319e0e8980..b8295d2cd605d 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 d807c5b8bb63b..d8e3bf5f30ab8 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 c19b71d1540b5..e922942e8acff 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) {