[SPIRV] Add vulkan memory option to DXC (microsoft#6848)

Some recent features require the vulkan memory model. We need a way to
support it in DXC. This commit adds an option to enable the Vulkan
memory model. The generated code will not change, but the compiler will
call the upgrade-memory-model pass in spirv-opt to generate the correct
code.

The expectation is that new code that requieres the vulkan memory model
can be written using the vulkan memory model, and then have the pass
update other references.

This change will undo the fix for microsoft#6066. The device scope in GLSL450 was
not a real device scope. It corresponds to QueueFamily scope in the
Vulkan memory model. We will start to use that scope for the atomic
operations in order to keep the behavoiur of the atomics the same
between the two memory models.

Haveing a true device scope atomic will not be possible with DXC. We
should be able to do something better when atomic operations are
implemented in clang.

Fixes microsoft#5784

Author: s-perron

external/SPIRV-Tools

@@ -396,6 +396,8 @@ def fspv_entrypoint_name_EQ : Joined<["-"], "fspv-entrypoint-name=">, Group<spir
   HelpText<"Specify the SPIR-V entry point name. Defaults to the HLSL entry point name.">;
 def fspv_enable_maximal_reconvergence: Flag<["-"], "fspv-enable-maximal-reconvergence">, Group<spirv_Group>, Flags<[CoreOption, DriverOption]>,
   HelpText<"Enables the MaximallyReconvergesKHR execution mode for this module.">;
+def fspv_use_vulkan_memory_model: Flag<["-"], "fspv-use-vulkan-memory-model">, Group<spirv_Group>, Flags<[CoreOption, DriverOption]>,
+  HelpText<"Generates SPIR-V modules that use the Vulkan memory model instead of GLSL450.">;
 def fvk_auto_shift_bindings: Flag<["-"], "fvk-auto-shift-bindings">, Group<spirv_Group>, Flags<[CoreOption, DriverOption]>,
   HelpText<"Apply fvk-*-shift to resources without an explicit register assignment.">;
 def Wno_vk_ignored_features : Joined<["-"], "Wno-vk-ignored-features">, Group<spirv_Group>, Flags<[CoreOption, DriverOption, HelpHidden]>,

include/dxc/Support/HLSLOptions.td

@@ -68,6 +68,7 @@ struct SpirvCodeGenOptions {
   bool fixFuncCallArguments;
   bool allowRWStructuredBufferArrays;
   bool enableMaximalReconvergence;
+  bool useVulkanMemoryModel;
   bool IEEEStrict;
   /// Maximum length in words for the OpString literal containing the shader
   /// source for DebugSource and DebugSourceContinued. If the source code length

include/dxc/Support/SPIRVOptions.h

@@ -1085,6 +1085,8 @@ int ReadDxcOpts(const OptTable *optionTable, unsigned flagsToInclude,
       Args.hasFlag(OPT_fvk_allow_rwstructuredbuffer_arrays, OPT_INVALID, false);
   opts.SpirvOptions.enableMaximalReconvergence =
       Args.hasFlag(OPT_fspv_enable_maximal_reconvergence, OPT_INVALID, false);
+  opts.SpirvOptions.useVulkanMemoryModel =
+      Args.hasFlag(OPT_fspv_use_vulkan_memory_model, OPT_INVALID, false);
 
   if (!handleVkShiftArgs(Args, OPT_fvk_b_shift, "b", &opts.SpirvOptions.bShift,
                          errors) ||

lib/DxcSupport/HLSLOptions.cpp

@@ -770,6 +770,9 @@ class SpirvBuilder {
   /// the feature used at the given source location.
   inline void requireCapability(spv::Capability, SourceLocation loc = {});
 
+  /// \brief Returns true if the module requires the given capability.
+  inline bool hasCapability(spv::Capability cap);
+
   /// \brief Adds an extension to the module under construction for translating
   /// the given target at the given source location.
   inline void requireExtension(llvm::StringRef extension, SourceLocation);
@@ -896,6 +899,11 @@ void SpirvBuilder::requireCapability(spv::Capability cap, SourceLocation loc) {
   }
 }
 
+bool SpirvBuilder::hasCapability(spv::Capability cap) {
+  SpirvCapability capability({}, cap);
+  return mod->hasCapability(capability);
+}
+
 void SpirvBuilder::requireExtension(llvm::StringRef ext, SourceLocation loc) {
   auto *extension = new (context) SpirvExtension(loc, ext);
   if (!mod->addExtension(extension))

tools/clang/include/clang/SPIRV/SpirvBuilder.h

@@ -102,6 +102,9 @@ class SpirvModule {
   // Returns false otherwise (e.g. if the capability already existed).
   bool addCapability(SpirvCapability *cap);
 
+  // Returns true if the capability is in the module.
+  bool hasCapability(SpirvCapability &cap);
+
   // Set the memory model of the module.
   void setMemoryModel(SpirvMemoryModel *model);
 

tools/clang/include/clang/SPIRV/SpirvModule.h

@@ -46,7 +46,7 @@ void CapabilityVisitor::addCapability(spv::Capability cap, SourceLocation loc) {
 void CapabilityVisitor::addCapabilityForType(const SpirvType *type,
                                              SourceLocation loc,
                                              spv::StorageClass sc) {
-  // Defent against instructions that do not have a return type.
+  // Defend against instructions that do not have a return type.
   if (!type)
     return;
 
@@ -827,8 +827,6 @@ void CapabilityVisitor::AddVulkanMemoryModelForVolatile(SpirvDecoration *decor,
                    "Volatile builtin variable in raytracing", loc);
     }
     addCapability(spv::Capability::VulkanMemoryModel, loc);
-    spvBuilder.setMemoryModel(spv::AddressingModel::Logical,
-                              spv::MemoryModel::VulkanKHR);
   }
 }
 
@@ -894,10 +892,6 @@ bool CapabilityVisitor::visit(SpirvModule *, Visitor::Phase phase) {
                                          {spv::Capability::RayTracingKHR});
   }
 
-  addExtensionAndCapabilitiesIfEnabled(
-      Extension::KHR_vulkan_memory_model,
-      {spv::Capability::VulkanMemoryModelDeviceScope});
-
   addExtensionAndCapabilitiesIfEnabled(
       Extension::NV_shader_subgroup_partitioned,
       {spv::Capability::GroupNonUniformPartitionedNV});

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -87,6 +87,12 @@ FeatureManager::FeatureManager(DiagnosticsEngine &de,
   } else {
     for (auto ext : opts.allowedExtensions)
       allowExtension(ext);
+
+    // The option to use the vulkan memory model implies the extension is
+    // available.
+    if (opts.useVulkanMemoryModel) {
+      allowExtension("SPV_KHR_vulkan_memory_model");
+    }
   }
 }
 

tools/clang/lib/SPIRV/FeatureManager.cpp

@@ -874,6 +874,10 @@ void SpirvEmitter::HandleTranslationUnit(ASTContext &context) {
   if (context.getDiagnostics().hasErrorOccurred())
     return;
 
+  if (!UpgradeToVulkanMemoryModelIfNeeded(&m)) {
+    return;
+  }
+
   // Check the existance of Texture and Sampler with
   // [[vk::combinedImageSampler]] for the same descriptor set and binding.
   auto resourceInfoForSampledImages =
@@ -14673,8 +14677,9 @@ SpirvEmitter::createFunctionScopeTempFromParameter(const ParmVarDecl *param) {
   return tempVar;
 }
 
-bool SpirvEmitter::spirvToolsFixupOpExtInst(std::vector<uint32_t> *mod,
-                                            std::string *messages) {
+bool SpirvEmitter::spirvToolsRunPass(std::vector<uint32_t> *mod,
+                                     spvtools::Optimizer::PassToken token,
+                                     std::string *messages) {
   spvtools::Optimizer optimizer(featureManager.getTargetEnv());
   optimizer.SetMessageConsumer(
       [messages](spv_message_level_t /*level*/, const char * /*source*/,
@@ -14691,33 +14696,28 @@ bool SpirvEmitter::spirvToolsFixupOpExtInst(std::vector<uint32_t> *mod,
   options.set_preserve_bindings(spirvOptions.preserveBindings);
   options.set_max_id_bound(spirvOptions.maxId);
 
-  optimizer.RegisterPass(
-      spvtools::CreateOpExtInstWithForwardReferenceFixupPass());
-
+  optimizer.RegisterPass(std::move(token));
   return optimizer.Run(mod->data(), mod->size(), mod, options);
 }
 
+bool SpirvEmitter::spirvToolsFixupOpExtInst(std::vector<uint32_t> *mod,
+                                            std::string *messages) {
+  spvtools::Optimizer::PassToken token =
+      spvtools::CreateOpExtInstWithForwardReferenceFixupPass();
+  return spirvToolsRunPass(mod, std::move(token), messages);
+}
+
 bool SpirvEmitter::spirvToolsTrimCapabilities(std::vector<uint32_t> *mod,
                                               std::string *messages) {
-  spvtools::Optimizer optimizer(featureManager.getTargetEnv());
-  optimizer.SetMessageConsumer(
-      [messages](spv_message_level_t /*level*/, const char * /*source*/,
-                 const spv_position_t & /*position*/,
-                 const char *message) { *messages += message; });
-
-  string::RawOstreamBuf printAllBuf(llvm::errs());
-  std::ostream printAllOS(&printAllBuf);
-  if (spirvOptions.printAll)
-    optimizer.SetPrintAll(&printAllOS);
-
-  spvtools::OptimizerOptions options;
-  options.set_run_validator(false);
-  options.set_preserve_bindings(spirvOptions.preserveBindings);
-  options.set_max_id_bound(spirvOptions.maxId);
-
-  optimizer.RegisterPass(spvtools::CreateTrimCapabilitiesPass());
+  spvtools::Optimizer::PassToken token = spvtools::CreateTrimCapabilitiesPass();
+  return spirvToolsRunPass(mod, std::move(token), messages);
+}
 
-  return optimizer.Run(mod->data(), mod->size(), mod, options);
+bool SpirvEmitter::spirvToolsUpgradeToVulkanMemoryModel(
+    std::vector<uint32_t> *mod, std::string *messages) {
+  spvtools::Optimizer::PassToken token =
+      spvtools::CreateUpgradeMemoryModelPass();
+  return spirvToolsRunPass(mod, std::move(token), messages);
 }
 
 bool SpirvEmitter::spirvToolsOptimize(std::vector<uint32_t> *mod,
@@ -15123,5 +15123,26 @@ SpirvEmitter::splatScalarToGenerate(QualType type, SpirvInstruction *scalar,
   return {};
 }
 
+bool SpirvEmitter::UpgradeToVulkanMemoryModelIfNeeded(
+    std::vector<uint32_t> *module) {
+  // DXC generates code assuming the vulkan memory model is not used. However,
+  // if a feature is used that requires the Vulkan memory model, then some code
+  // may need to be rewritten.
+  if (!spirvOptions.useVulkanMemoryModel &&
+      !spvBuilder.hasCapability(spv::Capability::VulkanMemoryModel))
+    return true;
+
+  std::string messages;
+  if (!spirvToolsUpgradeToVulkanMemoryModel(module, &messages)) {
+    emitFatalError("failed to use the vulkan memory model: %0", {}) << messages;
+    emitNote("please file a bug report on "
+             "https://github.com/Microsoft/DirectXShaderCompiler/issues "
+             "with source code if possible",
+             {});
+    return false;
+  }
+  return true;
+}
+
 } // end namespace spirv
 } // end namespace clang

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -36,6 +36,7 @@
 
 #include "ConstEvaluator.h"
 #include "DeclResultIdMapper.h"
+#include "spirv-tools/optimizer.hpp"
 
 namespace spvtools {
 namespace opt {
@@ -1200,6 +1201,13 @@ class SpirvEmitter : public ASTConsumer {
   /// Returns true on success and false otherwise.
   bool spirvToolsOptimize(std::vector<uint32_t> *mod, std::string *messages);
 
+  // \brief Runs the pass represented by the given pass token on the module.
+  // Returns true if the pass was successfully run. Any messages from the
+  // optimizer are returned in `messages`.
+  bool spirvToolsRunPass(std::vector<uint32_t> *mod,
+                         spvtools::Optimizer::PassToken token,
+                         std::string *messages);
+
   // \brief Calls SPIRV-Tools optimizer fix-opextinst-opcodes pass. This pass
   // fixes OpExtInst/OpExtInstWithForwardRefsKHR opcodes to use the correct one
   // depending of the presence of forward references.
@@ -1214,6 +1222,13 @@ class SpirvEmitter : public ASTConsumer {
   bool spirvToolsTrimCapabilities(std::vector<uint32_t> *mod,
                                   std::string *messages);
 
+  // \brief Runs the upgrade memory model pass using SPIRV-Tools's optimizer.
+  // This pass will modify the module, |mod|, so that it conforms to the Vulkan
+  // memory model instead of the GLSL450 memory model. Returns
+  // info/warning/error messages via |messages|.
+  bool spirvToolsUpgradeToVulkanMemoryModel(std::vector<uint32_t> *mod,
+                                            std::string *messages);
+
   /// \brief Helper function to run SPIRV-Tools optimizer's legalization passes.
   /// Runs the SPIRV-Tools legalization on the given SPIR-V module |mod|, and
   /// gets the info/warning/error messages via |messages|. If
@@ -1268,6 +1283,11 @@ class SpirvEmitter : public ASTConsumer {
                                           SpirvInstruction *scalar,
                                           SpirvLayoutRule rule);
 
+  /// Modifies the instruction in the code that use the GLSL450 memory module to
+  /// use the Vulkan memory model. This is done only if it has been requested or
+  /// the Vulkan memory model capability has been added to the module.
+  bool UpgradeToVulkanMemoryModelIfNeeded(std::vector<uint32_t> *module);
+
 public:
   /// \brief Wrapper method to create a fatal error message and report it
   /// in the diagnostic engine associated with this consumer.