[AArch64][SVE][InstCombine] Move last{a,b} before binop if one operand is a splat value

[mnadeem]

Summary

Move the last{a,b} operation to the vector operand of the binary instruction if
the binop's operand is a splat value. This essentially converts the binop
to a scalar operation.

Example:

// If x and/or y is a splat value:
lastX (binop (x, y)) --> binop(lastX(x), lastX(y))

[mnadeem]

Seeing this pattern internally, looking for comments about whether this is relevant to the community.

[david-arm]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

564

Is this always guaranteed to be safe? Could one of the lanes have previously generated an exception that we have now dropped?

llvm/test/Transforms/InstCombine/AArch64/sve-intrinsic-opts-lasta-lastb.ll

188

It might be worth having a FP binary operator test too here, since the m_BinOp match covers those as well.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

564

@david-arm I don't believe there's any need to worry about exceptions here as LLVM BinOps are defined to have no side effects.

575

Should this just be an else {? as by this point getSplatValue(LHS) should be producing a result that is known safe to use.

llvm/test/Transforms/InstCombine/AArch64/sve-intrinsic-opts-lasta-lastb.ll

188

I agree as looking at the code above I think these will highlight that the transform is going to unnecessarily drop fastmath flags. The same is possibly also true for things like sdiv's exact flag.

[mnadeem]

• Updated the code to make sure that the binary operator's flags are copied as well.
• Added more tests to check that the math flags are also copied.
• Added tests for add and fadd.

[mnadeem]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

575

Removed the use of getSplatValue() altogether.

[paulwalker-arm]

llvm/test/Transforms/InstCombine/AArch64/sve-intrinsic-opts-lasta-lastb.ll

278–348

Given you have div and div tests with and without flags, do these add/fadd tests testing anything new? If not then there's not much point keeping them.

278–348

Given you have div and fdiv tests with and without flags, do these add/fadd tests testing anything new? If not then there's not much point keeping them.

350

I was going to group this test with the above comment but looking at the output there is no transformation. Is this intended? I would have expected a scalar add.

[paulwalker-arm]

Sorry forgot to ask if you've used update_test_checks.py to generate the CHECK lines within sve-intrinsic-opts-lasta-lastb.ll like the other tests within the file. If not then please do so as otherwise the next person along will have unnecessary changes in their patch.

[mnadeem]

Fix formatting, remove/update tests according to comments.

[mnadeem]

llvm/test/Transforms/InstCombine/AArch64/sve-intrinsic-opts-lasta-lastb.ll

350

Removed that test as it was not relevant to this patch. The shuffle was pushed after the add by some other transformation.

[paulwalker-arm]

llvm/test/Transforms/InstCombine/AArch64/sve-intrinsic-opts-lasta-lastb.ll

287

don't

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

auto *ConvertFromSVBool =
                               {II.getType()}, {ConvertToSVBool});
 
   ConvertFromSVBool->takeName(&II);
   return IC.replaceInstUsesWith(II, ConvertFromSVBool);
 }
 
 static Optional<Instruction *> instCombineSVELast(InstCombiner &IC,
                                                   IntrinsicInst &II) {
+  IRBuilder<> Builder(II.getContext());
+  Builder.SetInsertPoint(&II);
   Value *Pg = II.getArgOperand(0);
   Value *Vec = II.getArgOperand(1);
-  bool IsAfter = II.getIntrinsicID() == Intrinsic::aarch64_sve_lasta;
+  auto IntrinsicID = II.getIntrinsicID();
+  bool IsAfter = IntrinsicID == Intrinsic::aarch64_sve_lasta;
 
   // lastX(splat(X)) --> X
   if (auto *SplatVal = getSplatValue(Vec))
     return IC.replaceInstUsesWith(II, SplatVal);
 
+  // If x and/or y is a splat value then:
+  // lastX (binop (x, y)) --> binop(lastX(x), lastX(y))
+  Value *LHS, *RHS;
+  if (match(Vec, m_OneUse(m_BinOp(m_Value(LHS), m_Value(RHS))))) {
+    if (isSplatValue(LHS) || isSplatValue(RHS)) {
+      auto *OldBinOp = cast<BinaryOperator>(Vec);
+      auto OpC = OldBinOp->getOpcode();
+      auto *NewLHS =
+          Builder.CreateIntrinsic(IntrinsicID, {Vec->getType()}, {Pg, LHS});
+      auto *NewRHS =
+          Builder.CreateIntrinsic(IntrinsicID, {Vec->getType()}, {Pg, RHS});
+      auto *NewBinOp = BinaryOperator::CreateWithCopiedFlags(
+          OpC, NewLHS, NewRHS, OldBinOp, OldBinOp->getName(), &II);
+      return IC.replaceInstUsesWith(II, NewBinOp);
+    }
+  }
+
   auto *C = dyn_cast<Constant>(Pg);
   if (IsAfter && C && C->isNullValue()) {
     // The intrinsic is extracting lane 0 so use an extract instead.
     auto *IdxTy = Type::getInt64Ty(II.getContext());
     auto *Extract = ExtractElementInst::Create(Vec, ConstantInt::get(IdxTy, 0));
     Extract->insertBefore(&II);
     Extract->takeName(&II);
     return IC.replaceInstUsesWith(II, Extract);

llvm/test/Transforms/InstCombine/AArch64/sve-intrinsic-opts-lasta-lastb.ll

;
   %pg = tail call <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 15)
   %e = tail call i8 @llvm.aarch64.sve.lastb.nxv16i8(<vscale x 16 x i1> %pg, <vscale x 16 x i8> %v)
   ret i8 %e
 }
 
 ; Return the splatted value irrespective of the predicate.
 define i8 @lasta_splat(<vscale x 16 x i1> %pg, i8 %a) #0 {
 ; OPT-LABEL: @lasta_splat(
-; OPT-NEXT:    ret i8 %a
+; OPT-NEXT:    ret i8 [[A:%.*]]
+;
   %splat_insert = insertelement <vscale x 16 x i8> poison, i8 %a, i32 0
   %splat = shufflevector <vscale x 16 x i8> %splat_insert, <vscale x 16 x i8> poison, <vscale x 16 x i32> zeroinitializer
   %last = tail call i8 @llvm.aarch64.sve.lasta.nxv16i8(<vscale x 16 x i1> %pg, <vscale x 16 x i8> %splat)
   ret i8 %last
 }
 
 define i8 @lastb_splat(<vscale x 16 x i1> %pg, i8 %a) #0 {
 ; OPT-LABEL: @lastb_splat(
-; OPT-NEXT:    ret i8 %a
+; OPT-NEXT:    ret i8 [[A:%.*]]
+;
   %splat_insert = insertelement <vscale x 16 x i8> poison, i8 %a, i32 0
   %splat = shufflevector <vscale x 16 x i8> %splat_insert, <vscale x 16 x i8> poison, <vscale x 16 x i32> zeroinitializer
   %last = tail call i8 @llvm.aarch64.sve.lastb.nxv16i8(<vscale x 16 x i1> %pg, <vscale x 16 x i8> %splat)
   ret i8 %last
 }
 
