[LV] Update logic for calculating register usage due to invariants

[sushgokh]

Summary

Previously, while calculating register usage due to invariants, it was assumed that invariant would always be part of widening instructions. This resulted in calculating vector register types for vectors which cant be legalized. e.g. For the attached SVE test case, it resulted in calculating register usage for type nxv4i128 resulting in legalization error(and crash).

An invariant might not always need a vector register. For e.g., invariant might just be used for iteration check.

This patch checks if the invariant is part of any widening instruction and considers register usage accordingly. Fixes issue 60493

[madhur13490]

I see the problem of calculating incorrect register usage. However, is it just limited to loop invariants?

[sushgokh]

In D143422#4108881, @madhur13490 wrote:

I see the problem of calculating incorrect register usage. However, is it just limited to loop invariants?

Yes

[sushgokh]

One thing I missed out is a check if the invariant is taking part in vector ops. If its indeed taking part in vector ops, we will need vector registers for it. But, if it doesnt, then scalar registers should be sufficient.

I need to rework on this patch

[sushgokh]

Hi,
Could someone please have a look at the patch?

[sushgokh]

any suggestions/comments on the patch ?

[sdesmalen]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6126

It seems that Inst can only ever be an Instruction, so this can be a cast. It would be even better to make LoopInvariants be a SmallVector of Instruction rather than Value (as a NFC precursory patch) , given that the loop that populates it skips non-Instruction values.

6128

If User is not a Instruction, the dyn_cast will result in a nullptr and the call to getParent() will fail.

6151

I think you can write the algorithm above a bit shorter using a lambda:

bool IsScalar = all_of(cast<Instruction>(Inst)->users(), [&](User *U) {
  auto *I = dyn_cast<Instruction>(U);
  return !I || TheLoop != LI->getLoopFor(I->getParent()) ||
         isScalarAfterVectorization(I, VFs[i]);
});      

ElementCount VF = IsScalar ? ElementCount::getFixed(1) : VFs[i];
unsigned  ClassID = TTI.getRegisterClassForType(VF.isVector(), Inst->getType());
Invariant[ClassID] += GetRegusage(Inst->getType(), VF);;

[sdesmalen]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6151

I think you can write the algorithm above a bit shorter using a lambda:

bool IsScalar = all_of(cast<Instruction>(Inst)->users(), [&](User *U) {
  auto *I = dyn_cast<Instruction>(U);
  return !I || TheLoop != LI->getLoopFor(I->getParent()) ||
         isScalarAfterVectorization(I, VFs[i]);
});      

ElementCount VF = IsScalar ? ElementCount::getFixed(1) : VFs[i];
unsigned  ClassID = TTI.getRegisterClassForType(VF.isVector(), Inst->getType());
Invariant[ClassID] += GetRegusage(Inst->getType(), VF);;

@sushgokh did you try the above suggestion?

When I try it locally, all the tests in your patch still pass.

[sushgokh]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6151

@sdesmalen Yes, I tried your suggestion and all the tests passed. But,my only intention was not to assume anything(as I have stated in the comment). But if you advise, will go with your suggestion

[sdesmalen]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6151

That case should already be covered because in getMaxLegalScalableVF it goes through all instructions in the loop to see if the element type is legal for scalable vectors.
So if there are any uses of this scalar value that are non-scalar, then the code in getMaxLegalScalableVF should have prevented the LV from considering a scalable VF for this loop.

Based on that, I believe my suggestion is safe.

[sushgokh]

@sdesmalen Have updated the patch. Thanks for the explanation

[sdesmalen]

Thanks for making the changes. I've just left a few more comments on the test.

llvm/test/Transforms/LoopVectorize/AArch64/reg-usage.ll

3

nit: ; Run:

3

If the test uses -debug-only, then it also needs REQUIRES: asserts

5

Can you add a more elaborate description here of what it's testing?

8

nit: Can you remove the dso_local and the local_unnamed_addr, as they're not needed.

29

Can this test be simplified a bit further? A lot of the instructions here (e.g. some the PHIs in the loop, the load %0 and subsequent GEP, etc) are not necessary for this test.

[sushgokh]

@sdesmalen I couldnt simplify the crashing test further. I hope this suffices

[fhahn]

LGTM, thanks!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6122–6123

nit: LoopInvariants stores Instruction* ,so there should be no need to use cast.

