[SPIR-V] Add capability trimming pass (#5518)

This new optimization passes determines the required capabilities for a
module, and strips unneeded capabilities.
This means if some dead-code uses a capability, it won't be added as a
requirement for the module anymore.

This interacts with `[[vk::ext_capabilities]]` and
[[vk::ext_extension]]` attributes, as the extension/capability they
declare could be stripped if not required.

Still markes as draft for now as there are inconsistencies with debug
instruction I need to figure out:
 - if 2 DebugScope are generated, SPIRV-Tools squash them into 1.
- seems like we started to generated the wrong ones, but test didn't saw
the 2 conflicting scopes, only checked the first one.

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -453,8 +453,6 @@ bool CapabilityVisitor::visit(SpirvImageOp *instr) {
                        instr->getStorageClass());
   if (instr->hasOffset() || instr->hasConstOffsets())
     addCapability(spv::Capability::ImageGatherExtended);
-  if (instr->hasMinLod())
-    addCapability(spv::Capability::MinLod);
   if (instr->isSparse())
     addCapability(spv::Capability::SparseResidency);
 
@@ -864,6 +862,13 @@ bool CapabilityVisitor::visit(SpirvModule *, Visitor::Phase phase) {
     addCapability(spv::Capability::Shader);
     addCapability(spv::Capability::Linkage);
   }
+
+  // SPIRV-Tools now has a pass to trim superfluous capabilities. This means we
+  // can remove most capability-selection logic from here, and just add
+  // capabilities by default. SPIRV-Tools will clean those up. Note: this pass
+  // supports only some capabilities. This list should be expanded to match the
+  // supported capabilities.
+  addCapability(spv::Capability::MinLod);
   return true;
 }
 

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -883,10 +883,10 @@ void SpirvEmitter::HandleTranslationUnit(ASTContext &context) {
       !dsetbindingsToCombineImageSampler.empty() ||
       spirvOptions.signaturePacking;
 
+  // Run legalization passes
   if (spirvOptions.codeGenHighLevel) {
     beforeHlslLegalization = needsLegalization;
   } else {
-    // Run legalization passes
     if (needsLegalization) {
       std::string messages;
       if (!spirvToolsLegalize(&m, &messages,
@@ -902,8 +902,8 @@ void SpirvEmitter::HandleTranslationUnit(ASTContext &context) {
       }
     }
 
-    // Run optimization passes
     if (theCompilerInstance.getCodeGenOpts().OptimizationLevel > 0) {
+      // Run optimization passes
       std::string messages;
       if (!spirvToolsOptimize(&m, &messages)) {
         emitFatalError("failed to optimize SPIR-V: %0", {}) << messages;
@@ -914,6 +914,25 @@ void SpirvEmitter::HandleTranslationUnit(ASTContext &context) {
         return;
       }
     }
+
+    // Trim unused capabilities.
+    // When optimizations are enabled, some optimization passes like DCE could
+    // make some capabilities useless. To avoid logic duplication between this
+    // pass, and DXC, DXC generates some capabilities unconditionally. This
+    // means we should run this pass, even when optimizations are disabled.
+    {
+      std::string messages;
+      if (!spirvToolsTrimCapabilities(&m, &messages)) {
+        emitFatalError("failed to trim capabilities: %0", {}) << messages;
+        emitNote("please file a bug report on "
+                 "https://github.com/Microsoft/DirectXShaderCompiler/issues "
+                 "with source code if possible",
+                 {});
+        return;
+      } else if (!messages.empty()) {
+        emitWarning("SPIR-V capability trimming: %0", {}) << messages;
+      }
+    }
   }
 
   // Validate the generated SPIR-V code
@@ -13998,6 +14017,28 @@ bool SpirvEmitter::spirvToolsValidate(std::vector<uint32_t> *mod,
   return tools.Validate(mod->data(), mod->size(), options);
 }
 
+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);
+
+  optimizer.RegisterPass(spvtools::CreateTrimCapabilitiesPass());
+
+  return optimizer.Run(mod->data(), mod->size(), mod, options);
+}
+
 bool SpirvEmitter::spirvToolsOptimize(std::vector<uint32_t> *mod,
                                       std::string *messages) {
   spvtools::Optimizer optimizer(featureManager.getTargetEnv());

tools/clang/lib/SPIRV/SpirvEmitter.h

@@ -1187,6 +1187,14 @@ class SpirvEmitter : public ASTConsumer {
   /// Returns true on success and false otherwise.
   bool spirvToolsOptimize(std::vector<uint32_t> *mod, std::string *messages);
 
+  // \brief Calls SPIRV-Tools optimizer's, but only with the capability trimming
+  // pass. Removes unused capabilities from the given SPIR-V module |mod|, and
+  // returns info/warning/error messages via |messages|. This pass doesn't trim
+  // all capabilities. To see the list of supported capabilities, check the pass
+  // headers.
+  bool spirvToolsTrimCapabilities(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

tools/clang/test/CodeGenSPIRV/bezier.domain.hlsl2spv

@@ -1,4 +1,4 @@
-// RUN: %dxc -T ds_6_0 -E BezierEvalDS %s | FileCheck %s
+// RUN: %dxc -T ds_6_0 -E BezierEvalDS %s -O0 | FileCheck %s
 
 #define MAX_POINTS 4
 
@@ -200,4 +200,4 @@ DS_OUTPUT BezierEvalDS( HS_CONSTANT_DATA_OUTPUT input,
 //      %Output = OpVariable %_ptr_Function_DS_OUTPUT Function
 //          %83 = OpLoad %DS_OUTPUT %Output
 //                OpReturnValue %83
-//                OpFunctionEnd
+//                OpFunctionEnd

tools/clang/test/CodeGenSPIRV/bezier.hull.hlsl2spv

@@ -1,4 +1,4 @@
-// RUN: %dxc -T hs_6_0 -E SubDToBezierHS %s | FileCheck %s
+// RUN: %dxc -T hs_6_0 -E SubDToBezierHS %s -O0 | FileCheck %s
 
 #define MAX_POINTS 3
 

tools/clang/test/CodeGenSPIRV/empty-struct-interface.vs.hlsl2spv

@@ -1,4 +1,4 @@
-// RUN: %dxc -T vs_6_0 -E main %s | FileCheck %s
+// RUN: %dxc -T vs_6_0 -E main %s -O0 | FileCheck %s
 
 // There is no interface variable for VSIn or VSOut empty structs (See OpEntryPoint below).
 

tools/clang/test/CodeGenSPIRV/passthru-cs.hlsl2spv

@@ -1,4 +1,4 @@
-// RUN: %dxc -T cs_6_0 -E main %s | FileCheck %s
+// RUN: %dxc -T cs_6_0 -E main %s -O0 | FileCheck %s
 
 // Source: https://msdn.microsoft.com/en-us/library/windows/desktop/ff476330(v=vs.85).aspx
 
@@ -98,4 +98,4 @@ void main( uint3 DTid : SV_DispatchThreadID )
 //          %43 = OpAccessChain %_ptr_Uniform_uint %BufferOut %uint_0 %41
 //                OpStore %43 %42
 //                OpReturn
-//                OpFunctionEnd
+//                OpFunctionEnd

tools/clang/test/CodeGenSPIRV/passthru-ps.hlsl2spv

@@ -1,4 +1,4 @@
-// RUN: %dxc -T ps_6_0 -E main %s | FileCheck %s
+// RUN: %dxc -T ps_6_0 -E main %s -O0 | FileCheck %s
 
 float4 main(float4 input: COLOR): SV_Target
 {
@@ -49,4 +49,4 @@ float4 main(float4 input: COLOR): SV_Target
 //    %bb_entry = OpLabel
 //          %19 = OpLoad %v4float %input
 //                OpReturnValue %19
-//                OpFunctionEnd
+//                OpFunctionEnd

tools/clang/test/CodeGenSPIRV/passthru-vs.hlsl2spv

@@ -1,4 +1,4 @@
-// RUN: %dxc -T vs_6_0 -E main %s | FileCheck %s
+// RUN: %dxc -T vs_6_0 -E main %s -O0 | FileCheck %s
 
 struct PSInput {
   float4 position : SV_Position;
@@ -86,4 +86,4 @@ PSInput main(float4 position: POSITION, float4 color: COLOR) {
 //                OpStore %35 %34
 //          %36 = OpLoad %PSInput %result
 //                OpReturnValue %36
-//                OpFunctionEnd
+//                OpFunctionEnd

tools/clang/test/CodeGenSPIRV_Lit/capability.trimmed.o3.hlsl

@@ -0,0 +1,13 @@
+// RUN: %dxc -T ps_6_0 -E main -O3 %s -spirv | FileCheck %s
+
+Texture1D   <float4> t : register(t1);
+SamplerState gSampler : register(s2);
+
+// CHECK-NOT: OpCapability MinLod
+
+float4 main(uint clamp : A) : SV_Target {
+    if (clamp < 0) {
+      return t.Sample(gSampler, 0.5f, 2, float(clamp));
+    }
+    return float4(0, 0, 0, 0);
+}

tools/clang/test/CodeGenSPIRV_Lit/capability.untrimmed.fcgl.hlsl

@@ -0,0 +1,12 @@
+// RUN: %dxc -T cs_6_0 -E main -fcgl %s -spirv | FileCheck %s
+
+// DXC generates by default all capabilities, and calls a specific pass
+// to trim unwanted capabilities. This is done to prevent code duplication
+// between the optimizer, and DXC.
+// -fcgl is used for debug purposed, and disable all passes, even
+// legalization. So if -fcgl is called, the capability should be present,
+// event if unused.
+
+// CHECK: OpCapability MinLod
+[numthreads(1, 1, 1)]
+void main() { }

tools/clang/test/CodeGenSPIRV_Lit/capability.untrimmed.o0.hlsl

@@ -0,0 +1,13 @@
+// RUN: %dxc -T ps_6_0 -E main -O0 %s -spirv | FileCheck %s
+
+Texture1D   <float4> t : register(t1);
+SamplerState gSampler : register(s2);
+
+// CHECK: OpCapability MinLod
+
+float4 main(uint clamp : A) : SV_Target {
+    if (clamp < 0) {
+      return t.Sample(gSampler, 0.5f, 2, float(clamp));
+    }
+    return float4(0, 0, 0, 0);
+}