[AArch64] Implement getArithmeticReductionCost

[mssimpso]

Summary

This patch provides an implementation of getArithmeticReductionCost for AArch64. We can specialize the cost of add reductions since they are computed using the 'addv' instruction.

[evandro]

Do you intend to consider AArch64TTIImpl::getMinMaxRdxCost() too?

[mssimpso]

In D44490#1037716, @evandro wrote:

Do you intend to consider AArch64TTIImpl::getMinMaxRdxCost() too?

Yes, I'm planning to implement getMinMaxRdxCost in a follow-on patch.

[samparker]

I'm happy to see this! LGTM with just one small comment, no need to re-review. Thanks!

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

911

You could perform the IsPairwise check at the beginning of the function to exit early.

[mssimpso]

In D44490#1038465, @samparker wrote:

I'm happy to see this! LGTM with just one small comment, no need to re-review. Thanks!

Thanks, Sam!

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

911

Yes, will do!

lib/Target/AArch64/AArch64TargetTransformInfo.h

public:
   unsigned getMaxPrefetchIterationsAhead();
 
   bool shouldExpandReduction(const IntrinsicInst *II) const {
     return false;
   }
 
   bool useReductionIntrinsic(unsigned Opcode, Type *Ty,
                              TTI::ReductionFlags Flags) const;
+
+  int getArithmeticReductionCost(unsigned Opcode, Type *Ty,
+                                 bool IsPairwiseForm);
   /// @}
 };
 
 } // end namespace llvm
 
 #endif // LLVM_LIB_TARGET_AARCH64_AARCH64TARGETTRANSFORMINFO_H

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

return (ScalarBits < 64) &&
            (ScalarBits * Ty->getVectorNumElements() >= 128);
   case Instruction::FCmp:
     return Flags.NoNaN;
   default:
     llvm_unreachable("Unhandled reduction opcode");
   }
   return false;
 }
+
+int AArch64TTIImpl::getArithmeticReductionCost(unsigned Opcode, Type *ValTy,
+                                               bool IsPairwise) {
+
+  std::pair<int, MVT> LT = TLI->getTypeLegalizationCost(DL, ValTy);
+  MVT MTy = LT.second;
+  int ISD = TLI->InstructionOpcodeToISD(Opcode);
+  assert(ISD && "Invalid opcode");
+
+  // Horizontal adds can use the 'addv' instruction. We model the cost of these
+  // instructions as normal vector adds. This is the only arithmetic vector
+  // reduction operation for which we have an instruction.
+  static const CostTblEntry CostTblNoPairwise[]{
+      {ISD::ADD, MVT::v8i8,  1},
+      {ISD::ADD, MVT::v16i8, 1},
+      {ISD::ADD, MVT::v4i16, 1},
+      {ISD::ADD, MVT::v8i16, 1},
+      {ISD::ADD, MVT::v4i32, 1},
+  };
+
+  if (!IsPairwise)
+    if (const auto *Entry = CostTableLookup(CostTblNoPairwise, ISD, MTy))
+      return LT.first * Entry->Cost;
+
+  return BaseT::getArithmeticReductionCost(Opcode, ValTy, IsPairwise);
+}

test/Analysis/CostModel/AArch64/vector-reduce.ll

 ; RUN: opt < %s -mtriple=aarch64--linux-gnu -cost-model -analyze | FileCheck %s --check-prefix=COST
 ; RUN: llc < %s -mtriple=aarch64--linux-gnu | FileCheck %s --check-prefix=CODE
 
 ; COST-LABEL: add.i8.v8i8
-; COST:       Found an estimated cost of 27 for instruction: %r = call i8 @llvm.experimental.vector.reduce.add.i8.v8i8(<8 x i8> %v)
+; COST:       Found an estimated cost of 1 for instruction: %r = call i8 @llvm.experimental.vector.reduce.add.i8.v8i8(<8 x i8> %v)
 ; CODE-LABEL: add.i8.v8i8
 ; CODE:       addv b0, v0.8b
 define i8 @add.i8.v8i8(<8 x i8> %v) {
   %r = call i8 @llvm.experimental.vector.reduce.add.i8.v8i8(<8 x i8> %v)
   ret i8 %r
 }
 
 ; COST-LABEL: add.i8.v16i8
-; COST:       Found an estimated cost of 53 for instruction: %r = call i8 @llvm.experimental.vector.reduce.add.i8.v16i8(<16 x i8> %v)
+; COST:       Found an estimated cost of 1 for instruction: %r = call i8 @llvm.experimental.vector.reduce.add.i8.v16i8(<16 x i8> %v)
 ; CODE-LABEL: add.i8.v16i8
 ; CODE:       addv b0, v0.16b
 define i8 @add.i8.v16i8(<16 x i8> %v) {
   %r = call i8 @llvm.experimental.vector.reduce.add.i8.v16i8(<16 x i8> %v)
   ret i8 %r
 }
 
 ; COST-LABEL: add.i16.v4i16
