Skip to content

Commit

Permalink
Support mixture of precompiled and non-precompiled modules (shader-sl…
Browse files Browse the repository at this point in the history
…ang#4860)

* Support mixture of precompiled and non-precompiled modules

This changes the implementation of precompile DXIL modules to
accept combinations of modules with precompiled DXIL, ones without,
and ones with a mixture of precompiled DXIL and Slang IR.

During precompilation, module IR is analyzed to find public functions
which appear to be capable of being compiled as HLSL, and those
functions are given a HLSLExport decoration, ensuring they are emitted
as HLSL and preserved in the precompiled DXIL blob. The IR for those
functions is then tagged with a new decoration AvailableInDXIL, which
marks that their implementation is present in the embedded DXIL blob.
The DXIL blob is attached to the IR as before, inside a EmbeddedDXIL
BlobLit instruction.

The logic that determines whether or not functions should be
precompiled to DXIL is a placeholder at this point, returning true
always. A subsequent change will add selection criteria.

During module linking, the full module IR is available, as well
as the optional EmbeddedDXIL blob. The IR for functions implemented
by the blob are tagged with AvailableInDXIL in the module IR.

After linking the IR for all modules to program level IR, the IR for
the functions marked AvailableInDXIL are deleted from the linked IR,
prior to emitting HLSL and compiling linking the result.

This change also changes the point of time when the module IR is
checked for EmbeddedDXIL blobs. Instead of happening at load time
as before, it happens during immediately before final linking, meaning
that the blob does not need to be independently stored with the module
separate from the IR as was done previously.

Work on shader-slang#4792

* Clean up debug prints

* Call isSimpleHLSLDataType stub

* Address feedback on precompiled dxil support

Allow for IR filtering both before and after linking.
Only mark AvailableInDXIL those functions which pass
both filtering stages. Functions are corrlated using
mangled function names.

Rather than delete functions entirely when linking with
libraries that include precompiled DXIL, instead convert
the IR function definitions to declarations by gutting
them, removing child blocks.

* Use artifact metadata and name list instead of linkedir hack

* Use String instead of UnownedStringSlice

* Update tests

* Renaming

* Minor edits

* Don't fully remove functions post-link

* Unexport before collecting metadata
  • Loading branch information
cheneym2 authored Aug 30, 2024
1 parent 87d3d4f commit ddf4a32
Show file tree
Hide file tree
Showing 26 changed files with 425 additions and 102 deletions.
4 changes: 1 addition & 3 deletions include/slang.h
Original file line number Diff line number Diff line change
Expand Up @@ -4929,9 +4929,7 @@ namespace slang
int targetIndex,
bool value) = 0;

virtual SLANG_NO_THROW void SLANG_MCALL setTargetEmbedDXIL(
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;
};
Expand Down
5 changes: 5 additions & 0 deletions source/compiler-core/slang-artifact-associated-impl.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -308,4 +308,9 @@ Slice<ShaderBindingRange> ArtifactPostEmitMetadata::getUsedBindingRanges()
return Slice<ShaderBindingRange>(m_usedBindings.getBuffer(), m_usedBindings.getCount());
}

Slice<String> ArtifactPostEmitMetadata::getExportedFunctionMangledNames()
{
return Slice<String>(m_exportedFunctionMangledNames.getBuffer(), m_exportedFunctionMangledNames.getCount());
}

} // namespace Slang
4 changes: 3 additions & 1 deletion source/compiler-core/slang-artifact-associated-impl.h
Original file line number Diff line number Diff line change
Expand Up @@ -155,13 +155,15 @@ class ArtifactPostEmitMetadata : public ComBaseObject, public IArtifactPostEmitM

// IArtifactPostEmitMetadata
SLANG_NO_THROW virtual Slice<ShaderBindingRange> SLANG_MCALL getUsedBindingRanges() SLANG_OVERRIDE;

SLANG_NO_THROW virtual Slice<String> SLANG_MCALL getExportedFunctionMangledNames() SLANG_OVERRIDE;

void* getInterface(const Guid& uuid);
void* getObject(const Guid& uuid);

static ComPtr<IArtifactPostEmitMetadata> create() { return ComPtr<IArtifactPostEmitMetadata>(new ThisType); }

