[llvm] [AMDGPU] Handle `s_setreg_imm32_b32` targeting `MODE` register (PR #174681)

[Shilei Tian via llvm-commits]

================
@@ -315,6 +338,82 @@ AMDGPULowerVGPREncoding::handleCoissue(MachineBasicBlock::instr_iterator I) {
   return Prev;
 }
 
+/// Convert mode value from S_SET_VGPR_MSB format to MODE register format.
+/// S_SET_VGPR_MSB uses: (src0[0-1], src1[2-3], src2[4-5], dst[6-7])
+/// MODE register uses:  (dst[0-1], src0[2-3], src1[4-5], src2[6-7])
+/// This is a left rotation by 2 bits on an 8-bit value.
+static int64_t convertModeToSetregFormat(int64_t Mode) {
+  assert(isUInt<8>(Mode) && "Mode expected to be 8-bit");
+  return llvm::rotl<uint8_t>(static_cast<uint8_t>(Mode), /*R=*/2);
+}
+
+bool AMDGPULowerVGPREncoding::updateSetregModeImm(MachineInstr &MI,
+                                                  int64_t ModeValue) {
+  assert(MI.getOpcode() == AMDGPU::S_SETREG_IMM32_B32);
+
+  // Convert from S_SET_VGPR_MSB format to MODE register format
+  int64_t SetregMode = convertModeToSetregFormat(ModeValue);
+
+  MachineOperand *ImmOp = TII->getNamedOperand(MI, AMDGPU::OpName::imm);
+  int64_t OldImm = ImmOp->getImm();
+  int64_t NewImm =
+      (OldImm & ~AMDGPU::Hwreg::VGPR_MSB_MASK) | (SetregMode << VGPRMSBShift);
+  ImmOp->setImm(NewImm);
+  return NewImm != OldImm;
+}
+
+bool AMDGPULowerVGPREncoding::handleSetregMode(MachineInstr &MI) {
+  using namespace AMDGPU::Hwreg;
+
+  assert(MI.getOpcode() == AMDGPU::S_SETREG_IMM32_B32 &&
+         "only S_SETREG_IMM32_B32 needs to be handled");
+
+  MachineOperand *SIMM16Op = TII->getNamedOperand(MI, AMDGPU::OpName::simm16);
+  assert(SIMM16Op && "SIMM16Op must be present");
+
+  auto [HwRegId, Offset, Size] = HwregEncoding::decode(SIMM16Op->getImm());
+  (void)Offset;
+  if (HwRegId != ID_MODE)
+    return false;
+
+  int64_t ModeValue = static_cast<int64_t>(CurrentMode);
+
+  // Case 1: Size <= 12 - the original instruction uses imm32[0:Size-1], so
+  // imm32[12:19] is unused. Safe to set imm32[12:19] to the correct VGPR
+  // MSBs.
+  if (Size <= VGPRMSBShift) {
+    // This instruction now acts as MostRecentModeSet so it can be updated if
+    // CurrentMode changes via piggybacking.
+    MostRecentModeSet = &MI;
----------------
shiltian wrote:

This is where we treat them the same.

https://github.com/llvm/llvm-project/pull/174681