-; COST:       Found an estimated cost of 13 for instruction: %r = call i16 @llvm.experimental.vector.reduce.add.i16.v4i16(<4 x i16> %v)
+; COST:       Found an estimated cost of 1 for instruction: %r = call i16 @llvm.experimental.vector.reduce.add.i16.v4i16(<4 x i16> %v)
 ; CODE-LABEL: add.i16.v4i16
 ; CODE:       addv h0, v0.4h
 define i16 @add.i16.v4i16(<4 x i16> %v) {
   %r = call i16 @llvm.experimental.vector.reduce.add.i16.v4i16(<4 x i16> %v)
   ret i16 %r
 }
 
 ; COST-LABEL: add.i16.v8i16
-; COST:       Found an estimated cost of 27 for instruction: %r = call i16 @llvm.experimental.vector.reduce.add.i16.v8i16(<8 x i16> %v)
+; COST:       Found an estimated cost of 1 for instruction: %r = call i16 @llvm.experimental.vector.reduce.add.i16.v8i16(<8 x i16> %v)
 ; CODE-LABEL: add.i16.v8i16
 ; CODE:       addv h0, v0.8h
 define i16 @add.i16.v8i16(<8 x i16> %v) {
   %r = call i16 @llvm.experimental.vector.reduce.add.i16.v8i16(<8 x i16> %v)
   ret i16 %r
 }
 
 ; COST-LABEL: add.i32.v4i32
-; COST:       Found an estimated cost of 13 for instruction: %r = call i32 @llvm.experimental.vector.reduce.add.i32.v4i32(<4 x i32> %v)
+; COST:       Found an estimated cost of 1 for instruction: %r = call i32 @llvm.experimental.vector.reduce.add.i32.v4i32(<4 x i32> %v)
 ; CODE-LABEL: add.i32.v4i32
 ; CODE:       addv s0, v0.4s
 define i32 @add.i32.v4i32(<4 x i32> %v) {
   %r = call i32 @llvm.experimental.vector.reduce.add.i32.v4i32(<4 x i32> %v)
   ret i32 %r
 }
 
 ; COST-LABEL: umin.i8.v8i8

test/Transforms/SLPVectorizer/AArch64/horizontal.ll

 ; CHECK: select <4 x i1>
 
 ; YAML:      --- !Passed
 ; YAML-NEXT: Pass:            slp-vectorizer
 ; YAML-NEXT: Name:            VectorizedHorizontalReduction
 ; YAML-NEXT: Function:        test_select
 ; YAML-NEXT: Args:
 ; YAML-NEXT:   - String:          'Vectorized horizontal reduction with cost '
-; YAML-NEXT:   - Cost:            '4'
+; YAML-NEXT:   - Cost:            '-8'
 ; YAML-NEXT:   - String:          ' and with tree size '
 ; YAML-NEXT:   - TreeSize:        '8'
 
 define i32 @test_select(i32* noalias nocapture readonly %blk1, i32* noalias nocapture readonly %blk2, i32 %lx, i32 %h) {
 entry:
   %cmp.22 = icmp sgt i32 %h, 0
   br i1 %cmp.22, label %for.body.lr.ph, label %for.end
 

 ; CHECK: mul nsw <4 x i32>
 
 ; YAML:      --- !Passed
 ; YAML-NEXT: Pass:            slp-vectorizer
 ; YAML-NEXT: Name:            VectorizedHorizontalReduction
 ; YAML-NEXT: Function:        reduction_with_br
 ; YAML-NEXT: Args:
 ; YAML-NEXT:   - String:          'Vectorized horizontal reduction with cost '
-; YAML-NEXT:   - Cost:            '1'
+; YAML-NEXT:   - Cost:            '-11'
 ; YAML-NEXT:   - String:          ' and with tree size '
 ; YAML-NEXT:   - TreeSize:        '3'
 
 entry:
   %cmp.16 = icmp sgt i32 %h, 0
   br i1 %cmp.16, label %for.body.lr.ph, label %for.end
 
 for.body.lr.ph:                                   ; preds = %entry

 ; CHECK: select <8 x i1>
 
 ; YAML:      --- !Passed
 ; YAML-NEXT: Pass:            slp-vectorizer
 ; YAML-NEXT: Name:            VectorizedHorizontalReduction
 ; YAML-NEXT: Function:        test_unrolled_select
 ; YAML-NEXT: Args:
 ; YAML-NEXT:   - String:          'Vectorized horizontal reduction with cost '
-; YAML-NEXT:   - Cost:            '-33'
+; YAML-NEXT:   - Cost:            '-47'
 ; YAML-NEXT:   - String:          ' and with tree size '
 ; YAML-NEXT:   - TreeSize:        '10'
 
 define i32 @test_unrolled_select(i8* noalias nocapture readonly %blk1, i8* noalias nocapture readonly %blk2, i32 %lx, i32 %h, i32 %lim) #0 {
 entry:
   %cmp.43 = icmp sgt i32 %h, 0
   br i1 %cmp.43, label %for.body.lr.ph, label %for.end