[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

6107–6108

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);;

6110

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.

6112

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

[sdesmalen]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6107–6108

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

6107–6108

@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

6107–6108

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

2 ↗

(On Diff #499724)

nit: ; Run:

2 ↗

(On Diff #499724)

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

4 ↗

(On Diff #499724)

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

7 ↗

(On Diff #499724)

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

28 ↗

(On Diff #499724)

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

6107–6108

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

6108

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

1 ↗

(On Diff #499921)

this is outdated?

41 ↗

(On Diff #499921)

nit: change the function to return void?

[paulkirth]

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

20–21 ↗

(On Diff #501024)

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

for (unsigned int i = 0, s = IdxToInstr.size(); i < s; ++i) {
 
     LLVM_DEBUG(dbgs() << "LV(REG): At #" << i << " Interval # "
                       << OpenIntervals.size() << '\n');
 
     // Add the current instruction to the list of open intervals.
     OpenIntervals.insert(I);
   }
 
-  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 =
-          TTI.getRegisterClassForType(VFs[i].isVector(), Inst->getType());
-      Invariant[ClassID] += Usage;
-    }
+    unsigned Usage = GetRegUsage(Inst->getType(), ElementCount::getFixed(1));
+    unsigned ClassID = TTI.getRegisterClassForType(false, Inst->getType());
 
-    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
-               << " registers\n";
-      }
-      dbgs() << "LV(REG): Found invariant usage: " << Invariant.size()
+    dbgs() << "LV(REG): Found invariant usage: " << Invariant.size()
+    for (const auto &pair : Invariant) {
+  });
+
+  RU.LoopInvariantRegs = Invariant;
+
+  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.
+
+    LLVM_DEBUG({
+      dbgs() << "LV(REG): VF = " << VFs[i] << '\n';
+      dbgs() << "LV(REG): Found max usage: " << MaxUsages[i].size()
              << " item\n";
-      for (const auto &pair : Invariant) {
+      for (const auto &pair : MaxUsages[i]) {
         dbgs() << "LV(REG): RegisterClass: "
                << TTI.getRegisterClassName(pair.first) << ", " << pair.second
                << " registers\n";
       }
     });
 
-    RU.LoopInvariantRegs = Invariant;
     RU.MaxLocalUsers = MaxUsages[i];
     RUs[i] = RU;
   }
 
   return RUs;
 }
 
 bool LoopVectorizationCostModel::useEmulatedMaskMemRefHack(Instruction *I,

llvm/test/Transforms/LoopVectorize/sve-invariant-legalization.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+;RUN: opt -mtriple arm64-linux -passes=loop-vectorize -mattr=+sve -disable-output %s | FileCheck %s --allow-empty
+
+@string = internal unnamed_addr constant [5 x i8] c"abcd\00", align 1
+define dso_local i32 @test(ptr %z) local_unnamed_addr {
+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
+}
+;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
+; CHECK: {{.*}}