[spirv] Update the SPIR-V codegen for HLSL discard. (#2895)

Generate OpDemoteToHelperInvocationEXT instruction only if
`-fspv-extension=SPV_EXT_demote_to_helper_invocation` is specified by
the user.

Fixes #2823.

tools/clang/include/clang/SPIRV/FeatureManager.h

@@ -59,19 +59,23 @@ public:
 
   /// Allows the given extension to be used in CodeGen.
   bool allowExtension(llvm::StringRef);
+
   /// Allows all extensions to be used in CodeGen.
   void allowAllKnownExtensions();
+
   /// Rqeusts the given extension for translating the given target feature at
   /// the given source location. Emits an error if the given extension is not
   /// permitted to use.
   bool requestExtension(Extension, llvm::StringRef target, SourceLocation);
 
   /// Translates extension name to symbol.
-  static Extension getExtensionSymbol(llvm::StringRef name);
+  Extension getExtensionSymbol(llvm::StringRef name);
+
   /// Translates extension symbol to name.
-  static const char *getExtensionName(Extension symbol);
+  const char *getExtensionName(Extension symbol);
+
   /// Returns true if the given extension is a KHR extension.
-  static bool isKHRExtension(llvm::StringRef name);
+  bool isKHRExtension(llvm::StringRef name);
 
   /// Returns the names of all known extensions as a string.
   std::string getKnownExtensions(const char *delimiter, const char *prefix = "",
@@ -92,9 +96,12 @@ public:
   bool isExtensionRequiredForTargetEnv(Extension);
 
   /// Returns true if the given extension is set in allowedExtensions
-  bool isExtensionEnabled(llvm::StringRef name);
+  bool isExtensionEnabled(Extension);
 
 private:
+  /// Returns whether codegen should allow usage of this extension by default.
+  bool enabledByDefault(Extension);
+
   /// \brief Wrapper method to create an error message and report it
   /// in the diagnostic engine associated with this object.
   template <unsigned N>

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -519,7 +519,7 @@ bool CapabilityVisitor::visit(SpirvEntryPoint *entryPoint) {
   case spv::ExecutionModel::AnyHitNV:
   case spv::ExecutionModel::MissNV:
   case spv::ExecutionModel::CallableNV:
-    if (featureManager.isExtensionEnabled("SPV_NV_ray_tracing")) {
+    if (featureManager.isExtensionEnabled(Extension::NV_ray_tracing)) {
       addCapability(spv::Capability::RayTracingNV);
       addExtension(Extension::NV_ray_tracing, "SPV_NV_ray_tracing", {});
     } else {

tools/clang/lib/SPIRV/FeatureManager.cpp

@@ -22,11 +22,8 @@ FeatureManager::FeatureManager(DiagnosticsEngine &de,
 
   if (opts.allowedExtensions.empty()) {
     // If no explicit extension control from command line, use the default mode:
-    // allowing all extensions.
-    // Special case : KHR_ray_tracing and NV_ray_tracing are mutually exclusive
-    // so enable only KHR extension by default
+    // allowing all extensions that are enabled by default.
     allowAllKnownExtensions();
-    allowedExtensions.reset(static_cast<unsigned>(Extension::NV_ray_tracing));
   } else {
     for (auto ext : opts.allowedExtensions)
       allowExtension(ext);
@@ -70,7 +67,13 @@ bool FeatureManager::allowExtension(llvm::StringRef name) {
   return true;
 }
 
-void FeatureManager::allowAllKnownExtensions() { allowedExtensions.set(); }
+void FeatureManager::allowAllKnownExtensions() {
+  allowedExtensions.set();
+  const auto numExtensions = static_cast<uint32_t>(Extension::Unknown);
+  for (uint32_t ext = 0; ext < numExtensions; ++ext)
+    if (!enabledByDefault(static_cast<Extension>(ext)))
+      allowedExtensions.reset(ext);
+}
 
 bool FeatureManager::requestExtension(Extension ext, llvm::StringRef target,
                                       SourceLocation srcLoc) {
@@ -230,14 +233,29 @@ bool FeatureManager::isExtensionRequiredForTargetEnv(Extension ext) {
   return required;
 }
 
-bool FeatureManager::isExtensionEnabled(llvm::StringRef name) {
+bool FeatureManager::isExtensionEnabled(Extension ext) {
   bool allowed = false;
-  Extension ext = getExtensionSymbol(name);
   if (ext != Extension::Unknown &&
       allowedExtensions.test(static_cast<unsigned>(ext)))
     allowed = true;
   return allowed;
 }
 
+bool FeatureManager::enabledByDefault(Extension ext) {
+  switch (ext) {
+    // KHR_ray_tracing and NV_ray_tracing are mutually exclusive so enable only
+    // KHR extension by default
+  case Extension::NV_ray_tracing:
+    return false;
+    // Enabling EXT_demote_to_helper_invocation changes the code generation
+    // behavior for the 'discard' statement. Therefore we will only enable it if
+    // the user explicitly asks for it.
+  case Extension::EXT_demote_to_helper_invocation:
+    return false;
+  default:
+    return true;
+  }
+}
+
 } // end namespace spirv
 } // end namespace clang

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -1422,14 +1422,27 @@ void SpirvEmitter::doDiscardStmt(const DiscardStmt *discardStmt) {
     return;
   }
 
-  // SPV_EXT_demote_to_helper_invocation SPIR-V extension provides a new
-  // instruction OpDemoteToHelperInvocationEXT allowing shaders to "demote" a
-  // fragment shader invocation to behave like a helper invocation for its
-  // duration. The demoted invocation will have no further side effects and will
-  // not output to the framebuffer, but remains active and can participate in
-  // computing derivatives and in subgroup operations. This is a better match
-  // for the "discard" instruction in HLSL.
-  spvBuilder.createDemoteToHelperInvocationEXT(discardStmt->getLoc());
+  if (featureManager.isExtensionEnabled(
+          Extension::EXT_demote_to_helper_invocation)) {
+    // SPV_EXT_demote_to_helper_invocation SPIR-V extension provides a new
+    // instruction OpDemoteToHelperInvocationEXT allowing shaders to "demote" a
+    // fragment shader invocation to behave like a helper invocation for its
+    // duration. The demoted invocation will have no further side effects and
+    // will not output to the framebuffer, but remains active and can
+    // participate in computing derivatives and in subgroup operations. This is
+    // a better match for the "discard" instruction in HLSL.
+    spvBuilder.createDemoteToHelperInvocationEXT(discardStmt->getLoc());
+  } else {
+    // Note: if/when the demote behavior becomes part of the core Vulkan spec,
+    // we should no longer generate OpKill for 'discard', and always generate
+    // the demote behavior.
+    spvBuilder.createKill(discardStmt->getLoc());
+    // Some statements that alter the control flow (break, continue, return, and
+    // discard), require creation of a new basic block to hold any statement
+    // that may follow them.
+    auto *newBB = spvBuilder.createBasicBlock();
+    spvBuilder.setInsertPoint(newBB);
+  }
 }
 
 void SpirvEmitter::doDoStmt(const DoStmt *theDoStmt,

tools/clang/test/CodeGenSPIRV/cf.discard.hlsl

@@ -1,30 +1,36 @@
 // Run: %dxc -T ps_6_0 -E main
 
-// According to the HLS spec, discard can only be called from a pixel shader.
-
-// CHECK: OpCapability DemoteToHelperInvocationEXT
-// CHECK: OpExtension "SPV_EXT_demote_to_helper_invocation"
+// Note: The proper way to translate HLSL 'discard' to SPIR-V is using the
+// 'OpDemoteToHelperInvocationEXT' instruction, which requires the
+// SPV_EXT_demote_to_helper_invocation extension.
+//
+// If you want that behavior, please use the following command line option:
+// '-fspv-extension=SPV_EXT_demote_to_helper_invocation'
+//
+// In the absence of this option DXC will generate 'OpKill' for 'discard'.
 
 void main() {
   int a, b;
   bool cond = true;
-  
+
   while(cond) {
 // CHECK: %while_body = OpLabel
     if(a==b) {
 // CHECK: %if_true = OpLabel
-// CHECK: OpDemoteToHelperInvocationEXT
-      discard;
-      break;
+// CHECK-NEXT: OpKill
+      {{discard;}}
+      discard;  // No SPIR-V should be emitted for this statement.
+      break;    // No SPIR-V should be emitted for this statement.
     } else {
-// CHECK: %if_false = OpLabel
+// CHECK-NEXT: %if_false = OpLabel
       ++a;
-// CHECK: OpDemoteToHelperInvocationEXT
+// CHECK: OpKill
       discard;
-      continue;
-      --b;
+      continue; // No SPIR-V should be emitted for this statement.
+      --b;      // No SPIR-V should be emitted for this statement.
     }
-// CHECK: %if_merge = OpLabel
+// CHECK-NEXT: %if_merge = OpLabel
+
   }
 // CHECK: %while_merge = OpLabel
 

tools/clang/test/CodeGenSPIRV/cf.discard.to-demote.hlsl

@@ -0,0 +1,31 @@
+// Run: %dxc -T ps_6_0 -E main -fspv-extension=SPV_EXT_demote_to_helper_invocation
+
+// According to the HLS spec, discard can only be called from a pixel shader.
+
+// CHECK: OpCapability DemoteToHelperInvocationEXT
+// CHECK: OpExtension "SPV_EXT_demote_to_helper_invocation"
+
+void main() {
+  int a, b;
+  bool cond = true;
+
+  while(cond) {
+// CHECK: %while_body = OpLabel
+    if(a==b) {
+// CHECK: %if_true = OpLabel
+// CHECK: OpDemoteToHelperInvocationEXT
+      discard;
+      break;
+    } else {
+// CHECK: %if_false = OpLabel
+      ++a;
+// CHECK: OpDemoteToHelperInvocationEXT
+      discard;
+      continue;
+      --b;
+    }
+// CHECK: %if_merge = OpLabel
+  }
+// CHECK: %while_merge = OpLabel
+
+}

tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -511,6 +511,7 @@ TEST_F(FileTest, BreakStmtMixed) { runFileTest("cf.break.mixed.hlsl"); }
 
 // For discard statement
 TEST_F(FileTest, Discard) { runFileTest("cf.discard.hlsl"); }
+TEST_F(FileTest, DiscardToDemote) { runFileTest("cf.discard.to-demote.hlsl"); }
 TEST_F(FileTest, DiscardCS) {
   // Using discard is only allowed in pixel shaders.
   runFileTest("cf.discard.cs.hlsl", Expect::Failure);