Adding patch and unittest to generalize ignorable induction casts in the LoopVectorizationLegality analysis.

[avery-laird]

Summary

This patch assumes that casts of the induction phi, from the same type, can be treated as the induction variable for the purposes of the legality analysis. This helps to make the cost model more accurate in certain cases, particularily with getMemInstScalarizationCost.

[bmahjour]

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

524

Isn't the subject cast available in the list returned by getCastInsts()?

528

Checking the source type of the cast to match the phi type seems redundant given that we begin by identifying casts that are users of the phi.

529

The comment in LoopVectorizationLegality.h indicates that this list is supposed to contain casts that are noop (possibly guaranteed by some runtime check). IIUC your use case may not have that guarantee. Looking through the code though, it seems we only use this list in cost-estimation, so perhaps it's fine...but if that's the case we should update the comment.

[avery-laird]

• determine whether getCastInsts() can be modified to have the desired behaviour directly (and implement if so)
• remove redundant type check
• if necessary, update comment in LoopVectorizationLegality.h to reflect changes

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

524

right now it is not, and that is the main issue. Ideally, getCastInsts() would return these types of casts. I'm currently looking into this further to see if there is a way to instead modify getCastInsts(). One issue with this, as you identified, is that getCastInsts() has stricter guarantees than required by the cost model.

[avery-laird]

After investigating further into whether this patch can be merged with getCastInsts(), I think there are some difficulties in doing so.

getCastInsts() looks for a specific pattern that I'm not sure fits this case. For example (from comments in IVDescriptor.cpp), given an IR sequence like this:

%for.body:
   %x = phi i64 [ 0, %ph ], [ %add, %for.body ]
   %casted_phi = and  %x, 2^n-1
   %add = add i64 %casted_phi, %step

The cast returned by getCastInsts() will be %casted_phi. The search is done by looking at the operands of %add, in this case finding %casted_phi, and testing whether PSE.getSCEV(%x) is equal to PSE.getSCEV(%casted_phi) under a predicate. Three differences are:

