[SeparateConstOffsetFromGEPPass] Added optional modification strategy

[eklepilkina]

Summary

This modification strategy tries to understand which GEP instrucions is profitable to modify for register pressure decreasing.

[anton-afanasyev]

Please rebase against precommited tests.

[eklepilkina]

Rebased against precommited tests

[reames]

At least for me, I need some context to be able to review this. What is the case which this improves in terms of codegen? And how common are such patterns? Keep in mind, I'm not terribly familiar with the pass here, so this may be pretty basic explanation. Do you have a bug with examples or analyze that lead to this change?

[eklepilkina]

Sorry, I had to provide the context at the beginning.

Now clang for RISC-V doesn't use offset addressing in generated assembly. Example from Dhrystone

 addiw   a0, s1, 5
 slli    a1, a0, 0x2
 add     a2, s4, a1
 sw      s2, 0(a2)
 addiw   a3, s1, 6
 slli    a3, a3, 0x2
 add     a3, a3, s4
 sw      s2, 0(a3)
 addiw   a3, s1, 35
 slli    a3, a3, 0x2
add     a3, a3, s4
sw      a0, 0(a3)

It's inefficient because we can use offsets.
Adding this pass allows to generate the next code

    addiw   a4, a2, 5
    slli    a5, a2, 2
    add a0, a0, a5
    sw  a3, 20(a0)
    sw  a3, 24(a0)
    sw  a4, 140(a0)
...

SeparateConstOffsetFromGEPPass is used to solve this problem in targets with limited addressing modes.
The changes inside pass was made because modification of all GEPs isn't profitable, seems that we need at least 2 GEPs and one value that was used for index that can be removed after modification. Otherwise we don't decrease register pressure.

[craig.topper]

This should be two patches, one changing the pass and one enabling for RISC-V.

llvm/lib/Target/RISCV/RISCVTargetMachine.cpp

171 ↗