6123

nit: should be able to use cast, as all users of instructions should be other instructions. Then there's also no need to check !I in the lambda.

llvm/test/Transforms/LoopVectorize/AArch64/reg-usage.ll

2

this is outdated?

42

nit: change the function to return void?

[paulkirth]

llvm/test/Transforms/LoopVectorize/AArch64/reg-usage.ll

21–22

I think this test fails when building w/ -LLVM_REVERSE_ITERATION=On. Would you mind addressing this?

For context, I was testing something else and ran into this failure. I bisected to this commit w/ reverse iteration enabled.

[sushgokh]

@paulkirth Trying to address the issue you have mentioned in D145472. Thanks

[paulkirth]

Thanks for the fix. That change looks good to me. Its unfortunate that we can express that ordering is unimportant in lit for some range of code. The only reason I found the issue here was that I had a similar problem when testing stack layout analysis, and never updated my build flags.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

LoopVectorizationCostModel::calculateRegisterUsage(ArrayRef<ElementCount> VFs) {
   SmallVector<Instruction *, 64> IdxToInstr;
   // Marks the end of each interval.
   IntervalMap EndPoint;
   // Saves the list of instruction indices that are used in the loop.
   SmallPtrSet<Instruction *, 8> Ends;
   // Saves the list of values that are used in the loop but are defined outside
   // the loop (not including non-instruction values such as arguments and
   // constants).
-  SmallPtrSet<Value *, 8> LoopInvariants;
+  SmallPtrSet<Instruction *, 8> LoopInvariants;
 
   for (BasicBlock *BB : make_range(DFS.beginRPO(), DFS.endRPO())) {
     for (Instruction &I : BB->instructionsWithoutDebug()) {
       IdxToInstr.push_back(&I);
 
       // Save the end location of each USE.
       for (Value *U : I.operands()) {
         auto *Instr = dyn_cast<Instruction>(U);
LoopVectorizationCostModel::calculateRegisterUsage(ArrayRef<ElementCount> VFs) {
   for (unsigned i = 0, e = VFs.size(); i < e; ++i) {
     // Note that elements in this SmallMapVector will be default constructed
     // as 0. So we can use "Invariant[ClassID] += n" in the code below even if
     // there is no previous entry for ClassID.
     SmallMapVector<unsigned, unsigned, 4> Invariant;
 
     for (auto *Inst : LoopInvariants) {
       // FIXME: The target might use more than one register for the type
       // even in the scalar case.
-      unsigned Usage =
-          VFs[i].isScalar() ? 1 : GetRegUsage(Inst->getType(), VFs[i]);
-      unsigned ClassID =
+      SmallVector<User *, 6> Users;
+      // get the users of Inst which form part of the loop
+      std::copy_if(Inst->user_begin(), Inst->user_end(),
+                   std::back_inserter(Users), [&](User *U) {
+                     Instruction *I = dyn_cast<Instruction>(U);
+                     return I && TheLoop == LI->getLoopFor(I->getParent());
+                   });
+      // We cant assume that if Inst user is scalar, there wont be any
+      // non-scalar user of the same Inst
+      auto nScalarInst =
+          std::count_if(Users.begin(), Users.end(), [&](User *U) {
+            Instruction *I = cast<Instruction>(U);
+            return isScalarAfterVectorization(I, VFs[i]);
+          });
+      auto nNonScalarInst = Users.size() - nScalarInst;
+      unsigned Usage, ClassID;
+      if (nScalarInst) {
+        Usage = GetRegUsage(Inst->getType(), ElementCount::getFixed(1));
+        ClassID = TTI.getRegisterClassForType(false, Inst->getType());
+        Invariant[ClassID] += Usage;
+      }
+
+      if (nNonScalarInst) {
+        Usage = GetRegUsage(Inst->getType(), VFs[i]);
+        ClassID =
             TTI.getRegisterClassForType(VFs[i].isVector(), Inst->getType());
         Invariant[ClassID] += Usage;
       }
+    }
 
     LLVM_DEBUG({
       dbgs() << "LV(REG): VF = " << VFs[i] << '\n';
       dbgs() << "LV(REG): Found max usage: " << MaxUsages[i].size()
              << " item\n";
       for (const auto &pair : MaxUsages[i]) {
         dbgs() << "LV(REG): RegisterClass: "
                << TTI.getRegisterClassName(pair.first) << ", " << pair.second

llvm/test/Transforms/LoopVectorize/AArch64/reg-usage.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+;RUN: opt -mtriple arm64-linux -passes=loop-vectorize -mattr=+sve -debug-only=loop-vectorize -disable-output <%s 2>&1 | FileCheck %s
+
+; Check that this doesn't crash while finding register usage
+
+@string = internal unnamed_addr constant [5 x i8] c"abcd\00", align 1
+define dso_local i32 @get_invariant_reg_usage(ptr %z) local_unnamed_addr {
+;CHECK: LV: Checking a loop in 'get_invariant_reg_usage'
+;CHECK: LV(REG): VF = vscale x 16
+;CHECK: LV(REG): Found max usage: 1 item
+;CHECK: LV(REG): RegisterClass: Generic::ScalarRC, 5 registers
+;CHECK: LV(REG): Found invariant usage: 1 item
+;CHECK: LV(REG): RegisterClass: Generic::ScalarRC, 3 registers
+
+L.entry:
+  %0 = load ptr, ptr %z, align 8
+  %1 = getelementptr i8, ptr %0, i64 16
+  %2 = load i128, ptr %1, align 16
+  %3 = icmp slt i128 %2, 1
+  br i1 %3, label %return, label %loopbody.preheader
+
+loopbody.preheader:        ;preds = %L.entry
+  %4 = load ptr, ptr %0, align 8
+  br label %loopbody
+
+loopbody:                  ;preds = %loopbody.preheader, %loopbody
+  %a = phi ptr [ %5, %loopbody ], [ %4, %loopbody.preheader ]
+  %b = phi ptr [ %6, %loopbody ], [ @string, %loopbody.preheader ]
+  %len_input = phi i128 [ %len, %loopbody ], [ %2, %loopbody.preheader ]
+  %len = add nsw i128 %len_input, -1
+  %5 = getelementptr i8, ptr %a, i64 1
+  %6 = getelementptr i8, ptr %b, i64 1
+  %7 = load i8, ptr %b, align 1
+  store i8 %7, ptr %a, align 4
+  %.not = icmp eq i128 %len, 0
+  br i1 %.not, label %return, label %loopbody
+
+return:                    ;preds = %loopexit, %L.entry
+  ret i32 undef
+}

llvm/test/Transforms/LoopVectorize/PowerPC/reg-usage.ll

 
 define void @double_(ptr nocapture %A, i32 %n) nounwind uwtable ssp {
 ;CHECK-LABEL: double_
 ;CHECK-PWR8: LV(REG): VF = 2
 ;CHECK-PWR8: LV(REG): Found max usage: 2 item
 ;CHECK-PWR8-NEXT: LV(REG): RegisterClass: PPC::GPRRC, 2 registers
 ;CHECK-PWR8-NEXT: LV(REG): RegisterClass: PPC::VSXRC, 5 registers
 ;CHECK-PWR8: LV(REG): Found invariant usage: 1 item
-;CHECK-PWR8-NEXT: LV(REG): RegisterClass: PPC::VSXRC, 1 registers
+;CHECK-PWR8-NEXT: LV(REG): RegisterClass: PPC::GPRRC, 1 registers
 
 ;CHECK-PWR9: LV(REG): VF = 1
 ;CHECK-PWR9: LV(REG): Found max usage: 2 item
 ;CHECK-PWR9-NEXT: LV(REG): RegisterClass: PPC::GPRRC, 2 registers
 ;CHECK-PWR9-NEXT: LV(REG): RegisterClass: PPC::VSXRC, 5 registers
 ;CHECK-PWR9: LV(REG): Found invariant usage: 1 item
 ;CHECK-PWR9-NEXT: LV(REG): RegisterClass: PPC::GPRRC, 1 registers
 

llvm/test/Transforms/LoopVectorize/RISCV/reg-usage.ll

 ; CHECK-SCALAR-NEXT: LV(REG): RegisterClass: RISCV::GPRRC, 2 registers
 ; CHECK-SCALAR-NEXT: LV(REG): RegisterClass: RISCV::FPRRC, 2 registers
 ; CHECK-SCALAR-NEXT: LV(REG): Found invariant usage: 1 item
 ; CHECK-SCALAR-NEXT: LV(REG): RegisterClass: RISCV::GPRRC, 1 registers
 ; CHECK-LMUL1:       LV(REG): Found max usage: 2 item
 ; CHECK-LMUL1-NEXT:  LV(REG): RegisterClass: RISCV::GPRRC, 2 registers
 ; CHECK-LMUL1-NEXT:  LV(REG): RegisterClass: RISCV::VRRC, 2 registers
 ; CHECK-LMUL1-NEXT:  LV(REG): Found invariant usage: 1 item
-; CHECK-LMUL1-NEXT:  LV(REG): RegisterClass: RISCV::VRRC, 2 registers
+; CHECK-LMUL1-NEXT:  LV(REG): RegisterClass: RISCV::GPRRC, 1 registers
 ; CHECK-LMUL2:       LV(REG): Found max usage: 2 item
 ; CHECK-LMUL2-NEXT:  LV(REG): RegisterClass: RISCV::GPRRC, 2 registers
 ; CHECK-LMUL2-NEXT:  LV(REG): RegisterClass: RISCV::VRRC, 4 registers
 ; CHECK-LMUL2-NEXT:  LV(REG): Found invariant usage: 1 item
-; CHECK-LMUL2-NEXT:  LV(REG): RegisterClass: RISCV::VRRC, 4 registers
+; CHECK-LMUL2-NEXT:  LV(REG): RegisterClass: RISCV::GPRRC, 1 registers
 ; CHECK-LMUL4:       LV(REG): Found max usage: 2 item
 ; CHECK-LMUL4-NEXT:  LV(REG): RegisterClass: RISCV::GPRRC, 2 registers
 ; CHECK-LMUL4-NEXT:  LV(REG): RegisterClass: RISCV::VRRC, 8 registers
 ; CHECK-LMUL4-NEXT:  LV(REG): Found invariant usage: 1 item
-; CHECK-LMUL4-NEXT:  LV(REG): RegisterClass: RISCV::VRRC, 8 registers
+; CHECK-LMUL4-NEXT:  LV(REG): RegisterClass: RISCV::GPRRC, 1 registers
 ; CHECK-LMUL8:       LV(REG): Found max usage: 2 item
 ; CHECK-LMUL8-NEXT:  LV(REG): RegisterClass: RISCV::GPRRC, 2 registers
 ; CHECK-LMUL8-NEXT:  LV(REG): RegisterClass: RISCV::VRRC, 16 registers
 ; CHECK-LMUL8-NEXT:  LV(REG): Found invariant usage: 1 item
-; CHECK-LMUL8-NEXT:  LV(REG): RegisterClass: RISCV::VRRC, 16 registers
+; CHECK-LMUL8-NEXT:  LV(REG): RegisterClass: RISCV::GPRRC, 1 registers
 
 entry:
   %conv = zext i32 %size to i64
   %cmp10.not = icmp eq i32 %size, 0
   br i1 %cmp10.not, label %for.cond.cleanup, label %for.body
 
 for.cond.cleanup:
   ret void

llvm/test/Transforms/LoopVectorize/RISCV/riscv-vector-reverse.ll

 ; CHECK-NEXT:  LV(REG): At #7 Interval # 3
 ; CHECK-NEXT:  LV(REG): At #9 Interval # 1
 ; CHECK-NEXT:  LV(REG): At #10 Interval # 2
 ; CHECK-NEXT:  LV(REG): VF = vscale x 4
 ; CHECK-NEXT:  LV(REG): Found max usage: 2 item
 ; CHECK-NEXT:  LV(REG): RegisterClass: RISCV::GPRRC, 3 registers
 ; CHECK-NEXT:  LV(REG): RegisterClass: RISCV::VRRC, 2 registers
 ; CHECK-NEXT:  LV(REG): Found invariant usage: 1 item
-; CHECK-NEXT:  LV(REG): RegisterClass: RISCV::VRRC, 4 registers
+; CHECK-NEXT:  LV(REG): RegisterClass: RISCV::GPRRC, 1 registers
 ; CHECK-NEXT:  LV: The target has 31 registers of RISCV::GPRRC register class
 ; CHECK-NEXT:  LV: The target has 32 registers of RISCV::VRRC register class
 ; CHECK-NEXT:  LV: Loop cost is 23
 ; CHECK-NEXT:  LV: IC is 2
 ; CHECK-NEXT:  LV: VF is vscale x 4
 ; CHECK-NEXT:  LV: Not Interleaving.
 ; CHECK-NEXT:  LV: Interleaving is not beneficial.
 ; CHECK-NEXT:  LV: Found a vectorizable loop (vscale x 4) in <stdin>

 ; CHECK-NEXT:  LV(REG): At #7 Interval # 3
 ; CHECK-NEXT:  LV(REG): At #9 Interval # 1
 ; CHECK-NEXT:  LV(REG): At #10 Interval # 2
 ; CHECK-NEXT:  LV(REG): VF = vscale x 4
 ; CHECK-NEXT:  LV(REG): Found max usage: 2 item
 ; CHECK-NEXT:  LV(REG): RegisterClass: RISCV::GPRRC, 3 registers
 ; CHECK-NEXT:  LV(REG): RegisterClass: RISCV::VRRC, 2 registers
 ; CHECK-NEXT:  LV(REG): Found invariant usage: 1 item
-; CHECK-NEXT:  LV(REG): RegisterClass: RISCV::VRRC, 4 registers
+; CHECK-NEXT:  LV(REG): RegisterClass: RISCV::GPRRC, 1 registers
 ; CHECK-NEXT:  LV: The target has 31 registers of RISCV::GPRRC register class
 ; CHECK-NEXT:  LV: The target has 32 registers of RISCV::VRRC register class
 ; CHECK-NEXT:  LV: Loop cost is 23
 ; CHECK-NEXT:  LV: IC is 2
 ; CHECK-NEXT:  LV: VF is vscale x 4
 ; CHECK-NEXT:  LV: Not Interleaving.
 ; CHECK-NEXT:  LV: Interleaving is not beneficial.
 ; CHECK-NEXT:  LV: Found a vectorizable loop (vscale x 4) in <stdin>

llvm/test/Transforms/LoopVectorize/X86/reg-usage-debug.ll

 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 ; CHECK: LV: Checking a loop in 'test_g'
 ; CHECK: LV(REG): Found max usage: 2 item
 ; CHECK-NEXT: LV(REG): RegisterClass: Generic::ScalarRC, 2 registers
 ; CHECK-NEXT: LV(REG): RegisterClass: Generic::VectorRC, 2 registers
 ; CHECK-NEXT: LV(REG): Found invariant usage: 1 item
-; CHECK-NEXT: LV(REG): RegisterClass: Generic::VectorRC, 2 registers
+; CHECK-NEXT: LV(REG): RegisterClass: Generic::ScalarRC, 1 registers
 
 define i32 @test_g(ptr nocapture readonly %a, i32 %n) local_unnamed_addr !dbg !6 {
 entry:
   tail call void @llvm.dbg.value(metadata ptr %a, i64 0, metadata !12, metadata !16), !dbg !17
   tail call void @llvm.dbg.value(metadata i32 %n, i64 0, metadata !13, metadata !16), !dbg !18
   tail call void @llvm.dbg.value(metadata i32 0, i64 0, metadata !15, metadata !16), !dbg !19
   tail call void @llvm.dbg.value(metadata i32 0, i64 0, metadata !14, metadata !16), !dbg !20
   tail call void @llvm.dbg.value(metadata i32 0, i64 0, metadata !15, metadata !16), !dbg !19
for.end:                                          ; preds = %for.end.loopexit, %entry
   ret i32 %r.0.lcssa, !dbg !38
 }
 
 ; CHECK: LV: Checking a loop in 'test'
 ; CHECK: LV(REG): Found max usage: 2 item
 ; CHECK-NEXT: LV(REG): RegisterClass: Generic::ScalarRC, 2 registers
 ; CHECK-NEXT: LV(REG): RegisterClass: Generic::VectorRC, 2 registers
 ; CHECK-NEXT: LV(REG): Found invariant usage: 1 item
-; CHECK-NEXT: LV(REG): RegisterClass: Generic::VectorRC, 2 registers
+; CHECK-NEXT: LV(REG): RegisterClass: Generic::ScalarRC, 1 registers
 
 define i32 @test(ptr nocapture readonly %a, i32 %n) local_unnamed_addr {
 entry:
   %cmp6 = icmp eq i32 %n, 0
   br i1 %cmp6, label %for.end, label %for.body.preheader
 
 for.body.preheader:                               ; preds = %entry
   %wide.trip.count = zext i32 %n to i64