[llvm] [WebAssembly][TTI] Avoid crash when costing scalable vector shifts (PR #212759)

[Anutosh Bhat via llvm-commits]

https://github.com/anutosh491 updated https://github.com/llvm/llvm-project/pull/212759

>From ff3e7d22d441b93e16756ac4d27f9e7220f98ec2 Mon Sep 17 00:00:00 2001
From: anutosh491 <andersonbhat491 at gmail.com>
Date: Wed, 29 Jul 2026 18:21:47 +0530
Subject: [PATCH] [WebAssembly][TTI] Avoid crash when costing scalable vector
 shifts

---
 .../WebAssemblyTargetTransformInfo.cpp        |  4 +-
 .../CostModel/WebAssembly/vector-shift.ll     | 42 +++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Analysis/CostModel/WebAssembly/vector-shift.ll

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyTargetTransformInfo.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyTargetTransformInfo.cpp
index 7d1136bb9beba..75c631055ae2d 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyTargetTransformInfo.cpp
+++ b/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);
diff --git a/llvm/test/Analysis/CostModel/WebAssembly/vector-shift.ll b/llvm/test/Analysis/CostModel/WebAssembly/vector-shift.ll
new file mode 100644
index 0000000000000..97f2a94b7b460
--- /dev/null
+++ b/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
+}