Initial -embed-spirv support (#4974)

* Initial -embed-spirv support Add support for SPIR-V precompilation using the framework established for DXIL. Work on #4883 * SLANG_UNUSED * Add linkage attributes to exported spirv functions * Combine DXIL and SPIRV paths * Whitespace fix * Merge remaining precompiled spirv/dxil paths * Change inst accessors to return codegentarget * Add unit test for precompiled spirv --------- Co-authored-by: Yong He <yonghe@outlook.com>
shader-slang · Sep 5, 2024 · 3136771 · 3136771
1 parent 33e8bfd
commit 3136771
Show file tree

Hide file tree

Showing 19 changed files with 214 additions and 109 deletions.
diff --git a/include/slang.h b/include/slang.h
@@ -956,8 +956,7 @@ extern "C"
  GenerateWholeProgram, // bool
  UseUpToDateBinaryModule, // bool, when set, will only load
  // precompiled modules if it is up-to-date with its source.
-
- EmbedDXIL, // bool
+ EmbedDownstreamIR, // bool
  ForceDXLayout, // bool
  CountOf,
  };
@@ -4929,9 +4928,9 @@ namespace slang
  int targetIndex,
  bool value) = 0;
 
- virtual SLANG_NO_THROW void SLANG_MCALL setEmbedDXIL(bool value) = 0;
-
  virtual SLANG_NO_THROW void SLANG_MCALL setTargetForceDXLayout(int targetIndex, bool value) = 0;
+
+ virtual SLANG_NO_THROW void SLANG_MCALL setTargetEmbedDownstreamIR(int targetIndex, bool value) = 0;
  };
 
  #define SLANG_UUID_ICompileRequest ICompileRequest::getTypeGuid()

diff --git a/source/slang/slang-compiler-tu.cpp b/source/slang/slang-compiler-tu.cpp
@@ -54,7 +54,7 @@ namespace Slang
  * Precompile the module for the given target.
  *
  * This function creates a target program and emits the precompiled blob as
- * an embedded blob in the module IR, e.g. DXIL.
+ * an embedded blob in the module IR, e.g. DXIL, SPIR-V.
  * Because the IR for the Slang Module may violate the restrictions of the
  * target language, the emitted target blob may not be able to include the
  * full module, but rather only the subset that can be precompiled. For
@@ -77,7 +77,7 @@ namespace Slang
  * in the linked IR which survived the additional pruning.
  *
  * Functions that are rejected after linking+legalization (inside
- * emitPrecompiled*):
+ * emitPrecompiledDownstreamIR):
  * - (DXIL) Functions that return or take a HLSLStructuredBufferType
  * - (DXIL) Functions that return or take a Matrix type
  *
@@ -86,17 +86,13 @@ namespace Slang
  * phase of filtering.
  *
  * The original module IR functions matching those are then marked with
- * "AvailableIn*" (e.g. AvailableInDXILDecoration) to indicate to future
+ * "AvailableInDownstreamIRDecoration" to indicate to future
  * module users which functions are present in the precompiled blob.
  */
  SLANG_NO_THROW SlangResult SLANG_MCALL Module::precompileForTarget(
  SlangCompileTarget target,
  slang::IBlob** outDiagnostics)
  {
- if (target != SLANG_DXIL)
- {
- return SLANG_FAIL;
- }
  CodeGenTarget targetEnum = CodeGenTarget(target);
 
  auto module = getIRModule();
@@ -130,10 +126,15 @@ namespace Slang
  {
  case CodeGenTarget::DXIL:
  tp.getOptionSet().add(CompilerOptionName::Profile, Profile::RawEnum::DX_Lib_6_6);
- tp.getOptionSet().add(CompilerOptionName::EmbedDXIL, true);
  break;
+ case CodeGenTarget::SPIRV:
+ break;
+ default:
+ return SLANG_FAIL;
  }
 
+ tp.getOptionSet().add(CompilerOptionName::EmbedDownstreamIR, true);
+
  CodeGenContext::EntryPointIndices entryPointIndices;
 
  entryPointIndices.setCount(entryPointCount);
@@ -166,7 +167,7 @@ namespace Slang
  }
 
  ComPtr<IArtifact> outArtifact;
- SlangResult res = codeGenContext.emitPrecompiledDXIL(outArtifact);
+ SlangResult res = codeGenContext.emitPrecompiledDownstreamIR(outArtifact);
 
  sink.getBlobIfNeeded(outDiagnostics);
  if (res != SLANG_OK)
