[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 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());
+  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 =
+        Builder.CreateVectorSplat(RetTy->getElementCount(), ScalarArg);
+    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.
   auto *DupXIntrI = dyn_cast<IntrinsicInst>(OpIndices);
   if (!DupXIntrI || DupXIntrI->getIntrinsicID() != Intrinsic::aarch64_sve_dup_x)
     return None;
 
   auto *SplatValue = dyn_cast<ConstantInt>(DupXIntrI->getOperand(0));
   if (!SplatValue ||
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

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -instcombine < %s | FileCheck %s
+
+target triple = "aarch64"
+
+define <vscale x 4 x i32> @uunpkhi_splat(i16 %a) #0 {
+; CHECK-LABEL: @uunpkhi_splat(
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i16 [[A:%.*]] to i32
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 4 x i32> poison, i32 [[TMP1]], i32 0
+; CHECK-NEXT:    [[UNPACK:%.*]] = shufflevector <vscale x 4 x i32> [[DOTSPLATINSERT]], <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.nxv4i32(<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-NEXT:    [[TMP1:%.*]] = zext i16 [[A:%.*]] to i32
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 4 x i32> poison, i32 [[TMP1]], i32 0
+; CHECK-NEXT:    [[UNPACK:%.*]] = shufflevector <vscale x 4 x i32> [[DOTSPLATINSERT]], <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.uunpklo.nxv4i32(<vscale x 8 x i16> %splat)
+  ret <vscale x 4 x i32> %unpack
+}
+
+define <vscale x 4 x i32> @sunpkhi_splat(i16 %a) #0 {
+; CHECK-LABEL: @sunpkhi_splat(
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i16 [[A:%.*]] to i32
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 4 x i32> poison, i32 [[TMP1]], i32 0
+; CHECK-NEXT:    [[UNPACK:%.*]] = shufflevector <vscale x 4 x i32> [[DOTSPLATINSERT]], <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.nxv4i32(<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-NEXT:    [[TMP1:%.*]] = sext i16 [[A:%.*]] to i32
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 4 x i32> poison, i32 [[TMP1]], i32 0
+; CHECK-NEXT:    [[UNPACK:%.*]] = shufflevector <vscale x 4 x i32> [[DOTSPLATINSERT]], <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.sunpklo.nxv4i32(<vscale x 8 x i16> %splat)
+  ret <vscale x 4 x i32> %unpack
+}
+
+declare <vscale x 4 x i32> @llvm.aarch64.sve.uunpkhi.nxv4i32(<vscale x 8 x i16>)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.uunpklo.nxv4i32(<vscale x 8 x i16>)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.sunpkhi.nxv4i32(<vscale x 8 x i16>)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.sunpklo.nxv4i32(<vscale x 8 x i16>)
+
+attributes #0 = { "target-features"="+sve" }