[AArch64][SVEIntrinsicOpts] Optimize tbl+dup into dup+extractelement

[junparser]

Summary

According to D99324, this patch try to convert pattern sve_tbl(vec, sve_dup_x(idx)) to sve_dup_x(extractelement(vec, idx)) which simplify lowering and isel in later phase. Try to solve the problem for both ACLE and auto-vectorisation.

TestPlan: check-llvm

[paulwalker-arm]

llvm/lib/Target/AArch64/SVEIntrinsicOpts.cpp

466–467

Can you use Builder.CreateVectorSplat here?

469–470

I'm never really sure what to do in these instances but do we need to worry about the original DupXIntrI? Given we know its operand is constant it's easily removed if this was the only use.

[junparser]

Address the comment.

[junparser]

llvm/lib/Target/AArch64/SVEIntrinsicOpts.cpp

466–467

done

469–470

replaceAllUsesWith and eraseFromParent usually appear together. we can ignore the Dup intrinsic here. However, we should remove it since it has no side effect, no use. Remove it can be safe.

[joechrisellis]

Hi, thank you for the patch! The optimisation seems sound to me -- just a few comments on the test side. 😄

llvm/test/CodeGen/AArch64/sve-tbl-dupx.ll

2

I don't think -dce is necessary here.

3

Worth adding:

; RUN: FileCheck --check-prefix=WARN --allow-empty %s <%t
...
; If this check fails please read test/CodeGen/AArch64/README for instructions on how to resolve it.
; WARN-NOT: warning

To make sure your 2>%t is used. This block can usually be found in SVE tests. 😄

89

A few of these are unused at the minute. You can remove:

llvm.aarch64.sve.dup.x.nxv8f16
llvm.aarch64.sve.dup.x.nxv2f32
llvm.aarch64.sve.dup.x.nxv4f32
llvm.aarch64.sve.dup.x.nxv2f64

[junparser]

Address comments

[junparser]

In D99412#2655132, @joechrisellis wrote:

Hi, thank you for the patch! The optimisation seems sound to me -- just a few comments on the test side. 😄

Thanks for review!

[joechrisellis]

LGTM. 😄

[paulwalker-arm]

llvm/test/CodeGen/AArch64/sve-tbl-dupx.ll

3

Not a problem for this patch but for information once D98856 lands this practice will no longer be needed as the key warning will error by default for opt/llc tests.

13

This shows what I meant by my previous comment in that all the tests check for this stray instruction and so it seem preferable to me for SVEIntrinsicOpts to remove it when it's knowingly no longer used.

[junparser]

llvm/test/CodeGen/AArch64/sve-tbl-dupx.ll

13

done, thanks for remind.

llvm/lib/Target/AArch64/SVEIntrinsicOpts.cpp

private:
 
   /// Operates at the function-scope. I.e., optimizations are applied local to
   /// the functions themselves.
   bool optimizeFunctions(SmallSetVector<Function *, 4> &Functions);
 
   static bool optimizeConvertFromSVBool(IntrinsicInst *I);
   static bool optimizePTest(IntrinsicInst *I);
   static bool optimizeVectorMul(IntrinsicInst *I);
+  static bool optimizeTBL(IntrinsicInst *I);
 
   static bool processPhiNode(IntrinsicInst *I);
 };
 } // end anonymous namespace
 
 void SVEIntrinsicOpts::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequired<DominatorTreeWrapperPass>();
   AU.setPreservesCFG();
if (OpMultiplierInst && OpMultiplierInst->use_empty())
       OpMultiplierInst->eraseFromParent();
     if (OpPredicateInst && OpPredicateInst->use_empty())
       OpPredicateInst->eraseFromParent();
   }
 
   return Changed;
 }
 
