[CodeGen] Update transformations to use poison for shufflevector/insertelem's initial vector elem

[aqjune]

Summary

This patch is a part of D93817 and makes transformations in CodeGen use poison for shufflevector/insertelem's initial vector element.

The change in CodeGenPrepare.cpp is fine because the mask of shufflevector should be always zero.
It doesn't touch the second element (which is poison).

The change in InterleavedAccessPass.cpp is also fine becauses the mask is of the form <a, a+m, a+2m, .., a+km> where a+km is smaller than
the size of the first vector operand.
This is guaranteed by the caller of replaceBinOpShuffles, which is lowerInterleavedLoad.
It calls isDeInterleaveMask and isDeInterleaveMaskOfFactor to check the mask is the desirable form.
isDeInterleaveMask has the check that a+km is smaller than the vector size.
To check my understanding, I added an assertion & added a test to show that this optimization doesn't fire in such case.

[spatel]

LGTM - other reviewers are probably more familiar with InterleavedAccess than me, so might want to wait for a 2nd opinion.
I made another assert suggestion for that code. Similar to the assert that you have proposed in this patch, we could add the asserts before this patch just to make the shuffle assumptions clearer (assuming those are correct assertions).

llvm/lib/CodeGen/InterleavedAccessPass.cpp

363–365

Could we also assert here (and/or in the loop above this one) that:

VecTy->getNumElements() < Shuffle->getShuffleMask().size()

?

Ie, the shuffle is always reducing the number of elements from the input vector operands?

[nikic]

llvm/lib/CodeGen/CodeGenPrepare.cpp

6722

Wondering if you can replace this insertelement + shufflevector by a CreateVectorSplat call.

llvm/lib/CodeGen/InterleavedAccessPass.cpp

405