+; Check that we move the lastb before the binary operation so that the new binary op is scalar.
+define i8 @lastb_binop_RHS_splat_sdiv(<vscale x 16 x i1> %pg, i8 %scalar, <vscale x 16 x i8> %vector) #0 {
+; OPT-LABEL: @lastb_binop_RHS_splat_sdiv(
+; OPT-NEXT:    [[TMP1:%.*]] = call i8 @llvm.aarch64.sve.lastb.nxv16i8(<vscale x 16 x i1> [[PG:%.*]], <vscale x 16 x i8> [[VECTOR:%.*]])
+; OPT-NEXT:    [[BINOP1:%.*]] = sdiv i8 [[TMP1]], [[SCALAR:%.*]]
+; OPT-NEXT:    ret i8 [[BINOP1]]
+;
+  %splat_insert = insertelement <vscale x 16 x i8> poison, i8 %scalar, i32 0
+  %splat = shufflevector <vscale x 16 x i8> %splat_insert, <vscale x 16 x i8> poison, <vscale x 16 x i32> zeroinitializer
+  %binop = sdiv <vscale x 16 x i8> %vector, %splat
+  %last = tail call i8 @llvm.aarch64.sve.lastb.nxv16i8(<vscale x 16 x i1> %pg, <vscale x 16 x i8> %binop)
+  ret i8 %last
+}
+
+define i8 @lastb_binop_RHS_splat_sdiv_exact(<vscale x 16 x i1> %pg, i8 %scalar, <vscale x 16 x i8> %vector) #0 {
+; OPT-LABEL: @lastb_binop_RHS_splat_sdiv_exact(
+; OPT-NEXT:    [[TMP1:%.*]] = call i8 @llvm.aarch64.sve.lastb.nxv16i8(<vscale x 16 x i1> [[PG:%.*]], <vscale x 16 x i8> [[VECTOR:%.*]])
+; OPT-NEXT:    [[BINOP1:%.*]] = sdiv exact i8 [[TMP1]], [[SCALAR:%.*]]
+; OPT-NEXT:    ret i8 [[BINOP1]]
+;
+  %splat_insert = insertelement <vscale x 16 x i8> poison, i8 %scalar, i32 0
+  %splat = shufflevector <vscale x 16 x i8> %splat_insert, <vscale x 16 x i8> poison, <vscale x 16 x i32> zeroinitializer
+  %binop = sdiv exact <vscale x 16 x i8> %vector, %splat
+  %last = tail call i8 @llvm.aarch64.sve.lastb.nxv16i8(<vscale x 16 x i1> %pg, <vscale x 16 x i8> %binop)
+  ret i8 %last
+}
+
+define float @lastb_binop_RHS_splat_fdiv_float_fast(<vscale x 4 x i1> %pg, float %scalar, <vscale x 4 x float> %vector) #0 {
+; OPT-LABEL: @lastb_binop_RHS_splat_fdiv_float_fast(
+; OPT-NEXT:    [[TMP1:%.*]] = call float @llvm.aarch64.sve.lastb.nxv4f32(<vscale x 4 x i1> [[PG:%.*]], <vscale x 4 x float> [[VECTOR:%.*]])
+; OPT-NEXT:    [[BINOP1:%.*]] = fdiv fast float [[TMP1]], [[SCALAR:%.*]]
+; OPT-NEXT:    ret float [[BINOP1]]
+;
+  %splat_insert = insertelement <vscale x 4 x float> poison, float %scalar, i32 0
+  %splat = shufflevector <vscale x 4 x float> %splat_insert, <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer
+  %binop = fdiv fast <vscale x 4 x float> %vector, %splat
+  %last = tail call float @llvm.aarch64.sve.lastb.nxv4f32(<vscale x 4 x i1> %pg, <vscale x 4 x float> %binop)
+  ret float %last
+}
+
+define float @lastb_binop_RHS_splat_fdiv_float(<vscale x 4 x i1> %pg, float %scalar, <vscale x 4 x float> %vector) #0 {
+; OPT-LABEL: @lastb_binop_RHS_splat_fdiv_float(
+; OPT-NEXT:    [[TMP1:%.*]] = call float @llvm.aarch64.sve.lastb.nxv4f32(<vscale x 4 x i1> [[PG:%.*]], <vscale x 4 x float> [[VECTOR:%.*]])
+; OPT-NEXT:    [[BINOP1:%.*]] = fdiv float [[TMP1]], [[SCALAR:%.*]]
+; OPT-NEXT:    ret float [[BINOP1]]
+;
+  %splat_insert = insertelement <vscale x 4 x float> poison, float %scalar, i32 0
+  %splat = shufflevector <vscale x 4 x float> %splat_insert, <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer
+  %binop = fdiv <vscale x 4 x float> %vector, %splat
+  %last = tail call float @llvm.aarch64.sve.lastb.nxv4f32(<vscale x 4 x i1> %pg, <vscale x 4 x float> %binop)
+  ret float %last
+}
+
+define i8 @lastb_binop_LHS_splat_sdiv(<vscale x 16 x i1> %pg, i8 %scalar, <vscale x 16 x i8> %vector) #0 {
+; OPT-LABEL: @lastb_binop_LHS_splat_sdiv(
+; OPT-NEXT:    [[TMP1:%.*]] = call i8 @llvm.aarch64.sve.lastb.nxv16i8(<vscale x 16 x i1> [[PG:%.*]], <vscale x 16 x i8> [[VECTOR:%.*]])
+; OPT-NEXT:    [[BINOP1:%.*]] = sdiv i8 [[SCALAR:%.*]], [[TMP1]]
+; OPT-NEXT:    ret i8 [[BINOP1]]
+;
+  %splat_insert = insertelement <vscale x 16 x i8> poison, i8 %scalar, i32 0
+  %splat = shufflevector <vscale x 16 x i8> %splat_insert, <vscale x 16 x i8> poison, <vscale x 16 x i32> zeroinitializer
+  %binop = sdiv <vscale x 16 x i8> %splat, %vector
+  %last = tail call i8 @llvm.aarch64.sve.lastb.nxv16i8(<vscale x 16 x i1> %pg, <vscale x 16 x i8> %binop)
+  ret i8 %last
+}
+
+define i8 @lastb_binop_LHS_splat_sdiv_exact(<vscale x 16 x i1> %pg, i8 %scalar, <vscale x 16 x i8> %vector) #0 {
+; OPT-LABEL: @lastb_binop_LHS_splat_sdiv_exact(
+; OPT-NEXT:    [[TMP1:%.*]] = call i8 @llvm.aarch64.sve.lastb.nxv16i8(<vscale x 16 x i1> [[PG:%.*]], <vscale x 16 x i8> [[VECTOR:%.*]])
+; OPT-NEXT:    [[BINOP1:%.*]] = sdiv exact i8 [[SCALAR:%.*]], [[TMP1]]
+; OPT-NEXT:    ret i8 [[BINOP1]]
+;
+  %splat_insert = insertelement <vscale x 16 x i8> poison, i8 %scalar, i32 0
+  %splat = shufflevector <vscale x 16 x i8> %splat_insert, <vscale x 16 x i8> poison, <vscale x 16 x i32> zeroinitializer
+  %binop = sdiv exact <vscale x 16 x i8> %splat, %vector
+  %last = tail call i8 @llvm.aarch64.sve.lastb.nxv16i8(<vscale x 16 x i1> %pg, <vscale x 16 x i8> %binop)
+  ret i8 %last
+}
+
+define float @lastb_binop_LHS_splat_fdiv_float_fast(<vscale x 4 x i1> %pg, float %scalar, <vscale x 4 x float> %vector) #0 {
+; OPT-LABEL: @lastb_binop_LHS_splat_fdiv_float_fast(
+; OPT-NEXT:    [[TMP1:%.*]] = call float @llvm.aarch64.sve.lastb.nxv4f32(<vscale x 4 x i1> [[PG:%.*]], <vscale x 4 x float> [[VECTOR:%.*]])
+; OPT-NEXT:    [[BINOP1:%.*]] = fdiv fast float [[SCALAR:%.*]], [[TMP1]]
+; OPT-NEXT:    ret float [[BINOP1]]
+;
+  %splat_insert = insertelement <vscale x 4 x float> poison, float %scalar, i32 0
+  %splat = shufflevector <vscale x 4 x float> %splat_insert, <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer
+  %binop = fdiv fast <vscale x 4 x float> %splat, %vector
+  %last = tail call float @llvm.aarch64.sve.lastb.nxv4f32(<vscale x 4 x i1> %pg, <vscale x 4 x float> %binop)
+  ret float %last
+}
+
+define float @lastb_binop_LHS_splat_fdiv_float(<vscale x 4 x i1> %pg, float %scalar, <vscale x 4 x float> %vector) #0 {
+; OPT-LABEL: @lastb_binop_LHS_splat_fdiv_float(
+; OPT-NEXT:    [[TMP1:%.*]] = call float @llvm.aarch64.sve.lastb.nxv4f32(<vscale x 4 x i1> [[PG:%.*]], <vscale x 4 x float> [[VECTOR:%.*]])
+; OPT-NEXT:    [[BINOP1:%.*]] = fdiv float [[SCALAR:%.*]], [[TMP1]]
+; OPT-NEXT:    ret float [[BINOP1]]
+;
+  %splat_insert = insertelement <vscale x 4 x float> poison, float %scalar, i32 0
+  %splat = shufflevector <vscale x 4 x float> %splat_insert, <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer
+  %binop = fdiv <vscale x 4 x float> %splat, %vector
+  %last = tail call float @llvm.aarch64.sve.lastb.nxv4f32(<vscale x 4 x i1> %pg, <vscale x 4 x float> %binop)
+  ret float %last
+}
+
+define i8 @lastb_binop_LHS_RHS_splat_sdiv(<vscale x 16 x i1> %pg, i8 %scalar1, i8 %scalar2) #0 {
+; OPT-LABEL: @lastb_binop_LHS_RHS_splat_sdiv(
+; OPT-NEXT:    [[BINOP1:%.*]] = sdiv i8 [[SCALAR1:%.*]], [[SCALAR2:%.*]]
+; OPT-NEXT:    ret i8 [[BINOP1]]
+;
+  %splat_insert = insertelement <vscale x 16 x i8> poison, i8 %scalar1, i8 0
+  %splat = shufflevector <vscale x 16 x i8> %splat_insert, <vscale x 16 x i8> poison, <vscale x 16 x i32> zeroinitializer
+  %splat_insert2 = insertelement <vscale x 16 x i8> poison, i8 %scalar2, i8 0
+  %splat2 = shufflevector <vscale x 16 x i8> %splat_insert2, <vscale x 16 x i8> poison, <vscale x 16 x i32> zeroinitializer
+  %binop = sdiv <vscale x 16 x i8> %splat, %splat2
+  %last = tail call i8 @llvm.aarch64.sve.lastb.nxv16i8(<vscale x 16 x i1> %pg, <vscale x 16 x i8> %binop)
+  ret i8 %last
+}
+
+; Check that we dont do anything as the binary op has multiple uses.
+define i8 @lastb_binop_nochange(<vscale x 16 x i1> %pg, i8 %scalar, <vscale x 16 x i8> %vector) #0 {
+; OPT-LABEL: @lastb_binop_nochange(
+; OPT-NEXT:    [[SPLAT_INSERT:%.*]] = insertelement <vscale x 16 x i8> poison, i8 [[SCALAR:%.*]], i32 0
+; OPT-NEXT:    [[SPLAT:%.*]] = shufflevector <vscale x 16 x i8> [[SPLAT_INSERT]], <vscale x 16 x i8> poison, <vscale x 16 x i32> zeroinitializer
+; OPT-NEXT:    [[BINOP:%.*]] = sdiv <vscale x 16 x i8> [[VECTOR:%.*]], [[SPLAT]]
+; OPT-NEXT:    [[LAST:%.*]] = tail call i8 @llvm.aarch64.sve.lastb.nxv16i8(<vscale x 16 x i1> [[PG:%.*]], <vscale x 16 x i8> [[BINOP]])
+; OPT-NEXT:    call void @use(<vscale x 16 x i8> [[BINOP]])
+; OPT-NEXT:    ret i8 [[LAST]]
+;
+  %splat_insert = insertelement <vscale x 16 x i8> poison, i8 %scalar, i32 0
+  %splat = shufflevector <vscale x 16 x i8> %splat_insert, <vscale x 16 x i8> poison, <vscale x 16 x i32> zeroinitializer
+  %binop = sdiv <vscale x 16 x i8> %vector, %splat
+  %last = tail call i8 @llvm.aarch64.sve.lastb.nxv16i8(<vscale x 16 x i1> %pg, <vscale x 16 x i8> %binop)
+  call void @use(<vscale x 16 x i8> %binop)
+  ret i8 %last
+}
+
+declare void @use(<vscale x 16 x i8>)
 declare <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32)
 declare i8 @llvm.aarch64.sve.lasta.nxv16i8(<vscale x 16 x i1>, <vscale x 16 x i8>)
 declare i8 @llvm.aarch64.sve.lastb.nxv16i8(<vscale x 16 x i1>, <vscale x 16 x i8>)
+declare float @llvm.aarch64.sve.lastb.nxv4f32(<vscale x 4 x i1>, <vscale x 4 x float>)
 
 attributes #0 = { "target-features"="+sve" }