[Vulkan] Added debug saving of Vulkan shaders, environment variable d…

…ocumentation. (#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 <elunderberg@octoml.ai>

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.

src/runtime/vulkan/vulkan_device.cc

@@ -24,6 +24,7 @@
 #include <unordered_map>
 #include <utility>
 
+#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

src/runtime/vulkan/vulkan_instance.cc

@@ -22,6 +22,7 @@
 #include <cstdlib>
 #include <utility>
 
+#include "../../support/utils.h"
 #include "vulkan_common.h"
 
 namespace tvm {
@@ -32,10 +33,7 @@ VulkanInstance::VulkanInstance() {
   const auto layers = []() {
     std::vector<const char*> 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<VkLayerProperties> inst_layer_prop(inst_layer_prop_count);

src/support/utils.h

@@ -37,6 +37,7 @@
 #include <algorithm>
 #include <array>
 #include <cctype>
+#include <cstdlib>
 #include <memory>
 #include <sstream>
 #include <string>
@@ -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_

src/target/spirv/build_vulkan.cc

@@ -26,8 +26,12 @@
 #include <libspirv.h>
 #include <tvm/tir/transform.h>
 
+#include <fstream>
+#include <sstream>
+
 #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<const char*>(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) {