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

[via llvm-commits]

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

>From 6b3556bccb40ac76033369fa7a5a6891660b9bea 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

Benchmarks RAUW and RUOW when tracking is enabled.
---
 llvm/benchmarks/SandboxIRBench.cpp | 32 +++++++++++++++++++++++-------
 1 file changed, 25 insertions(+), 7 deletions(-)

diff --git a/llvm/benchmarks/SandboxIRBench.cpp b/llvm/benchmarks/SandboxIRBench.cpp
index 6cff2eefdb19f2..ca2cab664f91e5 100644
--- a/llvm/benchmarks/SandboxIRBench.cpp
+++ b/llvm/benchmarks/SandboxIRBench.cpp
@@ -34,15 +34,19 @@ static std::unique_ptr<Module> parseIR(LLVMContext &C, const char *IR) {
 }
 
 enum class IR {
-  LLVM,
-  SBox,
+  LLVM,           ///> LLVM IR
+  SBoxNoTracking, ///> Sandbox IR with tracking disabled
+  SBoxTracking,   ///> Sandbox IR with tracking enabled
 };
 // Traits to get llvm::BasicBlock/sandboxir::BasicBlock from IR::LLVM/IR::SBox.
 template <IR IRTy> struct TypeSelect {};
 template <> struct TypeSelect<IR::LLVM> {
   using BasicBlock = llvm::BasicBlock;
 };
-template <> struct TypeSelect<IR::SBox> {
+template <> struct TypeSelect<IR::SBoxNoTracking> {
+  using BasicBlock = sandboxir::BasicBlock;
+};
+template <> struct TypeSelect<IR::SBoxTracking> {
   using BasicBlock = sandboxir::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::SBoxTracking)
+    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::SBoxTracking)
+    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,18 +186,21 @@ 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>);
-BENCHMARK(GetType<IR::SBox>);
+BENCHMARK(GetType<IR::SBoxNoTracking>);
 
 BENCHMARK(BBWalk<IR::LLVM>)->Args({1024});
-BENCHMARK(BBWalk<IR::SBox>)->Args({1024});
+BENCHMARK(BBWalk<IR::SBoxTracking>)->Args({1024});
 
 BENCHMARK(RAUW<IR::LLVM>)->Args({512});
-BENCHMARK(RAUW<IR::SBox>)->Args({512});
+BENCHMARK(RAUW<IR::SBoxNoTracking>)->Args({512});
+BENCHMARK(RAUW<IR::SBoxTracking>)->Args({512});
 
 BENCHMARK(RUOW<IR::LLVM>)->Args({4096});
-BENCHMARK(RUOW<IR::SBox>)->Args({4096});
+BENCHMARK(RUOW<IR::SBoxNoTracking>)->Args({4096});
+BENCHMARK(RUOW<IR::SBoxTracking>)->Args({4096});
 
 BENCHMARK_MAIN();