[SLP] Add a check before skipping inserting a trunc after z|sext vectorization tree root

[haicheng]

Summary

If the vectorization tree root operations are integer cast operations, D41948 does not consider them as candidates for truncation. However, sometimes the trunc is needed by instcombine to trigger type shrinking or prevent type extension.

One example is like this, when vectorizing

  %800 = zext <16 x i16> %799 to <16 x i32>
  %803 = zext <16 x i16> %802 to <16 x i32>
  %804 = sub nsw <16 x i32> %800, %803
  %805 = extractelement <16 x i32> %804, i32 0
  %806 = sext i32 %805 to i64
  %808 = extractelement <16 x i32> %804, i32 1
  %809 = sext i32 %808 to i64
...
  %850 = extractelement <16 x i32> %804, i32 15
  %851 = sext i32 %850 to i64

to

  %800 = zext <16 x i16> %799 to <16 x i32>
  %803 = zext <16 x i16> %802 to <16 x i32>
  %804 = sub nsw <16 x i32> %800, %803
  %805 = sext <16 x i32> %804 to <16 x i64>
  %806 = trunc <16 x i64> %805 to <16 x i32>
  %807 = extractelement <16 x i32> %806, i32 0
  %808 = extractelement <16 x i32> %806, i32 1
...
  %837 = extractelement <16 x i32> %806, i32 15

A trunc after the vectorized sext is needed to prevent the types of the first three IRs (2 zext and 1 sub) from being expanded to i64.

This patch checks if the sources of the vectorized cast operations can be changed to evaluate in different types. If so, the trunc is still added.

[ABataev]

Not sure that this should be handled in SLPVectorizer. It is the problem of the instcombiner that it is too optimistic sometimes, not the problem of the SLPVectorizer. InstCombiner should check that this kind of the expansion is good for the current platform

[haicheng]

Matt, what do you think?

[haicheng]

Abandon in favor of D44868.

lib/Transforms/Vectorize/SLPVectorizer.cpp

default:
     return false;
   }
 
   // Record the value that we can demote.
   ToDemote.push_back(V);
   return true;
 }
 
