diff --git a/llvm/include/llvm/ProfileData/PGOCtxProfReader.h b/llvm/include/llvm/ProfileData/PGOCtxProfReader.h index a19b3f51d642..28f05e9073a8 100644 --- a/llvm/include/llvm/ProfileData/PGOCtxProfReader.h +++ b/llvm/include/llvm/ProfileData/PGOCtxProfReader.h @@ -73,7 +73,8 @@ public: }; class PGOCtxProfileReader final { - BitstreamCursor &Cursor; + StringRef Magic; + BitstreamCursor Cursor; Expected advance(); Error readMetadata(); Error wrongValue(const Twine &); @@ -84,7 +85,9 @@ class PGOCtxProfileReader final { bool canReadContext(); public: - PGOCtxProfileReader(BitstreamCursor &Cursor) : Cursor(Cursor) {} + PGOCtxProfileReader(StringRef Buffer) + : Magic(Buffer.substr(0, PGOCtxProfileWriter::ContainerMagic.size())), + Cursor(Buffer.substr(PGOCtxProfileWriter::ContainerMagic.size())) {} Expected> loadContexts(); }; diff --git a/llvm/include/llvm/ProfileData/PGOCtxProfWriter.h b/llvm/include/llvm/ProfileData/PGOCtxProfWriter.h index ecee7a2cfb53..db9a0fd77f83 100644 --- a/llvm/include/llvm/ProfileData/PGOCtxProfWriter.h +++ b/llvm/include/llvm/ProfileData/PGOCtxProfWriter.h @@ -13,6 +13,7 @@ #ifndef LLVM_PROFILEDATA_PGOCTXPROFWRITER_H_ #define LLVM_PROFILEDATA_PGOCTXPROFWRITER_H_ +#include "llvm/ADT/StringExtras.h" #include "llvm/Bitstream/BitCodeEnums.h" #include "llvm/Bitstream/BitstreamWriter.h" #include "llvm/ProfileData/CtxInstrContextNode.h" @@ -68,15 +69,7 @@ class PGOCtxProfileWriter final { public: PGOCtxProfileWriter(raw_ostream &Out, - std::optional VersionOverride = std::nullopt) - : Writer(Out, 0) { - Writer.EnterSubblock(PGOCtxProfileBlockIDs::ProfileMetadataBlockID, - CodeLen); - const auto Version = VersionOverride ? *VersionOverride : CurrentVersion; - Writer.EmitRecord(PGOCtxProfileRecords::Version, - SmallVector({Version})); - } - + std::optional VersionOverride = std::nullopt); ~PGOCtxProfileWriter() { Writer.ExitBlock(); } void write(const ctx_profile::ContextNode &); @@ -85,6 +78,7 @@ public: static constexpr unsigned CodeLen = 2; static constexpr uint32_t CurrentVersion = 1; static constexpr unsigned VBREncodingBits = 6; + static constexpr StringRef ContainerMagic = "CTXP"; }; } // namespace llvm diff --git a/llvm/lib/ProfileData/PGOCtxProfReader.cpp b/llvm/lib/ProfileData/PGOCtxProfReader.cpp index 1b42d8c765f2..0a0e7db457fa 100644 --- a/llvm/lib/ProfileData/PGOCtxProfReader.cpp +++ b/llvm/lib/ProfileData/PGOCtxProfReader.cpp @@ -12,6 +12,7 @@ //===----------------------------------------------------------------------===// #include "llvm/ProfileData/PGOCtxProfReader.h" +#include "llvm/ADT/StringExtras.h" #include "llvm/Bitstream/BitCodeEnums.h" #include "llvm/Bitstream/BitstreamReader.h" #include "llvm/ProfileData/InstrProf.h" @@ -139,6 +140,20 @@ PGOCtxProfileReader::readContext(bool ExpectIndex) { } Error PGOCtxProfileReader::readMetadata() { + if (Magic.size() < PGOCtxProfileWriter::ContainerMagic.size() || + Magic != PGOCtxProfileWriter::ContainerMagic) + return make_error(instrprof_error::invalid_prof, + "Invalid magic"); + + BitstreamEntry Entry; + RET_ON_ERR(Cursor.advance().moveInto(Entry)); + if (Entry.Kind != BitstreamEntry::SubBlock || + Entry.ID != bitc::BLOCKINFO_BLOCK_ID) + return unsupported("Expected Block ID"); + // We don't need the blockinfo to read the rest, it's metadata usable for e.g. + // llvm-bcanalyzer. + RET_ON_ERR(Cursor.SkipBlock()); + EXPECT_OR_RET(Blk, advance()); if (Blk->Kind != BitstreamEntry::SubBlock) return unsupported("Expected Version record"); diff --git a/llvm/lib/ProfileData/PGOCtxProfWriter.cpp b/llvm/lib/ProfileData/PGOCtxProfWriter.cpp index 508179756446..74cd8763cc76 100644 --- a/llvm/lib/ProfileData/PGOCtxProfWriter.cpp +++ b/llvm/lib/ProfileData/PGOCtxProfWriter.cpp @@ -16,6 +16,40 @@ using namespace llvm; using namespace llvm::ctx_profile; +PGOCtxProfileWriter::PGOCtxProfileWriter( + raw_ostream &Out, std::optional VersionOverride) + : Writer(Out, 0) { + static_assert(ContainerMagic.size() == 4); + Out.write(ContainerMagic.data(), ContainerMagic.size()); + Writer.EnterBlockInfoBlock(); + { + auto DescribeBlock = [&](unsigned ID, StringRef Name) { + Writer.EmitRecord(bitc::BLOCKINFO_CODE_SETBID, + SmallVector{ID}); + Writer.EmitRecord(bitc::BLOCKINFO_CODE_BLOCKNAME, + llvm::arrayRefFromStringRef(Name)); + }; + SmallVector Data; + auto DescribeRecord = [&](unsigned RecordID, StringRef Name) { + Data.clear(); + Data.push_back(RecordID); + llvm::append_range(Data, Name); + Writer.EmitRecord(bitc::BLOCKINFO_CODE_SETRECORDNAME, Data); + }; + DescribeBlock(PGOCtxProfileBlockIDs::ProfileMetadataBlockID, "Metadata"); + DescribeRecord(PGOCtxProfileRecords::Version, "Version"); + DescribeBlock(PGOCtxProfileBlockIDs::ContextNodeBlockID, "Context"); + DescribeRecord(PGOCtxProfileRecords::Guid, "GUID"); + DescribeRecord(PGOCtxProfileRecords::CalleeIndex, "CalleeIndex"); + DescribeRecord(PGOCtxProfileRecords::Counters, "Counters"); + } + Writer.ExitBlock(); + Writer.EnterSubblock(PGOCtxProfileBlockIDs::ProfileMetadataBlockID, CodeLen); + const auto Version = VersionOverride ? *VersionOverride : CurrentVersion; + Writer.EmitRecord(PGOCtxProfileRecords::Version, + SmallVector({Version})); +} + void PGOCtxProfileWriter::writeCounters(const ContextNode &Node) { Writer.EmitCode(bitc::UNABBREV_RECORD); Writer.EmitVBR(PGOCtxProfileRecords::Counters, VBREncodingBits); diff --git a/llvm/unittests/ProfileData/PGOCtxProfReaderWriterTest.cpp b/llvm/unittests/ProfileData/PGOCtxProfReaderWriterTest.cpp index 6c6798ded00b..476f293780d8 100644 --- a/llvm/unittests/ProfileData/PGOCtxProfReaderWriterTest.cpp +++ b/llvm/unittests/ProfileData/PGOCtxProfReaderWriterTest.cpp @@ -6,7 +6,7 @@ // //===----------------------------------------------------------------------===// -#include "llvm/Bitstream/BitstreamReader.h" +#include "llvm/Bitcode/BitcodeAnalyzer.h" #include "llvm/ProfileData/CtxInstrContextNode.h" #include "llvm/ProfileData/PGOCtxProfReader.h" #include "llvm/ProfileData/PGOCtxProfWriter.h" @@ -106,8 +106,20 @@ TEST_F(PGOCtxProfRWTest, RoundTrip) { MemoryBuffer::getFile(ProfileFile.path()); ASSERT_TRUE(!!MB); ASSERT_NE(*MB, nullptr); - BitstreamCursor Cursor((*MB)->getBuffer()); - PGOCtxProfileReader Reader(Cursor); + + // Check it's analyzable by the BCAnalyzer + BitcodeAnalyzer BA((*MB)->getBuffer()); + std::string AnalyzerDump; + raw_string_ostream OS(AnalyzerDump); + BCDumpOptions Opts(OS); + + // As in, expect no error. + EXPECT_FALSE(BA.analyze(Opts)); + EXPECT_TRUE(AnalyzerDump.find("getBuffer()); auto Expected = Reader.loadContexts(); ASSERT_TRUE(!!Expected); auto &Ctxes = *Expected; @@ -143,8 +155,7 @@ TEST_F(PGOCtxProfRWTest, InvalidCounters) { auto MB = MemoryBuffer::getFile(ProfileFile.path()); ASSERT_TRUE(!!MB); ASSERT_NE(*MB, nullptr); - BitstreamCursor Cursor((*MB)->getBuffer()); - PGOCtxProfileReader Reader(Cursor); + PGOCtxProfileReader Reader((*MB)->getBuffer()); auto Expected = Reader.loadContexts(); EXPECT_FALSE(Expected); consumeError(Expected.takeError()); @@ -152,16 +163,14 @@ TEST_F(PGOCtxProfRWTest, InvalidCounters) { } TEST_F(PGOCtxProfRWTest, Empty) { - BitstreamCursor Cursor(""); - PGOCtxProfileReader Reader(Cursor); + PGOCtxProfileReader Reader(""); auto Expected = Reader.loadContexts(); EXPECT_FALSE(Expected); consumeError(Expected.takeError()); } TEST_F(PGOCtxProfRWTest, Invalid) { - BitstreamCursor Cursor("Surely this is not valid"); - PGOCtxProfileReader Reader(Cursor); + PGOCtxProfileReader Reader("Surely this is not valid"); auto Expected = Reader.loadContexts(); EXPECT_FALSE(Expected); consumeError(Expected.takeError()); @@ -182,8 +191,8 @@ TEST_F(PGOCtxProfRWTest, ValidButEmpty) { auto MB = MemoryBuffer::getFile(ProfileFile.path()); ASSERT_TRUE(!!MB); ASSERT_NE(*MB, nullptr); - BitstreamCursor Cursor((*MB)->getBuffer()); - PGOCtxProfileReader Reader(Cursor); + + PGOCtxProfileReader Reader((*MB)->getBuffer()); auto Expected = Reader.loadContexts(); EXPECT_TRUE(!!Expected); EXPECT_TRUE(Expected->empty()); @@ -204,8 +213,8 @@ TEST_F(PGOCtxProfRWTest, WrongVersion) { auto MB = MemoryBuffer::getFile(ProfileFile.path()); ASSERT_TRUE(!!MB); ASSERT_NE(*MB, nullptr); - BitstreamCursor Cursor((*MB)->getBuffer()); - PGOCtxProfileReader Reader(Cursor); + + PGOCtxProfileReader Reader((*MB)->getBuffer()); auto Expected = Reader.loadContexts(); EXPECT_FALSE(Expected); consumeError(Expected.takeError()); @@ -228,8 +237,7 @@ TEST_F(PGOCtxProfRWTest, DuplicateRoots) { auto MB = MemoryBuffer::getFile(ProfileFile.path()); ASSERT_TRUE(!!MB); ASSERT_NE(*MB, nullptr); - BitstreamCursor Cursor((*MB)->getBuffer()); - PGOCtxProfileReader Reader(Cursor); + PGOCtxProfileReader Reader((*MB)->getBuffer()); auto Expected = Reader.loadContexts(); EXPECT_FALSE(Expected); consumeError(Expected.takeError()); @@ -255,8 +263,7 @@ TEST_F(PGOCtxProfRWTest, DuplicateTargets) { auto MB = MemoryBuffer::getFile(ProfileFile.path()); ASSERT_TRUE(!!MB); ASSERT_NE(*MB, nullptr); - BitstreamCursor Cursor((*MB)->getBuffer()); - PGOCtxProfileReader Reader(Cursor); + PGOCtxProfileReader Reader((*MB)->getBuffer()); auto Expected = Reader.loadContexts(); EXPECT_FALSE(Expected); consumeError(Expected.takeError());