[LV] widenIntOrFpInduction. NFC.

[SjoerdMeijer]

Summary

This untangles some of the logic in widenIntOrFpInduction, which was quite necessary IMHO as the different cases was extremely difficult to follow and read. So I've removed/replaced this with more straigth-line code, which I would like to do first before making some other functional changes. Most of the times I directly commit NFC patches, but here I would like to get feedback as this is a little bit of a rewrite.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1812

Instead of implicitly setting Step in the lambdas, would it be possible to have the closures return the step/IV they create and take it as argument if required?

1826–1827

Comment seems out of place now.

1846–1847

The comment seems a bit out of place now, should it be around line 1910?

1847–1860

it might be good to make the VF, Part and Step explicit parameters here (and potentially other places) to make things a bit more explicit at the call sites.

[samparker]

+1 to Florian's suggestions for lambda parameters and some explicit data flow.

[SjoerdMeijer]

Thanks for taking a look @fhahn ! Comments addressed.

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1816

Seems like it may be good to pass in ID.getStep() as an explicit parameter, as it forms the basis of the step being created.

1823

https://llvm.org/docs/CodingStandards.html#don-t-use-else-after-a-return

1827

+1 about the first part of the original comment "If we haven't yet vectorized the induction variable..." seeming out of place here now.
But the remaining part that explains what this code is doing, would help here to describe what this lambda is for, including that comment about truncation.
Possibly also resurrecting the original comment about ScalarIV (changing "will be" to "is"):

// The scalar value to broadcast. This will be derived from the canonical
// induction variable.

1846–1847

+1
Some explanation what this lambda is for may be good here though. It basically creates the vector values from the scalar IV, in the absence of creating a vector IV.

1859

Is CreateScalarIVCode() really needed, instead of invoking

Value *ScalarIV = CreateScalarIV(Step);
CreateSplatIV(ScalarIV, Step);

directly, as done below before calling buildScalarSteps()?
(They are not folded together because CreateScalarIV() may change Step, right?)

1883

Can continue early-exiting by doing next

if (!NeedsScalarIV) {
  createVectorIntOrFpInductionPHI(ID, Step, EntryVal);
  return;
}

cleaning up the checks for NeedsScalarIV below.

[SjoerdMeijer]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1859

Ha, that's funny: no, we don't need it. Probably I stared too long at this code to notice this duplication and it is indeed the same as:

Value *ScalarIV = CreateScalarIV(Step);
CreateSplatIV(ScalarIV, Step);

which was actually also my intention. You're exactly right about:

They are not folded together because CreateScalarIV() may change Step, right?)

I wanted to have a one-liner:

CreateSplatIV(CreateScalarIV(Step), Step);

but that's indeed not possible because Step may change.

1883

Ah, many thanks. That cleans up things even more and this function is becoming a real beauty with just straight-line code.

[SjoerdMeijer]

Many thanks for reviewing!
I think this has all comments addressed.

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1821

There's one more ID.getStep() >> Step

1910

Not sure where this first "If an induction ... doesn't need to be widened" sentence belongs now, but it doesn't fit here.
The rest does fit before calling buildScalarStep(), which is also/firstly invoked above... better place it there?

[SjoerdMeijer]

• replaced ID.getStep() -> Step
• removed the sentence on line 1899 - 1900, and
• moved the comments on lines 1900 - 1904 to 1889.

[Ayal]

This looks good to me, thanks. The title should be more descriptive, e.g., "[LV] Refactor widenIntOrFpInduction, NFC".
Please wait a day or so to see if @fhahn or @samparker have further comments.

[SjoerdMeijer]

Will do, and many thanks for all your reviews!

And just fyi, I've uploaded work-in-progress patch D76838 that builds on top of this refactoring.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

assert((IV->getType()->isIntegerTy() || IV != OldInduction) &&
          "Primary induction variable must have an integer type");
 
   auto II = Legal->getInductionVars().find(IV);
   assert(II != Legal->getInductionVars().end() && "IV is not an induction");
 
   auto ID = II->second;
   assert(IV->getType() == ID.getStartValue()->getType() && "Types must match");
 
-  // The scalar value to broadcast. This will be derived from the canonical
-  // induction variable.
-  Value *ScalarIV = nullptr;
-
   // The value from the original loop to which we are mapping the new induction
   // variable.
   Instruction *EntryVal = Trunc ? cast<Instruction>(Trunc) : IV;
 
-  // True if we have vectorized the induction variable.
-  auto VectorizedIV = false;
-
-  // Determine if we want a scalar version of the induction variable. This is
-  // true if the induction variable itself is not widened, or if it has at
-  // least one user in the loop that is not widened.
-  auto NeedsScalarIV = VF > 1 && needsScalarInduction(EntryVal);
+  auto &DL = OrigLoop->getHeader()->getModule()->getDataLayout();
 
   // Generate code for the induction step. Note that induction steps are
   // required to be loop-invariant
+  auto CreateStepValue = [&]() -> Value * {
     assert(PSE.getSE()->isLoopInvariant(ID.getStep(), OrigLoop) &&
            "Induction step should be loop invariant");
-  auto &DL = OrigLoop->getHeader()->getModule()->getDataLayout();
-  Value *Step = nullptr;
     if (PSE.getSE()->isSCEVable(IV->getType())) {
       SCEVExpander Exp(*PSE.getSE(), DL, "induction");
-    Step = Exp.expandCodeFor(ID.getStep(), ID.getStep()->getType(),
+      return Exp.expandCodeFor(ID.getStep(), ID.getStep()->getType(),
                                LoopVectorPreHeader->getTerminator());
-  } else {
-    Step = cast<SCEVUnknown>(ID.getStep())->getValue();
-  }
-
-  // Try to create a new independent vector induction variable. If we can't
-  // create the phi node, we will splat the scalar induction variable in each
-  // loop iteration.
-  if (VF > 1 && !shouldScalarizeInstruction(EntryVal)) {
-    createVectorIntOrFpInductionPHI(ID, Step, EntryVal);
-    VectorizedIV = true;
-  }
+    } else
+      return cast<SCEVUnknown>(ID.getStep())->getValue();
+  };
 