1. All instructions in the def-use chain must be binary, but sext for example is unary
2. Instructions must update the incoming value of the phi node to be considered (in this case through %add's use of %casted_phi), but this patch should include instructions that don't update the induction Phi
3. Instructions need to be of the same type, but this patch should include instructions that may be casts from one type to another

In order to modify getCastInsts() to return the desired casts, at least these three things need to be changed. Additionally, by only changing InductionCastsToIgnore, just the cost model is affected, but getCastInsts() also directly affects code generation.

It seems the goal of getCastInsts() is generally different from what this patch is trying to do, but I am very interested in any feedback! Depending on that, I will update the patch with steps (2) and (3) identified in the plan.

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

bool LoopVectorizationLegality::canVectorizeOuterLoop() {
 
   return Result;
 }
 
 void LoopVectorizationLegality::addInductionPhi(
     PHINode *Phi, const InductionDescriptor &ID,
     SmallPtrSetImpl<Value *> &AllowedExit) {
   Inductions[Phi] = ID;
+  Type *PhiTy = Phi->getType();
 
   // In case this induction also comes with casts that we know we can ignore
   // in the vectorized loop body, record them here. All casts could be recorded
   // here for ignoring, but suffices to record only the first (as it is the
   // only one that may bw used outside the cast sequence).
   const SmallVectorImpl<Instruction *> &Casts = ID.getCastInsts();
   if (!Casts.empty())
     InductionCastsToIgnore.insert(*Casts.begin());
+  for (Value *O : cast<User>(Phi)->users())
+    if (auto *I = dyn_cast<Instruction>(O))
+      if (TheLoop->contains(I)) {
+        if (auto *Cast = dyn_cast<CastInst>(I))
+          if (Cast->getSrcTy() == PhiTy) {
+            InductionCastsToIgnore.insert(Cast);
+          }
+      }
 
-  Type *PhiTy = Phi->getType();
   const DataLayout &DL = Phi->getModule()->getDataLayout();
 
   // Get the widest type.
   if (!PhiTy->isFloatingPointTy()) {
     if (!WidestIndTy)
       WidestIndTy = convertPointerToIntegerType(DL, PhiTy);
     else
       WidestIndTy = getWiderType(DL, PhiTy, WidestIndTy);

llvm/unittests/Transforms/Vectorize/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Analysis
   Core
   Vectorize
   AsmParser
   )
 
 add_llvm_unittest(VectorizeTests
   VPlanDominatorTreeTest.cpp
   VPlanLoopInfoTest.cpp
   VPlanPredicatorTest.cpp
   VPlanTest.cpp
   VPlanHCFGTest.cpp
   VPlanSlpTest.cpp
+  LoopVectorizationLegalityTest.cpp
   )

llvm/unittests/Transforms/Vectorize/LoopVectorizationLegalityTest.cpp

+//===- LoopVectorizationLegality.cpp - Loop Legality unit tests -----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Vectorize/LoopVectorizationLegality.h"
+#include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/BasicAliasAnalysis.h"
+#include "llvm/Analysis/DemandedBits.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/ProfileSummaryInfo.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/VectorUtils.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+class LoopVectorizationLegalityTest : public testing::Test {
+protected:
+  TargetLibraryInfoImpl TLII;
+  TargetLibraryInfo TLI;
+  DataLayout DL;
+
+  std::unique_ptr<AAResults> AARes;
+  std::unique_ptr<BasicAAResult> BasicAA;
+  std::unique_ptr<LoopAccessInfo> LAI;
+  std::unique_ptr<PredicatedScalarEvolution> PSE;
+  std::unique_ptr<OptimizationRemarkEmitter> ORE;
+  std::unique_ptr<TargetTransformInfo> TTI;
+  std::unique_ptr<LoopVectorizationRequirements> R;
+  std::unique_ptr<LoopVectorizeHints> LVH;
+  std::unique_ptr<BlockFrequencyInfo> BFI;
+  std::unique_ptr<ProfileSummaryInfo> PSI;
+  std::unique_ptr<std::function<const LoopAccessInfo &(Loop &)>> GetLAA;
+  std::unique_ptr<InterleavedAccessInfo> IAI;
+  std::unique_ptr<DemandedBits> DB;
+  std::unique_ptr<LLVMContext> Ctx;
+  std::unique_ptr<Module> M;
+  std::unique_ptr<LoopInfo> LI;
+  std::unique_ptr<DominatorTree> DT;
+  std::unique_ptr<AssumptionCache> AC;
+  std::unique_ptr<ScalarEvolution> SE;
+
+  LoopVectorizationLegalityTest()
+      : TLII(), TLI(TLII),
+        DL("e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-"
+           "f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:"
+           "16:32:64-S128"),
+        Ctx(new LLVMContext) {}
+
+  LoopVectorizationLegality getLoopVectorizationLegality(Function *F) {
+    TTI.reset(new TargetTransformInfo(M->getDataLayout()));
+    ORE.reset(new OptimizationRemarkEmitter(F));
+    DT.reset(new DominatorTree(*F));
+    LI.reset(new LoopInfo(*DT));
+
+    BasicBlock *LoopHeader = F->getEntryBlock().getSingleSuccessor();
+    Loop *L = LI->getLoopFor(LoopHeader);
+
+    R.reset(new LoopVectorizationRequirements());
+    LVH.reset(new LoopVectorizeHints(L, false, *ORE));
+    BFI.reset(new BlockFrequencyInfo());
+    PSI.reset(new ProfileSummaryInfo(*M));
+
+    GetLAA.reset(new std::function<const LoopAccessInfo &(Loop &)>(
+        [&](Loop &L) -> const LoopAccessInfo & { return *LAI; }));
+
+    AC.reset(new AssumptionCache(*F));
+    DB.reset(new DemandedBits(*F, *AC, *DT));
+    SE.reset(new ScalarEvolution(*F, TLI, *AC, *DT, *LI));
+    BasicAA.reset(new BasicAAResult(DL, *F, TLI, *AC, &*DT));
+    AARes.reset(new AAResults(TLI));
+    AARes->addAAResult(*BasicAA);
+
+    PSE.reset(new PredicatedScalarEvolution(*SE, *L));
+    LAI.reset(new LoopAccessInfo(L, &*SE, &TLI, &*AARes, &*DT, &*LI));
+
+    LoopVectorizationLegality Legal(L, *PSE, &*DT, &*TTI, &TLI, &*AARes, F,
+                                    &*GetLAA, &*LI, &*ORE, &*R, &*LVH, &*DB,
+                                    &*AC, &*BFI, &*PSI);
+    return Legal;
+  }
+
+  Module &parseModule(const char *ModuleString) {
+    SMDiagnostic Err;
+    M = parseAssemblyString(ModuleString, Err, *Ctx);
+    EXPECT_TRUE(M);
+    return *M;
+  }
+};
+
+TEST_F(LoopVectorizationLegalityTest, IgnoreInductionCastTest) {
+  const char *ModuleString =
+      "@a = internal global [100 x double] zeroinitializer, align 16\n"
+      "@b = internal global [100 x double] zeroinitializer, align 16\n"
+      "define void @bar() {\n"
+      "entry:\n"
+      "  br label %for.body\n"
+      "for.body:\n"
+      "  %i.01 = phi i32 [ 0, %entry ], [ %inc, %for.inc ]\n"
+      "  %idxprom = sext i32 %i.01 to i64\n"
+      "  %arrayidx = getelementptr inbounds [100 x double], [100 x double]* "
+      "@a, i64 0, i64 %idxprom\n"
+      "  %0 = load double, double* %arrayidx, align 8\n"
+      "  %tobool = fcmp une double %0, 0.000000e+00\n"
+      "  br i1 %tobool, label %if.then, label %for.inc\n"
+      "if.then:\n"
+      "  %add = fadd double 2.000000e+00, 1.000000e+00\n"
+      "  %idxprom1 = sext i32 %i.01 to i64\n"
+      "  %arrayidx2 = getelementptr inbounds [100 x double], [100 x double]* "
+      "@b, i64 0, i64 %idxprom1\n"
+      "  store double %add, double* %arrayidx2, align 8\n"
+      "  br label %for.inc\n"
+      "for.inc:\n"
+      "  %inc = add nsw i32 %i.01, 1\n"
+      "  %cmp = icmp slt i32 %inc, 100\n"
+      "  br i1 %cmp, label %for.body, label %for.end\n"
+      "for.end:\n"
+      "  ret void\n"
+      "}";
+
+  parseModule(ModuleString);
+  Function *F = M->getFunction("bar");
+
+  LoopVectorizationLegality Legal = getLoopVectorizationLegality(F);
+  EXPECT_TRUE(Legal.canVectorize(false));
+
+  int count = 0;
+  bool found = false;
+  for (BasicBlock &BB : *F) {
+    for (Instruction &I : BB) {
+      Value *Ptr = getLoadStorePointerOperand(&I);
+      if (!Ptr)
+        continue;
+      // look for the 2nd memory op.
+      if (++count != 2)
+        continue;
+
+      auto *Gep = dyn_cast<GetElementPtrInst>(Ptr);
+      EXPECT_TRUE(Gep);
+
+      Value *Opd = Gep->getOperand(2);
+      EXPECT_TRUE(Opd);
+      EXPECT_TRUE(Legal.isInductionVariable(Opd));
+      found = true;
+      break;
+    }
+    if (found)
+      break;
+  }
+}
 No newline at end of file