[AMDGPU] Add asm comments if setreg changes MSBs (#185774)

Author: rampitec

llvm/lib/Target/AMDGPU/AMDGPUMCInstLower.cpp

@@ -436,11 +436,19 @@ void AMDGPUAsmPrinter::emitInstruction(const MachineInstr *MI) {
                              MF->getInfo<SIMachineFunctionInfo>(),
                              *OutStreamer);
 
-    if (isVerbose() && MI->getOpcode() == AMDGPU::S_SET_VGPR_MSB) {
-      unsigned V = MI->getOperand(0).getImm() & 0xff;
-      OutStreamer->AddComment(
-          " msbs: dst=" + Twine(V >> 6) + " src0=" + Twine(V & 3) +
-          " src1=" + Twine((V >> 2) & 3) + " src2=" + Twine((V >> 4) & 3));
+    if (isVerbose() && (MI->getOpcode() == AMDGPU::S_SET_VGPR_MSB ||
+                        (MI->getOpcode() == AMDGPU::S_SETREG_IMM32_B32 &&
+                         STI.has1024AddressableVGPRs()))) {
+      std::optional<unsigned> V;
+      if (MI->getOpcode() == AMDGPU::S_SETREG_IMM32_B32)
+        V = AMDGPU::convertSetRegImmToVgprMSBs(*MI,
+                                               STI.hasSetregVGPRMSBFixup());
+      else
+        V = MI->getOperand(0).getImm() & 0xff;
+      if (V.has_value())
+        OutStreamer->AddComment(
+            " msbs: dst=" + Twine(*V >> 6) + " src0=" + Twine(*V & 3) +
+            " src1=" + Twine((*V >> 2) & 3) + " src2=" + Twine((*V >> 4) & 3));
     }
 
     MCInst TmpInst;

llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp

@@ -3589,6 +3589,25 @@ MCRegister getVGPRWithMSBs(MCRegister Reg, unsigned MSBs,
   return RC->getRegister(Idx);
 }
 
+std::optional<unsigned> convertSetRegImmToVgprMSBs(const MachineInstr &MI,
+                                                   bool HasSetregVGPRMSBFixup) {
+  assert(MI.getOpcode() == AMDGPU::S_SETREG_IMM32_B32);
+
+  constexpr unsigned VGPRMSBShift =
+      llvm::countr_zero_constexpr<unsigned>(AMDGPU::Hwreg::DST_VGPR_MSB);
+
+  auto [HwRegId, Offset, Size] =
+      Hwreg::HwregEncoding::decode(MI.getOperand(1).getImm());
+  if (HwRegId != Hwreg::ID_MODE ||
+      (!HasSetregVGPRMSBFixup && (Offset + Size) <= VGPRMSBShift))
+    return {};
+  unsigned Imm = MI.getOperand(0).getImm();
+  Imm = ((Imm >> Offset) & Hwreg::VGPR_MSB_MASK) >> VGPRMSBShift;
+  if (!HasSetregVGPRMSBFixup)
+    Imm &= llvm::maskTrailingOnes<unsigned>(Size);
+  return llvm::rotr<uint8_t>(static_cast<uint8_t>(Imm), /*R=*/2);
+}
+
 std::pair<const AMDGPU::OpName *, const AMDGPU::OpName *>
 getVGPRLoweringOperandTables(const MCInstrDesc &Desc) {
   static const AMDGPU::OpName VOPOps[4] = {

llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.h

@@ -31,6 +31,7 @@ struct Align;
 class Argument;
 class Function;
 class GlobalValue;
+class MachineInstr;
 class MCInstrInfo;
 class MCRegisterClass;
 class MCRegisterInfo;
@@ -1990,6 +1991,12 @@ unsigned getVGPREncodingMSBs(MCRegister Reg, const MCRegisterInfo &MRI);
 MCRegister getVGPRWithMSBs(MCRegister Reg, unsigned MSBs,
                            const MCRegisterInfo &MRI);
 
+/// \returns VGPR MSBs encoded in a S_SETREG_IMM32_B32 \p MI if it sets
+/// it. If \p HasSetregVGPRMSBFixup is true then size of the ID_MODE mask is
+/// ignored.
+std::optional<unsigned> convertSetRegImmToVgprMSBs(const MachineInstr &MI,
+                                                   bool HasSetregVGPRMSBFixup);
+
 // Returns a table for the opcode with a given \p Desc to map the VGPR MSB
 // set by the S_SET_VGPR_MSB to one of 4 sources. In case of VOPD returns 2
 // maps, one for X and one for Y component.

llvm/test/CodeGen/AMDGPU/code-size-estimate.ll

@@ -581,7 +581,8 @@ define amdgpu_ps float @s_fmaak_f32(float inreg %x, float inreg %y) {
 ;
 ; GFX1250-LABEL: s_fmaak_f32:
 ; GFX1250:       ; %bb.0:
-; GFX1250-NEXT:    s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 25, 1), 1 ; encoding: [0x41,0x06,0x80,0xb9,0x01,0x00,0x00,0x00]
+; GFX1250-NEXT:    s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 25, 1), 1 ; msbs: dst=0 src0=0 src1=0 src2=0
+; GFX1250-NEXT:    ; encoding: [0x41,0x06,0x80,0xb9,0x01,0x00,0x00,0x00]
 ; GFX1250-NEXT:    s_fmaak_f32 s0, s0, s1, 0x43800000 ; encoding: [0x00,0x01,0x80,0xa2,0x00,0x00,0x80,0x43]
 ; GFX1250-NEXT:    s_delay_alu instid0(SALU_CYCLE_3) ; encoding: [0x0b,0x00,0x87,0xbf]
 ; GFX1250-NEXT:    v_mov_b32_e32 v0, s0 ; encoding: [0x00,0x02,0x00,0x7e]

llvm/test/CodeGen/AMDGPU/vgpr-setreg-mode-swar.mir

@@ -1,5 +1,6 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 6
 # RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx1250 -run-pass=amdgpu-lower-vgpr-encoding -o - %s | FileCheck %s
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx1250 -start-before=amdgpu-lower-vgpr-encoding -o - %s | FileCheck --check-prefix=ASM %s
 
 ---
 # Case 1a: Size < 12 (size=4), imm32[12:19]=0
@@ -8,6 +9,10 @@
 # vgpr256/257 (both MSB=1): S_SET_VGPR_MSB mode = (1 << 0) | (1 << 6) = 65
 # Setreg (Size=4 <= 12) resets the mode scope and clears bits[12:19] to 0.
 # No VGPR instruction follows, so bits[12:19] remain 0. Setreg imm = 5.
+
+# ASM-LABEL: {{^}}setreg_size_lt_12:
+# ASM: s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 0, 4), 5 ;  msbs: dst=0 src0=0 src1=0 src2=0
+
 name:            setreg_size_lt_12
 tracksRegLiveness: true
 body:             |
@@ -27,6 +32,10 @@ body:             |
 # Case 1b: Size == 12 (boundary), imm32[12:19]=0
 # Setreg (Size=12 <= 12) resets the mode scope and clears bits[12:19] to 0.
 # No VGPR instruction follows, so bits[12:19] remain 0. Setreg imm = 0xABC = 2748.
+
+# ASM-LABEL: {{^}}setreg_size_eq_12:
+# ASM: s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 0, 12), 0xabc ;  msbs: dst=0 src0=0 src1=0 src2=0
+
 name:            setreg_size_eq_12
 tracksRegLiveness: true
 body:             |
@@ -46,6 +55,10 @@ body:             |
 # Case 1c: Size <= 12 with existing non-zero bits in imm32[12:19]
 # imm32 = 0x23005 (bits 12:19 = 0x23). Setreg resets mode scope and clears
 # bits[12:19] to 0. No VGPR instruction follows, so result = 0x00005 = 5.
+
+# ASM-LABEL: {{^}}setreg_size_lt_12_nonzero_upper:
+# ASM: s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 0, 4), 5 ;  msbs: dst=0 src0=0 src1=0 src2=0
+
 name:            setreg_size_lt_12_nonzero_upper
 tracksRegLiveness: true
 body:             |
@@ -66,6 +79,10 @@ body:             |
 # Case 2: Size > 12 (size=16), imm32[12:19] already matches VGPR MSBs
 # vgpr256/257: S_SET_VGPR_MSB mode = 65, MODE register mode = 5
 # imm32 = 0x5ABC = 23228 (bits 12:19 = 5), no modification needed
+
+# ASM-LABEL: {{^}}setreg_size_gt_12_match:
+# ASM: _setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 0, 16), 0x5abc ;  msbs: dst=1 src0=1 src1=0 src2=0
+
 name:            setreg_size_gt_12_match
 tracksRegLiveness: true
 body:             |
@@ -86,6 +103,10 @@ body:             |
 # Case 3: Size > 12 (size=16), imm32[12:19] doesn't match VGPR MSBs
 # vgpr256/257: S_SET_VGPR_MSB mode = 65, MODE register mode = 5
 # imm32 = 0x23ABC = 146108 (bits 12:19 = 0x23 != 5), must insert s_set_vgpr_msb after
+
+# ASM-LABEL: {{^}}setreg_size_gt_12_mismatch:
+# ASM: s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 0, 16), 0x23abc ;  msbs: dst=3 src0=0 src1=2 src2=0
+
 name:            setreg_size_gt_12_mismatch
 tracksRegLiveness: true
 body:             |
@@ -108,6 +129,10 @@ body:             |
 # Case 4: Non-MODE hwreg should not be modified
 # This uses ID_STATUS=2 instead of ID_MODE=1
 # vgpr256/257: S_SET_VGPR_MSB mode = 65
+
+# ASM-LABEL: {{^}}setreg_non_mode_hwreg:
+# ASM: s_setreg_imm32_b32 hwreg(HW_REG_WAVE_STATUS, 0, 4), 5{{$}}
+
 name:            setreg_non_mode_hwreg
 tracksRegLiveness: true
 body:             |
@@ -127,6 +152,10 @@ body:             |
 # Case 5: Size <= 12 with VGPR MSBs already present in imm32[12:19]
 # imm32 = 0x5005 = 20485 (bits 12:19 = 5). Setreg resets mode scope and clears
 # bits[12:19] to 0, regardless of prior content. Result = 5.
+
+# ASM-LABEL: {{^}}setreg_size_lt_12_already_correct:
+# ASM: s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 0, 4), 5 ;  msbs: dst=0 src0=0 src1=0 src2=0
+
 name:            setreg_size_lt_12_already_correct
 tracksRegLiveness: true
 body:             |
@@ -147,6 +176,10 @@ body:             |
 # Case 6: Different VGPR MSB value (using different high VGPRs)
 # vgpr512/513 (both MSB=2): S_SET_VGPR_MSB mode = (2 << 0) | (2 << 6) = 130
 # Setreg resets mode scope and clears bits[12:19] to 0. No VGPR follows. Result = 5.
+
+# ASM-LABEL: {{^}}setreg_different_vgpr_msb:
+# ASM: s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 0, 4), 5 ;  msbs: dst=0 src0=0 src1=0 src2=0
+
 name:            setreg_different_vgpr_msb
 tracksRegLiveness: true
 body:             |
@@ -170,6 +203,10 @@ body:             |
 #   MODE register mode = (1 << 0) | (1 << 2) | (2 << 4) = 37 (dst=1, src0=1, src1=2)
 # Piggybacking updates setreg imm32[12:19] from 0 to 37.
 # Final setreg imm = 5 | (37 << 12) = 151557
+
+# ASM-LABEL: {{^}}setreg_size_le_12_piggyback_superset:
+# ASM: s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 0, 4), 0x25005 ;  msbs: dst=1 src0=1 src1=2 src2=0
+
 name:            setreg_size_le_12_piggyback_superset
 tracksRegLiveness: true
 body:             |
@@ -195,6 +232,10 @@ body:             |
 # The setreg (Size=4 <= 12) resets the mode scope. Its bits[12:19] are cleared to 0.
 # The second VGPR's setMode piggybacks mode = 65 into the setreg's bits[12:19],
 # giving imm32 = 5 | (5 << 12) = 20485 = 0x5005. No separate S_SET_VGPR_MSB needed.
+
+# ASM-LABEL: {{^}}setreg_size_le_12_then_different_vgpr:
+# ASM: s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 0, 4), 0x5005 ;  msbs: dst=1 src0=1 src1=0 src2=0
+
 name:            setreg_size_le_12_then_different_vgpr
 tracksRegLiveness: true
 body:             |
@@ -224,6 +265,10 @@ body:             |
 # Second VGPR (vgpr512/513): S_SET_VGPR_MSB mode = 130, MODE register mode = 10
 # Since MostRecentModeSet = nullptr, a new s_set_vgpr_msb is inserted.
 # New s_set_vgpr_msb imm = NewMode | (OldMode << 8) = 130 | (65 << 8) = 16770
+
+# ASM-LABEL: {{^}}setreg_size_gt_12_match_then_different_vgpr:
+# ASM: s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 0, 16), 0x5abc ;  msbs: dst=1 src0=1 src1=0 src2=0
+
 name:            setreg_size_gt_12_match_then_different_vgpr
 tracksRegLiveness: true
 body:             |