diff --git a/lib/DxilDia/DxilDiaDataSource.cpp b/lib/DxilDia/DxilDiaDataSource.cpp index bf422b19b..c7edef518 100644 --- a/lib/DxilDia/DxilDiaDataSource.cpp +++ b/lib/DxilDia/DxilDiaDataSource.cpp @@ -17,6 +17,9 @@ #include "dxc/Support/FileIOHelper.h" #include "dxc/Support/dxcapi.impl.h" +#include "llvm/Support/MSFileSystem.h" +#include "llvm/Support/FileSystem.h" + #include "llvm/Support/MemoryBuffer.h" #include "llvm/IR/DebugInfo.h" #include "llvm/IR/LLVMContext.h" @@ -73,6 +76,14 @@ STDMETHODIMP dxil_dia::DataSource::loadDataFromIStream(_In_ IStream *pInputIStre return E_FAIL; } + // Setup filesystem because bitcode reader might emit warning + ::llvm::sys::fs::MSFileSystem* msfPtr; + IFT(CreateMSFileSystemForDisk(&msfPtr)); + std::unique_ptr<::llvm::sys::fs::MSFileSystem> msf(msfPtr); + + ::llvm::sys::fs::AutoPerThreadSystem pts(msf.get()); + IFTLLVM(pts.error_code()); + CComPtr pIStream = pInputIStream; CComPtr pContainer; if (SUCCEEDED(hlsl::pdb::LoadDataFromStream(m_pMalloc, pInputIStream, &pContainer))) { @@ -117,9 +128,8 @@ STDMETHODIMP dxil_dia::DataSource::loadDataFromIStream(_In_ IStream *pInputIStre UINT32 BlobSize; const char *pBitcode = nullptr; hlsl::GetDxilProgramBitcode(pDxilProgramHeader, &pBitcode, &BlobSize); - UINT32 offset = (UINT32)(pBitcode - (const char *)pDxilProgramHeader); std::unique_ptr p = llvm::MemoryBuffer::getMemBuffer( - llvm::StringRef(pBitcode, bufferSize - offset), "data"); + llvm::StringRef(pBitcode, BlobSize), "data", false /* RequiresNullTerminator */); pEmbeddedBuffer.swap(p); pBitcodeBuffer = pEmbeddedBuffer.get(); } diff --git a/tools/clang/unittests/HLSL/CompilerTest.cpp b/tools/clang/unittests/HLSL/CompilerTest.cpp index 307e6b449..f9f9007f0 100644 --- a/tools/clang/unittests/HLSL/CompilerTest.cpp +++ b/tools/clang/unittests/HLSL/CompilerTest.cpp @@ -260,6 +260,8 @@ public: TEST_METHOD(DiaLoadDebugThenOK) TEST_METHOD(DiaTableIndexThenOK) TEST_METHOD(DiaLoadDebugSubrangeNegativeThenOK) + TEST_METHOD(DiaLoadRelocatedBitcode) + TEST_METHOD(DiaLoadBitcodePlusExtraData) TEST_METHOD(CodeGenFloatingPointEnvironment) TEST_METHOD(CodeGenInclude) @@ -2405,11 +2407,8 @@ TEST_F(CompilerTest, DiaLoadBadBitcodeThenFail) { VERIFY_FAILED(pDiaSource->loadDataFromIStream(pStream)); } -static void CompileTestAndLoadDiaSource(dxc::DxcDllSupport &dllSupport, const char *source, wchar_t *profile, IDiaDataSource **ppDataSource) { +static void CompileAndGetDebugPart(dxc::DxcDllSupport &dllSupport, const char *source, wchar_t *profile, IDxcBlob **ppDebugPart) { CComPtr pContainer; - CComPtr pDebugContent; - CComPtr pDiaSource; - CComPtr pStream; CComPtr pLib; CComPtr pReflection; UINT32 index; @@ -2422,7 +2421,17 @@ static void CompileTestAndLoadDiaSource(dxc::DxcDllSupport &dllSupport, const ch VERIFY_SUCCEEDED(dllSupport.CreateInstance(CLSID_DxcContainerReflection, &pReflection)); VERIFY_SUCCEEDED(pReflection->Load(pContainer)); VERIFY_SUCCEEDED(pReflection->FindFirstPartKind(hlsl::DFCC_ShaderDebugInfoDXIL, &index)); - VERIFY_SUCCEEDED(pReflection->GetPartContent(index, &pDebugContent)); + VERIFY_SUCCEEDED(pReflection->GetPartContent(index, ppDebugPart)); +} + +static void CompileTestAndLoadDiaSource(dxc::DxcDllSupport &dllSupport, const char *source, wchar_t *profile, IDiaDataSource **ppDataSource) { + CComPtr pDebugContent; + CComPtr pStream; + CComPtr pDiaSource; + CComPtr pLib; + VERIFY_SUCCEEDED(dllSupport.CreateInstance(CLSID_DxcLibrary, &pLib)); + + CompileAndGetDebugPart(dllSupport, source, profile, &pDebugContent); VERIFY_SUCCEEDED(pLib->CreateStreamFromBlobReadOnly(pDebugContent, &pStream)); VERIFY_SUCCEEDED(dllSupport.CreateInstance(CLSID_DxcDiaDataSource, &pDiaSource)); VERIFY_SUCCEEDED(pDiaSource->loadDataFromIStream(pStream)); @@ -2460,6 +2469,148 @@ TEST_F(CompilerTest, DiaLoadDebugSubrangeNegativeThenOK) { VERIFY_SUCCEEDED(pDiaDataSource->openSession(&pDiaSession)); } +TEST_F(CompilerTest, DiaLoadRelocatedBitcode) { + + static const char source[] = R"( + SamplerState samp0 : register(s0); + Texture2DArray tex0 : register(t0); + + float4 foo(Texture2DArray textures[], int idx, SamplerState samplerState, float3 uvw) { + return textures[NonUniformResourceIndex(idx)].Sample(samplerState, uvw); + } + + [RootSignature( "DescriptorTable(SRV(t0)), DescriptorTable(Sampler(s0)) " )] + float4 main(int index : INDEX, float3 uvw : TEXCOORD) : SV_Target { + Texture2DArray textures[] = { + tex0, + }; + return foo(textures, index, samp0, uvw); + } + )"; + + CComPtr pPart; + CComPtr pDiaSource; + CComPtr pStream; + + CComPtr pLib; + VERIFY_SUCCEEDED(m_dllSupport.CreateInstance(CLSID_DxcLibrary, &pLib)); + + CompileAndGetDebugPart(m_dllSupport, source, L"ps_6_0", &pPart); + const char *pPartData = (char *)pPart->GetBufferPointer(); + const size_t uPartSize = pPart->GetBufferSize(); + + // Get program header + const hlsl::DxilProgramHeader *programHeader = (hlsl::DxilProgramHeader *)pPartData; + + const char *pBitcode = nullptr; + uint32_t uBitcodeSize = 0; + hlsl::GetDxilProgramBitcode(programHeader, &pBitcode, &uBitcodeSize); + VERIFY_IS_TRUE(uBitcodeSize % sizeof(UINT32) == 0); + + size_t uNewGapSize = 4 * 10; // Size of some bytes between program header and bitcode + size_t uNewSuffixeBytes = 4 * 10; // Size of some random bytes after the program + + hlsl::DxilProgramHeader newProgramHeader = {}; + memcpy(&newProgramHeader, programHeader, sizeof(newProgramHeader)); + newProgramHeader.BitcodeHeader.BitcodeOffset = uNewGapSize + sizeof(newProgramHeader.BitcodeHeader); + + unsigned uNewSizeInBytes = sizeof(newProgramHeader) + uNewGapSize + uBitcodeSize + uNewSuffixeBytes; + VERIFY_IS_TRUE(uNewSizeInBytes % sizeof(UINT32) == 0); + newProgramHeader.SizeInUint32 = uNewSizeInBytes / sizeof(UINT32); + + llvm::SmallVector buffer; + llvm::raw_svector_ostream OS(buffer); + + // Write the header + OS.write((char *)&newProgramHeader, sizeof(newProgramHeader)); + + // Write some garbage between the header and the bitcode + for (unsigned i = 0; i < uNewGapSize; i++) { + OS.write(0xFF); + } + + // Write the actual bitcode + OS.write(pBitcode, uBitcodeSize); + + // Write some garbage after the bitcode + for (unsigned i = 0; i < uNewSuffixeBytes; i++) { + OS.write(0xFF); + } + OS.flush(); + + // Try to load this new program, make sure dia is still okay. + CComPtr pNewProgramBlob; + VERIFY_SUCCEEDED(pLib->CreateBlobWithEncodingFromPinned(buffer.data(), buffer.size(), CP_ACP, &pNewProgramBlob)); + + CComPtr pNewProgramStream; + VERIFY_SUCCEEDED(pLib->CreateStreamFromBlobReadOnly(pNewProgramBlob, &pNewProgramStream)); + + CComPtr pDiaDataSource; + VERIFY_SUCCEEDED(m_dllSupport.CreateInstance(CLSID_DxcDiaDataSource, &pDiaDataSource)); + + VERIFY_SUCCEEDED(pDiaDataSource->loadDataFromIStream(pNewProgramStream)); +} + +TEST_F(CompilerTest, DiaLoadBitcodePlusExtraData) { + // Test that dia doesn't crash when bitcode has unused extra data at the end + + static const char source[] = R"( + SamplerState samp0 : register(s0); + Texture2DArray tex0 : register(t0); + + float4 foo(Texture2DArray textures[], int idx, SamplerState samplerState, float3 uvw) { + return textures[NonUniformResourceIndex(idx)].Sample(samplerState, uvw); + } + + [RootSignature( "DescriptorTable(SRV(t0)), DescriptorTable(Sampler(s0)) " )] + float4 main(int index : INDEX, float3 uvw : TEXCOORD) : SV_Target { + Texture2DArray textures[] = { + tex0, + }; + return foo(textures, index, samp0, uvw); + } + )"; + + CComPtr pPart; + CComPtr pDiaSource; + CComPtr pStream; + + CComPtr pLib; + VERIFY_SUCCEEDED(m_dllSupport.CreateInstance(CLSID_DxcLibrary, &pLib)); + + CompileAndGetDebugPart(m_dllSupport, source, L"ps_6_0", &pPart); + const char *pPartData = (char *)pPart->GetBufferPointer(); + const size_t uPartSize = pPart->GetBufferSize(); + + // Get program header + const hlsl::DxilProgramHeader *programHeader = (hlsl::DxilProgramHeader *)pPartData; + + const char *pBitcode = nullptr; + uint32_t uBitcodeSize = 0; + hlsl::GetDxilProgramBitcode(programHeader, &pBitcode, &uBitcodeSize); + + llvm::SmallVector buffer; + llvm::raw_svector_ostream OS(buffer); + + // Write the bitcode + OS.write(pBitcode, uBitcodeSize); + for (unsigned i = 0; i < 12; i++) { + OS.write(0xFF); + } + OS.flush(); + + // Try to load this new program, make sure dia is still okay. + CComPtr pNewProgramBlob; + VERIFY_SUCCEEDED(pLib->CreateBlobWithEncodingFromPinned(buffer.data(), buffer.size(), CP_ACP, &pNewProgramBlob)); + + CComPtr pNewProgramStream; + VERIFY_SUCCEEDED(pLib->CreateStreamFromBlobReadOnly(pNewProgramBlob, &pNewProgramStream)); + + CComPtr pDiaDataSource; + VERIFY_SUCCEEDED(m_dllSupport.CreateInstance(CLSID_DxcDiaDataSource, &pDiaDataSource)); + + VERIFY_SUCCEEDED(pDiaDataSource->loadDataFromIStream(pNewProgramStream)); +} TEST_F(CompilerTest, DiaLoadDebugThenOK) { CompileTestAndLoadDia(m_dllSupport, nullptr);