[llvm] [SandboxVec] Replace hard-coded context save() with transaction-save pass (PR #127690)

Tue Feb 18 11:34:23 PST 2025

llvmbot wrote:




@llvm/pr-subscribers-vectorizers

Author: vporpo (vporpo)

<details>
<summary>Changes</summary>

This patch implements a small region pass that saves the context's state. The patch is now used in the default pipeline to save the context state instead of the hard-coded call to Context::save().

The concept behind this is that the passes themselves should not have to do the actual saving/restoring of the IR state, because that would make it challenging to reorder them in the pipeline. Having separate save/restore passes makes the transformation passes more composable as parts of arbitrary pipelines.

---
Full diff: https://github.com/llvm/llvm-project/pull/127690.diff


15 Files Affected:

- (added) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/TransactionSave.h (+28) 
- (modified) llvm/lib/Transforms/Vectorize/CMakeLists.txt (+1) 
- (modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def (+1) 
- (modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp (-2) 
- (added) llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/TransactionSave.cpp (+20) 
- (modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp (+1-1) 
- (modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerPassBuilder.cpp (+1) 
- (modified) llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll (+1-1) 
- (modified) llvm/test/Transforms/SandboxVectorizer/bottomup_seed_slice.ll (+1-1) 
- (modified) llvm/test/Transforms/SandboxVectorizer/bottomup_seed_slice_pow2.ll (+2-2) 
- (modified) llvm/test/Transforms/SandboxVectorizer/cross_bbs.ll (+1-1) 
- (modified) llvm/test/Transforms/SandboxVectorizer/pack.ll (+1-1) 
- (modified) llvm/test/Transforms/SandboxVectorizer/repeated_instrs.ll (+1-1) 
- (modified) llvm/test/Transforms/SandboxVectorizer/scheduler.ll (+1-1) 
- (modified) llvm/test/Transforms/SandboxVectorizer/special_opcodes.ll (+1-1) 


``````````diff

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/TransactionSave.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/TransactionSave.h
new file mode 100644
index 0000000000000..73aafabfd0b0c
--- /dev/null
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/TransactionSave.h
@@ -0,0 +1,28 @@
+//===- TransactionSave.h ----------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This is a region pass that simply calls Context::save() to save the IR state.
+//
+
+#ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_TRANSACTIONSAVE_H
+#define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_TRANSACTIONSAVE_H
+
+#include "llvm/SandboxIR/Pass.h"
+#include "llvm/SandboxIR/Region.h"
+
+namespace llvm::sandboxir {
+
+class TransactionSave : public RegionPass {
+public:
+  TransactionSave() : RegionPass("tr-save") {}
+  bool runOnRegion(Region &Rgn, const Analyses &A) final;
+};
+
+} // namespace llvm::sandboxir
+
+#endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_TRANSACTIONSAVE_H
diff --git a/llvm/lib/Transforms/Vectorize/CMakeLists.txt b/llvm/lib/Transforms/Vectorize/CMakeLists.txt
index 676d23e1ebdf0..38670ba304e53 100644
--- a/llvm/lib/Transforms/Vectorize/CMakeLists.txt
+++ b/llvm/lib/Transforms/Vectorize/CMakeLists.txt
@@ -11,6 +11,7 @@ add_llvm_component_library(LLVMVectorize
   SandboxVectorizer/Passes/RegionsFromMetadata.cpp
   SandboxVectorizer/Passes/SeedCollection.cpp
   SandboxVectorizer/Passes/TransactionAcceptOrRevert.cpp
+  SandboxVectorizer/Passes/TransactionSave.cpp
   SandboxVectorizer/SandboxVectorizer.cpp
   SandboxVectorizer/SandboxVectorizerPassBuilder.cpp
   SandboxVectorizer/Scheduler.cpp
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def
index 722c6f5db4192..f745073a1eba6 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def
@@ -19,6 +19,7 @@
 
 REGION_PASS("null", ::llvm::sandboxir::NullPass)
 REGION_PASS("print-instruction-count", ::llvm::sandboxir::PrintInstructionCount)
+REGION_PASS("tr-save", ::llvm::sandboxir::TransactionSave)
 REGION_PASS("tr-accept", ::llvm::sandboxir::TransactionAlwaysAccept)
 REGION_PASS("tr-accept-or-revert", ::llvm::sandboxir::TransactionAcceptOrRevert)
 REGION_PASS("bottom-up-vec", ::llvm::sandboxir::BottomUpVec)
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
index 0001c9bb7c7e4..f3b62e36e5115 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
@@ -82,8 +82,6 @@ bool SeedCollection::runOnFunction(Function &F, const Analyses &A) {
           // Create a region containing the seed slice.
           auto &Ctx = F.getContext();
           Region Rgn(Ctx, A.getTTI());
-          // TODO: Replace save() with a save pass in the pass pipeline.
-          Ctx.save();
           Rgn.setAux(SeedSlice);
           // Run the region pass pipeline.
           Change |= RPM.runOnRegion(Rgn, A);
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/TransactionSave.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/TransactionSave.cpp
new file mode 100644
index 0000000000000..8d39d971273b4
--- /dev/null
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/TransactionSave.cpp
@@ -0,0 +1,20 @@
+//===- TransactionSave.cpp - Save the IR state ----------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/TransactionSave.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/InstructionCost.h"
+
+namespace llvm::sandboxir {
+
+bool TransactionSave::runOnRegion(Region &Rgn, const Analyses &A) {
+  Rgn.getContext().save();
+  return false;
+}
+
+} // namespace llvm::sandboxir
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp
index 4f17aa213bbb0..5837cc16fcbac 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp
@@ -36,7 +36,7 @@ SandboxVectorizerPass::SandboxVectorizerPass() : FPM("fpm") {
     //         - Bottom-up Vectorizer pass that starts from a seed
     //         - Accept or revert IR state pass
     FPM.setPassPipeline(
-        "seed-collection<bottom-up-vec,tr-accept-or-revert>",
+        "seed-collection<tr-save,bottom-up-vec,tr-accept-or-revert>",
         sandboxir::SandboxVectorizerPassBuilder::createFunctionPass);
   } else {
     // Create the user-defined pipeline.
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerPassBuilder.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerPassBuilder.cpp
index e552f0570dd9d..389f9cc4cae7c 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerPassBuilder.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerPassBuilder.cpp
@@ -7,6 +7,7 @@
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/TransactionAcceptOrRevert.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/TransactionAlwaysAccept.h"
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/TransactionSave.h"
 
 namespace llvm::sandboxir {
 
diff --git a/llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll b/llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll
index ad23126b681fd..c076c0e849fa9 100644
--- a/llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll
+++ b/llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll
@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<bottom-up-vec,tr-accept>" %s -S | FileCheck %s
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<tr-save,bottom-up-vec,tr-accept>" %s -S | FileCheck %s
 
 define void @store_load(ptr %ptr) {
 ; CHECK-LABEL: define void @store_load(
diff --git a/llvm/test/Transforms/SandboxVectorizer/bottomup_seed_slice.ll b/llvm/test/Transforms/SandboxVectorizer/bottomup_seed_slice.ll
index f4fcc5742f8a7..38c0816504481 100644
--- a/llvm/test/Transforms/SandboxVectorizer/bottomup_seed_slice.ll
+++ b/llvm/test/Transforms/SandboxVectorizer/bottomup_seed_slice.ll
@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<bottom-up-vec,tr-accept>" %s -S | FileCheck %s
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<tr-save,bottom-up-vec,tr-accept>" %s -S | FileCheck %s
 
 
 declare void @foo()
diff --git a/llvm/test/Transforms/SandboxVectorizer/bottomup_seed_slice_pow2.ll b/llvm/test/Transforms/SandboxVectorizer/bottomup_seed_slice_pow2.ll
index 4218ca830dccf..f2eb124494b5e 100644
--- a/llvm/test/Transforms/SandboxVectorizer/bottomup_seed_slice_pow2.ll
+++ b/llvm/test/Transforms/SandboxVectorizer/bottomup_seed_slice_pow2.ll
@@ -1,6 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2=false -sbvec-passes="seed-collection<bottom-up-vec,tr-accept>" %s -S | FileCheck %s --check-prefix=POW2
-; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2=true -sbvec-passes="seed-collection<bottom-up-vec,tr-accept>" %s -S | FileCheck %s --check-prefix=NON-POW2
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2=false -sbvec-passes="seed-collection<tr-save,bottom-up-vec,tr-accept>" %s -S | FileCheck %s --check-prefix=POW2
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2=true -sbvec-passes="seed-collection<tr-save,bottom-up-vec,tr-accept>" %s -S | FileCheck %s --check-prefix=NON-POW2
 
 define void @pow2(ptr %ptr, float %val) {
 ; POW2-LABEL: define void @pow2(
diff --git a/llvm/test/Transforms/SandboxVectorizer/cross_bbs.ll b/llvm/test/Transforms/SandboxVectorizer/cross_bbs.ll
index 75d7cc8d9cefe..c559f294f9695 100644
--- a/llvm/test/Transforms/SandboxVectorizer/cross_bbs.ll
+++ b/llvm/test/Transforms/SandboxVectorizer/cross_bbs.ll
@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<bottom-up-vec,tr-accept>" %s -S | FileCheck %s
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<tr-save,bottom-up-vec,tr-accept>" %s -S | FileCheck %s
 
 define void @cross_bbs(ptr %ptr) {
 ; CHECK-LABEL: define void @cross_bbs(
diff --git a/llvm/test/Transforms/SandboxVectorizer/pack.ll b/llvm/test/Transforms/SandboxVectorizer/pack.ll
index c1e22d3dc2e73..c5abddb7fba01 100644
--- a/llvm/test/Transforms/SandboxVectorizer/pack.ll
+++ b/llvm/test/Transforms/SandboxVectorizer/pack.ll
@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<bottom-up-vec,tr-accept>" %s -S | FileCheck %s
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<tr-save,bottom-up-vec,tr-accept>" %s -S | FileCheck %s
 
 define void @pack_constants(ptr %ptr) {
 ; CHECK-LABEL: define void @pack_constants(
diff --git a/llvm/test/Transforms/SandboxVectorizer/repeated_instrs.ll b/llvm/test/Transforms/SandboxVectorizer/repeated_instrs.ll
index 3e4ef6787a563..081267da77e5f 100644
--- a/llvm/test/Transforms/SandboxVectorizer/repeated_instrs.ll
+++ b/llvm/test/Transforms/SandboxVectorizer/repeated_instrs.ll
@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<bottom-up-vec,tr-accept>" %s -S | FileCheck %s
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<tr-save,bottom-up-vec,tr-accept>" %s -S | FileCheck %s
 
 define i32 @repeated_splat(ptr %ptr, i32 %v) #0 {
 ; CHECK-LABEL: define i32 @repeated_splat(
diff --git a/llvm/test/Transforms/SandboxVectorizer/scheduler.ll b/llvm/test/Transforms/SandboxVectorizer/scheduler.ll
index fd4847da920d1..7741d8c64c8fc 100644
--- a/llvm/test/Transforms/SandboxVectorizer/scheduler.ll
+++ b/llvm/test/Transforms/SandboxVectorizer/scheduler.ll
@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<bottom-up-vec,tr-accept>" %s -S | FileCheck %s
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<tr-save,bottom-up-vec,tr-accept>" %s -S | FileCheck %s
 
 ; This used to crash because the newly added pack   instructions would not update
 ; the DAG and scheduler, leading to def-after-use.
diff --git a/llvm/test/Transforms/SandboxVectorizer/special_opcodes.ll b/llvm/test/Transforms/SandboxVectorizer/special_opcodes.ll
index 1aca7cf2a8bd8..edb8d615e0055 100644
--- a/llvm/test/Transforms/SandboxVectorizer/special_opcodes.ll
+++ b/llvm/test/Transforms/SandboxVectorizer/special_opcodes.ll
@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<bottom-up-vec,tr-accept>" %s -S | FileCheck %s
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -sbvec-allow-non-pow2 -sbvec-passes="seed-collection<tr-save,bottom-up-vec,tr-accept>" %s -S | FileCheck %s
 
 ; This file includes tests for opcodes that need special checks.
 

``````````

</details>


https://github.com/llvm/llvm-project/pull/127690