Why is this checking only the back() element? Better check all? assert(llvm::all_of(Mask, [&](int Idx) { return Idx < ... });

[aqjune]

llvm/lib/CodeGen/InterleavedAccessPass.cpp

363–365

Yes, I'll add one.

405

It was simply because of its cost. Actually I was wondering whether it is acceptable to have such formula in the assertion.

[aqjune]

Address comments

llvm/lib/CodeGen/CodeGenPrepare.cpp

bool CodeGenPrepare::optimizeSelectInst(SelectInst *SI) {
   CurInstIterator = StartBlock->end();
   return true;
 }
 
 /// Some targets only accept certain types for splat inputs. For example a VDUP
 /// in MVE takes a GPR (integer) register, and the instruction that incorporate
 /// a VDUP (such as a VADD qd, qm, rm) also require a gpr register.
 bool CodeGenPrepare::optimizeShuffleVectorInst(ShuffleVectorInst *SVI) {
+  // Accept shuf(insertelem(undef/poison, val, 0), undef/poison, <0,0,..>) only
   if (!match(SVI, m_Shuffle(m_InsertElt(m_Undef(), m_Value(), m_ZeroInt()),
                             m_Undef(), m_ZeroMask())))
     return false;
   Type *NewType = TLI->shouldConvertSplatType(SVI);
   if (!NewType)
     return false;
 
   auto *SVIVecType = cast<FixedVectorType>(SVI->getType());
   assert(!NewType->isVectorTy() && "Expected a scalar type!");
   assert(NewType->getScalarSizeInBits() == SVIVecType->getScalarSizeInBits() &&
          "Expected a type of the same size!");
   auto *NewVecType =
       FixedVectorType::get(NewType, SVIVecType->getNumElements());
 
   // Create a bitcast (shuffle (insert (bitcast(..))))
   IRBuilder<> Builder(SVI->getContext());
   Builder.SetInsertPoint(SVI);
   Value *BC1 = Builder.CreateBitCast(
       cast<Instruction>(SVI->getOperand(0))->getOperand(1), NewType);
-  Value *Insert = Builder.CreateInsertElement(UndefValue::get(NewVecType), BC1,
-                                              (uint64_t)0);
-  Value *Shuffle = Builder.CreateShuffleVector(Insert, SVI->getShuffleMask());
+  Value *Shuffle = Builder.CreateVectorSplat(NewVecType->getNumElements(), BC1);
   Value *BC2 = Builder.CreateBitCast(Shuffle, SVIVecType);
 
   SVI->replaceAllUsesWith(BC2);
   RecursivelyDeleteTriviallyDeadInstructions(
       SVI, TLInfo, nullptr, [&](Value *V) { removeAllAssertingVHReferences(V); });
 
   // Also hoist the bitcast up to its operand if it they are not in the same
   // block.

llvm/lib/CodeGen/InterleavedAccessPass.cpp

 // As interleaved accesses are difficult to identified in CodeGen (mainly
 // because the VECTOR_SHUFFLE DAG node is quite different from the shufflevector
 // IR), we identify and transform them to intrinsics in this pass so the
 // intrinsics can be easily matched into target specific instructions later in
 // CodeGen.
 //
 // E.g. An interleaved load (Factor = 2):
 //        %wide.vec = load <8 x i32>, <8 x i32>* %ptr
-//        %v0 = shuffle <8 x i32> %wide.vec, <8 x i32> undef, <0, 2, 4, 6>
-//        %v1 = shuffle <8 x i32> %wide.vec, <8 x i32> undef, <1, 3, 5, 7>
+//        %v0 = shuffle <8 x i32> %wide.vec, <8 x i32> poison, <0, 2, 4, 6>
+//        %v1 = shuffle <8 x i32> %wide.vec, <8 x i32> poison, <1, 3, 5, 7>
 //
 // It could be transformed into a ld2 intrinsic in AArch64 backend or a vld2
 // intrinsic in ARM backend.
 //
 // In X86, this can be further optimized into a set of target
 // specific loads followed by an optimized sequence of shuffles.
 //
 // E.g. An interleaved store (Factor = 3):
bool InterleavedAccess::lowerInterleavedLoad(
   // and factor as the first shufflevector.
   for (auto *Shuffle : Shuffles) {
     if (Shuffle->getType() != VecTy)
       return false;
     if (!isDeInterleaveMaskOfFactor(Shuffle->getShuffleMask(), Factor,
                                     Index))
       return false;
 
+    assert(Shuffle->getShuffleMask().size() <= NumLoadElements);
     Indices.push_back(Index);
   }
   for (auto *Shuffle : BinOpShuffles) {
     if (Shuffle->getType() != VecTy)
       return false;
     if (!isDeInterleaveMaskOfFactor(Shuffle->getShuffleMask(), Factor,
                                     Index))
       return false;
 
+    assert(Shuffle->getShuffleMask().size() <= NumLoadElements);
+
     if (cast<Instruction>(Shuffle->getOperand(0))->getOperand(0) == LI)
       Indices.push_back(Index);
     if (cast<Instruction>(Shuffle->getOperand(0))->getOperand(1) == LI)
       Indices.push_back(Index);
   }
 
   // Try and modify users of the load that are extractelement instructions to
   // use the shufflevector instructions instead of the load.
bool InterleavedAccess::lowerInterleavedLoad(
   return true;
 }
 
 bool InterleavedAccess::replaceBinOpShuffles(
     ArrayRef<ShuffleVectorInst *> BinOpShuffles,
     SmallVectorImpl<ShuffleVectorInst *> &Shuffles, LoadInst *LI) {
   for (auto *SVI : BinOpShuffles) {
     BinaryOperator *BI = cast<BinaryOperator>(SVI->getOperand(0));
+    Type *BIOp0Ty = BI->getOperand(0)->getType();
     ArrayRef<int> Mask = SVI->getShuffleMask();
-
-    auto *NewSVI1 = new ShuffleVectorInst(
-        BI->getOperand(0), UndefValue::get(BI->getOperand(0)->getType()), Mask,
-        SVI->getName(), SVI);
+    assert(all_of(Mask, [&](int Idx) {
+      return Idx < (int)cast<FixedVectorType>(BIOp0Ty)->getNumElements();
+    }));
+
+    auto *NewSVI1 =
+        new ShuffleVectorInst(BI->getOperand(0), PoisonValue::get(BIOp0Ty),
+                              Mask, SVI->getName(), SVI);
     auto *NewSVI2 = new ShuffleVectorInst(
-        BI->getOperand(1), UndefValue::get(BI->getOperand(1)->getType()), Mask,
+        BI->getOperand(1), PoisonValue::get(BI->getOperand(1)->getType()), Mask,
         SVI->getName(), SVI);
     Value *NewBI = BinaryOperator::Create(BI->getOpcode(), NewSVI1, NewSVI2,
                                           BI->getName(), SVI);
     SVI->replaceAllUsesWith(NewBI);
     LLVM_DEBUG(dbgs() << "  Replaced: " << *BI << "\n    And   : " << *SVI
                       << "\n  With    : " << *NewSVI1 << "\n    And   : "
                       << *NewSVI2 << "\n    And   : " << *NewBI << "\n");
     RecursivelyDeleteTriviallyDeadInstructions(SVI);

llvm/test/Transforms/InterleavedAccess/X86/interleave-load-extract-shuffle-changes.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -interleaved-access -S %s | FileCheck %s
 
 target datalayout = "e-m:o-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-apple-macosx10.15.0"
 
 ; No interleaved load instruction is generated, but the shuffle is moved just
 ; after the load.
 define <2 x double> @shuffle_binop_fol(<4 x double>* %ptr) {
 ; CHECK-LABEL: @shuffle_binop_fol(
 ; CHECK-NEXT:  vector.body.preheader:
 ; CHECK-NEXT:    [[WIDE_LOAD:%.*]] = load <4 x double>, <4 x double>* [[PTR:%.*]], align 8
-; CHECK-NEXT:    [[EXTRACTED1:%.*]] = shufflevector <4 x double> [[WIDE_LOAD]], <4 x double> undef, <2 x i32> <i32 0, i32 2>
-; CHECK-NEXT:    [[EXTRACTED2:%.*]] = shufflevector <4 x double> <double 1.000000e+00, double 1.000000e+00, double 1.000000e+00, double 1.000000e+00>, <4 x double> undef, <2 x i32> <i32 0, i32 2>
+; CHECK-NEXT:    [[EXTRACTED1:%.*]] = shufflevector <4 x double> [[WIDE_LOAD]], <4 x double> poison, <2 x i32> <i32 0, i32 2>
+; CHECK-NEXT:    [[EXTRACTED2:%.*]] = shufflevector <4 x double> <double 1.000000e+00, double 1.000000e+00, double 1.000000e+00, double 1.000000e+00>, <4 x double> poison, <2 x i32> <i32 0, i32 2>
 ; CHECK-NEXT:    [[FADD3:%.*]] = fadd <2 x double> [[EXTRACTED1]], [[EXTRACTED2]]
 ; CHECK-NEXT:    ret <2 x double> [[FADD3]]
 ;
 vector.body.preheader:
   %wide.load = load <4 x double>, <4 x double>* %ptr, align 8
   %fadd = fadd <4 x double> %wide.load, <double 1.0, double 1.0, double 1.0, double 1.0>
   %extracted = shufflevector <4 x double> %fadd, <4 x double> undef, <2 x i32> <i32 0, i32 2>
   ret <2 x double> %extracted
 }
 
+define <2 x double> @shuffle_binop_fol_oob(<4 x double>* %ptr) {
+; CHECK-LABEL: @shuffle_binop_fol_oob(
+; CHECK-NEXT:  vector.body.preheader:
+; CHECK-NEXT:    [[WIDE_LOAD:%.*]] = load <4 x double>, <4 x double>* [[PTR:%.*]], align 8
+; CHECK-NEXT:    [[FADD:%.*]] = fadd <4 x double> [[WIDE_LOAD]], <double 1.000000e+00, double 1.000000e+00, double 1.000000e+00, double 1.000000e+00>
+; CHECK-NEXT:    [[EXTRACTED:%.*]] = shufflevector <4 x double> [[FADD]], <4 x double> undef, <2 x i32> <i32 0, i32 4>
+; CHECK-NEXT:    ret <2 x double> [[EXTRACTED]]
+;
+vector.body.preheader:
+  %wide.load = load <4 x double>, <4 x double>* %ptr, align 8
+  %fadd = fadd <4 x double> %wide.load, <double 1.0, double 1.0, double 1.0, double 1.0>
+  %extracted = shufflevector <4 x double> %fadd, <4 x double> undef, <2 x i32> <i32 0, i32 4>
+  ret <2 x double> %extracted
+}
+
 ; No interleaved load instruction is generated, but the extractelement
 ; instructions are updated to use the shuffle instead of the load.
 define void @shuffle_extract(<4 x double>* %ptr, i1 %c) {
 ; CHECK-LABEL: @shuffle_extract(
 ; CHECK-NEXT:  vector.body.preheader:
 ; CHECK-NEXT:    [[WIDE_LOAD:%.*]] = load <4 x double>, <4 x double>* [[PTR:%.*]], align 8
 ; CHECK-NEXT:    [[EXTRACTED:%.*]] = shufflevector <4 x double> [[WIDE_LOAD]], <4 x double> undef, <2 x i32> <i32 0, i32 2>
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[IF_THEN:%.*]], label [[IF_MERGE:%.*]]