[llvm] [SandboxIR][Bench] Add tests with tracking enabled (PR #108273)

[via llvm-commits]

https://github.com/vporpo created https://github.com/llvm/llvm-project/pull/108273

Benchmarks RAUW and RUOW when tracking is enabled.

>From 0c36bc55f6d306492753fe2f80f0ff477879b2c8 Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Wed, 4 Sep 2024 14:56:07 -0700
Subject: [PATCH] [SandboxIR][Bench] Add tests with tracking enabled

---
 llvm/benchmarks/SandboxIRBench.cpp | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/llvm/benchmarks/SandboxIRBench.cpp b/llvm/benchmarks/SandboxIRBench.cpp
index 6cff2eefdb19f2..d2740b6ba1cfd6 100644
--- a/llvm/benchmarks/SandboxIRBench.cpp
+++ b/llvm/benchmarks/SandboxIRBench.cpp
@@ -34,8 +34,9 @@ static std::unique_ptr<Module> parseIR(LLVMContext &C, const char *IR) {
 }
 
 enum class IR {
-  LLVM,
-  SBox,
+  LLVM,   ///> LLVM IR
+  SBox,   ///> Sandbox IR with tracking disabled
+  SBoxTr, ///> Sandbox IR with tracking enabled
 };
 // Traits to get llvm::BasicBlock/sandboxir::BasicBlock from IR::LLVM/IR::SBox.
 template <IR IRTy> struct TypeSelect {};
@@ -45,6 +46,9 @@ template <> struct TypeSelect<IR::LLVM> {
 template <> struct TypeSelect<IR::SBox> {
   using BasicBlock = sandboxir::BasicBlock;
 };
+template <> struct TypeSelect<IR::SBoxTr> {
+  using BasicBlock = sandboxir::BasicBlock;
+};
 
 template <IR IRTy>
 static typename TypeSelect<IRTy>::BasicBlock *
@@ -59,12 +63,22 @@ genIR(std::unique_ptr<llvm::Module> &LLVMM, LLVMContext &LLVMCtx,
 
   sandboxir::Function *F = Ctx.createFunction(LLVMF);
   sandboxir::BasicBlock *BB = &*F->begin();
+  // Start tracking if we are testing with tracking enabled.
+  if constexpr (IRTy == IR::SBoxTr)
+    Ctx.save();
+
   if constexpr (IRTy == IR::LLVM)
     return LLVMBB;
   else
     return BB;
 }
 
+template <IR IRTy> static void finalize(sandboxir::Context &Ctx) {
+  // Accept changes if we are tracking.
+  if constexpr (IRTy == IR::SBoxTr)
+    Ctx.accept();
+}
+
 static std::string generateBBWalkIR(unsigned Size) {
   std::stringstream SS;
   SS << "define void @foo(i32 %v1, i32 %v2) {\n";
@@ -132,6 +146,7 @@ template <IR IRTy> static void RAUW(benchmark::State &State) {
     Def1->replaceAllUsesWith(Def2);
     Def2->replaceAllUsesWith(Def1);
   }
+  finalize<IRTy>(Ctx);
 }
 
 static std::string generateRUOWIR(unsigned NumOperands) {
@@ -171,6 +186,7 @@ template <IR IRTy> static void RUOW(benchmark::State &State) {
   auto *Call = &*It++;
   for (auto _ : State)
     Call->replaceUsesOfWith(Arg0, Arg1);
+  finalize<IRTy>(Ctx);
 }
 
 BENCHMARK(GetType<IR::LLVM>);
@@ -181,8 +197,10 @@ BENCHMARK(BBWalk<IR::SBox>)->Args({1024});
 
 BENCHMARK(RAUW<IR::LLVM>)->Args({512});
 BENCHMARK(RAUW<IR::SBox>)->Args({512});
+BENCHMARK(RAUW<IR::SBoxTr>)->Args({512});
 
 BENCHMARK(RUOW<IR::LLVM>)->Args({4096});
 BENCHMARK(RUOW<IR::SBox>)->Args({4096});
+BENCHMARK(RUOW<IR::SBoxTr>)->Args({4096});
 
 BENCHMARK_MAIN();