[CostModel] Add costs for llvm.experimental.vector.{extract,insert} intrinsics

[bsmith]

Summary

This patch add cost model support for the new llvm.experimental.vector.{extract,insert} intrinsics, using the existing getExtractSubvectorOverhead/getInsertSubvectorOverhead functions for shuffles.

Previously this case would throw an assertion.

Depends on D92094

[joechrisellis]

Hi @bsmith, thank you for the patch! 😄

I noticed b3e77c6d55853eea5f5c32ec8a3510c0b0e438e1 earlier:

commit b3e77c6d55853eea5f5c32ec8a3510c0b0e438e1
Author: Christopher Tetreault <[e-mail]>
Date:   Tue Jun 16 14:05:21 2020 -0700

    [SVE] Remove invalid calls to VectorType::getNumElements from BasicTTIImpl

    Summary:
    Most of these operations are reasonable for scalable vectors. Due to
    this, we have decided not to change the interface to specifically take
    FixedVectorType despite the fact that the current implementations make
    fixed width assumptions. Instead, we cast to FixedVectorType and assert
    in the body. If a developer makes some change in the future that causes
    one of these asserts to fire, they should either change their code or
    make the function they are trying to call handle scalable vectors.

    Reviewers: efriedma, samparker, RKSimon, craig.topper, sdesmalen, c-rhodes

    Reviewed By: efriedma

    Subscribers: tschuett, rkruppe, psnobl, llvm-commits

    Tags: #llvm

    Differential Revision: https://reviews.llvm.org/D81495

This patch changes the interface of both get{Extract,Insert}SubvectorOverhead (among others) to accepting FixedVectorTypes instead of VectorTypes, so this patch is partially reverting some of the changes there. Not sure if this is pertinent here, but perhaps it's useful for context.

llvm/include/llvm/CodeGen/BasicTTIImpl.h

118

llvm.experimental.vector.extract can also extract a scalable from a scalable -- should we also change this FixedVectorType to VectorType too?

I suppose this cost model breaks down a bit when we're extracting/inserting a scalable into a scalable, since we don't know how many elements we have to 'work on'. But I think that it might be possible for this code to crash in the same way as before if we don't address this case.

143

Same as above -- do we want to change this too?

llvm/test/Analysis/CostModel/AArch64/sve-getIntrinsicInstrCost-vec-insert-extract.ll

22

Related to the above comments, could we add some tests for inserting/extracting a scalable from a scalable? 😄

[bsmith]

In D93043#2448152, @joechrisellis wrote:

Hi @bsmith, thank you for the patch! 😄

I noticed b3e77c6d55853eea5f5c32ec8a3510c0b0e438e1 earlier:

commit b3e77c6d55853eea5f5c32ec8a3510c0b0e438e1
Author: Christopher Tetreault <[e-mail]>
Date:   Tue Jun 16 14:05:21 2020 -0700

    [SVE] Remove invalid calls to VectorType::getNumElements from BasicTTIImpl

    <snip>
    If a developer makes some change in the future that causes
    one of these asserts to fire, they should either change their code or
    make the function they are trying to call handle scalable vectors.
    <snip>

This patch changes the interface of both get{Extract,Insert}SubvectorOverhead (among others) to accepting FixedVectorTypes instead of VectorTypes, so this patch is partially reverting some of the changes there. Not sure if this is pertinent here, but perhaps it's useful for context.

Agreed, I see this patch as doing as the highlighted part of the commit message indicates, hence not really a reversion as such.

llvm/include/llvm/CodeGen/BasicTTIImpl.h

118

I think currently we don't generate these intrinsics as extracting or inserting scalable vectors into another scalable vector, and considering adding costs for that is much more complex it should be a different patch, if we do end up adding it.

That said, I agree that this case shouldn't crash, which it does currently. I will make those cases fall back to BaseT::getIntrinsicInstrCost(ICA, CostKind) as the other scalable vector cases do, as well as adding tests.

