[AArch64][SVE][InstCombine] Unpack of a splat vector -> Scalar extend

[mnadeem]

Summary

Replace vector unpack operation with a scalar extend operation.

unpack(splat(X)) --> splat(extend(X))

If we have both, unpkhi and unpklo, for the same vector then we may
save a register in some cases, e.g:

Hi = unpkhi (splat(X))
Lo = unpklo(splat(X)) 
 --> Hi = Lo = splat(extend(X))

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

796

Can this use Builder.CreateVectorSplat? or is getScalableSplat doing something special?

[mnadeem]

• Use Builder.CreateVectorSplat(...) instead of using custom code to create splat.

[mnadeem]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

796

Replaced with CreateVectorSplat.

[paulwalker-arm]

llvm/test/Transforms/InstCombine/AArch64/sve-intrinsic-opts-unpkhi-unpklo.ll

2

Please use utils/update_test_checks.py to generate the CHECK lines.

27

The test is called uunpklo_splat but the IR is calling uunpkhi.

32–50

This looks like a redundant test which is not testing anything that is not already covered by the separate unsigned hi/lo tests.

73

The test is called sunpklo_splat but the IR is calling sunpkhi.

78–96

As above, this looks like a redundant test which is not testing anything that is not already covered by the separate signed hi/lo tests.

100–103

These should end .nxv4i32 to match the return type.

[mnadeem]

Fixed tests according to the comments.

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

if (OpPredicate == DupPg) {
       OpMultiplicand->takeName(&II);
       return IC.replaceInstUsesWith(II, OpMultiplicand);
     }
   }
 
   return None;
 }
 
