[llvm] [GlobalISel] Add known bits and sign-bits handling for G_SPLAT_VECTOR (PR #140204)

[David Green via llvm-commits]

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

>From 0a22c65d7f18c68ba4157003cb48ff590d9ff996 Mon Sep 17 00:00:00 2001
From: David Green <david.green at arm.com>
Date: Tue, 20 May 2025 12:26:00 +0100
Subject: [PATCH] [GlobalISel] Add known bits and sign-bits handling for
 G_SPLAT_VECTOR

This forwards the KnownBits / SignBits from the scalar operand through to the
vector elements. The value is implicitly truncated if it is larger than the
vector element size.  The change in GISelValueTracking::computeNumSignBits
allows scalable sign-bits using a single element demanded elts, and brings it
in line with what is done in ValueTracking and SDAG (and
GISelValueTracking::getKnownBits). That was my main reason for adding this
opcode, to prevent scalable vector asserts.
---
 .../CodeGen/GlobalISel/GISelValueTracking.cpp | 19 +++++++++-
 .../GlobalISel/knownbits-sve-splat.mir        | 35 +++++++++++++++++++
 2 files changed, 53 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/AArch64/GlobalISel/knownbits-sve-splat.mir

diff --git a/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp b/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp
index cafd63437df92..99794d690a91c 100644
--- a/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp
@@ -223,6 +223,14 @@ void GISelValueTracking::computeKnownBitsImpl(Register R, KnownBits &Known,
     }
     break;
   }
+  case TargetOpcode::G_SPLAT_VECTOR: {
+    computeKnownBitsImpl(MI.getOperand(1).getReg(), Known, APInt(1, 1),
+                         Depth + 1);
+    // Implicitly truncate the bits to match the official semantics of
+    // G_SPLAT_VECTOR.
+    Known = Known.trunc(BitWidth);
+    break;
+  }
   case TargetOpcode::COPY:
   case TargetOpcode::G_PHI:
   case TargetOpcode::PHI: {
@@ -904,6 +912,15 @@ unsigned GISelValueTracking::computeNumSignBits(Register R,
     }
     break;
   }
+  case TargetOpcode::G_SPLAT_VECTOR: {
+    // Check if the sign bits of source go down as far as the truncated value.
+    Register Src = MI.getOperand(1).getReg();
+    unsigned NumSrcSignBits = computeNumSignBits(Src, APInt(1, 1), Depth + 1);
+    unsigned NumSrcBits = MRI.getType(Src).getSizeInBits();
+    if (NumSrcSignBits > (NumSrcBits - TyBits))
+      return NumSrcSignBits - (NumSrcBits - TyBits);
+    break;
+  }
   case TargetOpcode::G_INTRINSIC:
   case TargetOpcode::G_INTRINSIC_W_SIDE_EFFECTS:
   case TargetOpcode::G_INTRINSIC_CONVERGENT:
@@ -939,7 +956,7 @@ unsigned GISelValueTracking::computeNumSignBits(Register R,
 unsigned GISelValueTracking::computeNumSignBits(Register R, unsigned Depth) {
   LLT Ty = MRI.getType(R);
   APInt DemandedElts =
-      Ty.isVector() ? APInt::getAllOnes(Ty.getNumElements()) : APInt(1, 1);
+      Ty.isFixedVector() ? APInt::getAllOnes(Ty.getNumElements()) : APInt(1, 1);
   return computeNumSignBits(R, DemandedElts, Depth);
 }
 
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/knownbits-sve-splat.mir b/llvm/test/CodeGen/AArch64/GlobalISel/knownbits-sve-splat.mir
new file mode 100644
index 0000000000000..27c11ee1641e9
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/knownbits-sve-splat.mir
@@ -0,0 +1,35 @@
+# NOTE: Assertions have been autogenerated by utils/update_givaluetracking_test_checks.py UTC_ARGS: --version 5
+# RUN: llc -mtriple aarch64 -mattr=+sve -passes="print<gisel-value-tracking>" %s -o - 2>&1 | FileCheck %s
+
+---
+name:            Scalable
+body:             |
+  bb.1:
+  ; CHECK-LABEL: name: @Scalable
+  ; CHECK-NEXT: %0:_ KnownBits:0000000000001010 SignBits:12
+  ; CHECK-NEXT: %1:_ KnownBits:0000000000001010 SignBits:12
+    %0:_(s16) = G_CONSTANT i16 10
+    %1:_(<vscale x 8 x s16>) = G_SPLAT_VECTOR %0(s16)
+...
+---
+name:            Scalable_trunc
+body:             |
+  bb.1:
+  ; CHECK-LABEL: name: @Scalable_trunc
+  ; CHECK-NEXT: %0:_ KnownBits:00000000000000000000000000001010 SignBits:28
+  ; CHECK-NEXT: %1:_ KnownBits:0000000000001010 SignBits:12
+    %0:_(s32) = G_CONSTANT i32 10
+    %1:_(<vscale x 8 x s16>) = G_SPLAT_VECTOR %0(s32)
+...
+---
+name:            Scalable_signbits
+body:             |
+  bb.1:
+  ; CHECK-LABEL: name: @Scalable_signbits
+  ; CHECK-NEXT: %0:_ KnownBits:???????????????? SignBits:1
+  ; CHECK-NEXT: %1:_ KnownBits:???????????????????????????????? SignBits:17
+  ; CHECK-NEXT: %2:_ KnownBits:???????????????????????????????? SignBits:17
+    %0:_(s16) = COPY $h0
+    %1:_(s32) = G_SEXT %0(s16)
+    %2:_(<vscale x 4 x s32>) = G_SPLAT_VECTOR %1(s32)
+...