+/// Return true if we can extend/trunc the type of the specified value in
+/// instcombine.
+static bool canEvaluateInDifferentType(Value *V) {
+  // Constants can always be evaluated in a different type.
+  if (isa<Constant>(V))
+    return true;
+
+  auto *I = dyn_cast<Instruction>(V);
+  if (!I)
+    return false;
+  // Instcombine does't change the types of something that has multiple uses.
+  // SLP may remove extractelement, so we will continue checking if all the
+  // uses are the vector operands of extractelement insts.
+  if (!I->hasOneUse() && (llvm::any_of(I->users(), [=](User *U) -> bool {
+        return !isa<ExtractElementInst>(U) || U->getOperand(0) != I;
+      })))
+    return false;
+  switch (I->getOpcode()) {
+  // We can always change the types of a cast. 
+  case Instruction::Trunc:
+  case Instruction::ZExt:
+  case Instruction::SExt:
+   return true;
+  case Instruction::Add:
+  case Instruction::Sub:
+  case Instruction::Mul:
+  case Instruction::And:
+  case Instruction::Or:
+  case Instruction::Xor:
+    // These operators can all arbitrarily be extended/trunced if their inputs
+    // can.
+    return canEvaluateInDifferentType(I->getOperand(0)) &&
+           canEvaluateInDifferentType(I->getOperand(1));
+  case Instruction::Select:
+    // We can change a select if we can change both true/false operands.
+    return canEvaluateInDifferentType(I->getOperand(1)) &&
+           canEvaluateInDifferentType(I->getOperand(2));
+  case Instruction::PHI: {
+    // We can change a phi if we can change all operands.
+    PHINode *PN = cast<PHINode>(I);
+    for (Value *IncValue : PN->incoming_values())
+      if (!canEvaluateInDifferentType(IncValue))
+        return false;
+    return true;
+  }
+  case Instruction::ExtractElement:
+    return canEvaluateInDifferentType(I->getOperand(0));
+    return true;
+  default:
+    break;
+  }
+  return false;
+}
+
 void BoUpSLP::computeMinimumValueSizes() {
   // If there are no external uses, the expression tree must be rooted by a
   // store. We can't demote in-memory values, so there is nothing to do here.
   if (ExternalUses.empty())
     return;
 
   // We only attempt to truncate integer expressions.
   auto &TreeRoot = VectorizableTree[0].Scalars;
if (!Root->hasOneUse() || Expr.count(*Root->user_begin()))
       return;
 
   // Conservatively determine if we can actually truncate the roots of the
   // expression. Collect the values that can be demoted in ToDemote and
   // additional roots that require investigating in Roots.
   SmallVector<Value *, 32> ToDemote;
   SmallVector<Value *, 4> Roots;
   for (auto *Root : TreeRoot) {
-    // Do not include top zext/sext/trunc operations to those to be demoted, it
-    // produces noise cast<vect>, trunc <vect>, exctract <vect>, cast <extract>
-    // sequence.
+    // Do not include top zext/sext/trunc operations to those to be demoted if
+    // the sources of these cast IRs cannot be extended, it produces noise
+    // cast<vect>, trunc <vect>, exctract <vect>, cast <extract> sequence.
     if (isa<Constant>(Root))
       continue;
     auto *I = dyn_cast<Instruction>(Root);
     if (!I || !I->hasOneUse() || !Expr.count(I))
       return;
-    if (isa<ZExtInst>(I) || isa<SExtInst>(I))
+    if ((isa<ZExtInst>(I) || isa<SExtInst>(I)) &&
+        !canEvaluateInDifferentType(I->getOperand(0)))
       continue;
     if (auto *TI = dyn_cast<TruncInst>(I)) {
       Roots.push_back(TI->getOperand(0));
       continue;
     }
     if (!collectValuesToDemote(Root, Expr, ToDemote, Roots))
       return;
   }

test/Transforms/SLPVectorizer/AArch64/ext-trunc.ll

+; RUN: opt -S -slp-vectorizer -instcombine < %s | FileCheck %s
+
+target datalayout = "e-m:e-i32:64-i128:128-n32:64-S128"
+target triple = "aarch64--linux-gnu"
+
+declare void @foo(i64, i64, i64, i64)
+
+define void @test(<4 x i16> %a, <4 x i16> %b, i64* %p) {
+; Make sure types of sub and its sources are not extended.
+; CHECK-LABEL: @test(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[Z0:%.*]] = zext <4 x i16> [[A:%.*]] to <4 x i32>
+; CHECK-NEXT:    [[Z1:%.*]] = zext <4 x i16> [[B:%.*]] to <4 x i32>
+; CHECK-NEXT:    [[SUB:%.*]] = sub nsw <4 x i32> [[Z0]], [[Z1]]
+entry:
+  %z0 = zext <4 x i16> %a to <4 x i32>
+  %z1 = zext <4 x i16> %b to <4 x i32>
+  %sub0 = sub <4 x i32> %z0, %z1
+  %e0 = extractelement <4 x i32> %sub0, i32 0
+  %s0 = sext i32 %e0 to i64
+  %gep0 = getelementptr inbounds i64, i64* %p, i64 %s0
+  %load0 = load i64, i64* %gep0
+  %e1 = extractelement <4 x i32> %sub0, i32 1
+  %s1 = sext i32 %e1 to i64
+  %gep1 = getelementptr inbounds i64, i64* %p, i64 %s1
+  %load1 = load i64, i64* %gep1
+  %e2 = extractelement <4 x i32> %sub0, i32 2
+  %s2 = sext i32 %e2 to i64
+  %gep2 = getelementptr inbounds i64, i64* %p, i64 %s2
+  %load2 = load i64, i64* %gep2
+  %e3 = extractelement <4 x i32> %sub0, i32 3
+  %s3 = sext i32 %e3 to i64
+  %gep3 = getelementptr inbounds i64, i64* %p, i64 %s3
+  %load3 = load i64, i64* %gep3
+  call void @foo(i64 %load0, i64 %load1, i64 %load2, i64 %load3)
+  ret void
+}