+bool SVEIntrinsicOpts::optimizeTBL(IntrinsicInst *I) {
+  assert(I->getIntrinsicID() == Intrinsic::aarch64_sve_tbl &&
+         "Unexpected opcode");
+
+  auto *OpVal = I->getOperand(0);
+  auto *OpIndices = I->getOperand(1);
+  VectorType *VTy = cast<VectorType>(I->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 false;
+
+  auto *SplatValue = dyn_cast<ConstantInt>(DupXIntrI->getOperand(0));
+  if (!SplatValue ||
+      SplatValue->getValue().uge(VTy->getElementCount().getKnownMinValue()))
+    return false;
+
+  // Convert sve_tbl(OpVal sve_dup_x(SplatValue)) to
+  // splat_vector(extractelement(OpVal, SplatValue)) for further optimization.
+  LLVMContext &Ctx = I->getContext();
+  IRBuilder<> Builder(Ctx);
+  Builder.SetInsertPoint(I);
+  auto *Extract = Builder.CreateExtractElement(OpVal, SplatValue);
+  auto *VectorSplat =
+      Builder.CreateVectorSplat(VTy->getElementCount(), Extract);
+
+  I->replaceAllUsesWith(VectorSplat);
+  I->eraseFromParent();
+  return true;
+}
+
 bool SVEIntrinsicOpts::optimizeConvertFromSVBool(IntrinsicInst *I) {
   assert(I->getIntrinsicID() == Intrinsic::aarch64_sve_convert_from_svbool &&
          "Unexpected opcode");
 
   // If the reinterpret instruction operand is a PHI Node
   if (isa<PHINode>(I->getArgOperand(0)))
     return processPhiNode(I);
 
case Intrinsic::aarch64_sve_convert_from_svbool:
     return optimizeConvertFromSVBool(IntrI);
   case Intrinsic::aarch64_sve_fmul:
   case Intrinsic::aarch64_sve_mul:
     return optimizeVectorMul(IntrI);
   case Intrinsic::aarch64_sve_ptest_any:
   case Intrinsic::aarch64_sve_ptest_first:
   case Intrinsic::aarch64_sve_ptest_last:
     return optimizePTest(IntrI);
+  case Intrinsic::aarch64_sve_tbl:
+    return optimizeTBL(IntrI);
   default:
     return false;
   }
 
   return true;
 }
 
 bool SVEIntrinsicOpts::optimizeIntrinsicCalls(
for (auto &F : M.getFunctionList()) {
     switch (F.getIntrinsicID()) {
     case Intrinsic::aarch64_sve_convert_from_svbool:
     case Intrinsic::aarch64_sve_ptest_any:
     case Intrinsic::aarch64_sve_ptest_first:
     case Intrinsic::aarch64_sve_ptest_last:
     case Intrinsic::aarch64_sve_ptrue:
     case Intrinsic::aarch64_sve_mul:
     case Intrinsic::aarch64_sve_fmul:
+    case Intrinsic::aarch64_sve_tbl:
       for (User *U : F.users())
         Functions.insert(cast<Instruction>(U)->getFunction());
       break;
     default:
       break;
     }
   }
 
   if (!Functions.empty())
     Changed |= optimizeFunctions(Functions);
 
   return Changed;
 }

