Skip to content

Commit

Permalink
Improve invalid shader handling (#2880)
Browse files Browse the repository at this point in the history 
* Make BuildComputePipeline fail if the shader does not pass validation

Validation might fail e.g. because the shader entry point is missing.
This change allows the error to bubble up through the call stack until vkCreateComputePipelines, instead of killing the app later on with exit code 1.
BuildComputePipeline is converted to early-return style.

* Prevent trimSpirvDebugInfo from looping indefinitely on invalid input

An invalid SPIR-V binary might contain data that does not match any opCode, and yields a wordCount of zero, triggering an infinite loop.
trimSpirvDebugInfo now returns Expected<unsigned> and will give an error in such a situation.
Caller functions getShaderCode and getCodeSize (as well as their callers BuildShaderModule and getModuleData) have been updated to handle the error.
  • Loading branch information
@5p4k
5p4k authored Dec 14, 2023
1 parent c552fb3 commit 2fcce3b
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 37 deletions.
56 changes: 30 additions & 26 deletions llpc/context/llpcCompiler.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -575,7 +575,11 @@ Result Compiler::BuildShaderModule(const ShaderModuleBuildInfo *shaderInfo, Shad
return Result::ErrorInvalidPointer;
}

unsigned codeSize = ShaderModuleHelper::getCodeSize(shaderInfo);
auto codeSizeOrErr = ShaderModuleHelper::getCodeSize(shaderInfo);
if (Error err = codeSizeOrErr.takeError())
return errorToResult(std::move(err));

const unsigned codeSize = *codeSizeOrErr;
size_t allocSize = sizeof(ShaderModuleData) + codeSize;

ShaderModuleData moduleData = {};
Expand Down Expand Up @@ -2289,13 +2293,15 @@ Result Compiler::BuildComputePipeline(const ComputePipelineBuildInfo *pipelineIn
const bool buildUsingRelocatableElf = relocatableElfRequested && canUseRelocatableComputeShaderElf(pipelineInfo);

Result result = validatePipelineShaderInfo(&pipelineInfo->cs);
if (result != Result::Success)
return result;

MetroHash::Hash cacheHash = {};
MetroHash::Hash pipelineHash = {};
cacheHash = PipelineDumper::generateHashForComputePipeline(pipelineInfo, true);
pipelineHash = PipelineDumper::generateHashForComputePipeline(pipelineInfo, false);

if (result == Result::Success && EnableOuts()) {
if (EnableOuts()) {
const ShaderModuleData *moduleData = reinterpret_cast<const ShaderModuleData *>(pipelineInfo->cs.pModuleData);
auto moduleHash = reinterpret_cast<const MetroHash::Hash *>(&moduleData->hash[0]);
LLPC_OUTS("\n===============================================================================\n");
Expand All @@ -2310,8 +2316,7 @@ Result Compiler::BuildComputePipeline(const ComputePipelineBuildInfo *pipelineIn
LLPC_OUTS("\n");
}

if (result == Result::Success)
dumpCompilerOptions(pipelineDumpFile);
dumpCompilerOptions(pipelineDumpFile);

std::optional<CacheAccessor> cacheAccessor;
if (cl::CacheFullPipelines) {
Expand All @@ -2325,40 +2330,39 @@ Result Compiler::BuildComputePipeline(const ComputePipelineBuildInfo *pipelineIn
result = buildComputePipelineInternal(&computeContext, pipelineInfo, buildUsingRelocatableElf, &candidateElf,
&pipelineOut->stageCacheAccess);

if (result == Result::Success) {
elfBin.codeSize = candidateElf.size();
elfBin.pCode = candidateElf.data();
}
if (cacheAccessor && pipelineOut->pipelineCacheAccess == CacheAccessInfo::CacheNotChecked)
pipelineOut->pipelineCacheAccess = CacheAccessInfo::CacheMiss;

if (result != Result::Success) {
return result;
}
elfBin.codeSize = candidateElf.size();
elfBin.pCode = candidateElf.data();
} else {
LLPC_OUTS("Cache hit for compute pipeline.\n");
elfBin = cacheAccessor->getElfFromCache();
pipelineOut->pipelineCacheAccess = CacheAccessInfo::InternalCacheHit;
}

if (result == Result::Success) {
void *allocBuf = nullptr;
if (pipelineInfo->pfnOutputAlloc) {
allocBuf = pipelineInfo->pfnOutputAlloc(pipelineInfo->pInstance, pipelineInfo->pUserData, elfBin.codeSize);
if (allocBuf) {
uint8_t *code = static_cast<uint8_t *>(allocBuf);
memcpy(code, elfBin.pCode, elfBin.codeSize);
if (!pipelineInfo->pfnOutputAlloc) // Allocator is not specified
return Result::ErrorInvalidPointer;

pipelineOut->pipelineBin.codeSize = elfBin.codeSize;
pipelineOut->pipelineBin.pCode = code;
} else
result = Result::ErrorOutOfMemory;
} else {
// Allocator is not specified
result = Result::ErrorInvalidPointer;
}
}
void *const allocBuf =
pipelineInfo->pfnOutputAlloc(pipelineInfo->pInstance, pipelineInfo->pUserData, elfBin.codeSize);
if (!allocBuf)
return Result::ErrorOutOfMemory;

if (cacheAccessor && !cacheAccessor->isInCache() && result == Result::Success) {
uint8_t *code = static_cast<uint8_t *>(allocBuf);
memcpy(code, elfBin.pCode, elfBin.codeSize);

pipelineOut->pipelineBin.codeSize = elfBin.codeSize;
pipelineOut->pipelineBin.pCode = code;

if (cacheAccessor && !cacheAccessor->isInCache()) {
cacheAccessor->setElfInCache(elfBin);
}
return result;

return Result::Success;
}

// =====================================================================================================================
Expand Down
29 changes: 22 additions & 7 deletions llpc/util/llpcShaderModuleHelper.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
*/
#include "llpcShaderModuleHelper.h"
#include "llpcDebug.h"
#include "llpcError.h"
#include "llpcUtil.h"
#include "spirvExt.h"
#include "vkgcUtil.h"
Expand Down Expand Up @@ -233,8 +234,9 @@ ShaderModuleUsage ShaderModuleHelper::getShaderModuleUsageInfo(const BinaryData
//
// @param spvBin : SPIR-V binary code
// @param codeBuffer : The buffer in which to copy the shader code.
// @returns : The number of bytes written to trimSpvBin
unsigned ShaderModuleHelper::trimSpirvDebugInfo(const BinaryData *spvBin, llvm::MutableArrayRef<unsigned> codeBuffer) {
// @returns : The number of bytes written to trimSpvBin or an error if invalid data was encountered
Expected<unsigned> ShaderModuleHelper::trimSpirvDebugInfo(const BinaryData *spvBin,
llvm::MutableArrayRef<unsigned> codeBuffer) {
bool writeCode = !codeBuffer.empty();
assert(codeBuffer.empty() || codeBuffer.size() > sizeof(SpirvHeader));

Expand All @@ -260,6 +262,12 @@ unsigned ShaderModuleHelper::trimSpirvDebugInfo(const BinaryData *spvBin, llvm::
while (codePos < end) {
unsigned opCode = (codePos[0] & OpCodeMask);
unsigned wordCount = (codePos[0] >> WordCountShift);

if (wordCount == 0 || codePos + wordCount > end) {
LLPC_ERRS("Invalid SPIR-V binary\n");
return createResultError(Result::ErrorInvalidShader);
}

bool skip = false;
switch (opCode) {
case OpSource:
Expand Down Expand Up @@ -496,8 +504,12 @@ Result ShaderModuleHelper::getModuleData(const ShaderModuleBuildInfo *shaderInfo

if (moduleData.binType == BinaryType::Spirv) {
moduleData.usage = ShaderModuleHelper::getShaderModuleUsageInfo(&shaderBinary);
moduleData.binCode = getShaderCode(shaderInfo, codeBuffer);
moduleData.usage.isInternalRtShader = shaderInfo->options.pipelineOptions.internalRtShaders;
auto codeOrErr = getShaderCode(shaderInfo, codeBuffer);
if (Error err = codeOrErr.takeError())
return errorToResult(std::move(err));

moduleData.binCode = *codeOrErr;

// Calculate SPIR-V cache hash
Hash cacheHash = {};
Expand All @@ -520,13 +532,16 @@ Result ShaderModuleHelper::getModuleData(const ShaderModuleBuildInfo *shaderInfo
// @param shaderInfo : Shader module build info
// @param codeBuffer [out] : A buffer to hold the shader code.
// @return : The BinaryData for the shaderCode written to codeBuffer.
BinaryData ShaderModuleHelper::getShaderCode(const ShaderModuleBuildInfo *shaderInfo,
MutableArrayRef<unsigned int> &codeBuffer) {
Expected<BinaryData> ShaderModuleHelper::getShaderCode(const ShaderModuleBuildInfo *shaderInfo,
MutableArrayRef<unsigned int> &codeBuffer) {
BinaryData code;
const BinaryData &shaderBinary = shaderInfo->shaderBin;
bool trimDebugInfo = cl::TrimDebugInfo && !(shaderInfo->options.pipelineOptions.internalRtShaders);
if (trimDebugInfo) {
code.codeSize = trimSpirvDebugInfo(&shaderBinary, codeBuffer);
auto sizeOrErr = trimSpirvDebugInfo(&shaderBinary, codeBuffer);
if (Error err = sizeOrErr.takeError())
return err;
code.codeSize = *sizeOrErr;
} else {
assert(shaderBinary.codeSize <= codeBuffer.size() * sizeof(codeBuffer.front()));
memcpy(codeBuffer.data(), shaderBinary.pCode, shaderBinary.codeSize);
Expand All @@ -539,7 +554,7 @@ BinaryData ShaderModuleHelper::getShaderCode(const ShaderModuleBuildInfo *shader
// =====================================================================================================================
// @param shaderInfo : Shader module build info
// @return : The number of bytes need to hold the code for this shader module.
unsigned ShaderModuleHelper::getCodeSize(const ShaderModuleBuildInfo *shaderInfo) {
Expected<unsigned> ShaderModuleHelper::getCodeSize(const ShaderModuleBuildInfo *shaderInfo) {
const BinaryData &shaderBinary = shaderInfo->shaderBin;
bool trimDebugInfo = cl::TrimDebugInfo && !(shaderInfo->options.pipelineOptions.internalRtShaders);
if (!trimDebugInfo)
Expand Down
10 changes: 6 additions & 4 deletions llpc/util/llpcShaderModuleHelper.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
#pragma once
#include "llpc.h"
#include <llvm/ADT/ArrayRef.h>
#include <llvm/Support/Error.h>
#include <vector>

namespace Llpc {
Expand All @@ -56,7 +57,8 @@ class ShaderModuleHelper {
public:
static ShaderModuleUsage getShaderModuleUsageInfo(const BinaryData *spvBinCode);

static unsigned trimSpirvDebugInfo(const BinaryData *spvBin, llvm::MutableArrayRef<unsigned> codeBuffer);
static llvm::Expected<unsigned> trimSpirvDebugInfo(const BinaryData *spvBin,
llvm::MutableArrayRef<unsigned> codeBuffer);

static Result optimizeSpirv(const BinaryData *spirvBinIn, BinaryData *spirvBinOut);

Expand All @@ -70,9 +72,9 @@ class ShaderModuleHelper {
static Result getShaderBinaryType(BinaryData shaderBinary, BinaryType &binaryType);
static Result getModuleData(const ShaderModuleBuildInfo *shaderInfo, llvm::MutableArrayRef<unsigned> codeBuffer,
Vkgc::ShaderModuleData &moduleData);
static unsigned getCodeSize(const ShaderModuleBuildInfo *shaderInfo);
static BinaryData getShaderCode(const ShaderModuleBuildInfo *shaderInfo,
llvm::MutableArrayRef<unsigned int> &codeBuffer);
static llvm::Expected<unsigned> getCodeSize(const ShaderModuleBuildInfo *shaderInfo);
static llvm::Expected<BinaryData> getShaderCode(const ShaderModuleBuildInfo *shaderInfo,
llvm::MutableArrayRef<unsigned int> &codeBuffer);
};

} // namespace Llpc

0 comments on commit 2fcce3b

Please sign in to comment.