[llvm] [AMDGPU] Only extend division result where needed (PR #204422)

[via llvm-commits]

https://github.com/LU-JOHN updated https://github.com/llvm/llvm-project/pull/204422

>From 17a6a6c6e5ac0bac68e86ed78f3fcb1e438fbbbe Mon Sep 17 00:00:00 2001
From: John Lu <John.Lu at amd.com>
Date: Wed, 17 Jun 2026 14:26:34 -0500
Subject: [PATCH] Only extend where needed

Signed-off-by: John Lu <John.Lu at amd.com>
---
 llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
index ef96aaa2296b7..38f254de8ecc0 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp
@@ -1361,16 +1361,14 @@ Value *AMDGPUCodeGenPrepareImpl::shrinkDivRem64(IRBuilder<> &Builder,
   if (NumDivBits <= (IsSigned ? 23 : 22)) {
     Narrowed = expandDivRemToFloatImpl(Builder, I, Num, Den, NumDivBits, IsDiv,
                                        IsSigned);
+    if (Narrowed)
+      Narrowed = IsSigned ? Builder.CreateSExt(Narrowed, Num->getType())
+                          : Builder.CreateZExt(Narrowed, Num->getType());
   } else if (NumDivBits <= 32) {
     Narrowed = expandDivRem32(Builder, I, Num, Den);
   }
 
-  if (Narrowed) {
-    return IsSigned ? Builder.CreateSExt(Narrowed, Num->getType()) :
-                      Builder.CreateZExt(Narrowed, Num->getType());
-  }
-
-  return nullptr;
+  return Narrowed;
 }
 
 void AMDGPUCodeGenPrepareImpl::expandDivRem64(BinaryOperator &I) const {