[TTI] Support ScalableVectorType in getShuffleCost with SK_Broadcast kind

[junparser]

Summary

Since shufflevector supports scalable vector in broadcast mode, it
should be handled as well in getShuffleCost. This patch does this which fixes
assertion in https://github.com/llvm/llvm-project/issues/51550,

TestPlan: check-llvm

[spatel]

It's fine to have a vectorizer test, but it would be better to have a more direct test of the cost model itself.

I think something like this can check for the crash (we could add variations to see if the cost is actually accurate - see /llvm/test/Analysis/CostModel/AArch64/sve-shuffle-broadcast.ll for an existing test model):

; RUN: opt -cost-model -analyze -mtriple=riscv64 -mattr=+m,+experimental-v -scalable-vectorization=on  < %s | FileCheck %s

define void @broadcast() {
  %zero = shufflevector <vscale x 16 x i8> undef, <vscale x 16 x i8> undef, <vscale x 16 x i32> zeroinitializer
  ret void
}

llvm/include/llvm/CodeGen/BasicTTIImpl.h

875

Can we assert that the Type is Fixed unless the shuffle is a broadcast?

assert((ImprovedKind == TTI::SK_Broadcast || isa<FixedVectorType>(Tp)) && "Unexpected shuffle of scalable vector");

[junparser]

In D116362#3212562, @spatel wrote:

It's fine to have a vectorizer test, but it would be better to have a more direct test of the cost model itself.

I think something like this can check for the crash (we could add variations to see if the cost is actually accurate - see /llvm/test/Analysis/CostModel/AArch64/sve-shuffle-broadcast.ll for an existing test model):

; RUN: opt -cost-model -analyze -mtriple=riscv64 -mattr=+m,+experimental-v -scalable-vectorization=on  < %s | FileCheck %s

define void @broadcast() {
  %zero = shufflevector <vscale x 16 x i8> undef, <vscale x 16 x i8> undef, <vscale x 16 x i32> zeroinitializer
  ret void
}

good point, I missed this because of that rv does not have such target api. I'll add it.

llvm/include/llvm/CodeGen/BasicTTIImpl.h

875

yes，we can.

[junparser]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

875

hmm, we cannot assert here, SK_ExtractSubvector and such kind handle scalable vector as well.

[junparser]

Add testcase in cost model.

[dmgreen]

This looks like it would treat the default cost of a broadcast as scalarising the first n elements from a <vscale x n x iM> vector, which doesn't sound right. Unless we are assuming that the default vscale is 1.

I would suspect that the cost should either be invalid or to be "cheap" under assumption that the operation is always present. AArch64 has some cost overrides in its backend for certain types of scalable broadcasts.

[spatel]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

875

But we are casting to fixed vector type in the case statements below here - if scalable type is possible, then it will crash? Do you have an example of a scalable vector with any of the other shuffle kinds?

[junparser]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

885

@spatel， I mean Tp here maybe scalable vector

[junparser]

In D116362#3213897, @dmgreen wrote:

This looks like it would treat the default cost of a broadcast as scalarising the first n elements from a <vscale x n x iM> vector, which doesn't sound right. Unless we are assuming that the default vscale is 1.

Yes， This patch keeps same behavior as before (release mode).

I would suspect that the cost should either be invalid or to be "cheap" under assumption that the operation is always present. AArch64 has some cost overrides in its backend for certain types of scalable broadcasts.

without target override，we cannot compute scalable vector broadcasts here. So I think set invalid for scalabe vector here make sense.

[junparser]

address comments. Set invalid cost for scalable vector broadcast.

[paulwalker-arm]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

87–88

Not absolutely against this change, but I'm not a fan. This gives the impression this private helper function supports any type of vector, which is clearly not the case and hence why you return an invalid cost. Essentially most of the function implementations within this file specially cost the scalarisation path, which is just an invalid path for scalable vectors.

For this reason I think it's better to catch these cases as early in the call chain as possible. This function taking a FixedVectorType means that you're more likely to realise you're doing something wrong at compile time rather than runtime. Although I'll concede this is a bit tenuous given it currently only has a single use, but then would we want to make similar changes to every *Overhead function when it's more likely the caller is just at fault.

877

Based on the above personally I think getShuffleCost is wholly unsafe for scalable vector shuffles and I'd sooner see the "you probably didn't want to get here" code here. For example:

case TTI::SK_Broadcast:
  if (!isa<FixedVectorType>(Tp))
    return InstructionCost::getInvalid();
  return getBroadcastShuffleOverhead(cast<FixedVectorType>(Tp));

This for me makes it clearer that if you'd rather not return an invalid cost then you'll need to fix the target specific implementation of this function.

Presumably the other shuffle types below need the same treatment?

[junparser]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

877

Based on the above personally I think getShuffleCost is wholly unsafe for scalable vector shuffles and I'd sooner see the "you probably didn't want to get here" code here. For example:

