[llvm] [GlobalISel] Check the correct register in OneUse check. (PR #114763)

[David Green via llvm-commits]

https://github.com/davemgreen created https://github.com/llvm/llvm-project/pull/114763

This fixes a bug that started triggering after #111730, where we could remove a load with multiple uses. It looks like the match should be checking the other register.

  %SrcReg = load..
  %DstReg = sign_extend_inreg %SrcReg

>From 0654f979a72dfad909a53c40cbbad2d166445578 Mon Sep 17 00:00:00 2001
From: David Green <david.green at arm.com>
Date: Mon, 4 Nov 2024 08:37:49 +0000
Subject: [PATCH] [GlobalISel] Check the correct register in OneUse check.

This fixes a bug that started triggering after #111730, where we could remove a
load with multiple uses. It looks like the match should be checking the other
register.

  %SrcReg = load..
  %DstReg = sign_extend_inreg %SrcReg
---
 llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp |  2 +-
 llvm/test/CodeGen/AArch64/load.ll              | 14 ++++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
index 7c1bda2163b7a0..1f2baa3fa9c0f8 100644
--- a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
@@ -1049,7 +1049,7 @@ bool CombinerHelper::matchSextInRegOfLoad(
 
   Register SrcReg = MI.getOperand(1).getReg();
   auto *LoadDef = getOpcodeDef<GLoad>(SrcReg, MRI);
-  if (!LoadDef || !MRI.hasOneNonDBGUse(DstReg))
+  if (!LoadDef || !MRI.hasOneNonDBGUse(SrcReg))
     return false;
 
   uint64_t MemBits = LoadDef->getMemSizeInBits().getValue();
diff --git a/llvm/test/CodeGen/AArch64/load.ll b/llvm/test/CodeGen/AArch64/load.ll
index 543605a0a09296..3fa5d64a210e19 100644
--- a/llvm/test/CodeGen/AArch64/load.ll
+++ b/llvm/test/CodeGen/AArch64/load.ll
@@ -465,3 +465,17 @@ define <2 x fp128> @load_v2f128(ptr %p) {
   %a = load <2 x fp128>, ptr %p
   ret <2 x fp128> %a
 }
+
+define i32 @load_i8_s16_extrasuse(ptr %ptr, ptr %ptr2) {
+; CHECK-LABEL: load_i8_s16_extrasuse:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ldr w8, [x0]
+; CHECK-NEXT:    sxtb w0, w8
+; CHECK-NEXT:    str w8, [x1]
+; CHECK-NEXT:    ret
+  %a = load i32, ptr %ptr
+  %s = shl i32 %a, 24
+  %b = ashr i32 %s, 24
+  store i32 %a, ptr %ptr2
+  ret i32 %b
+}