[bsmith]

• Fallback to default cost when using llvm.experimental.vector.{extract,insert} intrinsics with two scalable vectors
• Add tests for above cases to ensure it no longer asserts

[joechrisellis]

LGTM, but I am curious if @ctetreau has an opinion on this change, since I think he has done some work in this area.

llvm/include/llvm/CodeGen/BasicTTIImpl.h

118

SGTM!

[CarolineConcatto]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

121

If this starts to accept scalable vector, by replacing FixedVectorType to VectorType, this could be a problem.
See this function unsigned VectorType::getNumElements() const; inDerivedTypes.h
Maybe you will need to use InstructionCost class or just assert at the moment when the type is scalable

[bsmith]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

121

Is adding another assert is really necessary? In my mind if someone changed SubVTy from FixedVectorType to VectorType they should also make this function support scalable vectors.

The fact that SubVTy is a FixedVectorType means there already is an assert and statement that this function doesn't support scalable vectors for SubVTy.

[CarolineConcatto]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

121

The problem here is also that you are using NumSubElts to compute the cost.
For scalable vector the number of elements is not know at compiler time, only when in run time.
So, this is cost is not entirely correct, because it computes only the minimum cost.
I don't know if you looked the function VectorType::getNumElements() , but it return getKnonwMinValue().
For this example <vscale x 16 x i32> it would be 16, but you still needs to know the value of vscale.
Because the total number of elements if vscale * 16, for this example.

[CarolineConcatto]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

121

But after saying all that I think this part is OK
Because you are still keeping FixedVectorType *SubVTy.

[bsmith]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

121

Correct, if this function started taking a scalable vector for SubVTy, the use of getNumElements would need to change. For now in that case it falls back to BaseT::getIntrinsicInstrCost(ICA, CostKind) before getting to this function, hence should be fine.

[CarolineConcatto]

Just in case I had a look at the other functions:
getIntrinsicInstrCost
getShuffleCost and
getVectorInstrCost
To see what would be their behave now that you change some of the types to accept scalable type.
It all looks good to me.

llvm/include/llvm/CodeGen/BasicTTIImpl.h

for (int i = 0, e = VTy->getNumElements(); i < e; ++i) {
       Cost += thisT()->getVectorInstrCost(Instruction::InsertElement, VTy, i);
       Cost += thisT()->getVectorInstrCost(Instruction::ExtractElement, VTy, i);
     }
     return Cost;
   }
 
   /// Estimate a cost of subvector extraction as a sequence of extract and
   /// insert operations.