(On Diff #436758)

Can we add a command line option to control this like AArch64 and PowerPC have?

llvm/test/Transforms/SeparateConstOffsetFromGEP/RISCV/split-gep.ll

2

This test doesn't exist in the repo. Where is the patch that adds it?

[eklepilkina]

llvm/test/Transforms/SeparateConstOffsetFromGEP/RISCV/split-gep.ll

2

I was told in the first comment to rebase on precommited tests. These tests are added as precommited in separate commit. Should I commit them?

[craig.topper]

llvm/test/Transforms/SeparateConstOffsetFromGEP/RISCV/split-gep.ll

2

Why is the script NOTE: Assertions have been autogenerated by utils/update_test_checks.py not in the pre-committed version?

[eklepilkina]

This should be two patches, one changing the pass and one enabling for RISC-V.

As far as I want to turn pass with enabled strategy, should I wait approve and merge of the accepting strategy and only after this create the second review? Or create series of patches as mentionedin documentation https://llvm.org/docs/Phabricator.html#creating-a-patch-series?

[craig.topper]

llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp

368

I think there you should be a std::move on PreviousIndices

371

rhs -> RHS

970–974

PossibleBase.size() == 0 -> PossibleBases.empty()

[craig.topper]

In D127727#3582087, @eklepilkina wrote:

This should be two patches, one changing the pass and one enabling for RISC-V.

As far as I want to turn pass with enabled strategy, should I wait approve and merge of the accepting strategy and only after this create the second review? Or create series of patches as mentionedin documentation https://llvm.org/docs/Phabricator.html#creating-a-patch-series?

You should create a series of patches.

[eklepilkina]

Separate part with pass modification

[eklepilkina]

• [SeparateConstOffsetFromGEP] Fix comparator for map with GEP bases

[eklepilkina]

Refactoring

[eklepilkina]

Fix review

[eklepilkina]

Ping

[eklepilkina]

• [SeparateConstOffsetFromGEP] Fix ignoring condition

[eklepilkina]

Gentle ping

[mkazantsev]

Some nits from me. If I may, some advice if you want to make progress here.

First, it seems that some pieces of this patch can be split out as seprate NFC refactorings. If so, please do. It should greatly reduce the code to look at, and smaller patches are generally easier to comprehend and review.

Second, the motivation of this patch is obscure. The structures and algorithms that you are using are not obvious. This either needs a detailed explanation in comments, or maybe incremental approach, when each step is understandable.

Third, the benefit isn't obvious either. In your tests, new IR is bigger than the old IR. This is not necessarily bad, but it requires explanation. If you have some particular scenario where the resulting assembly is better with this change, then it makes sense to provide an LLC test which shows it.

Fourth, this patch claims to improve register pressure. At the same time, it is done in a middle-end pass, which means its impact on register pressure on different platforms might be different. Was it actually tested on other platforms than RICSV? Or you think that the algorithm should be profitable regardless of the platform? Can I see any benefit from this patch in X86 for example?

llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp

203

Can this go as a separate NFC?

271–277

Please commit this reformatting separately.

280

\p CheckProfitability ?..

360

This requires more explanation. I could not figure what are indices, which of them is being optimized, and what is precedence in this context. Maybe write a detailed comment on what's going on here and what does this structure represent?

449

Canonize -> Canonicalize

450

I guess it should be "Returns true if a change was made, false otherwise".

547

Maybe rename InstructionsToTransform -> GEPsToTransform?

1083

Pls commit separately if it is needed.

1127

Rename as separate NFC?

1373

To me, this code structure looks counter-intuitive. Why do we print "Try to split GEP "... only when we check profitability, and do it silently when we don't? If possible, please restructure it like

if (CheckProfitability) {
  // Do all required profitability checks
}
// Do common transform logic uniformly

I'm not sure if it's possible here because of this post-processing. If not, then the transform part should be unified somehow else.

1380

More natural way would be

if (!CurrentChanged)
  continue;
for ...

1381

The complexity of this is SortedInstructionsList.size() * SortedInstructionsList.size() * sum(SortedInstructionsList[J]) if I'm reading this correctly. Looks very expensive. Is there a cheaper way of doing this? Imagine you have 10k instructions on your list. It will just be stuck forever.

llvm/test/Transforms/SeparateConstOffsetFromGEP/RISCV/split-gep.ll

86

Why is it a better code tha the old one?

286

This code is bigger than it used to be. Can you explain why is it better?

[eklepilkina]

Can I see any benefit from this patch in X86 for example?

This pass was written for targets with limited addressing mode, so it isn't added to X86 pipeline. It's used under the flag on Aarch64 and on RISCV we also suggest to turn off it by default, but this patch helps to make these optimization be useful more often, and remove some regressions that was found if turn these pass on on all test-suite. I'll provide test-suite results on Aarch64 platform with turned this pass with and without this patch. But yes, the main measurements were made for RISCV.

And if you mean that these changes should be done later in pipeline, there is the problem with the current instruction selection that can work only with one BB, so CodeGenPrepare pass need to sunk such GEPs with const to generate adddressing by offset, so I believe this pass was created as middle-end part.

llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp

1127

I don't really like the idea to rename in separate NFC patch, because renaming is connected with changes that were made and the old name wasn't suitable any more

1373

I understand your concerns, but I don't see a good solution here, because I don't want to make the unneeded actions for original version without checking profitability.

1381

Imagine you have 10k instructions on your list

I amn't sure we should optimize this case, because it's mostly impossible, because this list is always quite small. I'll think some more, but I amn't sure that the optimization here is more important than readability.

llvm/test/Transforms/SeparateConstOffsetFromGEP/RISCV/split-gep.ll

86

In assembly we use one more register to save the result of new generated GEP instruction, bt we have no profit because registers that are used by adds are also needed as far as these values are used in other instructions.

286

This code is bigger on IR, and it's so becuase of repeating sext opertaions, but sext isn't so critical in assembly, at the same time pass generates 2 new GEP instructions that are used as base and we need registers for them

[eklepilkina]

• [NFC][SeparateConstOffsetFromGEP] Small refactoring and reformatting
• [SeparateConstOffsetFromGEPPass] Added optional modification strategy
• Review fixes (part 1)

[eklepilkina]

• [SeparateConstOffsetFromGEPPass] Added optional modification strategy
• Review fixes (part 1)

[mkazantsev]

Found some bugs, some style proposals as well. The general point still holds. If the patch is purposed to reduce register pressure on some platform, please provide a test which shows that this actually happens. This can only be shown on a llc test.

llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp

247

This comment is obsolete now, Extract does not have these new parameters.

272

And if V is not binop, should it change?

362

This is usually called BasePointer in other parts of optimizer.

363

Shouldn't %b also be a part of it? Or where does it go?
Maybe more elaborate example on how there can be more than one previous index?

388

APInt? Just to make sure this doesn't overflow.

389

Naturaly I'd expect this to be SmallVector<const ConstantInt *>, but the code below suggests there might not be constants. Misleading name?

557

Use DenseMapInfo<Value *>::getTombstoneKey() and same above

561

Why PreviousIndices size but not contents?

709

No { }

710

undef and poison are constants but not ConstantInt. Are you OK with them?

1381

"Mostly impossible" means "possible". We generally bail out on non-linear algorithms with some thresholds. This could also be the case here.

llvm/test/Transforms/SeparateConstOffsetFromGEP/RISCV/split-gep.ll

286

Then please provide a llc test that demonstrates a positive change.

The fact that "sext isn't so critical" is a way not obvious to me. Filling the upper part of the registry may sometimes be an extra operation.

[eklepilkina]

Sorry for delay. Looked more on different benchmarks from test-suite during searching a good test case. There are such cases. But a deep exploration shows that SeparateConstOffsetFromGEP pass isn't the main reason, it produces better IR, but in some cases later passes can make it worse and cause worse asssembly code. So hacks I have made in the current pass as workarounds for these particular cases don't seem to be the proper decision. As far as this pass isn't the main reason of regressions we got, I decided to abandon this review.

Thank you very much for review and sorry to bother you.

llvm/include/llvm/Transforms/Scalar.h

 // calls such as sqrt.
 //
 FunctionPass *createPartiallyInlineLibCallsPass();
 
 //===----------------------------------------------------------------------===//
 //
 // SeparateConstOffsetFromGEP - Split GEPs for better CSE
 //
-FunctionPass *createSeparateConstOffsetFromGEPPass(bool LowerGEP = false);
+FunctionPass *
+createSeparateConstOffsetFromGEPPass(bool LowerGEP = false,
+                                     bool CheckProfitability = false);
 
 //===----------------------------------------------------------------------===//
 //
 // SpeculativeExecution - Aggressively hoist instructions to enable
 // speculative execution on targets where branches are expensive.
 //
 FunctionPass *createSpeculativeExecutionPass();
 

llvm/include/llvm/Transforms/Scalar/SeparateConstOffsetFromGEP.h

 
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 
 class SeparateConstOffsetFromGEPPass
     : public PassInfoMixin<SeparateConstOffsetFromGEPPass> {
   bool LowerGEP;
+  bool CheckProfitability;
 
 public:
-  SeparateConstOffsetFromGEPPass(bool LowerGEP = false) : LowerGEP(LowerGEP) {}
+  SeparateConstOffsetFromGEPPass(bool LowerGEP = false,
+                                 bool CheckProfitability = false)
+      : LowerGEP(LowerGEP), CheckProfitability(CheckProfitability) {}
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_SCALAR_SEPARATECONSTOFFSETFROMGEP_H

llvm/lib/Passes/PassBuilder.cpp

 Expected<bool> parseLoopExtractorPassOptions(StringRef Params) {
   return parseSinglePassOption(Params, "single", "LoopExtractor");
 }
 
 Expected<bool> parseLowerMatrixIntrinsicsPassOptions(StringRef Params) {
   return parseSinglePassOption(Params, "minimal", "LowerMatrixIntrinsics");
 }
 
+Expected<bool> parseSeparateConstOffsetFromGEPPassOptions(StringRef Params) {
+  return parseSinglePassOption(Params, "check-profit",
+                               "SeparateConstOffsetFromGEP");
+}
+
 Expected<AddressSanitizerOptions> parseASanPassOptions(StringRef Params) {
   AddressSanitizerOptions Result;
   while (!Params.empty()) {
     StringRef ParamName;
     std::tie(ParamName, Params) = Params.split(';');
 
     if (ParamName == "kernel") {
       Result.CompileKernel = true;

llvm/lib/Passes/PassRegistry.def

 FUNCTION_PASS("print-predicateinfo", PredicateInfoPrinterPass(dbgs()))
 FUNCTION_PASS("print-mustexecute", MustExecutePrinterPass(dbgs()))
 FUNCTION_PASS("print-memderefs", MemDerefPrinterPass(dbgs()))
 FUNCTION_PASS("reassociate", ReassociatePass())
 FUNCTION_PASS("redundant-dbg-inst-elim", RedundantDbgInstEliminationPass())
 FUNCTION_PASS("reg2mem", RegToMemPass())
 FUNCTION_PASS("scalarize-masked-mem-intrin", ScalarizeMaskedMemIntrinPass())
 FUNCTION_PASS("scalarizer", ScalarizerPass())
-FUNCTION_PASS("separate-const-offset-from-gep", SeparateConstOffsetFromGEPPass())
 FUNCTION_PASS("sccp", SCCPPass())
 FUNCTION_PASS("sink", SinkingPass())
 FUNCTION_PASS("slp-vectorizer", SLPVectorizerPass())
 FUNCTION_PASS("slsr", StraightLineStrengthReducePass())
 FUNCTION_PASS("speculative-execution", SpeculativeExecutionPass())
 FUNCTION_PASS("sroa", SROAPass())
 FUNCTION_PASS("strip-gc-relocates", StripGCRelocates())
 FUNCTION_PASS("structurizecfg", StructurizeCFGPass())
FUNCTION_PASS_WITH_PARAMS("simplifycfg",
                           "no-forward-switch-cond;forward-switch-cond;"
                           "no-switch-range-to-icmp;switch-range-to-icmp;"
                           "no-switch-to-lookup;switch-to-lookup;"
                           "no-keep-loops;keep-loops;"
                           "no-hoist-common-insts;hoist-common-insts;"
                           "no-sink-common-insts;sink-common-insts;"
                           "bonus-inst-threshold=N"
                           )
+FUNCTION_PASS_WITH_PARAMS("separate-const-offset-from-gep",
+                          "SeparateConstOffsetFromGEPPass",
+                          [](bool CheckProfitability) {
+                            return SeparateConstOffsetFromGEPPass(false, CheckProfitability);
+                          },
+                          parseSeparateConstOffsetFromGEPPassOptions,
+                          "check-profit")
 FUNCTION_PASS_WITH_PARAMS("loop-vectorize",
                           "LoopVectorizePass",
                            [](LoopVectorizeOptions Opts) {
                              return LoopVectorizePass(Opts);
                            },
                           parseLoopVectorizeOptions,
                           "no-interleave-forced-only;interleave-forced-only;"
                           "no-vectorize-forced-only;vectorize-forced-only")

llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp

 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include <cassert>
 #include <cstdint>
+#include <map>
 #include <string>
 
 using namespace llvm;
 using namespace llvm::PatternMatch;
 
 #define DEBUG_TYPE "separate-const-offset-from-gep"
 
 static cl::opt<bool> DisableSeparateConstOffsetFromGEP(
     "disable-separate-const-offset-from-gep", cl::init(false),
     cl::desc("Do not separate the constant offset from a GEP instruction"),
     cl::Hidden);
 
 // Setting this flag may emit false positives when the input module already
 // contains dead instructions. Therefore, we set it only in unit tests that are
public:
   /// \p GEP The given GEP
   /// \p UserChainTail Outputs the tail of UserChain so that we can
   ///                  garbage-collect unused instructions in UserChain.
   static Value *Extract(Value *Idx, GetElementPtrInst *GEP,
                         User *&UserChainTail, const DominatorTree *DT);
 
   /// Looks for a constant offset from the given GEP index without extracting
   /// it. It returns the numeric value of the extracted constant offset (0 if
   /// failed). The meaning of the arguments are the same as Extract.
   static int64_t Find(Value *Idx, GetElementPtrInst *GEP,
-                      const DominatorTree *DT);
+                      const DominatorTree *DT, Value *&NonConstantBaseValue,
+                      bool CheckProfitability = false);
 
 private:
   ConstantOffsetExtractor(Instruction *InsertionPt, const DominatorTree *DT)
       : IP(InsertionPt), DL(InsertionPt->getModule()->getDataLayout()), DT(DT) {
   }
 
   /// Searches the expression that computes V for a non-zero constant C s.t.
   /// V can be reassociated into the form V' + C. If the searching is
   /// successful, returns C and update UserChain as a def-use chain from C to V;
   /// otherwise, UserChain is empty.
   ///
   /// \p V                    The given expression
   /// \p SignExtended         Whether V will be sign-extended in the computation
   ///                         of the GEP index
   /// \p ZeroExtended         Whether V will be zero-extended in the computation
   ///                         of the GEP index
   /// \p NonNegative          Whether V is guaranteed to be non-negative. For
   ///                         example, an index of an inbounds GEP is guaranteed
   ///                         to be non-negative. Levaraging this, we can better
   ///                         split inbounds GEPs.
-  APInt find(Value *V, bool SignExtended, bool ZeroExtended, bool NonNegative);
+  /// \p NonConstantBaseValue The second non-constant operand if V is binary
+  ///                         operator.
+  /// \p CheckProfitability   Check the possible profit of optimization and
+  ///                         search only instructions that can be optimized
+  ///                         without extra operations.
+  APInt find(Value *V, bool SignExtended, bool ZeroExtended, bool NonNegative,
+             Value *&NonConstantBaseValue, bool CheckProfitability = false);
 
   /// A helper function to look into both operands of a binary operator.
   APInt findInEitherOperand(BinaryOperator *BO, bool SignExtended,
-                            bool ZeroExtended);
+                            bool ZeroExtended, Value *&NonConstantBaseValue,
+                            bool CheckProfitability = false);
 
   /// After finding the constant offset C from the GEP index I, we build a new
   /// index I' s.t. I' + C = I. This function builds and returns the new
   /// index I' according to UserChain produced by function "find".
   ///
   /// The building conceptually takes two steps:
   /// 1) iteratively distribute s/zext towards the leaves of the expression tree
   /// that computes I
private:
 
   /// Insertion position of cloned instructions.
   Instruction *IP;
 
   const DataLayout &DL;
   const DominatorTree *DT;
 };
 
+/// GEPBaseInfo - structure contains information about possible common base for
+/// GEP instructions.
+/// In case if we have the next GEP instruction
+///
+///  %add = add nsw i64 %i, 5
+///  getelementptr inbounds [50 x i64], [50 x i64]* %array, i64 %a, i64 %b
+///
+/// GEPPointer - %array
+/// PreviousIndices - [%a]
+/// NonConstantBaseValue - %i
+///
+/// Structure stores information to calculate GEP instructions with same base.
+/// Needed to count how many GEP instructions can be optimized using one new
+/// GEP instruction.
+struct GEPBaseInfo {
+  /// Pointer used in GEP instruction.
+  const Value *GEPPointer;
+  /// Indexes that precede index that can be optimized.
+  SmallVector<const Value *> PreviousIndices;
+  /// Non constant value that will be used in new base GEP.
+  const Value *NonConstantBaseValue;
+
+  GEPBaseInfo(const Value *GEPPointer,
+              SmallVector<const Value *> PreviousIndices,
+              const Value *NonConstantBaseValue)
+      : GEPPointer(GEPPointer), PreviousIndices(PreviousIndices),
+        NonConstantBaseValue(NonConstantBaseValue) {}
+};
+
+/// GEPInfo - structure contains basic information about GEP instruction
+/// needed for their modification.
+struct GEPInfo {
+  GetElementPtrInst *GEPInstruction;
+  int64_t AccumulativeByteOffset;
+  SmallVector<const Value *> ConstantIndices;
+
+  GEPInfo(GetElementPtrInst *GEPInstruction, int64_t AccumulativeByteOffset,
+          SmallVector<const Value *> &&Indices)
+      : GEPInstruction(GEPInstruction),
+        AccumulativeByteOffset(AccumulativeByteOffset),
+        ConstantIndices(Indices) {}
+};
+
 /// A pass that tries to split every GEP in the function into a variadic
 /// base and a constant offset. It is a FunctionPass because searching for the
 /// constant offset may inspect other basic blocks.
 class SeparateConstOffsetFromGEPLegacyPass : public FunctionPass {
 public:
   static char ID;
 
-  SeparateConstOffsetFromGEPLegacyPass(bool LowerGEP = false)
-      : FunctionPass(ID), LowerGEP(LowerGEP) {
+  SeparateConstOffsetFromGEPLegacyPass(bool LowerGEP = false,
+                                       bool CheckProfitability = false)
+      : FunctionPass(ID), LowerGEP(LowerGEP),
+        CheckProfitability(CheckProfitability) {
     initializeSeparateConstOffsetFromGEPLegacyPassPass(
         *PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addRequired<ScalarEvolutionWrapperPass>();
     AU.addRequired<TargetTransformInfoWrapperPass>();
     AU.addRequired<LoopInfoWrapperPass>();
     AU.setPreservesCFG();
     AU.addRequired<TargetLibraryInfoWrapperPass>();
   }
 
   bool runOnFunction(Function &F) override;
 
 private:
   bool LowerGEP;
+  bool CheckProfitability;
 };
 
 /// A pass that tries to split every GEP in the function into a variadic
 /// base and a constant offset. It is a FunctionPass because searching for the
 /// constant offset may inspect other basic blocks.
 class SeparateConstOffsetFromGEP {
 public:
   SeparateConstOffsetFromGEP(
       DominatorTree *DT, ScalarEvolution *SE, LoopInfo *LI,
       TargetLibraryInfo *TLI,
-      function_ref<TargetTransformInfo &(Function &)> GetTTI, bool LowerGEP)
-      : DT(DT), SE(SE), LI(LI), TLI(TLI), GetTTI(GetTTI), LowerGEP(LowerGEP) {}
+      function_ref<TargetTransformInfo &(Function &)> GetTTI, bool LowerGEP,
+      bool CheckProfitability)
+      : DT(DT), SE(SE), LI(LI), TLI(TLI), GetTTI(GetTTI), LowerGEP(LowerGEP),
+        CheckProfitability(CheckProfitability) {}
 
   bool run(Function &F);
 
 private:
   /// Tries to split the given GEP into a variadic base and a constant offset,
   /// and returns true if the splitting succeeds.
-  bool splitGEP(GetElementPtrInst *GEP);
+  bool splitGEP(GetElementPtrInst *GEP, int64_t AccumulativeByteOffset);
+
+  /// Canonicalize GEP if needed and collect information to decide if GEP
+  /// modification is useful.
+  /// Returns true if a change was made, false otherwise.
+  bool preprocessGEP(GetElementPtrInst *GEP);
 
   /// Lower a GEP with multiple indices into multiple GEPs with a single index.
   /// Function splitGEP already split the original GEP into a variadic part and
   /// a constant offset (i.e., AccumulativeByteOffset). This function lowers the
   /// variadic part into a set of GEPs with a single index and applies
   /// AccumulativeByteOffset to it.
   /// \p Variadic                  The variadic part of the original GEP.
   /// \p AccumulativeByteOffset    The constant offset.
   void lowerToSingleIndexGEPs(GetElementPtrInst *Variadic,
                               int64_t AccumulativeByteOffset);
 
   /// Lower a GEP with multiple indices into ptrtoint+arithmetics+inttoptr form.
   /// Function splitGEP already split the original GEP into a variadic part and
   /// a constant offset (i.e., AccumulativeByteOffset). This function lowers the
   /// variadic part into a set of arithmetic operations and applies
   /// AccumulativeByteOffset to it.
   /// \p Variadic                  The variadic part of the original GEP.
   /// \p AccumulativeByteOffset    The constant offset.
   void lowerToArithmetics(GetElementPtrInst *Variadic,
                           int64_t AccumulativeByteOffset);
 
   /// Finds the constant offset within each index and accumulates them. If
   /// LowerGEP is true, it finds in indices of both sequential and structure
-  /// types, otherwise it only finds in sequential indices. The output
-  /// NeedsExtraction indicates whether we successfully find a non-zero constant
-  /// offset.
-  int64_t accumulateByteOffset(GetElementPtrInst *GEP, bool &NeedsExtraction);
+  /// types, otherwise it only finds in sequential indices.
+  void accumulateByteOffset(GetElementPtrInst *GEP);
 
   /// Canonicalize array indices to pointer-size integers. This helps to
   /// simplify the logic of splitting a GEP. For example, if a + b is a
   /// pointer-size integer, we have
   ///   gep base, a + b = gep (gep base, a), b
   /// However, this equality may not hold if the size of a + b is smaller than
   /// the pointer size, because LLVM conceptually sign-extends GEP indices to
   /// pointer size before computing the address
private:
   TargetLibraryInfo *TLI;
   // Retrieved lazily since not always used.
   function_ref<TargetTransformInfo &(Function &)> GetTTI;
 
   /// Whether to lower a GEP with multiple indices into arithmetic operations or
   /// multiple GEPs with a single index.
   bool LowerGEP;
 
+  /// Check the possible profit of optimization to reduce register pressure
+  /// or modify all possible GEPs.
+  bool CheckProfitability;
+
   DenseMap<const SCEV *, SmallVector<Instruction *, 2>> DominatingAdds;
   DenseMap<const SCEV *, SmallVector<Instruction *, 2>> DominatingSubs;
+
+  /// GEP instructions chosen for transformation
+  DenseMap<GEPBaseInfo, SmallVector<GEPInfo>> GEPsToTransform;
 };
 
 } // end anonymous namespace
 
+template <> struct llvm::DenseMapInfo<GEPBaseInfo> {
+  static inline GEPBaseInfo getEmptyKey() {
+    return GEPBaseInfo(nullptr, SmallVector<const Value *>(), nullptr);
+  }
+  static inline GEPBaseInfo getTombstoneKey() {
+    return GEPBaseInfo((Value *)(-1), SmallVector<const Value *>(),
+                       (Value *)(-1));
+  }
+  static unsigned getHashValue(const GEPBaseInfo &Val) {
+    return llvm::hash_combine(Val.GEPPointer, Val.PreviousIndices.size(),
+                              Val.NonConstantBaseValue);
+  }
+  static bool isEqual(const GEPBaseInfo &LHS, const GEPBaseInfo &RHS) {
+    return LHS.GEPPointer == RHS.GEPPointer &&
+           LHS.NonConstantBaseValue == RHS.NonConstantBaseValue &&
+           LHS.PreviousIndices == RHS.PreviousIndices;
+  }
+};
+
 char SeparateConstOffsetFromGEPLegacyPass::ID = 0;
 
 INITIALIZE_PASS_BEGIN(
     SeparateConstOffsetFromGEPLegacyPass, "separate-const-offset-from-gep",
     "Split GEPs to a variadic base and a constant offset for better CSE", false,
     false)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(ScalarEvolutionWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
 INITIALIZE_PASS_END(
     SeparateConstOffsetFromGEPLegacyPass, "separate-const-offset-from-gep",
     "Split GEPs to a variadic base and a constant offset for better CSE", false,
     false)
 
-FunctionPass *llvm::createSeparateConstOffsetFromGEPPass(bool LowerGEP) {
-  return new SeparateConstOffsetFromGEPLegacyPass(LowerGEP);
+FunctionPass *
+llvm::createSeparateConstOffsetFromGEPPass(bool LowerGEP,
+                                           bool CheckProfitability) {
+  return new SeparateConstOffsetFromGEPLegacyPass(LowerGEP, CheckProfitability);
 }
 
 bool ConstantOffsetExtractor::CanTraceInto(bool SignExtended,
                                             bool ZeroExtended,
                                             BinaryOperator *BO,
                                             bool NonNegative) {
   // We only consider ADD, SUB and OR, because a non-zero constant found in
   // expressions composed of these operations can be easily hoisted as a
if (ZeroExtended && !BO->hasNoUnsignedWrap())
       return false;
   }
 
   return true;
 }
 
 APInt ConstantOffsetExtractor::findInEitherOperand(BinaryOperator *BO,
                                                    bool SignExtended,
-                                                   bool ZeroExtended) {
+                                                   bool ZeroExtended,
+                                                   Value *&NonConstantBaseValue,
+                                                   bool CheckProfitability) {
   // Save off the current height of the chain, in case we need to restore it.
   size_t ChainLength = UserChain.size();
 
   // BO being non-negative does not shed light on whether its operands are
   // non-negative. Clear the NonNegative flag here.
-  APInt ConstantOffset = find(BO->getOperand(0), SignExtended, ZeroExtended,
-                              /* NonNegative */ false);
+  APInt ConstantOffset =
+      find(BO->getOperand(0), SignExtended, ZeroExtended,
+           /* NonNegative */ false, NonConstantBaseValue, CheckProfitability);
+
+  // Only sub and add instructions don't need adding extra instructions.
+  if (CheckProfitability && (BO->getOpcode() != Instruction::Sub &&
+                             BO->getOpcode() != Instruction::Add)) {
+    NonConstantBaseValue = nullptr;
+    ConstantOffset = 0;
+    UserChain.resize(ChainLength);
+    return ConstantOffset;
+  }
+
   // If we found a constant offset in the left operand, stop and return that.
   // This shortcut might cause us to miss opportunities of combining the
   // constant offsets in both operands, e.g., (a + 4) + (b + 5) => (a + b) + 9.
   // However, such cases are probably already handled by -instcombine,
   // given this pass runs after the standard optimizations.
-  if (ConstantOffset != 0) return ConstantOffset;
+  if (ConstantOffset != 0) {
+    if (!isa<ConstantInt>(BO->getOperand(1))) {
+      NonConstantBaseValue = BO->getOperand(1);
+    }
+    return ConstantOffset;
+  }
 
   // Reset the chain back to where it was when we started exploring this node,
   // since visiting the LHS didn't pan out.
   UserChain.resize(ChainLength);
 
-  ConstantOffset = find(BO->getOperand(1), SignExtended, ZeroExtended,
-                        /* NonNegative */ false);
+  ConstantOffset =
+      find(BO->getOperand(1), SignExtended, ZeroExtended,
+           /* NonNegative */ false, NonConstantBaseValue, CheckProfitability);
   // If U is a sub operator, negate the constant offset found in the right
   // operand.
   if (BO->getOpcode() == Instruction::Sub)
     ConstantOffset = -ConstantOffset;
 
   // If RHS wasn't a suitable candidate either, reset the chain again.
   if (ConstantOffset == 0)
     UserChain.resize(ChainLength);
 
+  if (!isa<ConstantInt>(BO->getOperand(0))) {
+    NonConstantBaseValue = BO->getOperand(0);
+  }
+
   return ConstantOffset;
 }
 
 APInt ConstantOffsetExtractor::find(Value *V, bool SignExtended,
-                                    bool ZeroExtended, bool NonNegative) {
+                                    bool ZeroExtended, bool NonNegative,
+                                    Value *&NonConstantBaseValue,
+                                    bool CheckProfitability) {
   // TODO(jingyue): We could trace into integer/pointer casts, such as
   // inttoptr, ptrtoint, bitcast, and addrspacecast. We choose to handle only
   // integers because it gives good enough results for our benchmarks.
   unsigned BitWidth = cast<IntegerType>(V->getType())->getBitWidth();
 
   // We cannot do much with Values that are not a User, such as an Argument.
   User *U = dyn_cast<User>(V);
   if (U == nullptr) return APInt(BitWidth, 0);
 
   APInt ConstantOffset(BitWidth, 0);
   if (ConstantInt *CI = dyn_cast<ConstantInt>(V)) {
     // Hooray, we found it!
     ConstantOffset = CI->getValue();
   } else if (BinaryOperator *BO = dyn_cast<BinaryOperator>(V)) {
     // Trace into subexpressions for more hoisting opportunities.
     if (CanTraceInto(SignExtended, ZeroExtended, BO, NonNegative))
-      ConstantOffset = findInEitherOperand(BO, SignExtended, ZeroExtended);
-  } else if (isa<TruncInst>(V)) {
       ConstantOffset =
-        find(U->getOperand(0), SignExtended, ZeroExtended, NonNegative)
+          findInEitherOperand(BO, SignExtended, ZeroExtended,
+                              NonConstantBaseValue, CheckProfitability);
+  } else if (isa<TruncInst>(V)) {
+    ConstantOffset = find(U->getOperand(0), SignExtended, ZeroExtended,
+                          NonNegative, NonConstantBaseValue, CheckProfitability)
                          .trunc(BitWidth);
   } else if (isa<SExtInst>(V)) {
-    ConstantOffset = find(U->getOperand(0), /* SignExtended */ true,
-                          ZeroExtended, NonNegative).sext(BitWidth);
+    ConstantOffset =
+        find(U->getOperand(0), /* SignExtended */ true, ZeroExtended,
+             NonNegative, NonConstantBaseValue, CheckProfitability)
+            .sext(BitWidth);
   } else if (isa<ZExtInst>(V)) {
     // As an optimization, we can clear the SignExtended flag because
     // sext(zext(a)) = zext(a). Verified in @sext_zext in split-gep.ll.
     //
     // Clear the NonNegative flag, because zext(a) >= 0 does not imply a >= 0.
-    ConstantOffset =
-        find(U->getOperand(0), /* SignExtended */ false,
-             /* ZeroExtended */ true, /* NonNegative */ false).zext(BitWidth);
+    ConstantOffset = find(U->getOperand(0), /* SignExtended */ false,
+                          /* ZeroExtended */ true, /* NonNegative */ false,
+                          NonConstantBaseValue, CheckProfitability)
+                         .zext(BitWidth);
   }
 
   // If we found a non-zero constant offset, add it to the path for
   // rebuildWithoutConstOffset. Zero is a valid constant offset, but doesn't
   // help this optimization.
   if (ConstantOffset != 0)
     UserChain.push_back(U);
   return ConstantOffset;
Value *ConstantOffsetExtractor::removeConstOffset(unsigned ChainIndex) {
   NewBO->takeName(BO);
   return NewBO;
 }
 
 Value *ConstantOffsetExtractor::Extract(Value *Idx, GetElementPtrInst *GEP,
                                         User *&UserChainTail,
                                         const DominatorTree *DT) {
   ConstantOffsetExtractor Extractor(GEP, DT);
+  Value *NonConstantBaseValue = nullptr;
   // Find a non-zero constant offset first.
   APInt ConstantOffset =
       Extractor.find(Idx, /* SignExtended */ false, /* ZeroExtended */ false,
-                     GEP->isInBounds());
+                     GEP->isInBounds(), NonConstantBaseValue);
   if (ConstantOffset == 0) {
     UserChainTail = nullptr;
     return nullptr;
   }
   // Separates the constant offset from the GEP index.
   Value *IdxWithoutConstOffset = Extractor.rebuildWithoutConstOffset();
   UserChainTail = Extractor.UserChain.back();
   return IdxWithoutConstOffset;
 }
 
 int64_t ConstantOffsetExtractor::Find(Value *Idx, GetElementPtrInst *GEP,
-                                      const DominatorTree *DT) {
+                                      const DominatorTree *DT,
+                                      Value *&NonConstantBaseValue,
+                                      bool CheckProfitability) {
   // If Idx is an index of an inbound GEP, Idx is guaranteed to be non-negative.
   return ConstantOffsetExtractor(GEP, DT)
       .find(Idx, /* SignExtended */ false, /* ZeroExtended */ false,
-            GEP->isInBounds())
+            GEP->isInBounds(), NonConstantBaseValue, CheckProfitability)
       .getSExtValue();
 }
 
 bool SeparateConstOffsetFromGEP::canonicalizeArrayIndicesToPointerSize(
     GetElementPtrInst *GEP) {
   bool Changed = false;
   Type *IntPtrTy = DL->getIntPtrType(GEP->getType());
   gep_type_iterator GTI = gep_type_begin(*GEP);
   for (User::op_iterator I = GEP->op_begin() + 1, E = GEP->op_end();
        I != E; ++I, ++GTI) {
     // Skip struct member indices which must be i32.
     if (GTI.isSequential()) {
       if ((*I)->getType() != IntPtrTy) {
         *I = CastInst::CreateIntegerCast(*I, IntPtrTy, true, "idxprom", GEP);
         Changed = true;
       }
     }
   }
   return Changed;
 }
 
-int64_t
-SeparateConstOffsetFromGEP::accumulateByteOffset(GetElementPtrInst *GEP,
-                                                 bool &NeedsExtraction) {
-  NeedsExtraction = false;
+void SeparateConstOffsetFromGEP::accumulateByteOffset(GetElementPtrInst *GEP) {
   int64_t AccumulativeByteOffset = 0;
   gep_type_iterator GTI = gep_type_begin(*GEP);
+  SmallVector<const Value *> ConstantIndices;
+  SmallVector<GEPBaseInfo, 2> PossibleBases;
+
   for (unsigned I = 1, E = GEP->getNumOperands(); I != E; ++I, ++GTI) {
+    Value *NonConstantBaseValue = nullptr;
     if (GTI.isSequential()) {
       // Tries to extract a constant offset from this GEP index.
-      int64_t ConstantOffset =
-          ConstantOffsetExtractor::Find(GEP->getOperand(I), GEP, DT);
+      int64_t ConstantOffset = ConstantOffsetExtractor::Find(
+          GEP->getOperand(I), GEP, DT, NonConstantBaseValue,
+          CheckProfitability);
       if (ConstantOffset != 0) {
-        NeedsExtraction = true;
+        if (CheckProfitability || PossibleBases.empty()) {
+          PossibleBases.emplace_back(
+              GEP->getPointerOperand(),
+              SmallVector<const Value *, 4>(GEP->idx_begin(),
+                                            GEP->idx_begin() + I - 1),
+              NonConstantBaseValue);
+        }
+
         // A GEP may have multiple indices.  We accumulate the extracted
         // constant offset to a byte offset, and later offset the remainder of
         // the original GEP with this byte offset.
         AccumulativeByteOffset +=
             ConstantOffset * DL->getTypeAllocSize(GTI.getIndexedType());
+        ConstantIndices.push_back(GEP->getOperand(I));
       }
     } else if (LowerGEP) {
       StructType *StTy = GTI.getStructType();
       uint64_t Field = cast<ConstantInt>(GEP->getOperand(I))->getZExtValue();
       // Skip field 0 as the offset is always 0.
       if (Field != 0) {
-        NeedsExtraction = true;
+        if (CheckProfitability || PossibleBases.empty()) {
+          PossibleBases.emplace_back(GEP->getPointerOperand(),
+                                     SmallVector<const Value *>(),
+                                     NonConstantBaseValue);
+        }
         AccumulativeByteOffset +=
             DL->getStructLayout(StTy)->getElementOffset(Field);
       }
     }
   }
-  return AccumulativeByteOffset;
+  TargetTransformInfo &TTI = GetTTI(*GEP->getFunction());
+
+  // If LowerGEP is disabled, before really splitting the GEP, check whether the
+  // backend supports the addressing mode we are about to produce. If no, this
+  // splitting probably won't be beneficial.
+  // If LowerGEP is enabled, even the extracted constant offset can not match
+  // the addressing mode, we can still do optimizations to other lowered parts
+  // of variable indices. Therefore, we don't check for addressing modes in that
+  // case.
+  if (!LowerGEP) {
+    unsigned AddrSpace = GEP->getPointerAddressSpace();
+    if (!TTI.isLegalAddressingMode(GEP->getResultElementType(),
+                                   /*BaseGV=*/nullptr, AccumulativeByteOffset,
+                                   /*HasBaseReg=*/true, /*Scale=*/0,
+                                   AddrSpace)) {
+      LLVM_DEBUG(
+          dbgs()
+          << "Don't optimize. The backend doesn't support the addressing mode\n");
+      return;
+    }
+  }
+  for (const GEPBaseInfo &Base : PossibleBases) {
+    if (GEPsToTransform.find(Base) == GEPsToTransform.end()) {
+      GEPsToTransform[Base] = SmallVector<GEPInfo>();
+    }
+    GEPsToTransform[Base].emplace_back(GEP, AccumulativeByteOffset,
+                                               std::move(ConstantIndices));
+  }
 }
 
 void SeparateConstOffsetFromGEP::lowerToSingleIndexGEPs(
     GetElementPtrInst *Variadic, int64_t AccumulativeByteOffset) {
   IRBuilder<> Builder(Variadic);
   Type *IntPtrTy = DL->getIntPtrType(Variadic->getType());
 
   Type *I8PtrTy =
if (ResultPtr->getType() != Variadic->getType())
     ResultPtr = Builder.CreateBitCast(ResultPtr, Variadic->getType());
 
   Variadic->replaceAllUsesWith(ResultPtr);
   Variadic->eraseFromParent();
 }
 
 void
 SeparateConstOffsetFromGEP::lowerToArithmetics(GetElementPtrInst *Variadic,
                                                int64_t AccumulativeByteOffset) {
   IRBuilder<> Builder(Variadic);
   Type *IntPtrTy = DL->getIntPtrType(Variadic->getType());
 
   Value *ResultPtr = Builder.CreatePtrToInt(Variadic->getOperand(0), IntPtrTy);
   gep_type_iterator GTI = gep_type_begin(*Variadic);
   // Create ADD/SHL/MUL arithmetic operations for each sequential indices. We
   // don't create arithmetics for structure indices, as they are accumulated
   // in the constant offset index.
ResultPtr = Builder.CreateAdd(
         ResultPtr, ConstantInt::get(IntPtrTy, AccumulativeByteOffset));
   }
 
   ResultPtr = Builder.CreateIntToPtr(ResultPtr, Variadic->getType());
   Variadic->replaceAllUsesWith(ResultPtr);
   Variadic->eraseFromParent();
 }
 
-bool SeparateConstOffsetFromGEP::splitGEP(GetElementPtrInst *GEP) {
+bool SeparateConstOffsetFromGEP::preprocessGEP(GetElementPtrInst *GEP) {
   // Skip vector GEPs.
   if (GEP->getType()->isVectorTy())
     return false;
 
   // The backend can already nicely handle the case where all indices are
   // constant.
   if (GEP->hasAllConstantIndices())
     return false;
 
   bool Changed = canonicalizeArrayIndicesToPointerSize(GEP);
 
-  bool NeedsExtraction;
-  int64_t AccumulativeByteOffset = accumulateByteOffset(GEP, NeedsExtraction);
+  accumulateByteOffset(GEP);
 
-  if (!NeedsExtraction)
-    return Changed;
+}
 
+bool SeparateConstOffsetFromGEP::splitGEP(GetElementPtrInst *GEP,
+                                          int64_t AccumulativeByteOffset) {
   TargetTransformInfo &TTI = GetTTI(*GEP->getFunction());
 
-  // If LowerGEP is disabled, before really splitting the GEP, check whether the
-  // backend supports the addressing mode we are about to produce. If no, this
-  // splitting probably won't be beneficial.
-  // If LowerGEP is enabled, even the extracted constant offset can not match
-  // the addressing mode, we can still do optimizations to other lowered parts
-  // of variable indices. Therefore, we don't check for addressing modes in that
-  // case.
-  if (!LowerGEP) {
-    unsigned AddrSpace = GEP->getPointerAddressSpace();
-    if (!TTI.isLegalAddressingMode(GEP->getResultElementType(),
-                                   /*BaseGV=*/nullptr, AccumulativeByteOffset,
-                                   /*HasBaseReg=*/true, /*Scale=*/0,
-                                   AddrSpace)) {
-      return Changed;
-    }
-  }
-
   // Remove the constant offset in each sequential index. The resultant GEP
   // computes the variadic base.
   // Notice that we don't remove struct field indices here. If LowerGEP is
   // disabled, a structure index is not accumulated and we still use the old
   // one. If LowerGEP is enabled, a structure index is accumulated in the
   // constant offset. LowerToSingleIndexGEPs or lowerToArithmetics will later
   // handle the constant offset and won't need a new structure index.
   gep_type_iterator GTI = gep_type_begin(*GEP);
if (skipFunction(F))
     return false;
   auto *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
   auto *SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
   auto *LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
   auto *TLI = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
   auto GetTTI = [this](Function &F) -> TargetTransformInfo & {
     return this->getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
   };
-  SeparateConstOffsetFromGEP Impl(DT, SE, LI, TLI, GetTTI, LowerGEP);
+  SeparateConstOffsetFromGEP Impl(DT, SE, LI, TLI, GetTTI, LowerGEP,
+                                  CheckProfitability);
   return Impl.run(F);
 }
 
 bool SeparateConstOffsetFromGEP::run(Function &F) {
   if (DisableSeparateConstOffsetFromGEP)
     return false;
 
   DL = &F.getParent()->getDataLayout();
   bool Changed = false;
 
-  LLVM_DEBUG(dbgs() << "========= Function " << F.getName() << " =========\n");
+  auto OnlyUsedInGEP = [](const GEPInfo &GEPInfo) {
+    bool OnlyUsedInGEP = GEPInfo.ConstantIndices.empty();
+    for (const Value *Index : GEPInfo.ConstantIndices) {
+      uint64_t NumUses = Index->getNumUses();
+      // In case of cast instruction check usages of both result and original
+      // value.
+      if (isa<SExtInst>(Index) || isa<ZExtInst>(Index) ||
+          isa<TruncInst>(Index)) {
+        NumUses += cast<Instruction>(Index)->getOperand(0)->getNumUses() - 1;
+      }
+      OnlyUsedInGEP |= NumUses == 1;
+    }
+    return OnlyUsedInGEP;
+  };
 
+  LLVM_DEBUG(dbgs() << "========= Function " << F.getName() << " =========\n");
   for (BasicBlock &B : F) {
+    GEPsToTransform.clear();
     if (!DT->isReachableFromEntry(&B))
       continue;
 
     for (Instruction &I : llvm::make_early_inc_range(B))
       if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(&I))
-        Changed |= splitGEP(GEP);
-    // No need to split GEP ConstantExprs because all its indices are constant
-    // already.
+        Changed |= preprocessGEP(GEP);
+
+    if (!CheckProfitability) {
+      for (const auto &GEPInfoPair : GEPsToTransform) {
+        for (const auto &GEPInfo : GEPInfoPair.second) {
+          LLVM_DEBUG(dbgs() << "Try to split GEP " << *GEPInfo.GEPInstruction
+                            << "\n");
+          Changed |=
+              splitGEP(GEPInfo.GEPInstruction, GEPInfo.AccumulativeByteOffset);
+          // No need to split GEP ConstantExprs because all its indices are
+          // constant already.
+        }
+      }
+    } else {
+      // As far as one instruction can be optimized using different base, choose
+      // the best base option based on possible effect for decreasing register
+      // pressure. Sort all found bases in decreasing order of possible effect.
+      SmallVector<std::pair<SmallVector<GEPInfo>, unsigned>>
+          SortedInstructionsList;
+      for (const auto &GEPInfoPair : GEPsToTransform) {
+        unsigned DeadValuesNumber = count_if(GEPInfoPair.second, OnlyUsedInGEP);
+        if (DeadValuesNumber > 0) {
+          SortedInstructionsList.emplace_back(GEPInfoPair.second,
+                                              DeadValuesNumber);
+        }
+      }
+      sort(SortedInstructionsList, [&OnlyUsedInGEP](auto &LHS, auto &RHS) {
+        return LHS.second > RHS.second || LHS.first.size() > RHS.first.size();
+      });
+
+      // Optimize all chosen GEPs
+      for (unsigned I = 0; I < SortedInstructionsList.size(); I++) {
+        auto DetailedInfoList = SortedInstructionsList[I].first;
+        if (DetailedInfoList.size() > 1 &&
+            any_of(DetailedInfoList, OnlyUsedInGEP)) {
+          for (const auto &GEPInfo : DetailedInfoList) {
+            LLVM_DEBUG(dbgs() << "Try to split GEP " << *GEPInfo.GEPInstruction
+                              << "\n");
+            bool CurrentChanged = splitGEP(GEPInfo.GEPInstruction,
+                                           GEPInfo.AccumulativeByteOffset);
+            Changed |= CurrentChanged;
+            // If GEP is already optimized remove it from lists connected with
+            // other bases.
+            if (!CurrentChanged)
+              continue;
+            for (unsigned J = I + 1; J < SortedInstructionsList.size(); J++) {
+              auto RemoveIt = remove_if(SortedInstructionsList[J].first,
+                                        [&GEPInfo](const struct GEPInfo &Info) {
+                                          return Info.GEPInstruction ==
+                                                 GEPInfo.GEPInstruction;
+                                        });
+              SortedInstructionsList[J].first.erase(
+                  RemoveIt, SortedInstructionsList[J].first.end());
+            }
+          }
+        }
+      }
+    }
   }
 
   Changed |= reuniteExts(F);
 
   if (VerifyNoDeadCode)
     verifyNoDeadCode(F);
 
   return Changed;

 SeparateConstOffsetFromGEPPass::run(Function &F, FunctionAnalysisManager &AM) {
   auto *DT = &AM.getResult<DominatorTreeAnalysis>(F);
   auto *SE = &AM.getResult<ScalarEvolutionAnalysis>(F);
   auto *LI = &AM.getResult<LoopAnalysis>(F);
   auto *TLI = &AM.getResult<TargetLibraryAnalysis>(F);
   auto GetTTI = [&AM](Function &F) -> TargetTransformInfo & {
     return AM.getResult<TargetIRAnalysis>(F);
   };
-  SeparateConstOffsetFromGEP Impl(DT, SE, LI, TLI, GetTTI, LowerGEP);
+  SeparateConstOffsetFromGEP Impl(DT, SE, LI, TLI, GetTTI, LowerGEP,
+                                  CheckProfitability);
   if (!Impl.run(F))
     return PreservedAnalyses::all();
   PreservedAnalyses PA;
   PA.preserveSet<CFGAnalyses>();
   return PA;
 }

llvm/test/Transforms/SeparateConstOffsetFromGEP/RISCV/split-gep.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt < %s -mtriple=riscv64-unknown-elf -passes='separate-const-offset-from-gep,early-cse' \
+; RUN: opt < %s -mtriple=riscv64-unknown-elf -passes='separate-const-offset-from-gep<check-profit>,early-cse' \
 ; RUN:       -S | FileCheck %s
 
 ; Several tests for -separate-const-offset-from-gep. The transformation
 ; heavily relies on TargetTransformInfo, so we put these tests under
 ; target-specific folders.
 
+; Positive test
 ; Simple case when GEPs should be optimized.
 define i64 @test1(i64* %array, i64 %i, i64 %j)  {
 ; CHECK-LABEL: @test1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[ADD:%.*]] = add nsw i64 [[I:%.*]], 5
 ; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i64, i64* [[ARRAY:%.*]], i64 [[I]]
 ; CHECK-NEXT:    [[GEP4:%.*]] = getelementptr inbounds i64, i64* [[TMP0]], i64 5
 ; CHECK-NEXT:    store i64 [[J:%.*]], i64* [[GEP4]], align 4
entry:
   %gep2 = getelementptr inbounds i64, i64* %array, i64 %add2
   store i64 %j, i64* %gep2
   %add3 = add nsw i64 %i, 35
   %gep3 = getelementptr inbounds i64, i64* %array, i64 %add3
   store i64 %add, i64* %gep3
   ret i64 undef
 }
 
+; Positive test
 ; Optimize GEPs when there sext instructions are needed to cast index value to expected type.
 define i32 @test2(i32* %array, i32 %i, i32 %j) {
 ; CHECK-LABEL: @test2(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[I:%.*]], 5
 ; CHECK-NEXT:    [[TMP0:%.*]] = sext i32 [[I]] to i64
 ; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i32, i32* [[ARRAY:%.*]], i64 [[TMP0]]
 ; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr inbounds i32, i32* [[TMP1]], i64 5
entry:
   store i32 %j, i32* %gep5
   %add6 = add nsw i32 %i, 35
   %sext7 = sext i32 %add6 to i64
   %gep8 = getelementptr inbounds i32, i32* %array, i64 %sext7
   store i32 %add, i32* %gep8
   ret i32 undef
 }
 
+; Negative test
 ; No need to modify because all values are also used in other expressions.
 ; Modification doesn't decrease register pressure.
 define i32 @test3(i32* %array, i32 %i) {
 ; CHECK-LABEL: @test3(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[I:%.*]], 5
-; CHECK-NEXT:    [[TMP0:%.*]] = sext i32 [[I]] to i64
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i32, i32* [[ARRAY:%.*]], i64 [[TMP0]]
-; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr inbounds i32, i32* [[TMP1]], i64 5
-; CHECK-NEXT:    store i32 [[ADD]], i32* [[GEP2]], align 4
+; CHECK-NEXT:    [[SEXT:%.*]] = sext i32 [[ADD]] to i64
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, i32* [[ARRAY:%.*]], i64 [[SEXT]]
+; CHECK-NEXT:    store i32 [[ADD]], i32* [[GEP]], align 4
 ; CHECK-NEXT:    [[ADD3:%.*]] = add nsw i32 [[I]], 6
-; CHECK-NEXT:    [[GEP54:%.*]] = getelementptr inbounds i32, i32* [[TMP1]], i64 6
-; CHECK-NEXT:    store i32 [[ADD3]], i32* [[GEP54]], align 4
+; CHECK-NEXT:    [[SEXT4:%.*]] = sext i32 [[ADD3]] to i64
+; CHECK-NEXT:    [[GEP5:%.*]] = getelementptr inbounds i32, i32* [[ARRAY]], i64 [[SEXT4]]
+; CHECK-NEXT:    store i32 [[ADD3]], i32* [[GEP5]], align 4
 ; CHECK-NEXT:    [[ADD6:%.*]] = add nsw i32 [[I]], 35
-; CHECK-NEXT:    [[GEP86:%.*]] = getelementptr inbounds i32, i32* [[TMP1]], i64 35
-; CHECK-NEXT:    store i32 [[ADD6]], i32* [[GEP86]], align 4
+; CHECK-NEXT:    [[SEXT7:%.*]] = sext i32 [[ADD6]] to i64
+; CHECK-NEXT:    [[GEP8:%.*]] = getelementptr inbounds i32, i32* [[ARRAY]], i64 [[SEXT7]]
+; CHECK-NEXT:    store i32 [[ADD6]], i32* [[GEP8]], align 4
 ; CHECK-NEXT:    ret i32 undef
 ;
 entry:
   %add = add nsw i32 %i, 5
   %sext = sext i32 %add to i64
   %gep = getelementptr inbounds i32, i32* %array, i64 %sext
   store i32 %add, i32* %gep
   %add3 = add nsw i32 %i, 6
   %sext4 = sext i32 %add3 to i64
   %gep5 = getelementptr inbounds i32, i32* %array, i64 %sext4
   store i32 %add3, i32* %gep5
   %add6 = add nsw i32 %i, 35
   %sext7 = sext i32 %add6 to i64
   %gep8 = getelementptr inbounds i32, i32* %array, i64 %sext7
   store i32 %add6, i32* %gep8
   ret i32 undef
 }
 
+; Positive test
 ; Optimized GEPs for multidimensional array with same base
 define i32 @test4([50 x i32]* %array2, i32 %i) {
 ; CHECK-LABEL: @test4(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[I:%.*]], 5
 ; CHECK-NEXT:    [[TMP0:%.*]] = sext i32 [[I]] to i64
 ; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr [50 x i32], [50 x i32]* [[ARRAY2:%.*]], i64 [[TMP0]], i64 [[TMP0]]
 ; CHECK-NEXT:    [[GEP3:%.*]] = getelementptr inbounds i32, i32* [[TMP1]], i64 255
entry:
   store i32 %i, i32* %gep5
   %add6 = add nsw i32 %i, 35
   %sext7 = sext i32 %add6 to i64
   %gep8 = getelementptr inbounds [50 x i32], [50 x i32]* %array2, i64 %sext7, i64 %j
   store i32 %i, i32* %gep8
   ret i32 undef
 }
 
+; Negative test
 ; No need to optimize GEPs, because there is critical amount with non-constant offsets.
 define i64 @test6(i64* %array, i64 %i, i64 %j) {
 ; CHECK-LABEL: @test6(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[ADD:%.*]] = add nsw i64 [[I:%.*]], 5
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i64, i64* [[ARRAY:%.*]], i64 [[J:%.*]]
 ; CHECK-NEXT:    store i64 [[ADD]], i64* [[GEP]], align 4
-; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i64, i64* [[ARRAY]], i64 [[I]]
-; CHECK-NEXT:    [[GEP52:%.*]] = getelementptr inbounds i64, i64* [[TMP0]], i64 6
-; CHECK-NEXT:    store i64 [[I]], i64* [[GEP52]], align 4
-; CHECK-NEXT:    store i64 [[I]], i64* [[TMP0]], align 4
+; CHECK-NEXT:    [[ADD3:%.*]] = add nsw i64 [[I]], 6
+; CHECK-NEXT:    [[GEP5:%.*]] = getelementptr inbounds i64, i64* [[ARRAY]], i64 [[ADD3]]
+; CHECK-NEXT:    store i64 [[I]], i64* [[GEP5]], align 4
+; CHECK-NEXT:    [[GEP8:%.*]] = getelementptr inbounds i64, i64* [[ARRAY]], i64 [[I]]
+; CHECK-NEXT:    store i64 [[I]], i64* [[GEP8]], align 4
 ; CHECK-NEXT:    ret i64 undef
 ;
 entry:
   %add = add nsw i64 %i, 5
   %gep = getelementptr inbounds i64, i64* %array, i64 %j
   store i64 %add, i64* %gep
   %add3 = add nsw i64 %i, 6
   %gep5 = getelementptr inbounds i64, i64* %array, i64 %add3
   store i64 %i, i64* %gep5
   %add6 = add nsw i64 %i, 35
   %gep8 = getelementptr inbounds i64, i64* %array, i64 %i
   store i64 %i, i64* %gep8
   ret i64 undef
 }
 
+; Negative test
 ; No need to optimize GEPs, because the base variable is different.
 define i32 @test7(i32* %array, i32 %i, i32 %j, i32 %k) {
 ; CHECK-LABEL: @test7(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[I:%.*]], 5
-; CHECK-NEXT:    [[TMP0:%.*]] = sext i32 [[I]] to i64
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i32, i32* [[ARRAY:%.*]], i64 [[TMP0]]
-; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr inbounds i32, i32* [[TMP1]], i64 5
-; CHECK-NEXT:    store i32 [[ADD]], i32* [[GEP2]], align 4
-; CHECK-NEXT:    [[TMP2:%.*]] = sext i32 [[K:%.*]] to i64
-; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr i32, i32* [[ARRAY]], i64 [[TMP2]]
-; CHECK-NEXT:    [[GEP54:%.*]] = getelementptr inbounds i32, i32* [[TMP3]], i64 6
-; CHECK-NEXT:    store i32 [[I]], i32* [[GEP54]], align 4
-; CHECK-NEXT:    [[TMP4:%.*]] = sext i32 [[J:%.*]] to i64
-; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr i32, i32* [[ARRAY]], i64 [[TMP4]]
-; CHECK-NEXT:    [[GEP86:%.*]] = getelementptr inbounds i32, i32* [[TMP5]], i64 35
-; CHECK-NEXT:    store i32 [[I]], i32* [[GEP86]], align 4
+; CHECK-NEXT:    [[SEXT:%.*]] = sext i32 [[ADD]] to i64
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, i32* [[ARRAY:%.*]], i64 [[SEXT]]
+; CHECK-NEXT:    store i32 [[ADD]], i32* [[GEP]], align 4
+; CHECK-NEXT:    [[ADD3:%.*]] = add nsw i32 [[K:%.*]], 6
+; CHECK-NEXT:    [[SEXT4:%.*]] = sext i32 [[ADD3]] to i64
+; CHECK-NEXT:    [[GEP5:%.*]] = getelementptr inbounds i32, i32* [[ARRAY]], i64 [[SEXT4]]
+; CHECK-NEXT:    store i32 [[I]], i32* [[GEP5]], align 4
+; CHECK-NEXT:    [[ADD6:%.*]] = add nsw i32 [[J:%.*]], 35
+; CHECK-NEXT:    [[SEXT7:%.*]] = sext i32 [[ADD6]] to i64
+; CHECK-NEXT:    [[GEP8:%.*]] = getelementptr inbounds i32, i32* [[ARRAY]], i64 [[SEXT7]]
+; CHECK-NEXT:    store i32 [[I]], i32* [[GEP8]], align 4
 ; CHECK-NEXT:    ret i32 undef
 ;
 entry:
   %add = add nsw i32 %i, 5
   %sext = sext i32 %add to i64
   %gep = getelementptr inbounds i32, i32* %array, i64 %sext
   store i32 %add, i32* %gep
   %add3 = add nsw i32 %k, 6
   %sext4 = sext i32 %add3 to i64
   %gep5 = getelementptr inbounds i32, i32* %array, i64 %sext4
   store i32 %i, i32* %gep5
   %add6 = add nsw i32 %j, 35
   %sext7 = sext i32 %add6 to i64
   %gep8 = getelementptr inbounds i32, i32* %array, i64 %sext7
   store i32 %i, i32* %gep8
   ret i32 undef
 }
 
+; Negative test
 ; No need to optimize GEPs, because the base of GEP instructions is different.
 define i32 @test8(i32* %array, i32* %array2, i32* %array3, i32 %i) {
 ; CHECK-LABEL: @test8(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[I:%.*]], 5
-; CHECK-NEXT:    [[TMP0:%.*]] = sext i32 [[I]] to i64
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i32, i32* [[ARRAY:%.*]], i64 [[TMP0]]
-; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr inbounds i32, i32* [[TMP1]], i64 5
-; CHECK-NEXT:    store i32 [[ADD]], i32* [[GEP2]], align 4
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr i32, i32* [[ARRAY2:%.*]], i64 [[TMP0]]
-; CHECK-NEXT:    [[GEP54:%.*]] = getelementptr inbounds i32, i32* [[TMP2]], i64 6
-; CHECK-NEXT:    store i32 [[I]], i32* [[GEP54]], align 4
-; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr i32, i32* [[ARRAY3:%.*]], i64 [[TMP0]]
-; CHECK-NEXT:    [[GEP86:%.*]] = getelementptr inbounds i32, i32* [[TMP3]], i64 35
-; CHECK-NEXT:    store i32 [[I]], i32* [[GEP86]], align 4
+; CHECK-NEXT:    [[SEXT:%.*]] = sext i32 [[ADD]] to i64
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, i32* [[ARRAY:%.*]], i64 [[SEXT]]
+; CHECK-NEXT:    store i32 [[ADD]], i32* [[GEP]], align 4
+; CHECK-NEXT:    [[ADD3:%.*]] = add nsw i32 [[I]], 6
+; CHECK-NEXT:    [[SEXT4:%.*]] = sext i32 [[ADD3]] to i64
+; CHECK-NEXT:    [[GEP5:%.*]] = getelementptr inbounds i32, i32* [[ARRAY2:%.*]], i64 [[SEXT4]]
+; CHECK-NEXT:    store i32 [[I]], i32* [[GEP5]], align 4
+; CHECK-NEXT:    [[ADD6:%.*]] = add nsw i32 [[I]], 35
+; CHECK-NEXT:    [[SEXT7:%.*]] = sext i32 [[ADD6]] to i64
+; CHECK-NEXT:    [[GEP8:%.*]] = getelementptr inbounds i32, i32* [[ARRAY3:%.*]], i64 [[SEXT7]]
+; CHECK-NEXT:    store i32 [[I]], i32* [[GEP8]], align 4
 ; CHECK-NEXT:    ret i32 undef
 ;
 entry:
   %add = add nsw i32 %i, 5
   %sext = sext i32 %add to i64
   %gep = getelementptr inbounds i32, i32* %array, i64 %sext
   store i32 %add, i32* %gep
   %add3 = add nsw i32 %i, 6
   %sext4 = sext i32 %add3 to i64
   %gep5 = getelementptr inbounds i32, i32* %array2, i64 %sext4
   store i32 %i, i32* %gep5
   %add6 = add nsw i32 %i, 35
   %sext7 = sext i32 %add6 to i64
   %gep8 = getelementptr inbounds i32, i32* %array3, i64 %sext7
   store i32 %i, i32* %gep8
   ret i32 undef
 }
 
+; Negative test
 ; No need to optimize GEPs of multidimensional array, because the base of GEP instructions is different.
 define i32 @test9([50 x i32]* %array, i32 %i) {
 ; CHECK-LABEL: @test9(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[I:%.*]], 5
-; CHECK-NEXT:    [[TMP0:%.*]] = sext i32 [[I]] to i64
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr [50 x i32], [50 x i32]* [[ARRAY:%.*]], i64 0, i64 [[TMP0]]
-; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr inbounds i32, i32* [[TMP1]], i64 5
-; CHECK-NEXT:    store i32 [[ADD]], i32* [[GEP2]], align 4
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr [50 x i32], [50 x i32]* [[ARRAY]], i64 [[TMP0]], i64 [[TMP0]]
-; CHECK-NEXT:    [[GEP54:%.*]] = getelementptr inbounds i32, i32* [[TMP2]], i64 6
-; CHECK-NEXT:    store i32 [[I]], i32* [[GEP54]], align 4
-; CHECK-NEXT:    [[GEP87:%.*]] = getelementptr inbounds i32, i32* [[TMP2]], i64 335
-; CHECK-NEXT:    store i32 [[I]], i32* [[GEP87]], align 4
+; CHECK-NEXT:    [[SEXT:%.*]] = sext i32 [[ADD]] to i64
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [50 x i32], [50 x i32]* [[ARRAY:%.*]], i64 0, i64 [[SEXT]]
+; CHECK-NEXT:    store i32 [[ADD]], i32* [[GEP]], align 4
+; CHECK-NEXT:    [[ADD3:%.*]] = add nsw i32 [[I]], 6
+; CHECK-NEXT:    [[SEXT4:%.*]] = sext i32 [[ADD3]] to i64
+; CHECK-NEXT:    [[INT:%.*]] = sext i32 [[I]] to i64
+; CHECK-NEXT:    [[GEP5:%.*]] = getelementptr inbounds [50 x i32], [50 x i32]* [[ARRAY]], i64 [[INT]], i64 [[SEXT4]]
+; CHECK-NEXT:    store i32 [[I]], i32* [[GEP5]], align 4
+; CHECK-NEXT:    [[ADD6:%.*]] = add nsw i32 [[I]], 35
+; CHECK-NEXT:    [[SEXT7:%.*]] = sext i32 [[ADD6]] to i64
+; CHECK-NEXT:    [[GEP8:%.*]] = getelementptr inbounds [50 x i32], [50 x i32]* [[ARRAY]], i64 [[SEXT4]], i64 [[SEXT7]]
+; CHECK-NEXT:    store i32 [[I]], i32* [[GEP8]], align 4
 ; CHECK-NEXT:    ret i32 undef
 ;
 entry:
   %add = add nsw i32 %i, 5
   %sext = sext i32 %add to i64
   %gep = getelementptr inbounds [50 x i32], [50 x i32]* %array, i64 0, i64 %sext
   store i32 %add, i32* %gep
   %add3 = add nsw i32 %i, 6