[Driver][SPIR-V] Use consistent tools to convert between text and binary form (#120266)

Currently we produce SPIR-V text with `spirv-dis` but assemble it with
`llvm-spirv`. The SPIR-V text format is different between the tools so
the assemble fails. Use `spirv-as` for assembly as it uses the same
format.

---------

Signed-off-by: Sarnie, Nick <[email protected]>

Author: sarnex

clang/lib/Driver/ToolChains/SPIRV.cpp

@@ -26,8 +26,8 @@ void SPIRV::constructTranslateCommand(Compilation &C, const Tool &T,
   llvm::opt::ArgStringList CmdArgs(Args);
   CmdArgs.push_back(Input.getFilename());
 
-  if (Input.getType() == types::TY_PP_Asm)
-    CmdArgs.push_back("-to-binary");
+  assert(Input.getType() != types::TY_PP_Asm && "Unexpected input type");
+
   if (Output.getType() == types::TY_PP_Asm)
     CmdArgs.push_back("--spirv-tools-dis");
 
@@ -46,6 +46,31 @@ void SPIRV::constructTranslateCommand(Compilation &C, const Tool &T,
                                          Exec, CmdArgs, Input, Output));
 }
 
+void SPIRV::constructAssembleCommand(Compilation &C, const Tool &T,
+                                     const JobAction &JA,
+                                     const InputInfo &Output,
+                                     const InputInfo &Input,
+                                     const llvm::opt::ArgStringList &Args) {
+  llvm::opt::ArgStringList CmdArgs(Args);
+  CmdArgs.push_back(Input.getFilename());
+
+  assert(Input.getType() == types::TY_PP_Asm && "Unexpected input type");
+
+  CmdArgs.append({"-o", Output.getFilename()});
+
+  // Try to find "spirv-as-<LLVM_VERSION_MAJOR>". Otherwise, fall back to
+  // plain "spirv-as".
+  using namespace std::string_literals;
+  auto VersionedTool = "spirv-as-"s + std::to_string(LLVM_VERSION_MAJOR);
+  std::string ExeCand = T.getToolChain().GetProgramPath(VersionedTool.c_str());
+  if (!llvm::sys::fs::can_execute(ExeCand))
+    ExeCand = T.getToolChain().GetProgramPath("spirv-as");
+
+  const char *Exec = C.getArgs().MakeArgString(ExeCand);
+  C.addCommand(std::make_unique<Command>(JA, T, ResponseFileSupport::None(),
+                                         Exec, CmdArgs, Input, Output));
+}
+
 void SPIRV::Translator::ConstructJob(Compilation &C, const JobAction &JA,
                                      const InputInfo &Output,
                                      const InputInfoList &Inputs,
@@ -57,12 +82,29 @@ void SPIRV::Translator::ConstructJob(Compilation &C, const JobAction &JA,
   constructTranslateCommand(C, *this, JA, Output, Inputs[0], {});
 }
 
+void SPIRV::Assembler::ConstructJob(Compilation &C, const JobAction &JA,
+                                    const InputInfo &Output,
+                                    const InputInfoList &Inputs,
+                                    const ArgList &Args,
+                                    const char *AssembleOutput) const {
+  claimNoWarnArgs(Args);
+  if (Inputs.size() != 1)
+    llvm_unreachable("Invalid number of input files.");
+  constructAssembleCommand(C, *this, JA, Output, Inputs[0], {});
+}
+
 clang::driver::Tool *SPIRVToolChain::getTranslator() const {
   if (!Translator)
     Translator = std::make_unique<SPIRV::Translator>(*this);
   return Translator.get();
 }
 
+clang::driver::Tool *SPIRVToolChain::getAssembler() const {
+  if (!Assembler)
+    Assembler = std::make_unique<SPIRV::Assembler>(*this);
+  return Assembler.get();
+}
+
 clang::driver::Tool *SPIRVToolChain::SelectTool(const JobAction &JA) const {
   Action::ActionClass AC = JA.getKind();
   return SPIRVToolChain::getTool(AC);
@@ -73,8 +115,9 @@ clang::driver::Tool *SPIRVToolChain::getTool(Action::ActionClass AC) const {
   default:
     break;
   case Action::BackendJobClass:
-  case Action::AssembleJobClass:
     return SPIRVToolChain::getTranslator();
+  case Action::AssembleJobClass:
+    return SPIRVToolChain::getAssembler();
   }
   return ToolChain::getTool(AC);
 }

clang/lib/Driver/ToolChains/SPIRV.h

@@ -22,6 +22,11 @@ void constructTranslateCommand(Compilation &C, const Tool &T,
                                const InputInfo &Input,
                                const llvm::opt::ArgStringList &Args);
 
+void constructAssembleCommand(Compilation &C, const Tool &T,
+                              const JobAction &JA, const InputInfo &Output,
+                              const InputInfo &Input,
+                              const llvm::opt::ArgStringList &Args);
+
 class LLVM_LIBRARY_VISIBILITY Translator : public Tool {
 public:
   Translator(const ToolChain &TC)
@@ -47,13 +52,25 @@ class LLVM_LIBRARY_VISIBILITY Linker final : public Tool {
                     const char *LinkingOutput) const override;
 };
 
+class LLVM_LIBRARY_VISIBILITY Assembler final : public Tool {
+public:
+  Assembler(const ToolChain &TC) : Tool("SPIRV::Assembler", "spirv-as", TC) {}
+  bool hasIntegratedAssembler() const override { return false; }
+  bool hasIntegratedCPP() const override { return false; }
+  void ConstructJob(Compilation &C, const JobAction &JA,
+                    const InputInfo &Output, const InputInfoList &Inputs,
+                    const llvm::opt::ArgList &TCArgs,
+                    const char *AssembleOutput) const override;
+};
+
 } // namespace SPIRV
 } // namespace tools
 
 namespace toolchains {
 
 class LLVM_LIBRARY_VISIBILITY SPIRVToolChain : public ToolChain {
   mutable std::unique_ptr<Tool> Translator;
+  mutable std::unique_ptr<Tool> Assembler;
 
 public:
   SPIRVToolChain(const Driver &D, const llvm::Triple &Triple,
@@ -81,6 +98,8 @@ class LLVM_LIBRARY_VISIBILITY SPIRVToolChain : public ToolChain {
 
 private:
   clang::driver::Tool *getTranslator() const;
+  clang::driver::Tool *getAssembler() const;
+
   bool NativeLLVMSupport;
 };
 

clang/test/Driver/spirv-toolchain.cl

@@ -43,7 +43,7 @@
 // Check assembly input -> object output
 // RUN: %clang -### --target=spirv64 -x assembler -c %s 2>&1 | FileCheck --check-prefix=ASM %s
 // RUN: %clang -### --target=spirv32 -x assembler -c %s 2>&1 | FileCheck --check-prefix=ASM %s
-// ASM: {{llvm-spirv.*"}} {{".*"}} "-to-binary" "-o" {{".*o"}}
+// ASM: {{spirv-as.*"}} {{".*"}} "-o" {{".*o"}}
 
 //-----------------------------------------------------------------------------
 // Check --save-temps.
@@ -56,7 +56,7 @@
 // TMP-SAME: "-o" [[BC:".*bc"]]
 // TMP-SAME: [[I]]
 // TMP: {{llvm-spirv.*"}} [[BC]] "--spirv-tools-dis" "-o" [[S:".*s"]]
-// TMP: {{llvm-spirv.*"}} [[S]] "-to-binary" "-o" {{".*o"}}
+// TMP: {{spirv-as.*"}} [[S]] "-o" {{".*o"}}
 
 //-----------------------------------------------------------------------------
 // Check linking when multiple input files are passed.