diff --git a/llvm/include/llvm/SandboxIR/Context.h b/llvm/include/llvm/SandboxIR/Context.h index f2056de87cb94..b0d6f8335d9e0 100644 --- a/llvm/include/llvm/SandboxIR/Context.h +++ b/llvm/include/llvm/SandboxIR/Context.h @@ -44,11 +44,12 @@ class Context { protected: LLVMContext &LLVMCtx; - friend class Type; // For LLVMCtx. - friend class PointerType; // For LLVMCtx. - friend class IntegerType; // For LLVMCtx. - friend class StructType; // For LLVMCtx. - friend class Region; // For LLVMCtx. + friend class Type; // For LLVMCtx. + friend class PointerType; // For LLVMCtx. + friend class IntegerType; // For LLVMCtx. + friend class StructType; // For LLVMCtx. + friend class Region; // For LLVMCtx. + friend class IRSnapshotChecker; // To snapshot LLVMModuleToModuleMap. Tracker IRTracker; diff --git a/llvm/include/llvm/SandboxIR/Instruction.h b/llvm/include/llvm/SandboxIR/Instruction.h index d9642365908d2..2a59d72e28552 100644 --- a/llvm/include/llvm/SandboxIR/Instruction.h +++ b/llvm/include/llvm/SandboxIR/Instruction.h @@ -11,6 +11,7 @@ #include "llvm/IR/IRBuilder.h" #include "llvm/IR/Instructions.h" +#include "llvm/IR/Module.h" #include "llvm/IR/PatternMatch.h" #include "llvm/SandboxIR/BasicBlock.h" #include "llvm/SandboxIR/Constant.h" diff --git a/llvm/include/llvm/SandboxIR/Tracker.h b/llvm/include/llvm/SandboxIR/Tracker.h index dab20eb809ba0..9a031f3270837 100644 --- a/llvm/include/llvm/SandboxIR/Tracker.h +++ b/llvm/include/llvm/SandboxIR/Tracker.h @@ -42,13 +42,12 @@ #include "llvm/ADT/PointerUnion.h" #include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StableHashing.h" #include "llvm/IR/IRBuilder.h" #include "llvm/IR/Instruction.h" -#include "llvm/IR/Module.h" #include "llvm/SandboxIR/Use.h" #include "llvm/Support/Debug.h" #include -#include namespace llvm::sandboxir { @@ -64,9 +63,56 @@ class SwitchInst; class ConstantInt; class ShuffleVectorInst; class CmpInst; -class Module; class GlobalVariable; +#ifndef NDEBUG + +/// A class that saves hashes and textual IR snapshots of functions in a +/// SandboxIR Context, and does hash comparison when `expectNoDiff` is called. +/// If hashes differ, it prints textual IR for both old and new versions to +/// aid debugging. +/// +/// This is used as an additional debug check when reverting changes to +/// SandboxIR, to verify the reverted state matches the initial state. +class IRSnapshotChecker { + Context &Ctx; + + // A snapshot of textual IR for a function, with a hash for quick comparison. + struct FunctionSnapshot { + llvm::stable_hash Hash; + std::string TextualIR; + }; + + // A snapshot for each llvm::Function found in every module in the SandboxIR + // Context. In practice there will always be one module, but sandbox IR + // save/restore ops work at the Context level, so we must take the full state + // into account. + using ContextSnapshot = DenseMap; + + ContextSnapshot OrigContextSnapshot; + + // Dumps to a string the textual IR for a single Function. + std::string dumpIR(const llvm::Function &F) const; + + // Returns a snapshot of all the modules in the sandbox IR context. + ContextSnapshot takeSnapshot() const; + + // Compares two snapshots and returns true if they differ. + bool diff(const ContextSnapshot &Orig, const ContextSnapshot &Curr) const; + +public: + IRSnapshotChecker(Context &Ctx) : Ctx(Ctx) {} + + /// Saves a snapshot of the current state. If there was any previous snapshot, + /// it will be replaced with the new one. + void save(); + + /// Checks current state against saved state, crashes if different. + void expectNoDiff(); +}; + +#endif // NDEBUG + /// The base class for IR Change classes. class IRChangeBase { protected: @@ -405,6 +451,10 @@ class Tracker { TrackerState State = TrackerState::Disabled; Context &Ctx; +#ifndef NDEBUG + IRSnapshotChecker SnapshotChecker; +#endif + public: #ifndef NDEBUG /// Helps catch bugs where we are creating new change objects while in the @@ -412,7 +462,15 @@ class Tracker { bool InMiddleOfCreatingChange = false; #endif // NDEBUG - explicit Tracker(Context &Ctx) : Ctx(Ctx) {} + explicit Tracker(Context &Ctx) + : Ctx(Ctx) +#ifndef NDEBUG + , + SnapshotChecker(Ctx) +#endif + { + } + ~Tracker(); Context &getContext() const { return Ctx; } /// Record \p Change and take ownership. This is the main function used to diff --git a/llvm/lib/SandboxIR/Tracker.cpp b/llvm/lib/SandboxIR/Tracker.cpp index d35e3ba84990f..27ed37aa9bdd3 100644 --- a/llvm/lib/SandboxIR/Tracker.cpp +++ b/llvm/lib/SandboxIR/Tracker.cpp @@ -10,12 +10,75 @@ #include "llvm/ADT/STLExtras.h" #include "llvm/IR/BasicBlock.h" #include "llvm/IR/Instruction.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/StructuralHash.h" #include "llvm/SandboxIR/Instruction.h" #include using namespace llvm::sandboxir; #ifndef NDEBUG + +std::string IRSnapshotChecker::dumpIR(const llvm::Function &F) const { + std::string Result; + raw_string_ostream SS(Result); + F.print(SS, /*AssemblyAnnotationWriter=*/nullptr); + return Result; +} + +IRSnapshotChecker::ContextSnapshot IRSnapshotChecker::takeSnapshot() const { + ContextSnapshot Result; + for (const auto &Entry : Ctx.LLVMModuleToModuleMap) + for (const auto &F : *Entry.first) { + FunctionSnapshot Snapshot; + Snapshot.Hash = StructuralHash(F, /*DetailedHash=*/true); + Snapshot.TextualIR = dumpIR(F); + Result[&F] = Snapshot; + } + return Result; +} + +bool IRSnapshotChecker::diff(const ContextSnapshot &Orig, + const ContextSnapshot &Curr) const { + bool DifferenceFound = false; + for (const auto &[F, OrigFS] : Orig) { + auto CurrFSIt = Curr.find(F); + if (CurrFSIt == Curr.end()) { + DifferenceFound = true; + dbgs() << "Function " << F->getName() << " not found in current IR.\n"; + dbgs() << OrigFS.TextualIR << "\n"; + continue; + } + const FunctionSnapshot &CurrFS = CurrFSIt->second; + if (OrigFS.Hash != CurrFS.Hash) { + DifferenceFound = true; + dbgs() << "Found IR difference in Function " << F->getName() << "\n"; + dbgs() << "Original:\n" << OrigFS.TextualIR << "\n"; + dbgs() << "Current:\n" << CurrFS.TextualIR << "\n"; + } + } + // Check that Curr doesn't contain any new functions. + for (const auto &[F, CurrFS] : Curr) { + if (!Orig.contains(F)) { + DifferenceFound = true; + dbgs() << "Function " << F->getName() + << " found in current IR but not in original snapshot.\n"; + dbgs() << CurrFS.TextualIR << "\n"; + } + } + return DifferenceFound; +} + +void IRSnapshotChecker::save() { OrigContextSnapshot = takeSnapshot(); } + +void IRSnapshotChecker::expectNoDiff() { + ContextSnapshot CurrContextSnapshot = takeSnapshot(); + if (diff(OrigContextSnapshot, CurrContextSnapshot)) { + llvm_unreachable( + "Original and current IR differ! Probably a checkpointing bug."); + } +} + void UseSet::dump() const { dump(dbgs()); dbgs() << "\n"; @@ -275,7 +338,12 @@ void CmpSwapOperands::dump() const { } #endif -void Tracker::save() { State = TrackerState::Record; } +void Tracker::save() { + State = TrackerState::Record; +#if !defined(NDEBUG) && defined(EXPENSIVE_CHECKS) + SnapshotChecker.save(); +#endif +} void Tracker::revert() { assert(State == TrackerState::Record && "Forgot to save()!"); @@ -283,6 +351,9 @@ void Tracker::revert() { for (auto &Change : reverse(Changes)) Change->revert(*this); Changes.clear(); +#if !defined(NDEBUG) && defined(EXPENSIVE_CHECKS) + SnapshotChecker.expectNoDiff(); +#endif } void Tracker::accept() { diff --git a/llvm/unittests/SandboxIR/TrackerTest.cpp b/llvm/unittests/SandboxIR/TrackerTest.cpp index 4f2cfa6b06ecd..4eedab124bfa0 100644 --- a/llvm/unittests/SandboxIR/TrackerTest.cpp +++ b/llvm/unittests/SandboxIR/TrackerTest.cpp @@ -1844,3 +1844,71 @@ define void @foo(i32 %arg, float %farg) { Ctx.revert(); EXPECT_FALSE(FAdd->getFastMathFlags() != OrigFMF); } + +// IRSnapshotChecker is only defined in debug mode. +#ifndef NDEBUG + +TEST_F(TrackerTest, IRSnapshotCheckerNoChanges) { + parseIR(C, R"IR( +define i32 @foo(i32 %arg) { + %add0 = add i32 %arg, %arg + ret i32 %add0 +} +)IR"); + Function &LLVMF = *M->getFunction("foo"); + sandboxir::Context Ctx(C); + + [[maybe_unused]] auto *F = Ctx.createFunction(&LLVMF); + sandboxir::IRSnapshotChecker Checker(Ctx); + Checker.save(); + Checker.expectNoDiff(); +} + +TEST_F(TrackerTest, IRSnapshotCheckerDiesWithUnexpectedChanges) { + parseIR(C, R"IR( +define i32 @foo(i32 %arg) { + %add0 = add i32 %arg, %arg + %add1 = add i32 %add0, %arg + ret i32 %add1 +} +)IR"); + Function &LLVMF = *M->getFunction("foo"); + sandboxir::Context Ctx(C); + + auto *F = Ctx.createFunction(&LLVMF); + auto *BB = &*F->begin(); + auto It = BB->begin(); + sandboxir::Instruction *Add0 = &*It++; + sandboxir::Instruction *Add1 = &*It++; + sandboxir::IRSnapshotChecker Checker(Ctx); + Checker.save(); + Add1->setOperand(1, Add0); + EXPECT_DEATH(Checker.expectNoDiff(), "Found IR difference"); +} + +TEST_F(TrackerTest, IRSnapshotCheckerSaveMultipleTimes) { + parseIR(C, R"IR( +define i32 @foo(i32 %arg) { + %add0 = add i32 %arg, %arg + %add1 = add i32 %add0, %arg + ret i32 %add1 +} +)IR"); + Function &LLVMF = *M->getFunction("foo"); + sandboxir::Context Ctx(C); + + auto *F = Ctx.createFunction(&LLVMF); + auto *BB = &*F->begin(); + auto It = BB->begin(); + sandboxir::Instruction *Add0 = &*It++; + sandboxir::Instruction *Add1 = &*It++; + sandboxir::IRSnapshotChecker Checker(Ctx); + Checker.save(); + Add1->setOperand(1, Add0); + // Now IR differs from the last snapshot. Let's take a new snapshot. + Checker.save(); + // The new snapshot should have replaced the old one, so this should succeed. + Checker.expectNoDiff(); +} + +#endif // NDEBUG