@@ -183,7 +184,8 @@ namespace Slang
  for (const auto& mangledName : metadata->getExportedFunctionMangledNames())
  {
  auto moduleInst = nameToFunction[mangledName];
- builder.addDecoration(moduleInst, kIROp_AvailableInDXILDecoration);
+ builder.addDecoration(moduleInst, kIROp_AvailableInDownstreamIRDecoration,
+ builder.getIntValue(builder.getIntType(), (int)targetReq->getTarget()));
  auto moduleDec = moduleInst->findDecoration<IRDownstreamModuleExportDecoration>();
  moduleDec->removeAndDeallocate();
  }
@@ -207,12 +209,7 @@ namespace Slang
  // Add the precompiled blob to the module
  builder.setInsertInto(module);
 
- switch (targetReq->getTarget())
- {
- case CodeGenTarget::DXIL:
- builder.emitEmbeddedDXIL(blob);
- break;
- }
+ builder.emitEmbeddedDownstreamIR(targetReq->getTarget(), blob);
 
  return SLANG_OK;
  }

diff --git a/source/slang/slang-compiler.cpp b/source/slang/slang-compiler.cpp
@@ -767,9 +767,9 @@ namespace Slang
 # pragma warning(pop)
 #endif
 
- SlangResult CodeGenContext::emitPrecompiledDXIL(ComPtr<IArtifact>& outArtifact)
+ SlangResult CodeGenContext::emitPrecompiledDownstreamIR(ComPtr<IArtifact>& outArtifact)
  {
- return emitWithDownstreamForEntryPoints(outArtifact);
+ return _emitEntryPoints(outArtifact);
  }
 
  String GetHLSLProfileName(Profile profile)