+static Value *getScalableSplat(Value *V, unsigned NumElts,
+                               IRBuilder<> &Builder) {
+  Value *Poison =
+      PoisonValue::get(ScalableVectorType::get(V->getType(), NumElts));
+  Value *Zeros = ConstantAggregateZero::get(
+      ScalableVectorType::get(Builder.getInt32Ty(), NumElts));
+  Value *NewVal = Builder.CreateInsertElement(
+      Poison, V, ConstantInt::get(Builder.getInt32Ty(), 0));
+  NewVal = Builder.CreateShuffleVector(NewVal, Poison, Zeros);
+  return NewVal;
+}
+
+static Optional<Instruction *> instCombineSVEUnpack(InstCombiner &IC,
+                                                    IntrinsicInst &II) {
+  IRBuilder<> Builder(II.getContext());
+  Builder.SetInsertPoint(&II);
+  Value *UnpackArg = II.getArgOperand(0);
+  auto *RetTy = cast<ScalableVectorType>(II.getType());
+  auto RetNumElements = RetTy->getMinNumElements();
+  bool isSigned = II.getIntrinsicID() == Intrinsic::aarch64_sve_sunpkhi ||
+                  II.getIntrinsicID() == Intrinsic::aarch64_sve_sunpklo;
+
+  // Hi = uunpkhi(splat(X)) --> Hi = splat(extend(X))
+  // Lo = uunpklo(splat(X)) --> Lo = splat(extend(X))
+  if (auto *ScalarArg = getSplatValue(UnpackArg)) {
+    ScalarArg =
+        Builder.CreateIntCast(ScalarArg, RetTy->getScalarType(), isSigned);
+    Value *NewVal = getScalableSplat(ScalarArg, RetNumElements, Builder);
+    NewVal->takeName(&II);
+    return IC.replaceInstUsesWith(II, NewVal);
+  }
+
+  return None;
+}
 static Optional<Instruction *> instCombineSVETBL(InstCombiner &IC,
                                                  IntrinsicInst &II) {
   auto *OpVal = II.getOperand(0);
   auto *OpIndices = II.getOperand(1);
   VectorType *VTy = cast<VectorType>(II.getType());
 
   // Check whether OpIndices is an aarch64_sve_dup_x intrinsic call with
   // constant splat value < minimal element count of result.
AArch64TTIImpl::instCombineIntrinsic(InstCombiner &IC,
   case Intrinsic::aarch64_sve_ptest_first:
   case Intrinsic::aarch64_sve_ptest_last:
     return instCombineSVEPTest(IC, II);
   case Intrinsic::aarch64_sve_mul:
   case Intrinsic::aarch64_sve_fmul:
     return instCombineSVEVectorMul(IC, II);
   case Intrinsic::aarch64_sve_tbl:
     return instCombineSVETBL(IC, II);
+  case Intrinsic::aarch64_sve_uunpkhi:
+  case Intrinsic::aarch64_sve_uunpklo:
+  case Intrinsic::aarch64_sve_sunpkhi:
+  case Intrinsic::aarch64_sve_sunpklo:
+    return instCombineSVEUnpack(IC, II);
   }
 
   return None;
 }
 
 bool AArch64TTIImpl::isWideningInstruction(Type *DstTy, unsigned Opcode,
                                            ArrayRef<const Value *> Args) {
 

llvm/test/Transforms/InstCombine/AArch64/sve-intrinsic-opts-unpkhi-unpklo.ll

+; RUN: opt -S -instcombine < %s | FileCheck %s
+
+target triple = "aarch64-unknown-linux-gnu"
+
+define <vscale x 4 x i32> @uunpkhi_splat(i16 %a) #0 {
+; CHECK-LABEL: @uunpkhi_splat(
+; CHECK: %1 = zext i16 %a to i32
+; CHECK: %2 = insertelement <vscale x 4 x i32> poison, i32 %1, i32 0
+; CHECK: %unpack = shufflevector <vscale x 4 x i32> %2, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:   ret <vscale x 4 x i32> %unpack
+  %splat_insert = insertelement <vscale x 8 x i16> poison, i16 %a, i32 0
+  %splat = shufflevector <vscale x 8 x i16> %splat_insert, <vscale x 8 x i16> poison, <vscale x 8 x i32> zeroinitializer
+  %unpack = tail call <vscale x 4 x i32> @llvm.aarch64.sve.uunpkhi.nxv8i32
+               (<vscale x 8 x i16> %splat)
+  ret <vscale x 4 x i32> %unpack
+}
+
+define <vscale x 4 x i32> @uunpklo_splat(i16 %a) #0 {
+; CHECK-LABEL: @uunpklo_splat(
+; CHECK: %1 = zext i16 %a to i32
+; CHECK: %2 = insertelement <vscale x 4 x i32> poison, i32 %1, i32 0
+; CHECK: %unpack = shufflevector <vscale x 4 x i32> %2, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:   ret <vscale x 4 x i32> %unpack
+  %splat_insert = insertelement <vscale x 8 x i16> poison, i16 %a, i32 0
+  %splat = shufflevector <vscale x 8 x i16> %splat_insert, <vscale x 8 x i16> poison, <vscale x 8 x i32> zeroinitializer
+  %unpack = tail call <vscale x 4 x i32> @llvm.aarch64.sve.uunpkhi.nxv8i32
+               (<vscale x 8 x i16> %splat)
+  ret <vscale x 4 x i32> %unpack
+}
+
+define void @uunpk_splat(i16 %a) #0 {
+; CHECK-LABEL: @uunpk_splat(
+; CHECK: %1 = zext i16 %a to i32
+; CHECK: %2 = insertelement <vscale x 4 x i32> poison, i32 %1, i32 0
+; CHECK: %unpackhi = shufflevector <vscale x 4 x i32> %2, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK: %3 = zext i16 %a to i32
+; CHECK: %4 = insertelement <vscale x 4 x i32> poison, i32 %3, i32 0
+; CHECK: %unpacklo = shufflevector <vscale x 4 x i32> %4, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK: call void @use(<vscale x 4 x i32> %unpackhi, <vscale x 4 x i32> %unpacklo)
+; CHECK-NEXT:   ret void
+  %splat_insert = insertelement <vscale x 8 x i16> poison, i16 %a, i32 0
+  %splat = shufflevector <vscale x 8 x i16> %splat_insert, <vscale x 8 x i16> poison, <vscale x 8 x i32> zeroinitializer
+  %unpackhi = tail call <vscale x 4 x i32> @llvm.aarch64.sve.uunpkhi.nxv8i32
+               (<vscale x 8 x i16> %splat)
+  %unpacklo = tail call <vscale x 4 x i32> @llvm.aarch64.sve.uunpklo.nxv8i32
+               (<vscale x 8 x i16> %splat)
+  call void @use(<vscale x 4 x i32> %unpackhi, <vscale x 4 x i32> %unpacklo)
+  ret void
+}
+
+define <vscale x 4 x i32> @sunpkhi_splat(i16 %a) #0 {
+; CHECK-LABEL: @sunpkhi_splat(
+; CHECK: %1 = sext i16 %a to i32
+; CHECK: %2 = insertelement <vscale x 4 x i32> poison, i32 %1, i32 0
+; CHECK: %unpack = shufflevector <vscale x 4 x i32> %2, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:   ret <vscale x 4 x i32> %unpack
+  %splat_insert = insertelement <vscale x 8 x i16> poison, i16 %a, i32 0
+  %splat = shufflevector <vscale x 8 x i16> %splat_insert, <vscale x 8 x i16> poison, <vscale x 8 x i32> zeroinitializer
+  %unpack = tail call <vscale x 4 x i32> @llvm.aarch64.sve.sunpkhi.nxv8i32
+               (<vscale x 8 x i16> %splat)
+  ret <vscale x 4 x i32> %unpack
+}
+
+define <vscale x 4 x i32> @sunpklo_splat(i16 %a) #0 {
+; CHECK-LABEL: @sunpklo_splat(
+; CHECK: %1 = sext i16 %a to i32
+; CHECK: %2 = insertelement <vscale x 4 x i32> poison, i32 %1, i32 0
+; CHECK: %unpack = shufflevector <vscale x 4 x i32> %2, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:   ret <vscale x 4 x i32> %unpack
+  %splat_insert = insertelement <vscale x 8 x i16> poison, i16 %a, i32 0
+  %splat = shufflevector <vscale x 8 x i16> %splat_insert, <vscale x 8 x i16> poison, <vscale x 8 x i32> zeroinitializer
+  %unpack = tail call <vscale x 4 x i32> @llvm.aarch64.sve.sunpkhi.nxv8i32
+               (<vscale x 8 x i16> %splat)
+  ret <vscale x 4 x i32> %unpack
+}
+
+define void @sunpk_splat(i16 %a) #0 {
+; CHECK-LABEL: @sunpk_splat(
+; CHECK: %1 = sext i16 %a to i32
+; CHECK: %2 = insertelement <vscale x 4 x i32> poison, i32 %1, i32 0
+; CHECK: %unpackhi = shufflevector <vscale x 4 x i32> %2, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK: %3 = sext i16 %a to i32
+; CHECK: %4 = insertelement <vscale x 4 x i32> poison, i32 %3, i32 0
+; CHECK: %unpacklo = shufflevector <vscale x 4 x i32> %4, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK: call void @use(<vscale x 4 x i32> %unpackhi, <vscale x 4 x i32> %unpacklo)
+; CHECK-NEXT:   ret void
+  %splat_insert = insertelement <vscale x 8 x i16> poison, i16 %a, i32 0
+  %splat = shufflevector <vscale x 8 x i16> %splat_insert, <vscale x 8 x i16> poison, <vscale x 8 x i32> zeroinitializer
+  %unpackhi = tail call <vscale x 4 x i32> @llvm.aarch64.sve.sunpkhi.nxv8i32
+               (<vscale x 8 x i16> %splat)
+  %unpacklo = tail call <vscale x 4 x i32> @llvm.aarch64.sve.sunpklo.nxv8i32
+               (<vscale x 8 x i16> %splat)
+  call void @use(<vscale x 4 x i32> %unpackhi, <vscale x 4 x i32> %unpacklo)
+  ret void
+}
+
+declare void @use(<vscale x 4 x i32> , <vscale x 4 x i32>)
+
+declare <vscale x 4 x i32> @llvm.aarch64.sve.uunpkhi.nxv8i32(<vscale x 8 x i16>)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.uunpklo.nxv8i32(<vscale x 8 x i16>)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.sunpkhi.nxv8i32(<vscale x 8 x i16>)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.sunpklo.nxv8i32(<vscale x 8 x i16>)
+
+attributes #0 = { "target-features"="+sve" }