[spir-v] Add option to print SPIR-V between passes (#4503)

This debug/development option matches the `--print-all` flag in
spirv-opt or `--print-after-all`/`--print-before-all` in llvm's opt.

The new option can be enabled with `-fspv-print-all`.

Author: kuhar

include/dxc/Support/HLSLOptions.td

@@ -367,6 +367,8 @@ def Wno_vk_ignored_features : Joined<["-"], "Wno-vk-ignored-features">, Group<sp
   HelpText<"Do not emit warnings for ingored features resulting from no Vulkan support">;
 def Wno_vk_emulated_features : Joined<["-"], "Wno-vk-emulated-features">, Group<spirv_Group>, Flags<[CoreOption, DriverOption, HelpHidden]>,
   HelpText<"Do not emit warnings for emulated features resulting from no direct mapping">;
+def fspv_print_all: Flag<["-"], "fspv-print-all">, Group<spirv_Group>, Flags<[CoreOption, DriverOption]>,
+  HelpText<"Print the SPIR-V module before each pass and after the last one. Useful for debugging SPIR-V legalization and optimization passes.">;
 def Oconfig : CommaJoined<["-"], "Oconfig=">, Group<spirv_Group>, Flags<[CoreOption]>,
   HelpText<"Specify a comma-separated list of SPIRV-Tools passes to customize optimization configuration (see http://khr.io/hlsl2spirv#optimization)">;
 // SPIRV Change Ends

include/dxc/Support/SPIRVOptions.h

@@ -89,6 +89,8 @@ struct SpirvCodeGenOptions {
 
   bool signaturePacking; ///< Whether signature packing is enabled or not
 
+  bool printAll; // Dump SPIR-V module before each pass and after the last one.
+
   // String representation of all command line options.
   std::string clOptions;
 };

lib/DxcSupport/HLSLOptions.cpp

@@ -986,6 +986,8 @@ int ReadDxcOpts(const OptTable *optionTable, unsigned flagsToInclude,
     opts.SpirvOptions.allowedExtensions.push_back(A->getValue());
   }
 
+  opts.SpirvOptions.printAll = Args.hasFlag(OPT_fspv_print_all, OPT_INVALID, false);
+
   opts.SpirvOptions.debugInfoFile = opts.SpirvOptions.debugInfoSource = false;
   opts.SpirvOptions.debugInfoLine = opts.SpirvOptions.debugInfoTool = false;
   opts.SpirvOptions.debugInfoRich = false;
@@ -1081,6 +1083,7 @@ int ReadDxcOpts(const OptTable *optionTable, unsigned flagsToInclude,
       Args.hasFlag(OPT_fspv_reduce_load_size, OPT_INVALID, false) ||
       Args.hasFlag(OPT_fspv_reflect, OPT_INVALID, false) ||
       Args.hasFlag(OPT_fspv_fix_func_call_arguments, OPT_INVALID, false) ||
+      Args.hasFlag(OPT_fspv_print_all, OPT_INVALID, false) ||
       Args.hasFlag(OPT_Wno_vk_ignored_features, OPT_INVALID, false) ||
       Args.hasFlag(OPT_Wno_vk_emulated_features, OPT_INVALID, false) ||
       Args.hasFlag(OPT_fvk_auto_shift_bindings, OPT_INVALID, false) ||

tools/clang/include/clang/SPIRV/String.h

@@ -9,11 +9,13 @@
 #ifndef LLVM_CLANG_SPIRV_STRING_H
 #define LLVM_CLANG_SPIRV_STRING_H
 
-#include <string>
-#include <vector>
-
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <streambuf>
+#include <string>
+#include <vector>
 
 namespace clang {
 namespace spirv {
@@ -28,6 +30,36 @@ std::vector<uint32_t> encodeSPIRVString(llvm::StringRef strChars);
 /// It follows the SPIR-V string encoding requirements.
 std::string decodeSPIRVString(llvm::ArrayRef<uint32_t> strWords);
 
+/// \brief Stream buffer implementation that writes to `llvm::raw_ostream`.
+/// Intended to be used with APIs that write to `std::ostream`.
+///
+/// Sample:
+/// ```
+/// RawOstreamBuf buf(llvm::errs());
+/// std::ostream os(&buf);
+/// some_print(os);
+/// ```
+class RawOstreamBuf : public std::streambuf {
+public:
+  RawOstreamBuf(llvm::raw_ostream &os) : os(os) {}
+  using char_type = std::streambuf::char_type;
+  using int_type = std::streambuf::int_type;
+
+protected:
+  std::streamsize xsputn(const char_type *s, std::streamsize count) override {
+    os << llvm::StringRef(s, count);
+    return count;
+  }
+
+  int_type overflow(int_type c) override {
+    os << char_type(c);
+    return 0;
+  }
+
+private:
+  llvm::raw_ostream &os;
+};
+
 } // end namespace string
 } // end namespace spirv
 } // end namespace clang

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -13032,12 +13032,16 @@ bool SpirvEmitter::spirvToolsValidate(std::vector<uint32_t> *mod,
 bool SpirvEmitter::spirvToolsOptimize(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);
 
@@ -13073,6 +13077,11 @@ bool SpirvEmitter::spirvToolsLegalize(std::vector<uint32_t> *mod,
                  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);
   // Add interface variable SROA if the signature packing is enabled.

tools/clang/test/CodeGenSPIRV/fspv-print-all.hlsl

@@ -0,0 +1,17 @@
+// RUN: %dxc -T ps_6_0 -E main -spirv -fspv-print-all
+
+// We expect the whole module to be printed multiple times, but we cannot check
+// that here as stderr contents do not appear as test diagnostics. Instead,
+// just confirm that the file compiles fine with `-fspv-print-all` enabled.
+
+// CHECK: OpEntryPoint Fragment %main
+
+struct Input
+{
+  float4 color : COLOR;
+};
+
+float4 main(Input input) : SV_TARGET
+{
+  return input.color;
+}

tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -3126,4 +3126,6 @@ float4 PSMain(float4 color : COLOR) : SV_TARGET { return color; }
 
 TEST_F(FileTest, RenameEntrypoint) { runFileTest("fspv-entrypoint-name.hlsl"); }
 
+TEST_F(FileTest, PrintAll) { runFileTest("fspv-print-all.hlsl"); }
+
 } // namespace

tools/clang/unittests/SPIRV/StringTest.cpp

@@ -7,9 +7,11 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "gmock/gmock.h"
 #include "clang/SPIRV/String.h"
+#include "llvm/Support/raw_ostream.h"
+#include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include <ostream>
 
 namespace {
 
@@ -79,7 +81,34 @@ TEST(String, EncodeAndDecodeString) {
 
   EXPECT_EQ(str, result);
 }
-
 // TODO: Add more ModuleBuilder tests
 
+TEST(String, RawOstreamBuf) {
+  // Set up the following output stream structure:
+  //  os -> buf -> rawOS -> underlyingBuffer
+  //
+  // Check that the contents written to `os` appear in `underlyingBuffer`.
+  std::string underlingBuffer;
+  llvm::raw_string_ostream rawOS(underlingBuffer);
+  string::RawOstreamBuf buf(rawOS);
+  std::ostream os(&buf);
+
+  // Flushes both buffers.
+  auto flush = [&rawOS, &os] {
+    os.flush();
+    rawOS.flush();
+  };
+
+  EXPECT_EQ(underlingBuffer, "");
+  os << "test";
+  flush();
+  EXPECT_EQ(underlingBuffer, "test");
+  os << 13 << 37 << "\n";
+  flush();
+  EXPECT_EQ(underlingBuffer, "test1337\n");
+  os << ' ';
+  flush();
+  EXPECT_EQ(underlingBuffer, "test1337\n ");
+}
+
 } // anonymous namespace