[llvm] [AMDGPU] Generalize extractSubregFromImm (PR #208765)

Jay Foad via llvm-commits llvm-commits at lists.llvm.org
Tue Jul 14 05:01:18 PDT 2026


https://github.com/jayfoad updated https://github.com/llvm/llvm-project/pull/208765

>From f7e1a439962110b042c286a9a5d9a81220a4b19b Mon Sep 17 00:00:00 2001
From: Jay Foad <jay.foad at amd.com>
Date: Wed, 28 May 2025 13:10:46 +0100
Subject: [PATCH 1/2] [AMDGPU] Generalize extractSubregFromImm

Simplify the code and potentially handle more subreg indices, if any are
defined in the future.
---
 llvm/lib/Target/AMDGPU/SIFoldOperands.cpp | 33 ++++++++---------
 llvm/lib/Target/AMDGPU/SIInstrInfo.cpp    | 43 ++++++++---------------
 llvm/lib/Target/AMDGPU/SIInstrInfo.h      |  3 +-
 3 files changed, 32 insertions(+), 47 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp b/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp
index 3fdbb74eb3342..4c37d9cf8b107 100644
--- a/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp
+++ b/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp
@@ -103,9 +103,9 @@ struct FoldableDef {
   /// Return the effective immediate value defined by this instruction, after
   /// application of any subregister extracts which may exist between the use
   /// and def instruction.
-  std::optional<int64_t> getEffectiveImmVal() const {
+  int64_t getEffectiveImmVal(const SIInstrInfo &TII) const {
     assert(isImm());
-    return SIInstrInfo::extractSubregFromImm(ImmToFold, DefSubReg);
+    return TII.extractSubregFromImm(ImmToFold, DefSubReg);
   }
 
   /// Check if it is legal to fold this effective value into \p MI's \p OpNo
@@ -114,13 +114,11 @@ struct FoldableDef {
                       unsigned OpIdx) const {
     switch (Kind) {
     case MachineOperand::MO_Immediate: {
-      std::optional<int64_t> ImmToFold = getEffectiveImmVal();
-      if (!ImmToFold)
-        return false;
+      int64_t ImmToFold = getEffectiveImmVal(TII);
 
       // TODO: Should verify the subregister index is supported by the class
       // TODO: Avoid the temporary MachineOperand
-      MachineOperand TmpOp = MachineOperand::CreateImm(*ImmToFold);
+      MachineOperand TmpOp = MachineOperand::CreateImm(ImmToFold);
       return TII.isOperandLegal(MI, OpIdx, &TmpOp);
     }
     case MachineOperand::MO_FrameIndex: {
@@ -611,7 +609,7 @@ bool SIFoldOperandsImpl::updateOperand(FoldCandidate &Fold) const {
 
   std::optional<int64_t> ImmVal;
   if (Fold.isImm())
-    ImmVal = Fold.Def.getEffectiveImmVal();
+    ImmVal = Fold.Def.getEffectiveImmVal(*TII);
 
   if (ImmVal && canUseImmWithOpSel(Fold.UseMI, Fold.UseOpNo, *ImmVal)) {
     if (tryFoldImmWithOpSel(Fold.UseMI, Fold.UseOpNo, *ImmVal))
@@ -790,9 +788,9 @@ static bool isPKF32InstrReplicatesLower32BitsOfScalarOperand(
 // literal) and replicates the bits to both channels. Therefore, if the hi and
 // lo are not same, we can't fold it.
 static bool checkImmOpForPKF32InstrReplicatesLower32BitsOfScalarOperand(
-    const FoldableDef &OpToFold) {
+    const FoldableDef &OpToFold, const SIInstrInfo &TII) {
   assert(OpToFold.isImm() && "Expected immediate operand");
-  uint64_t ImmVal = OpToFold.getEffectiveImmVal().value();
+  uint64_t ImmVal = OpToFold.getEffectiveImmVal(TII);
   uint32_t Lo = Lo_32(ImmVal);
   uint32_t Hi = Hi_32(ImmVal);
   return Lo == Hi;
@@ -839,8 +837,8 @@ bool SIFoldOperandsImpl::tryAddToFoldList(
 
   bool IsLegal = OpToFold.isOperandLegal(*TII, *MI, OpNo);
   if (!IsLegal && OpToFold.isImm()) {
-    if (std::optional<int64_t> ImmVal = OpToFold.getEffectiveImmVal())
-      IsLegal = canUseImmWithOpSel(MI, OpNo, *ImmVal);
+    int64_t ImmVal = OpToFold.getEffectiveImmVal(*TII);
+    IsLegal = canUseImmWithOpSel(MI, OpNo, ImmVal);
   }
 
   if (!IsLegal) {
@@ -955,7 +953,8 @@ bool SIFoldOperandsImpl::tryAddToFoldList(
   // src0 or src1.
   if (OpToFold.isImm() &&
       isPKF32InstrReplicatesLower32BitsOfScalarOperand(ST, MI, OpNo) &&
-      !checkImmOpForPKF32InstrReplicatesLower32BitsOfScalarOperand(OpToFold))
+      !checkImmOpForPKF32InstrReplicatesLower32BitsOfScalarOperand(OpToFold,
+                                                                   *TII))
     return false;
 
   appendFoldCandidate(FoldList, MI, OpNo, OpToFold);
@@ -1181,7 +1180,8 @@ bool SIFoldOperandsImpl::tryToFoldACImm(
 
   if (OpToFold.isImm() && OpToFold.isOperandLegal(*TII, *UseMI, UseOpIdx)) {
     if (isPKF32InstrReplicatesLower32BitsOfScalarOperand(ST, UseMI, UseOpIdx) &&
-        !checkImmOpForPKF32InstrReplicatesLower32BitsOfScalarOperand(OpToFold))
+        !checkImmOpForPKF32InstrReplicatesLower32BitsOfScalarOperand(OpToFold,
+                                                                     *TII))
       return false;
     appendFoldCandidate(FoldList, UseMI, UseOpIdx, OpToFold);
     return true;
@@ -1348,7 +1348,7 @@ bool SIFoldOperandsImpl::foldOperand(
         if (MovOp == AMDGPU::AV_MOV_B32_IMM_PSEUDO &&
             (!OpToFold.isImm() ||
              !TII->isImmOperandLegal(MovDesc, SrcIdx,
-                                     *OpToFold.getEffectiveImmVal())))
+                                     OpToFold.getEffectiveImmVal(*TII))))
           break;
 
         if (!MRI->constrainRegClass(SrcReg, MovSrcRC))
@@ -1362,7 +1362,8 @@ bool SIFoldOperandsImpl::foldOperand(
         // only matters for these concrete cases.
         // TODO: Handle non-imm case if it's useful.
         if (!OpToFold.isImm() ||
-            !TII->isImmOperandLegal(MovDesc, 1, *OpToFold.getEffectiveImmVal()))
+            !TII->isImmOperandLegal(MovDesc, 1,
+                                    OpToFold.getEffectiveImmVal(*TII)))
           break;
       }
 
@@ -1469,7 +1470,7 @@ bool SIFoldOperandsImpl::foldOperand(
 
         if (OpToFold.isImm()) {
           UseMI->getOperand(1).ChangeToImmediate(
-              *OpToFold.getEffectiveImmVal());
+              OpToFold.getEffectiveImmVal(*TII));
         } else if (OpToFold.isFI())
           UseMI->getOperand(1).ChangeToFrameIndex(OpToFold.getFI());
         else {
diff --git a/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp b/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp
index e945db946c158..419d533d37d18 100644
--- a/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp
@@ -3477,28 +3477,14 @@ void SIInstrInfo::mutateAndCleanupImplicit(MachineInstr &MI,
     MI.removeOperand(I);
 }
 
-std::optional<int64_t> SIInstrInfo::extractSubregFromImm(int64_t Imm,
-                                                         unsigned SubRegIndex) {
-  switch (SubRegIndex) {
-  case AMDGPU::NoSubRegister:
+int64_t SIInstrInfo::extractSubregFromImm(int64_t Imm,
+                                          unsigned SubRegIndex) const {
+  if (SubRegIndex == AMDGPU::NoSubRegister)
     return Imm;
-  case AMDGPU::sub0:
-    return SignExtend64<32>(Imm);
-  case AMDGPU::sub1:
-    return SignExtend64<32>(Imm >> 32);
-  case AMDGPU::lo16:
-    return SignExtend64<16>(Imm);
-  case AMDGPU::hi16:
-    return SignExtend64<16>(Imm >> 16);
-  case AMDGPU::sub1_lo16:
-    return SignExtend64<16>(Imm >> 32);
-  case AMDGPU::sub1_hi16:
-    return SignExtend64<16>(Imm >> 48);
-  default:
-    return std::nullopt;
-  }
-
-  llvm_unreachable("covered subregister switch");
+  assert(RI.getSubRegIdxSize(SubRegIndex) > 0);
+  assert(RI.getSubRegIdxOffset(SubRegIndex) >= 0);
+  return SignExtend64(Imm >> RI.getSubRegIdxOffset(SubRegIndex),
+                      RI.getSubRegIdxSize(SubRegIndex));
 }
 
 static unsigned getNewFMAAKInst(const GCNSubtarget &ST, unsigned Opc) {
@@ -3624,7 +3610,7 @@ bool SIInstrInfo::foldImmediate(MachineInstr &UseMI, MachineInstr &DefMI,
     MCRegister MovDstPhysReg =
         DstReg.isPhysical() ? DstReg.asMCReg() : MCRegister();
 
-    std::optional<int64_t> SubRegImm = extractSubregFromImm(Imm, UseSubReg);
+    int64_t SubRegImm = extractSubregFromImm(Imm, UseSubReg);
 
     // TODO: Try to fold with AMDGPU::V_MOV_B16_t16_e64
     for (unsigned MovOp :
@@ -3670,7 +3656,7 @@ bool SIInstrInfo::foldImmediate(MachineInstr &UseMI, MachineInstr &DefMI,
       // FIXME: isImmOperandLegal should have form that doesn't require existing
       // MachineInstr or MachineOperand
       if (!RI.opCanUseLiteralConstant(OpInfo.OperandType) &&
-          !isInlineConstant(*SubRegImm, OpInfo.OperandType))
+          !isInlineConstant(SubRegImm, OpInfo.OperandType))
         break;
 
       NewOpc = MovOp;
@@ -3689,7 +3675,7 @@ bool SIInstrInfo::foldImmediate(MachineInstr &UseMI, MachineInstr &DefMI,
 
     const MCInstrDesc &NewMCID = get(NewOpc);
     UseMI.setDesc(NewMCID);
-    UseMI.getOperand(1).ChangeToImmediate(*SubRegImm);
+    UseMI.getOperand(1).ChangeToImmediate(SubRegImm);
     UseMI.addImplicitDefUseOperands(*MF);
     return true;
   }
@@ -3771,7 +3757,7 @@ bool SIInstrInfo::foldImmediate(MachineInstr &UseMI, MachineInstr &DefMI,
       if (pseudoToMCOpcode(NewOpc) == -1)
         return false;
 
-      const std::optional<int64_t> SubRegImm = extractSubregFromImm(
+      int64_t SubRegImm = extractSubregFromImm(
           Imm, RegSrc == Src1 ? Src0->getSubReg() : Src1->getSubReg());
 
       // FIXME: This would be a lot easier if we could return a new instruction
@@ -3790,7 +3776,7 @@ bool SIInstrInfo::foldImmediate(MachineInstr &UseMI, MachineInstr &DefMI,
         UseMI.untieRegOperand(
             AMDGPU::getNamedOperandIdx(Opc, AMDGPU::OpName::src2));
 
-      Src1->ChangeToImmediate(*SubRegImm);
+      Src1->ChangeToImmediate(SubRegImm);
 
       removeModOperands(UseMI);
       UseMI.setDesc(get(NewOpc));
@@ -3865,11 +3851,10 @@ bool SIInstrInfo::foldImmediate(MachineInstr &UseMI, MachineInstr &DefMI,
         UseMI.untieRegOperand(
             AMDGPU::getNamedOperandIdx(Opc, AMDGPU::OpName::src2));
 
-      const std::optional<int64_t> SubRegImm =
-          extractSubregFromImm(Imm, Src2->getSubReg());
+      int64_t SubRegImm = extractSubregFromImm(Imm, Src2->getSubReg());
 
       // ChangingToImmediate adds Src2 back to the instruction.
-      Src2->ChangeToImmediate(*SubRegImm);
+      Src2->ChangeToImmediate(SubRegImm);
 
       // These come before src2.
       removeModOperands(UseMI);
diff --git a/llvm/lib/Target/AMDGPU/SIInstrInfo.h b/llvm/lib/Target/AMDGPU/SIInstrInfo.h
index 8e15b7b45b609..93ec831f3ff2c 100644
--- a/llvm/lib/Target/AMDGPU/SIInstrInfo.h
+++ b/llvm/lib/Target/AMDGPU/SIInstrInfo.h
@@ -459,8 +459,7 @@ class SIInstrInfo final : public AMDGPUGenInstrInfo {
   /// e.g. %imm = S_MOV_B64 K[0:63]
   ///      USE %imm.sub1
   /// This will return K[32:63]
-  static std::optional<int64_t> extractSubregFromImm(int64_t ImmVal,
-                                                     unsigned SubRegIndex);
+  int64_t extractSubregFromImm(int64_t ImmVal, unsigned SubRegIndex) const;
 
   bool foldImmediate(MachineInstr &UseMI, MachineInstr &DefMI, Register Reg,
                      MachineRegisterInfo *MRI) const final;

>From 9d3e6c89dc78a10de4c679ab307a9cec6aeda496 Mon Sep 17 00:00:00 2001
From: Jay Foad <jay.foad at amd.com>
Date: Tue, 14 Jul 2026 13:00:44 +0100
Subject: [PATCH 2/2] Tweak asserts

---
 llvm/lib/Target/AMDGPU/SIInstrInfo.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp b/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp
index 419d533d37d18..921bddd0ac01d 100644
--- a/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp
@@ -3481,8 +3481,8 @@ int64_t SIInstrInfo::extractSubregFromImm(int64_t Imm,
                                           unsigned SubRegIndex) const {
   if (SubRegIndex == AMDGPU::NoSubRegister)
     return Imm;
-  assert(RI.getSubRegIdxSize(SubRegIndex) > 0);
-  assert(RI.getSubRegIdxOffset(SubRegIndex) >= 0);
+  assert(RI.getSubRegIdxSize(SubRegIndex) != -1u);
+  assert(RI.getSubRegIdxOffset(SubRegIndex) != -1u);
   return SignExtend64(Imm >> RI.getSubRegIdxOffset(SubRegIndex),
                       RI.getSubRegIdxSize(SubRegIndex));
 }



More information about the llvm-commits mailing list