-  // If we haven't yet vectorized the induction variable, or if we will create
-  // a scalar one, we need to define the scalar induction variable and step
-  // values. If we were given a truncation type, truncate the canonical
-  // induction variable and step. Otherwise, derive these values from the
-  // induction descriptor.
-  if (!VectorizedIV || NeedsScalarIV) {
-    ScalarIV = Induction;
+  auto CreateScalarIV = [&](Value *&Step) -> Value * {
+    Value *ScalarIV = Induction;
     if (IV != OldInduction) {
       ScalarIV = IV->getType()->isIntegerTy()
                      ? Builder.CreateSExtOrTrunc(Induction, IV->getType())
                      : Builder.CreateCast(Instruction::SIToFP, Induction,
                                           IV->getType());
       ScalarIV = emitTransformedIndex(Builder, ScalarIV, PSE.getSE(), DL, ID);
       ScalarIV->setName("offset.idx");
     }
     if (Trunc) {
       auto *TruncType = cast<IntegerType>(Trunc->getType());
       assert(Step->getType()->isIntegerTy() &&
              "Truncation requires an integer step");
       ScalarIV = Builder.CreateTrunc(ScalarIV, TruncType);
       Step = Builder.CreateTrunc(Step, TruncType);
     }
+    return ScalarIV;
+  };
+
+  auto CreateSplatIV = [&](Value *ScalarIV, Value *Step) {
+    Value *Broadcasted = getBroadcastInstrs(ScalarIV);
+    for (unsigned Part = 0; Part < UF; ++Part) {
+      Value *EntryPart =
+          getStepVector(Broadcasted, VF * Part, Step, ID.getInductionOpcode());
+      VectorLoopValueMap.setVectorValue(EntryVal, Part, EntryPart);
+      if (Trunc)
+        addMetadata(EntryPart, Trunc);
+      recordVectorLoopValueForInductionCast(ID, EntryVal, EntryPart, Part);
     }
+  };
 
-  // If we haven't yet vectorized the induction variable, splat the scalar
-  // induction variable, and build the necessary step vectors.
-  // TODO: Don't do it unless the vectorized IV is really required.
-  if (!VectorizedIV) {
+  auto CreateScalarIVCode = [&](Value *Step) {
+    Value *ScalarIV = CreateScalarIV(Step);
     Value *Broadcasted = getBroadcastInstrs(ScalarIV);
     for (unsigned Part = 0; Part < UF; ++Part) {
       Value *EntryPart =
           getStepVector(Broadcasted, VF * Part, Step, ID.getInductionOpcode());
       VectorLoopValueMap.setVectorValue(EntryVal, Part, EntryPart);
       if (Trunc)
         addMetadata(EntryPart, Trunc);
       recordVectorLoopValueForInductionCast(ID, EntryVal, EntryPart, Part);
     }
+  };
+
+  // Now do the actual transformations, and start with creating the step value.
+  Value *Step = CreateStepValue();
+  if (VF <= 1) {
+    CreateScalarIVCode(Step);
+    return;
+  }
+
+  // Determine if we want a scalar version of the induction variable. This is
+  // true if the induction variable itself is not widened, or if it has at
+  // least one user in the loop that is not widened.
+  auto NeedsScalarIV = needsScalarInduction(EntryVal);
+
+  // Try to create a new independent vector induction variable. If we can't
+  // create the phi node, we will splat the scalar induction variable in each
+  // loop iteration.
+  if (!shouldScalarizeInstruction(EntryVal)) {
+    createVectorIntOrFpInductionPHI(ID, Step, EntryVal);
+    // If we will create a scalar IV, we need to define the scalar induction
+    // variable and step values. If we were given a truncation type, truncate
+    // the canonical induction variable and step. Otherwise, derive these
+    // values from the induction descriptor.
+    if (NeedsScalarIV) {
+      Value *ScalarIV = CreateScalarIV(Step);
+      buildScalarSteps(ScalarIV, Step, EntryVal, ID);
+    }
+    return;
   }
 
+  // We are done here if we don't need a scalar IV.
+  if (!NeedsScalarIV)
+    return;
+
+  // If we haven't yet vectorized the induction variable, splat the scalar
+  // induction variable, and build the necessary step vectors.
+  // TODO: Don't do it unless the vectorized IV is really required.
+  Value *ScalarIV = CreateScalarIV(Step);
+  CreateSplatIV(ScalarIV, Step);
+
   // If an induction variable is only used for counting loop iterations or
   // calculating addresses, it doesn't need to be widened. Create scalar steps
   // that can be used by instructions we will later scalarize. Note that the
   // addition of the scalar steps will not increase the number of instructions
   // in the loop in the common case prior to InstCombine. We will be trading
   // one vector extract for each scalar step.
-  if (NeedsScalarIV)
-    buildScalarSteps(ScalarIV, Step, EntryVal, ID);
 }
 
 Value *InnerLoopVectorizer::getStepVector(Value *Val, int StartIdx, Value *Step,
                                           Instruction::BinaryOps BinOp) {
   // Create and check the types.
   assert(Val->getType()->isVectorTy() && "Must be a vector");
   int VLen = Val->getType()->getVectorNumElements();