llvm/test/CodeGen/AArch64/sve-tbl-dupx.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -aarch64-sve-intrinsic-opts -dce < %s 2>%t | FileCheck %s
+; op2 = tbl(op1 dup_x(idx)) -> op2 = dup_x(extractelement(op1, idx))
+
+define <vscale x 16 x i8> @dup_ext_i8(<vscale x 16 x i8> %data) {
+; CHECK-LABEL: @dup_ext_i8(
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <vscale x 16 x i8> [[DATA:%.*]], i8 1
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 16 x i8> poison, i8 [[TMP1]], i32 0
+; CHECK-NEXT:    [[DOTSPLAT:%.*]] = shufflevector <vscale x 16 x i8> [[DOTSPLATINSERT]], <vscale x 16 x i8> poison, <vscale x 16 x i32> zeroinitializer
+; CHECK-NEXT:    ret <vscale x 16 x i8> [[DOTSPLAT]]
+;
+  %tmp = call <vscale x 16 x i8> @llvm.aarch64.sve.dup.x.nxv16i8(i8 1)
+  %out = call <vscale x 16 x i8> @llvm.aarch64.sve.tbl.nxv16i8(<vscale x 16 x i8>  %data, <vscale x 16 x i8> %tmp)
+  ret <vscale x 16 x i8> %out
+}
+
+define <vscale x 8 x i16> @dup_ext_i16(<vscale x 8 x i16> %data) {
+; CHECK-LABEL: @dup_ext_i16(
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <vscale x 8 x i16> [[DATA:%.*]], i16 1
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 8 x i16> poison, i16 [[TMP1]], i32 0
+; CHECK-NEXT:    [[DOTSPLAT:%.*]] = shufflevector <vscale x 8 x i16> [[DOTSPLATINSERT]], <vscale x 8 x i16> poison, <vscale x 8 x i32> zeroinitializer
+; CHECK-NEXT:    ret <vscale x 8 x i16> [[DOTSPLAT]]
+;
+  %tmp = call <vscale x 8 x i16> @llvm.aarch64.sve.dup.x.nxv8i16(i16 1)
+  %out = call <vscale x 8 x i16> @llvm.aarch64.sve.tbl.nxv8i16(<vscale x 8 x i16>  %data, <vscale x 8 x i16> %tmp)
+  ret <vscale x 8 x i16> %out
+}
+
+define <vscale x 4 x i32> @dup_ext_i32(<vscale x 4 x i32> %data) {
+; CHECK-LABEL: @dup_ext_i32(
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <vscale x 4 x i32> [[DATA:%.*]], i32 1
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 4 x i32> poison, i32 [[TMP1]], i32 0
+; CHECK-NEXT:    [[DOTSPLAT:%.*]] = 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> [[DOTSPLAT]]
+;
+  %tmp = call <vscale x 4 x i32> @llvm.aarch64.sve.dup.x.nxv4i32(i32 1)
+  %out = call <vscale x 4 x i32> @llvm.aarch64.sve.tbl.nxv4i32(<vscale x 4 x i32>  %data, <vscale x 4 x i32> %tmp)
+  ret <vscale x 4 x i32> %out
+}
+
+define <vscale x 2 x i64> @dup_ext_i64(<vscale x 2 x i64> %data) {
+; CHECK-LABEL: @dup_ext_i64(
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <vscale x 2 x i64> [[DATA:%.*]], i64 1
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 2 x i64> poison, i64 [[TMP1]], i32 0
+; CHECK-NEXT:    [[DOTSPLAT:%.*]] = shufflevector <vscale x 2 x i64> [[DOTSPLATINSERT]], <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT:    ret <vscale x 2 x i64> [[DOTSPLAT]]
+;
+  %tmp = call <vscale x 2 x i64> @llvm.aarch64.sve.dup.x.nxv2i64(i64 1)
+  %out = call <vscale x 2 x i64> @llvm.aarch64.sve.tbl.nxv2i64(<vscale x 2 x i64>  %data, <vscale x 2 x i64> %tmp)
+  ret <vscale x 2 x i64> %out
+}
+
+define <vscale x 8 x half> @dup_ext_f16(<vscale x 8 x half> %data) {
+; CHECK-LABEL: @dup_ext_f16(
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <vscale x 8 x half> [[DATA:%.*]], i16 1
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 8 x half> poison, half [[TMP1]], i32 0
+; CHECK-NEXT:    [[DOTSPLAT:%.*]] = shufflevector <vscale x 8 x half> [[DOTSPLATINSERT]], <vscale x 8 x half> poison, <vscale x 8 x i32> zeroinitializer
+; CHECK-NEXT:    ret <vscale x 8 x half> [[DOTSPLAT]]
+;
+  %tmp = call <vscale x 8 x i16> @llvm.aarch64.sve.dup.x.nxv8i16(i16 1)
+  %out = call <vscale x 8 x half> @llvm.aarch64.sve.tbl.nxv8f16(<vscale x 8 x half>  %data, <vscale x 8 x i16> %tmp)
+  ret <vscale x 8 x half> %out
+}
+
+define <vscale x 4 x float> @dup_ext_f32(<vscale x 4 x float> %data) {
+; CHECK-LABEL: @dup_ext_f32(
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <vscale x 4 x float> [[DATA:%.*]], i32 1
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 4 x float> poison, float [[TMP1]], i32 0
+; CHECK-NEXT:    [[DOTSPLAT:%.*]] = shufflevector <vscale x 4 x float> [[DOTSPLATINSERT]], <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:    ret <vscale x 4 x float> [[DOTSPLAT]]
+;
+  %tmp = call <vscale x 4 x i32> @llvm.aarch64.sve.dup.x.nxv4i32(i32 1)
+  %out = call <vscale x 4 x float> @llvm.aarch64.sve.tbl.nxv4f32(<vscale x 4 x float>  %data, <vscale x 4 x i32> %tmp)
+  ret <vscale x 4 x float> %out
+}
+
+define <vscale x 2 x double> @dup_ext_f64(<vscale x 2 x double> %data) {
+; CHECK-LABEL: @dup_ext_f64(
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <vscale x 2 x double> [[DATA:%.*]], i64 1
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 2 x double> poison, double [[TMP1]], i32 0
+; CHECK-NEXT:    [[DOTSPLAT:%.*]] = shufflevector <vscale x 2 x double> [[DOTSPLATINSERT]], <vscale x 2 x double> poison, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT:    ret <vscale x 2 x double> [[DOTSPLAT]]
+;
+  %tmp = call <vscale x 2 x i64> @llvm.aarch64.sve.dup.x.nxv2i64(i64 1)
+  %out = call <vscale x 2 x double> @llvm.aarch64.sve.tbl.nxv2f64(<vscale x 2 x double> %data, <vscale x 2 x i64> %tmp)
+  ret <vscale x 2 x double> %out
+}
+
+declare <vscale x 16 x i8> @llvm.aarch64.sve.dup.x.nxv16i8( i8)
+declare <vscale x 8 x i16> @llvm.aarch64.sve.dup.x.nxv8i16(i16)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.dup.x.nxv4i32(i32)
+declare <vscale x 2 x i64> @llvm.aarch64.sve.dup.x.nxv2i64(i64)
+declare <vscale x 8 x half> @llvm.aarch64.sve.dup.x.nxv8f16(half)
+declare <vscale x 2 x float> @llvm.aarch64.sve.dup.x.nxv2f32(float)
+declare <vscale x 4 x float> @llvm.aarch64.sve.dup.x.nxv4f32(float)
+declare <vscale x 2 x double> @llvm.aarch64.sve.dup.x.nxv2f64(double)
+declare <vscale x 16 x i8> @llvm.aarch64.sve.tbl.nxv16i8( <vscale x 16 x i8>, <vscale x 16 x i8>)
+declare <vscale x 8 x i16> @llvm.aarch64.sve.tbl.nxv8i16( <vscale x 8 x i16>, <vscale x 8 x i16>)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.tbl.nxv4i32( <vscale x 4 x i32>, <vscale x 4 x i32>)
+declare <vscale x 2 x i64> @llvm.aarch64.sve.tbl.nxv2i64( <vscale x 2 x i64>, <vscale x 2 x i64>)
+declare <vscale x 8 x half> @llvm.aarch64.sve.tbl.nxv8f16( <vscale x 8 x half>, <vscale x 8 x i16>)
+declare <vscale x 4 x float> @llvm.aarch64.sve.tbl.nxv4f32( <vscale x 4 x float>, <vscale x 4 x i32>)
+declare <vscale x 2 x double> @llvm.aarch64.sve.tbl.nxv2f64( <vscale x 2 x double>, <vscale x 2 x i64>)