[llvm] [AArch64] Guard against getRegisterBitWidth returning zero in vector instr cost. (PR #117749)

[David Green via llvm-commits]

https://github.com/davemgreen updated https://github.com/llvm/llvm-project/pull/117749

>From f649f012639bebbe14a902556725d272a0681705 Mon Sep 17 00:00:00 2001
From: David Green <david.green at arm.com>
Date: Fri, 29 Nov 2024 01:51:36 +0000
Subject: [PATCH] [AArch64] Guard against getRegisterBitWidth returning zero in
 vector instr cost.

If the getRegisterBitWidth is zero (such as in sme streaming functions), then
we could hit a crash from using % RegWidth.

It took a while to figure out what was going wrong so there are a few other
minor cleanups here too.
---
 .../AArch64/AArch64TargetTransformInfo.cpp     |  2 +-
 .../AArch64/extract_float_streaming.ll         | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Analysis/CostModel/AArch64/extract_float_streaming.ll

diff --git a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
index a415ab4c7f7dbc..d1536a276a9040 100644
--- a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
@@ -3259,7 +3259,7 @@ InstructionCost AArch64TTIImpl::getVectorInstrCostHelper(
       auto RegWidth =
           getRegisterBitWidth(TargetTransformInfo::RGK_FixedWidthVector)
               .getFixedValue();
-      return (Idx == 0 || (Idx * EltSz) % RegWidth == 0);
+      return Idx == 0 || (RegWidth != 0 && (Idx * EltSz) % RegWidth == 0);
     };
 
     // Check if the type constraints on input vector type and result scalar type
diff --git a/llvm/test/Analysis/CostModel/AArch64/extract_float_streaming.ll b/llvm/test/Analysis/CostModel/AArch64/extract_float_streaming.ll
new file mode 100644
index 00000000000000..84502abceed3b0
--- /dev/null
+++ b/llvm/test/Analysis/CostModel/AArch64/extract_float_streaming.ll
@@ -0,0 +1,18 @@
+; NOTE: Assertions have been autogenerated by utils/update_analyze_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes="print<cost-model>" 2>&1 -disable-output -mtriple=aarch64-unknown-linux -mattr=+sme | FileCheck %s
+
+define double @extract_case7(<4 x double> %a) "aarch64_pstate_sm_enabled" {
+; CHECK-LABEL: 'extract_case7'
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 2 for instruction: %0 = extractelement <4 x double> %a, i32 1
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: %1 = extractelement <4 x double> %a, i32 2
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 2 for instruction: %res = fmul double %0, %1
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: ret double %res
+;
+entry:
+  %1 = extractelement <4 x double> %a, i32 1
+  %2 = extractelement <4 x double> %a, i32 2
+  %res = fmul double %1, %2
+  ret double %res
+}
+
+declare void @foo(double)