-  unsigned getExtractSubvectorOverhead(FixedVectorType *VTy, int Index,
+  unsigned getExtractSubvectorOverhead(VectorType *VTy, int Index,
                                        FixedVectorType *SubVTy) {
     assert(VTy && SubVTy &&
            "Can only extract subvectors from vectors");
     int NumSubElts = SubVTy->getNumElements();
-    assert((Index + NumSubElts) <= (int)VTy->getNumElements() &&
+    assert((!isa<FixedVectorType>(VTy) ||
+            (Index + NumSubElts) <=
+                (int)cast<FixedVectorType>(VTy)->getNumElements()) &&
            "SK_ExtractSubvector index out of range");
 
     unsigned Cost = 0;
     // Subvector extraction cost is equal to the cost of extracting element from
     // the source type plus the cost of inserting them into the result vector
     // type.
     for (int i = 0; i != NumSubElts; ++i) {
       Cost += thisT()->getVectorInstrCost(Instruction::ExtractElement, VTy,
                                           i + Index);
       Cost +=
           thisT()->getVectorInstrCost(Instruction::InsertElement, SubVTy, i);
     }
     return Cost;
   }
 
   /// Estimate a cost of subvector insertion as a sequence of extract and
   /// insert operations.
-  unsigned getInsertSubvectorOverhead(FixedVectorType *VTy, int Index,
+  unsigned getInsertSubvectorOverhead(VectorType *VTy, int Index,
                                       FixedVectorType *SubVTy) {
     assert(VTy && SubVTy &&
            "Can only insert subvectors into vectors");
     int NumSubElts = SubVTy->getNumElements();
-    assert((Index + NumSubElts) <= (int)VTy->getNumElements() &&
+    assert((!isa<FixedVectorType>(VTy) ||
+            (Index + NumSubElts) <=
+                (int)cast<FixedVectorType>(VTy)->getNumElements()) &&
            "SK_InsertSubvector index out of range");
 
     unsigned Cost = 0;
     // Subvector insertion cost is equal to the cost of extracting element from
     // the source type plus the cost of inserting them into the result vector
     // type.
     for (int i = 0; i != NumSubElts; ++i) {
       Cost +=
case TTI::SK_Broadcast:
       return getBroadcastShuffleOverhead(cast<FixedVectorType>(Tp));
     case TTI::SK_Select:
     case TTI::SK_Reverse:
     case TTI::SK_Transpose:
     case TTI::SK_PermuteSingleSrc:
     case TTI::SK_PermuteTwoSrc:
       return getPermuteShuffleOverhead(cast<FixedVectorType>(Tp));
     case TTI::SK_ExtractSubvector:
-      return getExtractSubvectorOverhead(cast<FixedVectorType>(Tp), Index,
+      return getExtractSubvectorOverhead(Tp, Index,
                                          cast<FixedVectorType>(SubTp));
     case TTI::SK_InsertSubvector:
-      return getInsertSubvectorOverhead(cast<FixedVectorType>(Tp), Index,
+      return getInsertSubvectorOverhead(Tp, Index,
                                         cast<FixedVectorType>(SubTp));
     }
     llvm_unreachable("Unknown TTI::ShuffleKind");
   }
 
   unsigned getCastInstrCost(unsigned Opcode, Type *Dst, Type *Src,
                             TTI::CastContextHint CCH,
                             TTI::TargetCostKind CostKind,
case Intrinsic::masked_gather: {
         return BaseT::getIntrinsicInstrCost(ICA, CostKind);
       assert(VF.isScalar() && "Can't vectorize types here.");
       const Value *Mask = Args[2];
       bool VarMask = !isa<Constant>(Mask);
       Align Alignment = cast<ConstantInt>(Args[1])->getAlignValue();
       return thisT()->getGatherScatterOpCost(Instruction::Load, RetTy, Args[0],
                                              VarMask, Alignment, CostKind, I);
     }
+    case Intrinsic::experimental_vector_extract: {
+      // FIXME: Handle case where a scalable vector is extracted from a scalable
+      // vector
+      if (isa<ScalableVectorType>(RetTy))
+        return BaseT::getIntrinsicInstrCost(ICA, CostKind);
+      unsigned Index = cast<ConstantInt>(Args[1])->getZExtValue();
+      return thisT()->getShuffleCost(TTI::SK_ExtractSubvector,
+                                     cast<VectorType>(Args[0]->getType()),
+                                     Index, cast<VectorType>(RetTy));
+    }
+    case Intrinsic::experimental_vector_insert: {
+      // FIXME: Handle case where a scalable vector is inserted into a scalable
+      // vector
+      if (isa<ScalableVectorType>(Args[1]->getType()))
+        return BaseT::getIntrinsicInstrCost(ICA, CostKind);
+      unsigned Index = cast<ConstantInt>(Args[2])->getZExtValue();
+      return thisT()->getShuffleCost(
+          TTI::SK_InsertSubvector, cast<VectorType>(Args[0]->getType()), Index,
+          cast<VectorType>(Args[1]->getType()));
+    }
     case Intrinsic::vector_reduce_add:
     case Intrinsic::vector_reduce_mul:
     case Intrinsic::vector_reduce_and:
     case Intrinsic::vector_reduce_or:
     case Intrinsic::vector_reduce_xor:
     case Intrinsic::vector_reduce_smax:
     case Intrinsic::vector_reduce_smin:
     case Intrinsic::vector_reduce_fmax:

llvm/test/Analysis/CostModel/AArch64/sve-getIntrinsicInstrCost-vec-insert-extract.ll

+; RUN: opt < %s -cost-model -analyze -mtriple=aarch64-linux-gnu -mattr=+sve | FileCheck %s
+
+define <16 x i32> @extract_cost(<vscale x 4 x i32> %vec) {
+; CHECK-LABEL: 'extract_cost'
+; CHECK-NEXT: Cost Model: Found an estimated cost of 72 for instruction:   %ret = call <16 x i32> @llvm.experimental.vector.extract.v16i32.nxv4i32(<vscale x 4 x i32> %vec, i64 0)
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <16 x i32> %ret
+
+  %ret = call <16 x i32> @llvm.experimental.vector.extract.v16i32.nxv4i32(<vscale x 4 x i32> %vec, i64 0)
+  ret <16 x i32> %ret
+}
+
+define <vscale x 4 x i32> @insert_cost(<vscale x 4 x i32> %vec, <16 x i32> %subVec) {
+; CHECK-LABEL: 'insert_cost'
+; CHECK-NEXT: Cost Model: Found an estimated cost of 72 for instruction:   %ret = call <vscale x 4 x i32> @llvm.experimental.vector.insert.nxv4i32.v16i32(<vscale x 4 x i32> %vec, <16 x i32> %subVec, i64 0)
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 4 x i32> %ret
+
+  %ret = call <vscale x 4 x i32> @llvm.experimental.vector.insert.nxv4i32.v16i32(<vscale x 4 x i32> %vec, <16 x i32> %subVec, i64 0)
+  ret <vscale x 4 x i32> %ret
+}
+
+define <vscale x 4 x i32> @extract_cost_scalable(<vscale x 16 x i32> %vec) {
+; CHECK-LABEL: 'extract_cost_scalable'
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %ret = call <vscale x 4 x i32> @llvm.experimental.vector.extract.nxv4i32.nxv16i32(<vscale x 16 x i32> %vec, i64 0)
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 4 x i32> %ret
+
+  %ret = call <vscale x 4 x i32> @llvm.experimental.vector.extract.nxv4i32.nxv16i32(<vscale x 16 x i32> %vec, i64 0)
+  ret <vscale x 4 x i32> %ret
+}
+
+define <vscale x 16 x i32> @insert_cost_scalable(<vscale x 16 x i32> %vec, <vscale x 4 x i32> %subVec) {
+; CHECK-LABEL: 'insert_cost_scalable'
+; CHECK-NEXT: Cost Model: Found an estimated cost of 1 for instruction:   %ret = call <vscale x 16 x i32> @llvm.experimental.vector.insert.nxv16i32.nxv4i32(<vscale x 16 x i32> %vec, <vscale x 4 x i32> %subVec, i64 0)
+; CHECK-NEXT: Cost Model: Found an estimated cost of 0 for instruction:   ret <vscale x 16 x i32> %ret
+
+  %ret = call <vscale x 16 x i32> @llvm.experimental.vector.insert.nxv16i32.nxv4i32(<vscale x 16 x i32> %vec, <vscale x 4 x i32> %subVec, i64 0)
+  ret <vscale x 16 x i32> %ret
+}
+
+declare <16 x i32> @llvm.experimental.vector.extract.v16i32.nxv4i32(<vscale x 4 x i32>, i64)
+declare <vscale x 4 x i32> @llvm.experimental.vector.insert.nxv4i32.v16i32(<vscale x 4 x i32>, <16 x i32>, i64)
+declare <vscale x 4 x i32> @llvm.experimental.vector.extract.nxv4i32.nxv16i32(<vscale x 16 x i32>, i64)
+declare <vscale x 16 x i32> @llvm.experimental.vector.insert.nxv16i32.nxv4i32(<vscale x 16 x i32>, <vscale x 4 x i32>, i64)