Revert "[LoopVectorize] Simplify scalar cost calculation in getInstructionCost."

[fhahn]

Summary

This reverts 6998f8ae2d14e096aff33968f226587b5c1a193a.

Unfortunately another issue surfaced where the underlying assumption for
the patch does not hold and we hit the assertion for an Add instruction.

In pr55096-scalarize-add.ll, there's an add that will get scalarized (as
it is an induction update that has scalarized users) and there will be
multiple copies after vectorization.

The way we currently compute the costs for instructions involved in
induction updates is not really accurate, but I don't think there's an
easy way to ensure scalarized instructions will have single copy after
vectorization.

Fixes #55096.

[david-arm]

Hi @fhahn, I think I remember when doing this patch these asserts were put in to deal with reviewer requests, but in my original patch I didn't have any asserts at all. I worry that reverting the patch doesn't solve the underlying problem, which is that the cost calculations are sometimes incorrect and difficult to follow. As a compromise could we just remove the assert, but keep the simpler (and easier to understand) logic? This patch has been in the codebase for over a year now without problems - personally it feels more appropriate to fix the assert rather than revert. What do you think?

[mkazantsev]

Ping? :)

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

 }
 
 InstructionCost
 LoopVectorizationCostModel::getInstructionCost(Instruction *I, ElementCount VF,
                                                Type *&VectorTy) {
   Type *RetTy = I->getType();
   if (canTruncateToMinimalBitwidth(I, VF))
     RetTy = IntegerType::get(RetTy->getContext(), MinBWs[I]);
+  VectorTy = isScalarAfterVectorization(I, VF) ? RetTy : ToVectorTy(RetTy, VF);
   auto SE = PSE.getSE();
   TTI::TargetCostKind CostKind = TTI::TCK_RecipThroughput;
 
-  auto hasSingleCopyAfterVectorization = [this](Instruction *I,
-                                                ElementCount VF) -> bool {
-    if (VF.isScalar())
-      return true;
-
-    auto Scalarized = InstsToScalarize.find(VF);
-    assert(Scalarized != InstsToScalarize.end() &&
-           "VF not yet analyzed for scalarization profitability");
-    return !Scalarized->second.count(I) &&
-           llvm::all_of(I->users(), [&](User *U) {
-             auto *UI = cast<Instruction>(U);
-             return !Scalarized->second.count(UI);
-           });
-  };
-  (void) hasSingleCopyAfterVectorization;
-
-  if (isScalarAfterVectorization(I, VF)) {
-    // With the exception of GEPs and PHIs, after scalarization there should
-    // only be one copy of the instruction generated in the loop. This is
-    // because the VF is either 1, or any instructions that need scalarizing
-    // have already been dealt with by the the time we get here. As a result,
-    // it means we don't have to multiply the instruction cost by VF.
-    assert(I->getOpcode() == Instruction::GetElementPtr ||
-           I->getOpcode() == Instruction::PHI ||
-           (I->getOpcode() == Instruction::BitCast &&
-            I->getType()->isPointerTy()) ||
-           hasSingleCopyAfterVectorization(I, VF));
-    VectorTy = RetTy;
-  } else
-    VectorTy = ToVectorTy(RetTy, VF);
-
   // TODO: We need to estimate the cost of intrinsic calls.
   switch (I->getOpcode()) {
   case Instruction::GetElementPtr:
     // We mark this instruction as zero-cost because the cost of GEPs in
     // vectorized code depends on whether the corresponding memory instruction
     // is scalarized or not. Therefore, we handle GEPs with the memory
     // instruction cost.
     return 0;
case Instruction::Xor: {
     Value *Op2 = I->getOperand(1);
     TargetTransformInfo::OperandValueProperties Op2VP;
     TargetTransformInfo::OperandValueKind Op2VK =
         TTI.getOperandInfo(Op2, Op2VP);
     if (Op2VK == TargetTransformInfo::OK_AnyValue && Legal->isUniform(Op2))
       Op2VK = TargetTransformInfo::OK_UniformValue;
 
     SmallVector<const Value *, 4> Operands(I->operand_values());
-    return TTI.getArithmeticInstrCost(
-        I->getOpcode(), VectorTy, CostKind, TargetTransformInfo::OK_AnyValue,
-        Op2VK, TargetTransformInfo::OP_None, Op2VP, Operands, I);
+    unsigned N = isScalarAfterVectorization(I, VF) ? VF.getKnownMinValue() : 1;
+    return N * TTI.getArithmeticInstrCost(I->getOpcode(), VectorTy, CostKind,
+                                          TargetTransformInfo::OK_AnyValue,
+                                          Op2VK, TargetTransformInfo::OP_None,
+                                          Op2VP, Operands, I);
   }
   case Instruction::FNeg: {
-    return TTI.getArithmeticInstrCost(
-        I->getOpcode(), VectorTy, CostKind, TargetTransformInfo::OK_AnyValue,
-        TargetTransformInfo::OK_AnyValue, TargetTransformInfo::OP_None,
-        TargetTransformInfo::OP_None, I->getOperand(0), I);
+    unsigned N = isScalarAfterVectorization(I, VF) ? VF.getKnownMinValue() : 1;
+    return N * TTI.getArithmeticInstrCost(I->getOpcode(), VectorTy, CostKind,
+                                          TargetTransformInfo::OK_AnyValue,
+                                          TargetTransformInfo::OK_AnyValue,
+                                          TargetTransformInfo::OP_None,
+                                          TargetTransformInfo::OP_None,
+                                          I->getOperand(0), I);
   }
   case Instruction::Select: {
     SelectInst *SI = cast<SelectInst>(I);
     const SCEV *CondSCEV = SE->getSCEV(SI->getCondition());
     bool ScalarCond = (SE->isLoopInvariant(CondSCEV, TheLoop));
 
     const Value *Op0, *Op1;
     using namespace llvm::PatternMatch;
if (Width.isVector()) {
       assert(Decision != CM_Unknown &&
              "CM decision should be taken at this point");
       if (Decision == CM_Scalarize)
         Width = ElementCount::getFixed(1);
     }
     VectorTy = ToVectorTy(getLoadStoreType(I), Width);
     return getMemoryInstructionCost(I, VF);
   }
-  case Instruction::BitCast:
-    if (I->getType()->isPointerTy())
-      return 0;
-    LLVM_FALLTHROUGH;
   case Instruction::ZExt:
   case Instruction::SExt:
   case Instruction::FPToUI:
   case Instruction::FPToSI:
   case Instruction::FPExt:
   case Instruction::PtrToInt:
   case Instruction::IntToPtr:
   case Instruction::SIToFP:
   case Instruction::UIToFP:
   case Instruction::Trunc:
-  case Instruction::FPTrunc: {
+  case Instruction::FPTrunc:
+  case Instruction::BitCast: {
     // Computes the CastContextHint from a Load/Store instruction.
     auto ComputeCCH = [&](Instruction *I) -> TTI::CastContextHint {
       assert((isa<LoadInst>(I) || isa<StoreInst>(I)) &&
              "Expected a load or a store!");
 
       if (VF.isScalar() || !TheLoop->contains(I))
         return TTI::CastContextHint::Normal;
 
if (canTruncateToMinimalBitwidth(I, VF)) {
             largestIntegerVectorType(ToVectorTy(I->getType(), VF), MinVecTy);
       } else if (Opcode == Instruction::ZExt || Opcode == Instruction::SExt) {
         SrcVecTy = largestIntegerVectorType(SrcVecTy, MinVecTy);
         VectorTy =
             smallestIntegerVectorType(ToVectorTy(I->getType(), VF), MinVecTy);
       }
     }
 
-    return TTI.getCastInstrCost(Opcode, VectorTy, SrcVecTy, CCH, CostKind, I);
+    unsigned N;
+    if (isScalarAfterVectorization(I, VF)) {
+      assert(!VF.isScalable() && "VF is assumed to be non scalable");
+      N = VF.getKnownMinValue();
+    } else
+      N = 1;
+    return N *
+           TTI.getCastInstrCost(Opcode, VectorTy, SrcVecTy, CCH, CostKind, I);
   }
   case Instruction::Call: {
     if (RecurrenceDescriptor::isFMulAddIntrinsic(I))
       if (auto RedCost = getReductionPatternCost(I, VF, VectorTy, CostKind))
         return *RedCost;
     bool NeedToScalarize;
     CallInst *CI = cast<CallInst>(I);
     InstructionCost CallCost = getVectorCallCost(CI, VF, NeedToScalarize);
     if (getVectorIntrinsicIDForCall(CI, TLI)) {
       InstructionCost IntrinsicCost = getVectorIntrinsicCost(CI, VF);
       return std::min(CallCost, IntrinsicCost);
     }
     return CallCost;
   }
   case Instruction::ExtractValue:
     return TTI.getInstructionCost(I, TTI::TCK_RecipThroughput);
   case Instruction::Alloca:
     // We cannot easily widen alloca to a scalable alloca, as
     // the result would need to be a vector of pointers.
     if (VF.isScalable())
       return InstructionCost::getInvalid();
     LLVM_FALLTHROUGH;
   default:
-    // This opcode is unknown. Assume that it is the same as 'mul'.
-    return TTI.getArithmeticInstrCost(Instruction::Mul, VectorTy, CostKind);
+    // The cost of executing VF copies of the scalar instruction. This opcode
+    // is unknown. Assume that it is the same as 'mul'.
+    return VF.getKnownMinValue() * TTI.getArithmeticInstrCost(
+                                       Instruction::Mul, VectorTy, CostKind) +
+           getScalarizationOverhead(I, VF);
   } // end of switch.
 }
 
 char LoopVectorize::ID = 0;
 
 static const char lv_name[] = "Loop Vectorization";
 
 INITIALIZE_PASS_BEGIN(LoopVectorize, LV_NAME, lv_name, false, false)

llvm/test/Transforms/LoopVectorize/AArch64/no_vector_instructions.ll

 ; REQUIRES: asserts
 ; RUN: opt < %s -loop-vectorize -force-vector-interleave=1 -S -debug-only=loop-vectorize 2>&1 | FileCheck %s
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64--linux-gnu"
 
 ; CHECK-LABEL: all_scalar
 ; CHECK:       LV: Found scalar instruction: %i.next = add nuw nsw i64 %i, 2
-; CHECK:       LV: Found an estimated cost of 1 for VF 2 For instruction: %i.next = add nuw nsw i64 %i, 2
+; CHECK:       LV: Found an estimated cost of 2 for VF 2 For instruction: %i.next = add nuw nsw i64 %i, 2
 ; CHECK:       LV: Not considering vector loop of width 2 because it will not generate any vector instructions
 ;
 define void @all_scalar(i64* %a, i64 %n) {
 entry:
   br label %for.body
 
 for.body:
   %i = phi i64 [ 0, %entry ], [ %i.next, %for.body ]

llvm/test/Transforms/LoopVectorize/X86/pr55096-scalarize-add.ll

 ; RUN: opt -passes=loop-vectorize -force-vector-width=2 -force-vector-interleave=1 -S %s | FileCheck %s
 
-; REQUIRES: asserts
-; XFAIL: *
-
 target triple = "x86_64-apple-macosx"
 
-; CHECK: vector.body
-
 define void @test_pr55096(i64 %c, ptr %p) {
+; CHECK-LABEL: @test_pr55096(
+; CHECK:       vector.body:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, %vector.ph ], [ [[INDEX_NEXT:%.*]], %pred.store.continue3 ]
+; CHECK-NEXT:    [[VEC_IND:%.*]] = phi <2 x i64> [ <i64 122, i64 123>, %vector.ph ], [ [[VEC_IND_NEXT:%.*]], %pred.store.continue3 ]
+; CHECK-NEXT:    [[TMP0:%.*]] = trunc i64 [[INDEX]] to i16
+; CHECK-NEXT:    [[TMP1:%.*]] = mul i16 [[TMP0]], 2008
+; CHECK-NEXT:    [[OFFSET_IDX:%.*]] = add i16 6229, [[TMP1]]
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult <2 x i64> [[VEC_IND]], {{.*}}
+; CHECK-NEXT:    [[TMP3:%.*]] = xor <2 x i1> [[TMP2]], <i1 true, i1 true>
+; CHECK-NEXT:    [[TMP4:%.*]] = extractelement <2 x i1> [[TMP3]], i32 0
+; CHECK-NEXT:    br i1 [[TMP4]], label %pred.store.if, label %pred.store.continue
+; CHECK:       pred.store.if:
+; CHECK-NEXT:    [[TMP5:%.*]] = add i16 [[OFFSET_IDX]], 0
+; CHECK-NEXT:    [[TMP6:%.*]] = add i16 [[TMP5]], 2008
+; CHECK-NEXT:    [[TMP7:%.*]] = udiv i16 4943, [[TMP6]]
+; CHECK-NEXT:    [[TMP8:%.*]] = getelementptr inbounds i16, ptr [[P:%.*]], i16 [[TMP7]]
+; CHECK-NEXT:    store i16 0, ptr [[TMP8]], align 2
+; CHECK-NEXT:    br label %pred.store.continue
+; CHECK:       pred.store.continue:
+; CHECK-NEXT:    [[TMP9:%.*]] = phi i16 [ poison, %vector.body ], [ [[TMP7]], %pred.store.if ]
+; CHECK-NEXT:    [[TMP10:%.*]] = extractelement <2 x i1> [[TMP3]], i32 1
+; CHECK-NEXT:    br i1 [[TMP10]], label %pred.store.if2, label %pred.store.continue3
+; CHECK:       pred.store.if2:
+; CHECK-NEXT:    [[TMP11:%.*]] = add i16 [[OFFSET_IDX]], 2008
+; CHECK-NEXT:    [[TMP12:%.*]] = add i16 [[TMP11]], 2008
+; CHECK-NEXT:    [[TMP13:%.*]] = udiv i16 4943, [[TMP12]]
+; CHECK-NEXT:    [[TMP14:%.*]] = getelementptr inbounds i16, ptr [[P]], i16 [[TMP13]]
+; CHECK-NEXT:    store i16 0, ptr [[TMP14]], align 2
+; CHECK-NEXT:    br label %pred.store.continue3
+; CHECK:       pred.store.continue3:
+; CHECK-NEXT:    [[TMP15:%.*]] = phi i16 [ poison, %pred.store.continue ], [ [[TMP13]], %pred.store.if2 ]
+; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 2
+; CHECK-NEXT:    [[VEC_IND_NEXT]] = add <2 x i64> [[VEC_IND]], <i64 2, i64 2>
+; CHECK-NEXT:    [[TMP16:%.*]] = icmp eq i64 [[INDEX_NEXT]], 340
+; CHECK-NEXT:    br i1 [[TMP16]], label %middle.block, label %vector.body
+;
 entry:
   br label %loop.header
 
 loop.header:
   %iv.1 = phi i64 [ 122, %entry ], [ %iv.1.next, %loop.latch ]
   %iv.2 = phi i16 [ 6229, %entry ], [ %iv.2.next, %loop.latch ]
   %iv.2.next = add i16 %iv.2, 2008
   %cmp = icmp ult i64 %iv.1, %c