@@ -1257,10 +1257,7 @@ namespace Slang
  {
  CodeGenContext sourceCodeGenContext(this, sourceTarget, extensionTracker);
 
- if (target == CodeGenTarget::DXILAssembly || target == CodeGenTarget::DXIL)
- {
- sourceCodeGenContext.removeAvailableInDXIL = true;
- }
+ sourceCodeGenContext.removeAvailableInDownstreamIR = true;
 
  SLANG_RETURN_ON_FAIL(sourceCodeGenContext.emitEntryPointsSource(sourceArtifact));
  sourceCodeGenContext.maybeDumpIntermediate(sourceArtifact);
@@ -1561,20 +1558,23 @@ namespace Slang
  // Load embedded precompiled libraries from IR into library artifacts
  program->enumerateIRModules([&](IRModule* irModule)
  {
- for (auto inst : irModule->getModuleInst()->getChildren())
+ for (auto globalInst : irModule->getModuleInst()->getChildren())
  {
  if (target == CodeGenTarget::DXILAssembly || target == CodeGenTarget::DXIL)
  {
- if (inst->getOp() == kIROp_EmbeddedDXIL)
+ if (auto inst = as<IREmbeddedDownstreamIR>(globalInst))
  {
- auto slice = static_cast<IRBlobLit*>(inst->getOperand(0))->getStringSlice();
- ArtifactDesc desc = ArtifactDescUtil::makeDescForCompileTarget(SLANG_DXIL);
- desc.kind = ArtifactKind::Library;
+ if (inst->getTarget() == CodeGenTarget::DXIL)
+ {
+ auto slice = inst->getBlob()->getStringSlice();
+ ArtifactDesc desc = ArtifactDescUtil::makeDescForCompileTarget(SLANG_DXIL);
+ desc.kind = ArtifactKind::Library;
 
- auto library = ArtifactUtil::createArtifact(desc);
+  auto library = ArtifactUtil::createArtifact(desc);
 
- library->addRepresentationUnknown(StringBlob::create(slice));
- libraries.add(library);
+ library->addRepresentationUnknown(StringBlob::create(slice));
+ libraries.add(library);
+ }
  }
  }
  }
@@ -2030,7 +2030,7 @@ namespace Slang
  { 
  if (auto artifact = targetProgram->getExistingWholeProgramResult())
  {
- if (!targetProgram->getOptionSet().getBoolOption(CompilerOptionName::EmbedDXIL))
+ if (!targetProgram->getOptionSet().getBoolOption(CompilerOptionName::EmbedDownstreamIR))
  {
  artifacts.add(ComPtr<IArtifact>(artifact));
  }
@@ -2285,6 +2285,9 @@ namespace Slang
  auto linkage = getLinkage();
  for (auto targetReq : linkage->targets)
  {
+ if (targetReq->getOptionSet().getBoolOption(CompilerOptionName::EmbedDownstreamIR))
+ continue;
+
  auto targetProgram = program->getTargetProgram(targetReq);
  generateOutput(targetProgram);
  }

diff --git a/source/slang/slang-compiler.h b/source/slang/slang-compiler.h
@@ -2733,13 +2733,13 @@ namespace Slang
 
  SlangResult emitEntryPoints(ComPtr<IArtifact>& outArtifact);
 
- SlangResult emitPrecompiledDXIL(ComPtr<IArtifact>& outArtifact);
+ SlangResult emitPrecompiledDownstreamIR(ComPtr<IArtifact>& outArtifact);
 
  void maybeDumpIntermediate(IArtifact* artifact);
 
- // Used to cause instructions available in precompiled DXIL to be
+ // Used to cause instructions available in precompiled blobs to be
  // removed between IR linking and target source generation.
- bool removeAvailableInDXIL = false;
+ bool removeAvailableInDownstreamIR = false;
 
  protected:
  CodeGenTarget m_targetFormat = CodeGenTarget::Unknown;
@@ -2816,7 +2816,7 @@ namespace Slang
  virtual SLANG_NO_THROW void SLANG_MCALL setTargetForceGLSLScalarBufferLayout(int targetIndex, bool value) SLANG_OVERRIDE;
  virtual SLANG_NO_THROW void SLANG_MCALL setTargetForceDXLayout(int targetIndex, bool value) SLANG_OVERRIDE;
  virtual SLANG_NO_THROW void SLANG_MCALL setTargetGenerateWholeProgram(int targetIndex, bool value) SLANG_OVERRIDE;
- virtual SLANG_NO_THROW void SLANG_MCALL setEmbedDXIL(bool value) SLANG_OVERRIDE;
+ virtual SLANG_NO_THROW void SLANG_MCALL setTargetEmbedDownstreamIR(int targetIndex, bool value) SLANG_OVERRIDE;
  virtual SLANG_NO_THROW void SLANG_MCALL setMatrixLayoutMode(SlangMatrixLayoutMode mode) SLANG_OVERRIDE;
  virtual SLANG_NO_THROW void SLANG_MCALL setDebugInfoLevel(SlangDebugInfoLevel level) SLANG_OVERRIDE;
  virtual SLANG_NO_THROW void SLANG_MCALL setOptimizationLevel(SlangOptimizationLevel level) SLANG_OVERRIDE;

diff --git a/source/slang/slang-emit-spirv.cpp b/source/slang/slang-emit-spirv.cpp
@@ -3933,6 +3933,17 @@ struct SPIRVEmitContext
  }
  break;
  }
+ case kIROp_DownstreamModuleExportDecoration:
+ {
+ requireSPIRVCapability(SpvCapabilityLinkage);
+ auto name = decoration->getParent()->findDecoration<IRExportDecoration>()->getMangledName();
+ emitInst(getSection(SpvLogicalSectionID::Annotations),
+ decoration,
+ SpvOpDecorate,
+ dstID,
+ SpvDecorationLinkageAttributes, name, SpvLinkageTypeExport);
+ break;
+ }
  // ...
  }
 
@@ -6847,12 +6858,27 @@ SlangResult emitSPIRVFromIR(
 #endif
 
  auto shouldPreserveParams = codeGenContext->getTargetProgram()->getOptionSet().getBoolOption(CompilerOptionName::PreserveParameters);
+ auto generateWholeProgram = codeGenContext->getTargetProgram()->getOptionSet().getBoolOption(CompilerOptionName::GenerateWholeProgram);
  for (auto inst : irModule->getGlobalInsts())
  {
  if (as<IRDebugSource>(inst))
+ {
  context.ensureInst(inst);
+ }
  if (shouldPreserveParams && as<IRGlobalParam>(inst))
+ {
  context.ensureInst(inst);
+ }
+ if (generateWholeProgram)
+ {
+ if (auto func = as<IRFunc>(inst))
+ {
+ if (func->findDecoration<IRDownstreamModuleExportDecoration>())
+ {
+ context.ensureInst(inst);
+ }
+ }
+ }
  }
 
  // Emit source language info.

diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
@@ -416,31 +416,47 @@ bool checkStaticAssert(IRInst* inst, DiagnosticSink* sink)
  return false;
 }
 
-static void unexportNonEmbeddableDXIL(IRModule* irModule)
+static void unexportNonEmbeddableIR(CodeGenTarget target, IRModule* irModule)
 {
  for (auto inst : irModule->getGlobalInsts())
  {
  if (inst->getOp() == kIROp_Func)
  {
- // DXIL does not permit HLSLStructureBufferType in exported functions
- // or sadly Matrices (https://github.com/shader-slang/slang/issues/4880)
- auto type = as<IRFuncType>(inst->getFullType());
- auto argCount = type->getOperandCount();
- for (UInt aa = 0; aa < argCount; ++aa)
+ bool remove = false;
+ if (target == CodeGenTarget::HLSL)
  {
- auto operand = type->getOperand(aa);
- if (operand->getOp() == kIROp_HLSLStructuredBufferType ||
- operand->getOp() == kIROp_MatrixType)
+ // DXIL does not permit HLSLStructureBufferType in exported functions
+ // or sadly Matrices (https://github.com/shader-slang/slang/issues/4880)
+ auto type = as<IRFuncType>(inst->getFullType());
+ auto argCount = type->getOperandCount();
+ for (UInt aa = 0; aa < argCount; ++aa)
  {
- if (auto dec = inst->findDecoration<IRPublicDecoration>())
+ auto operand = type->getOperand(aa);
+ if (operand->getOp() == kIROp_HLSLStructuredBufferType ||
+ operand->getOp() == kIROp_MatrixType)
  {
- dec->removeAndDeallocate();
+ remove = true;
+ break;
  }
- if (auto dec = inst->findDecoration<IRDownstreamModuleExportDecoration>())
- {
- dec->removeAndDeallocate();
- }
- break;
+ }
+ }
+ else if (target == CodeGenTarget::SPIRV)
+ {
+ // SPIR-V does not allow exporting entry points
+ if (inst->findDecoration<IREntryPointDecoration>())
+ {
+ remove = true;
+ }
+ }
+ if (remove)
+ {
+ if (auto dec = inst->findDecoration<IRPublicDecoration>())
+ {
+ dec->removeAndDeallocate();
+ }
+ if (auto dec = inst->findDecoration<IRDownstreamModuleExportDecoration>())
+ {
+ dec->removeAndDeallocate();
  }
  }
  }
@@ -778,9 +794,9 @@ Result linkAndOptimizeIR(
  break;
  }
 
- if (codeGenContext->removeAvailableInDXIL)
+ if (codeGenContext->removeAvailableInDownstreamIR)
  {
- removeAvailableInDownstreamModuleDecorations(irModule);
+ removeAvailableInDownstreamModuleDecorations(target, irModule);
  }
 
  if (targetProgram->getOptionSet().shouldRunNonEssentialValidation())
@@ -1490,11 +1506,9 @@ Result linkAndOptimizeIR(
  auto metadata = new ArtifactPostEmitMetadata;
  outLinkedIR.metadata = metadata;
 
- if (targetProgram->getOptionSet().getBoolOption(CompilerOptionName::EmbedDXIL))
+ if (targetProgram->getOptionSet().getBoolOption(CompilerOptionName::EmbedDownstreamIR))
  {
- // We need to make sure that we don't try to export any functions that can't
- // be part of a DXIL library interface, eg. with resources.
- unexportNonEmbeddableDXIL(irModule);
+ unexportNonEmbeddableIR(target, irModule);
  }
 
  collectMetadata(irModule, *metadata);

diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
@@ -806,7 +806,7 @@ INST_RANGE(BindingQuery, GetRegisterIndex, GetRegisterSpace)
  INST(NumThreadsDecoration, numThreads, 3, 0)
  INST(WaveSizeDecoration, waveSize, 1, 0)
 
- INST(AvailableInDXILDecoration, availableInDXIL, 0, 0)
+ INST(AvailableInDownstreamIRDecoration, availableInDownstreamIR, 1, 0)
 
  // Added to IRParam parameters to an entry point
  /* GeometryInputPrimitiveTypeDecoration */
@@ -1244,7 +1244,7 @@ INST(DebugVar, DebugVar, 4, 0)
 INST(DebugValue, DebugValue, 2, 0)
 
 /* Embedded Precompiled Libraries */
-INST(EmbeddedDXIL, EmbeddedDXIL, 1, 0)
+INST(EmbeddedDownstreamIR, EmbeddedDownstreamIR, 2, 0)
 
 /* Inline assembly */