[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

362

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

6725

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

llvm/lib/CodeGen/InterleavedAccessPass.cpp

401

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

362

Yes, I'll add one.

401

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

 //===- CodeGenPrepare.cpp - Prepare a function for code generation --------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
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,
+  Value *Insert = Builder.CreateInsertElement(PoisonValue::get(NewVecType), BC1,
                                               (uint64_t)0);
   Value *Shuffle = Builder.CreateShuffleVector(Insert, SVI->getShuffleMask());
   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

 //===- InterleavedAccessPass.cpp ------------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the Interleaved Access pass, which identifies
 // interleaved memory accesses and transforms them into target specific
 // intrinsics.
 //
 // An interleaved load reads data from memory into several vectors, with
 // DE-interleaving the data on a factor. An interleaved store writes several
 // vectors to memory with RE-interleaving the data on a factor.
 //
 // 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):
for (auto *Shuffle : Shuffles) {
     Indices.push_back(Index);
   }
   for (auto *Shuffle : BinOpShuffles) {
     if (Shuffle->getType() != VecTy)
       return false;
     if (!isDeInterleaveMaskOfFactor(Shuffle->getShuffleMask(), Factor,
                                     Index))
       return false;
 
     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::replaceBinOpShuffles(
     ArrayRef<ShuffleVectorInst *> BinOpShuffles,
     SmallVectorImpl<ShuffleVectorInst *> &Shuffles, LoadInst *LI) {
   for (auto *SVI : BinOpShuffles) {
     BinaryOperator *BI = cast<BinaryOperator>(SVI->getOperand(0));
     ArrayRef<int> Mask = SVI->getShuffleMask();
+    assert(Mask.back() < 0 ||
+           (unsigned)Mask.back() <
+               cast<FixedVectorType>(BI->getOperand(0)->getType())
+                   ->getNumElements());
 
     auto *NewSVI1 = new ShuffleVectorInst(
-        BI->getOperand(0), UndefValue::get(BI->getOperand(0)->getType()), Mask,
+        BI->getOperand(0), PoisonValue::get(BI->getOperand(0)->getType()), 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:%.*]]