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

…ocumentation.

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

@@ -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) {