case TTI::SK_Broadcast:
  if (!isa<FixedVectorType>(Tp))
    return InstructionCost::getInvalid();
  return getBroadcastShuffleOverhead(cast<FixedVectorType>(Tp));

This for me makes it clearer that if you'd rather not return an invalid cost then you'll need to fix the target specific implementation of this function.

make sense to me.

Presumably the other shuffle types below need the same treatment?

Yes, we also need handle intrinsic like vector.reverse with scalable vector.

[junparser]

address paul's comment.

[junparser]

update testcase.

[sdesmalen]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

877–879

nit: is it worth writing this as:

if (auto *FVT = dyn_cast<FixedVectorType>(Tp))
  return getBroadcastShuffleOverhead(FVT);
return InstructionCost::getInvalid()

884

(same suggestion here)

llvm/test/Analysis/CostModel/RISCV/rvv-shuffle.ll

4 ↗

(On Diff #397189)

nit: -scalable-vectorization=off|on is an option to the LoopVectorizer, not the costmodel, so it won't have any effect when combined with -cost-model -analyze

[junparser]

address comment.

[rogfer01]

llvm/test/Analysis/CostModel/RISCV/rvv-shuffle.ll

14 ↗

(On Diff #397494)

I'm curious, why does this one have cost zero and all others are invalid?

Maybe we're getting the min element count somewhere and assuming this is a scalar boxed in a fixed vector of length 1?

llvm/include/llvm/CodeGen/BasicTTIImpl.h

 //===- BasicTTIImpl.h -------------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
private:
   using BaseT = TargetTransformInfoImplCRTPBase<T>;
   using TTI = TargetTransformInfo;
 
   /// Helper function to access this as a T.
   T *thisT() { return static_cast<T *>(this); }
 
   /// Estimate a cost of Broadcast as an extract and sequence of insert
   /// operations.
-  InstructionCost getBroadcastShuffleOverhead(FixedVectorType *VTy) {
+  InstructionCost getBroadcastShuffleOverhead(VectorType *VTy) {
     InstructionCost Cost = 0;
     // Broadcast cost is equal to the cost of extracting the zero'th element
     // plus the cost of inserting it into every element of the result vector.
     Cost += thisT()->getVectorInstrCost(Instruction::ExtractElement, VTy, 0);
 
-    for (int i = 0, e = VTy->getNumElements(); i < e; ++i) {
+    for (int i = 0, e = VTy->getElementCount().getKnownMinValue(); i < e; ++i) {
       Cost += thisT()->getVectorInstrCost(Instruction::InsertElement, VTy, i);
     }
     return Cost;
   }
 
   /// Estimate a cost of shuffle as a sequence of extract and insert
   /// operations.
   InstructionCost getPermuteShuffleOverhead(FixedVectorType *VTy) {
case TTI::SK_Splice:
       break;
     }
     return Kind;
   }
 
   InstructionCost getShuffleCost(TTI::ShuffleKind Kind, VectorType *Tp,
                                  ArrayRef<int> Mask, int Index,
                                  VectorType *SubTp) {
-
     switch (improveShuffleKindFromMask(Kind, Mask)) {
     case TTI::SK_Broadcast:
-      return getBroadcastShuffleOverhead(cast<FixedVectorType>(Tp));
+      return getBroadcastShuffleOverhead(Tp);
     case TTI::SK_Select:
     case TTI::SK_Splice:
     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(Tp, Index,
                                          cast<FixedVectorType>(SubTp));
     case TTI::SK_InsertSubvector:
       return getInsertSubvectorOverhead(Tp, Index,
                                         cast<FixedVectorType>(SubTp));
     }
     llvm_unreachable("Unknown TTI::ShuffleKind");
   }

llvm/test/Analysis/CostModel/RISCV/rvv-shuffle-broadcast.ll

+; NOTE: Assertions have been autogenerated by utils/update_analyze_test_checks.py
+; Check getShuffleCost for SK_BroadCast with scalable vector
+
+; RUN: opt -cost-model -analyze -mtriple=riscv64 -mattr=+m,+experimental-v -scalable-vectorization=on  < %s | FileCheck %s
+
+define void  @broadcast() {
+; CHECK-LABEL: 'broadcast'
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 9 for instruction: %zero = shufflevector <vscale x 8 x i8> undef, <vscale x 8 x i8> undef, <vscale x 8 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 17 for instruction: %1 = shufflevector <vscale x 16 x i8> undef, <vscale x 16 x i8> undef, <vscale x 16 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 5 for instruction: %2 = shufflevector <vscale x 4 x i16> undef, <vscale x 4 x i16> undef, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 9 for instruction: %3 = shufflevector <vscale x 8 x i16> undef, <vscale x 8 x i16> undef, <vscale x 8 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 3 for instruction: %4 = shufflevector <vscale x 2 x i32> undef, <vscale x 2 x i32> undef, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 5 for instruction: %5 = shufflevector <vscale x 4 x i32> undef, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: %6 = shufflevector <vscale x 1 x i64> undef, <vscale x 1 x i64> undef, <vscale x 1 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 3 for instruction: %7 = shufflevector <vscale x 2 x i64> undef, <vscale x 2 x i64> undef, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 17 for instruction: %8 = shufflevector <vscale x 16 x i1> undef, <vscale x 16 x i1> undef, <vscale x 16 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 9 for instruction: %9 = shufflevector <vscale x 8 x i1> undef, <vscale x 8 x i1> undef, <vscale x 8 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 5 for instruction: %10 = shufflevector <vscale x 4 x i1> undef, <vscale x 4 x i1> undef, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 3 for instruction: %11 = shufflevector <vscale x 2 x i1> undef, <vscale x 2 x i1> undef, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: ret void
+;
+  %zero = shufflevector <vscale x 8 x i8> undef, <vscale x 8 x i8> undef, <vscale x 8 x i32> zeroinitializer
+  %1 = shufflevector <vscale x 16 x i8> undef, <vscale x 16 x i8> undef, <vscale x 16 x i32> zeroinitializer
+  %2 = shufflevector <vscale x 4 x i16> undef, <vscale x 4 x i16> undef, <vscale x 4 x i32> zeroinitializer
+  %3 = shufflevector <vscale x 8 x i16> undef, <vscale x 8 x i16> undef, <vscale x 8 x i32> zeroinitializer
+  %4 = shufflevector <vscale x 2 x i32> undef, <vscale x 2 x i32> undef, <vscale x 2 x i32> zeroinitializer
+  %5 = shufflevector <vscale x 4 x i32> undef, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+  %6 = shufflevector <vscale x 1 x i64> undef, <vscale x 1 x i64> undef, <vscale x 1 x i32> zeroinitializer
+  %7 = shufflevector <vscale x 2 x i64> undef, <vscale x 2 x i64> undef, <vscale x 2 x i32> zeroinitializer
+  %8 = shufflevector <vscale x 16 x i1> undef, <vscale x 16 x i1> undef, <vscale x 16 x i32> zeroinitializer
+  %9 = shufflevector <vscale x 8 x i1> undef, <vscale x 8 x i1> undef, <vscale x 8 x i32> zeroinitializer
+  %10 = shufflevector <vscale x 4 x i1> undef, <vscale x 4 x i1> undef, <vscale x 4 x i32> zeroinitializer
+  %11 = shufflevector <vscale x 2 x i1> undef, <vscale x 2 x i1> undef, <vscale x 2 x i32> zeroinitializer
+  ret void
+}
+

llvm/test/Transforms/LoopVectorize/RISCV/scalable-vectorizer-shufflecost.ll

+; REQUIRES: asserts
+; RUN: opt -mtriple=riscv64 -mattr=+m,+experimental-v -loop-vectorize \
+; RUN: -scalable-vectorization=on -S < %s 2>&1 | FileCheck %s
+
+@a = global i32 0, align 4
+@b = global i32* null, align 8
+
+; CHECK: <vscale x 2 x i32>
+define void @test(i32 signext %d) {
+entry:
+  %call = tail call noalias nonnull dereferenceable(4) i8* @_Znwm(i64 4)
+  %c = bitcast i8* %call to i32*
+  %cmp12 = icmp sgt i32 %d, 0
+  %0 = load i32, i32* @a, align 4
+  %tobool.not14 = icmp eq i32 %0, 0
+  br i1 %tobool.not14, label %for.end4, label %for.cond1.preheader.preheader
+
+for.cond1.preheader.preheader:                    ; preds = %entry
+  %wide.trip.count = zext i32 %d to i64
+  %1 = load i32*, i32** @b, align 8
+  br label %for.cond1.preheader
+
+for.cond1.preheader:                              ; preds = %for.cond1.preheader.backedge, %for.cond1.preheader.preheader
+  br i1 %cmp12, label %for.body2, label %for.cond1.preheader.backedge
+
+for.cond1.preheader.backedge:                     ; preds = %for.body2, %for.cond1.preheader
+  br label %for.cond1.preheader
+
+for.body2:                                        ; preds = %for.cond1.preheader, %for.body2
+  %indvars.iv = phi i64 [ %indvars.iv.next, %for.body2 ], [ 0, %for.cond1.preheader ]
+  %2 = load i32, i32* %1, align 4
+  %arrayidx3 = getelementptr inbounds i32, i32* %c, i64 %indvars.iv
+  store i32 %2, i32* %arrayidx3, align 4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond.not = icmp eq i64 %indvars.iv.next, %wide.trip.count
+  br i1 %exitcond.not, label %for.cond1.preheader.backedge, label %for.body2
+
+for.end4:                                         ; preds = %entry
+  ret void
+}
+
+declare nonnull i8* @_Znwm(i64)