List<ShaderBindingRange> m_usedBindings;
List<String> m_exportedFunctionMangledNames;
};

} // namespace Slang
Expand Down
3 changes: 3 additions & 0 deletions source/compiler-core/slang-artifact-associated.h
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,9 @@ class IArtifactPostEmitMetadata : public ICastable

/// Get the binding ranges
SLANG_NO_THROW virtual Slice<ShaderBindingRange> SLANG_MCALL getUsedBindingRanges() = 0;

/// Get the list of functions that were exported in the linked IR
SLANG_NO_THROW virtual Slice<String> SLANG_MCALL getExportedFunctionMangledNames() = 0;
};

} // namespace Slang
Expand Down
138 changes: 135 additions & 3 deletions source/slang/slang-compiler-tu.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,91 @@
#include "../core/slang-basic.h"
#include "slang-compiler.h"
#include "slang-ir-insts.h"
#include "slang-ir-util.h"
#include "slang-capability.h"

namespace Slang
{
// Only attempt to precompile functions:
// 1) With function bodies (not just empty decls)
// 2) Not marked with unsafeForceInlineDecoration
// 3) Have a simple HLSL data type as the return or parameter type
static bool attemptPrecompiledExport(IRInst* inst)
{
if (inst->getOp() != kIROp_Func)
{
return false;
}

// Skip functions with no body
bool hasBody = false;
for (auto child : inst->getChildren())
{
if (child->getOp() == kIROp_Block)
{
hasBody = true;
break;
}
}
if (!hasBody)
{
return false;
}

// Skip functions marked with unsafeForceInlineDecoration
if (inst->findDecoration<IRUnsafeForceInlineEarlyDecoration>())
{
return false;
}

// Skip non-simple HLSL data types, filters out generics
if (!isSimpleHLSLDataType(inst))
{
return false;
}

return true;
}

/*
* 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.
* 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
* example, DXIL libraries do not allow resources like structured buffers
* to appear in the library interface. Also, no target languages allow
* generics to be precompiled.
*
* Some restrictions can be enforced up front before linking, but some are
* done during target generation in between IR linking+legalization and
* target source emission.
*
* Functions which can be rejected up front:
* - Functions with no body
* - Functions marked with unsafeForceInlineDecoration
* - Functions that define or use generics
*
* The functions not rejected up front are marked with
* DownstreamModuleExportDecoration which indicates functions we're trying to
* export for precompilation, and this also helps to identify the functions
* in the linked IR which survived the additional pruning.
*
* Functions that are rejected after linking+legalization (inside
* emitPrecompiled*):
* - (DXIL) Functions that return or take a HLSLStructuredBufferType
* - (DXIL) Functions that return or take a Matrix type
*
* emitPrecompiled* produces the output artifact containing target language
* blob, and as metadata, the list of functions which survived the second
* phase of filtering.
*
* The original module IR functions matching those are then marked with
* "AvailableIn*" (e.g. AvailableInDXILDecoration) 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)
Expand All @@ -20,6 +101,7 @@ namespace Slang

auto module = getIRModule();
auto linkage = getLinkage();
auto builder = IRBuilder(module);

DiagnosticSink sink(linkage->getSourceManager(), Lexer::sourceLocationLexer);
applySettingsToDiagnosticSink(&sink, &sink, linkage->m_optionSet);
Expand Down Expand Up @@ -48,6 +130,7 @@ namespace Slang
{
case CodeGenTarget::DXIL:
tp.getOptionSet().add(CompilerOptionName::Profile, Profile::RawEnum::DX_Lib_6_6);
tp.getOptionSet().add(CompilerOptionName::EmbedDXIL, true);
break;
}

Expand All @@ -59,20 +142,69 @@ namespace Slang
CodeGenContext::Shared sharedCodeGenContext(&tp, entryPointIndices, &sink, nullptr);
CodeGenContext codeGenContext(&sharedCodeGenContext);

// Mark all public functions as exported, ensure there's at least one. Store a mapping
// of function name to IRInst* for later reference. After linking is done, we'll scan
// the linked result to see which functions survived the pruning and are included in the
// precompiled blob.
Dictionary<String, IRInst*> nameToFunction;
bool hasAtLeastOneFunction = false;
for (auto inst : module->getGlobalInsts())
{
if (attemptPrecompiledExport(inst))
{
hasAtLeastOneFunction = true;
builder.addDecoration(inst, kIROp_DownstreamModuleExportDecoration);
nameToFunction[inst->findDecoration<IRExportDecoration>()->getMangledName()] = inst;
}
}

// Bail if there are no functions to export. That's not treated as an error
// because it's possible that the module just doesn't have any simple HLSL.
if (!hasAtLeastOneFunction)
{
return SLANG_OK;
}

ComPtr<IArtifact> outArtifact;
SlangResult res = codeGenContext.emitTranslationUnit(outArtifact);
SlangResult res = codeGenContext.emitPrecompiledDXIL(outArtifact);

sink.getBlobIfNeeded(outDiagnostics);

if (res != SLANG_OK)
{
return res;
}

auto metadata = findAssociatedRepresentation<IArtifactPostEmitMetadata>(outArtifact);
if (!metadata)
{
return SLANG_E_NOT_AVAILABLE;
}

for (const auto& mangledName : metadata->getExportedFunctionMangledNames())
{
auto moduleInst = nameToFunction[mangledName];
builder.addDecoration(moduleInst, kIROp_AvailableInDXILDecoration);
auto moduleDec = moduleInst->findDecoration<IRDownstreamModuleExportDecoration>();
moduleDec->removeAndDeallocate();
}

// Finally, clean up the transient export decorations left over in the module. These are
// represent functions that were pruned from the IR after linking, before target generation.
for (auto moduleInst : module->getGlobalInsts())
{
if (moduleInst->getOp() == kIROp_Func)
{
if (auto dec = moduleInst->findDecoration<IRDownstreamModuleExportDecoration>())
{
dec->removeAndDeallocate();
}
}
}

ISlangBlob* blob;
outArtifact->loadBlob(ArtifactKeep::Yes, &blob);

auto builder = IRBuilder(module);
// Add the precompiled blob to the module
builder.setInsertInto(module);

switch (targetReq->getTarget())
Expand Down
56 changes: 23 additions & 33 deletions source/slang/slang-compiler.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -350,7 +350,7 @@ namespace Slang
Profile Profile::lookUp(UnownedStringSlice const& name)
{
#define PROFILE(TAG, NAME, STAGE, VERSION) if(name == UnownedTerminatedStringSlice(#NAME)) return Profile::TAG;
#define PROFILE_ALIAS(TAG, DEF, NAME) if(name == UnownedTerminatedStringSlice(#NAME)) return Profile::TAG;
#define PROFILE_ALIAS(TAG, DEF, NAME) if(name == UnownedTerminatedStringSlice(#NAME)) return Profile::TAG;
#include "slang-profile-defs.h"

return Profile::Unknown;
Expand Down Expand Up @@ -767,7 +767,7 @@ namespace Slang
# pragma warning(pop)
#endif

SlangResult CodeGenContext::emitTranslationUnit(ComPtr<IArtifact>& outArtifact)
SlangResult CodeGenContext::emitPrecompiledDXIL(ComPtr<IArtifact>& outArtifact)
{
return emitWithDownstreamForEntryPoints(outArtifact);
}
Expand Down Expand Up @@ -1094,23 +1094,6 @@ namespace Slang
return SLANG_OK;
}

bool CodeGenContext::isPrecompiled()
{
auto program = getProgram();

bool allPrecompiled = true;
program->enumerateIRModules([&](IRModule* irModule)
{
// TODO: Conditionalize this on target
if (!irModule->precompiledDXIL)
{
allPrecompiled = false;
}
});

return allPrecompiled;
}

SlangResult CodeGenContext::emitWithDownstreamForEntryPoints(ComPtr<IArtifact>& outArtifact)
{
outArtifact.setNull();
Expand Down Expand Up @@ -1272,15 +1255,17 @@ namespace Slang
}
else
{
if (!isPrecompiled())
CodeGenContext sourceCodeGenContext(this, sourceTarget, extensionTracker);

if (target == CodeGenTarget::DXILAssembly || target == CodeGenTarget::DXIL)
{
CodeGenContext sourceCodeGenContext(this, sourceTarget, extensionTracker);
sourceCodeGenContext.removeAvailableInDXIL = true;
}

SLANG_RETURN_ON_FAIL(sourceCodeGenContext.emitEntryPointsSource(sourceArtifact));
sourceCodeGenContext.maybeDumpIntermediate(sourceArtifact);
SLANG_RETURN_ON_FAIL(sourceCodeGenContext.emitEntryPointsSource(sourceArtifact));
sourceCodeGenContext.maybeDumpIntermediate(sourceArtifact);

sourceLanguage = (SourceLanguage)TypeConvertUtil::getSourceLanguageFromTarget((SlangCompileTarget)sourceTarget);
}
sourceLanguage = (SourceLanguage)TypeConvertUtil::getSourceLanguageFromTarget((SlangCompileTarget)sourceTarget);
}

if (sourceArtifact)
Expand Down Expand Up @@ -1571,24 +1556,29 @@ namespace Slang
libraries.addRange(linkage->m_libModules.getBuffer(), linkage->m_libModules.getCount());
}

if (isPrecompiled())
auto program = getProgram();

// Load embedded precompiled libraries from IR into library artifacts
program->enumerateIRModules([&](IRModule* irModule)
{
auto program = getProgram();
program->enumerateIRModules([&](IRModule* irModule)
for (auto inst : irModule->getModuleInst()->getChildren())
{
if (target == CodeGenTarget::DXILAssembly || target == CodeGenTarget::DXIL)
{
// TODO: conditionalize on target
if (irModule->precompiledDXIL)
if (inst->getOp() == kIROp_EmbeddedDXIL)
{
auto slice = static_cast<IRBlobLit*>(inst->getOperand(0))->getStringSlice();
ArtifactDesc desc = ArtifactDescUtil::makeDescForCompileTarget(SLANG_DXIL);
desc.kind = ArtifactKind::Library;

auto library = ArtifactUtil::createArtifact(desc);

library->addRepresentationUnknown(irModule->precompiledDXIL);
library->addRepresentationUnknown(StringBlob::create(slice));
libraries.add(library);
}
});
}
}
}
});

options.compilerSpecificArguments = allocator.allocate(compilerSpecificArguments);
options.requiredCapabilityVersions = SliceUtil::asSlice(requiredCapabilityVersions);
Expand Down
14 changes: 6 additions & 8 deletions source/slang/slang-compiler.h
Original file line number Diff line number Diff line change
Expand Up @@ -555,7 +555,6 @@ namespace Slang
/// and parsing via Slang reflection, and is not recommended for future APIs to use.
///
Scope* _getOrCreateScopeForLegacyLookup(ASTBuilder* astBuilder);

protected:
ComponentType(Linkage* linkage);

Expand Down Expand Up @@ -2732,10 +2731,14 @@ namespace Slang

SlangResult emitEntryPoints(ComPtr<IArtifact>& outArtifact);

SlangResult emitTranslationUnit(ComPtr<IArtifact>& outArtifact);
SlangResult emitPrecompiledDXIL(ComPtr<IArtifact>& outArtifact);

void maybeDumpIntermediate(IArtifact* artifact);

// Used to cause instructions available in precompiled DXIL to be
// removed between IR linking and target source generation.
bool removeAvailableInDXIL = false;

protected:
CodeGenTarget m_targetFormat = CodeGenTarget::Unknown;
ExtensionTracker* m_extensionTracker = nullptr;
Expand Down Expand Up @@ -2772,11 +2775,6 @@ namespace Slang


SlangResult _emitEntryPoints(ComPtr<IArtifact>& outArtifact);

/* Checks if all modules in the target program are already compiled to the
target language, indicating that a pass-through linking using the
downstream compiler is viable.*/
bool isPrecompiled();
private:
Shared* m_shared = nullptr;
};
Expand Down Expand Up @@ -2816,7 +2814,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 setTargetEmbedDXIL(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 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;
Expand Down
Loading

0 comments on commit ddf4a32

Please sign in to comment.