[WebAssembly][TTI] Avoid crash when costing scalable vector shifts

llvm/lib/Target/WebAssembly/WebAssemblyTargetTransformInfo.cpp

@@ -82,7 +82,7 @@ InstructionCost WebAssemblyTTIImpl::getArithmeticInstrCost(
       BasicTTIImplBase<WebAssemblyTTIImpl>::getArithmeticInstrCost(
           Opcode, Ty, CostKind, Op1Info, Op2Info);
 
-  if (auto *VTy = dyn_cast<VectorType>(Ty)) {
+  if (auto *VTy = dyn_cast<FixedVectorType>(Ty)) {
     switch (Opcode) {
     case Instruction::LShr:
     case Instruction::AShr:
@@ -92,7 +92,7 @@ InstructionCost WebAssemblyTTIImpl::getArithmeticInstrCost(
       // approximation.
       if (!Op2Info.isUniform())
         Cost =
-            cast<FixedVectorType>(VTy)->getNumElements() *
+            VTy->getNumElements() *
             (TargetTransformInfo::TCC_Basic +
              getArithmeticInstrCost(Opcode, VTy->getElementType(), CostKind) +
              TargetTransformInfo::TCC_Basic);

llvm/test/Analysis/CostModel/WebAssembly/vector-shift.ll

@@ -0,0 +1,42 @@
+; NOTE: Assertions have been autogenerated by utils/update_analyze_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -mtriple=wasm32-unknown-unknown -mattr=+simd128 \
+; RUN:   -passes='print<cost-model>' -disable-output < %s 2>&1 | FileCheck %s
+
+define void @fixed_non_uniform(<4 x i32> %x, <4 x i32> %amount) {
+; CHECK-LABEL: 'fixed_non_uniform'
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 12 for instruction: %shl = shl <4 x i32> %x, %amount
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 12 for instruction: %lshr = lshr <4 x i32> %x, %amount
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 12 for instruction: %ashr = ashr <4 x i32> %x, %amount
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: ret void
+;
+  %shl = shl <4 x i32> %x, %amount
+  %lshr = lshr <4 x i32> %x, %amount
+  %ashr = ashr <4 x i32> %x, %amount
+  ret void
+}
+
+define void @scalable_uniform(<vscale x 4 x i32> %x) {
+; CHECK-LABEL: 'scalable_uniform'
+; CHECK-NEXT:  Cost Model: Invalid cost for instruction: %shl = shl <vscale x 4 x i32> %x, splat (i32 1)
+; CHECK-NEXT:  Cost Model: Invalid cost for instruction: %lshr = lshr <vscale x 4 x i32> %x, splat (i32 1)
+; CHECK-NEXT:  Cost Model: Invalid cost for instruction: %ashr = ashr <vscale x 4 x i32> %x, splat (i32 1)
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: ret void
+;
+  %shl = shl <vscale x 4 x i32> %x, splat (i32 1)
+  %lshr = lshr <vscale x 4 x i32> %x, splat (i32 1)
+  %ashr = ashr <vscale x 4 x i32> %x, splat (i32 1)
+  ret void
+}
+
+define void @scalable_non_uniform(<vscale x 4 x i32> %x, <vscale x 4 x i32> %amount) {
+; CHECK-LABEL: 'scalable_non_uniform'
+; CHECK-NEXT:  Cost Model: Invalid cost for instruction: %shl = shl <vscale x 4 x i32> %x, %amount
+; CHECK-NEXT:  Cost Model: Invalid cost for instruction: %lshr = lshr <vscale x 4 x i32> %x, %amount
+; CHECK-NEXT:  Cost Model: Invalid cost for instruction: %ashr = ashr <vscale x 4 x i32> %x, %amount
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: ret void
+;
+  %shl = shl <vscale x 4 x i32> %x, %amount
+  %lshr = lshr <vscale x 4 x i32> %x, %amount
+  %ashr = ashr <vscale x 4 x i32> %x, %amount
+  ret void
+}