[comgr][hotswap] preserve VGPR-MSB mode when splitting gfx1250 WMMAs

The WMMA split pass rewrites gfx1250 B0-only WMMAs into pairs of narrower
WMMAs that also run on A0. On gfx1250, VGPR operands past v255 are reached
through a persistent VGPR-MSB bank mode, so a split whose upper half crosses
v255 (or a K-split whose second half reuses dst as src2) must execute that
half under a different bank.

Make the pass VGPR-MSB aware:

- Recover the live VGPR-MSB mode at each WMMA via a whole-function CFG fixed
  point, fail-closed: an ambiguous or unprovable mode declines the split.
- Bracket the upper half with s_set_vgpr_msb (switch before, restore after),
  and emit nothing when no bank change is needed.
- Fail closed on every matched-but-unsplittable path (RequiredPatchFailed)
  instead of silently leaving an A0-illegal opcode in .text.

Use the same control-flow surface as the rewrite so the mode proof cannot be
seeded from a wrong incoming mode:

- Decline the analysis when collectDirectBranchTargets() reports an unresolved
  call target, since such a call may enter any function at an interior offset.
- Resolve absolute-immediate and PC-materialized s_swap_pc_i64 / s_set_pc_i64
  targets in the interior-entry pre-pass so a cross-function interior entry
  declines the split rather than seeding the symbol-start mode.
- Seed declared kernel-descriptor entries so an interior KD-entry block is
  analyzed from the real entry mode rather than misread as unreachable.

Add lit coverage for the bank-crossing, fail-closed, and loop-carried mode
cases, plus interior cross-function s_swap_pc_i64 (materialized and absolute),
an interior s_set_pc_i64 jump, and an interior kernel-descriptor entry.

Author: YashDeshpande25

amd/comgr/src/comgr-hotswap-b0a0.cpp

@@ -989,23 +989,18 @@ struct PcMaterializedCallInfo {
 /// by llvm/test/MC/AMDGPU/gfx1250_asm_salu_lit64.s. Stop at the first
 /// overlapping definition or control-flow boundary, so any variation remains
 /// unresolved and follows the existing fail-closed policy.
-static std::optional<PcMaterializedCallInfo>
-matchPcMaterializedCall(ArrayRef<InternalDecodedInst> Decoded, size_t CallIndex,
-                        const LLVMState &LS, uint64_t TextAddr) {
-  const InternalDecodedInst &Call = Decoded[CallIndex];
-  if (!Call.DecodeSucceeded || Call.Inst.getOpcode() != LS.SSwapPcI64Opcode ||
-      Call.Inst.getNumOperands() < 2 || !Call.Inst.getOperand(0).isReg() ||
-      !Call.Inst.getOperand(0).getReg())
-    return std::nullopt;
-  const MCOperand &TargetOperand =
-      Call.Inst.getOperand(Call.Inst.getNumOperands() - 1);
-  if (!TargetOperand.isReg() || !TargetOperand.getReg())
-    return std::nullopt;
-  MCRegister TargetRegister(TargetOperand.getReg());
-
+// Shared get-PC/add resolver behind matchPcMaterializedCall (s_swap_pc_i64
+// calls) and the WMMA split pass's s_set_pc_i64 jump handling. The backward
+// scan and fail-closed policy are identical for both transfer kinds; only the
+// terminating opcode and the return-register semantics differ, which the
+// callers handle.
+std::optional<MaterializedPcSequence>
+resolveMaterializedPcTarget(ArrayRef<InternalDecodedInst> Decoded,
+                            size_t TransferIndex, MCRegister TargetRegister,
+                            const LLVMState &LS, uint64_t TextAddr) {
   std::optional<size_t> AddIndex;
   int64_t AddImmediate = 0;
-  for (size_t I = CallIndex; I != 0;) {
+  for (size_t I = TransferIndex; I != 0;) {
     --I;
     const InternalDecodedInst &Candidate = Decoded[I];
     if (!Candidate.DecodeSucceeded || isControlFlowBoundary(Candidate, LS))
@@ -1041,23 +1036,43 @@ matchPcMaterializedCall(ArrayRef<InternalDecodedInst> Decoded, size_t CallIndex,
       return std::nullopt;
 
     std::optional<uint64_t> GetPcAddress = checkedAddUint64(
-        TextAddr, Candidate.Offset, "PC-materialized call instruction");
+        TextAddr, Candidate.Offset, "PC-materialized transfer instruction");
     if (!GetPcAddress)
       return std::nullopt;
     std::optional<uint64_t> PcValue = checkedAddUint64(
-        *GetPcAddress, Candidate.Size, "PC-materialized call PC value");
+        *GetPcAddress, Candidate.Size, "PC-materialized transfer PC value");
     if (!PcValue)
       return std::nullopt;
     // s_add_nc_u64 uses modulo-2^64 arithmetic. Casting the signed MC
     // immediate to uint64_t and adding it reproduces both positive and
     // negative literals, including INT64_MIN, without signed overflow.
-    return PcMaterializedCallInfo{
-        *PcValue + static_cast<uint64_t>(AddImmediate), Candidate.Offset,
-        Call.Offset, MCRegister(Call.Inst.getOperand(0).getReg())};
+    return MaterializedPcSequence{
+        *PcValue + static_cast<uint64_t>(AddImmediate), Candidate.Offset};
   }
   return std::nullopt;
 }
 
+static std::optional<PcMaterializedCallInfo>
+matchPcMaterializedCall(ArrayRef<InternalDecodedInst> Decoded, size_t CallIndex,
+                        const LLVMState &LS, uint64_t TextAddr) {
+  const InternalDecodedInst &Call = Decoded[CallIndex];
+  if (!Call.DecodeSucceeded || Call.Inst.getOpcode() != LS.SSwapPcI64Opcode ||
+      Call.Inst.getNumOperands() < 2 || !Call.Inst.getOperand(0).isReg() ||
+      !Call.Inst.getOperand(0).getReg())
+    return std::nullopt;
+  const MCOperand &TargetOperand =
+      Call.Inst.getOperand(Call.Inst.getNumOperands() - 1);
+  if (!TargetOperand.isReg() || !TargetOperand.getReg())
+    return std::nullopt;
+  std::optional<MaterializedPcSequence> Sequence = resolveMaterializedPcTarget(
+      Decoded, CallIndex, MCRegister(TargetOperand.getReg()), LS, TextAddr);
+  if (!Sequence)
+    return std::nullopt;
+  return PcMaterializedCallInfo{Sequence->Target, Sequence->SequenceStart,
+                                Call.Offset,
+                                MCRegister(Call.Inst.getOperand(0).getReg())};
+}
+
 struct KnownCallSite {
   size_t InstIndex = 0;
   uint64_t Target = 0;
@@ -2273,10 +2288,11 @@ static std::optional<uint32_t> applyGfx1250B0toA0Rules(
       *PoolVAddr, Elf.textAddr(), "trampoline pool base offset");
   if (!PoolBaseOffset)
     return std::nullopt;
-  PatchContext Ctx{
-      Config,      Decoded,           Text,         TextSize, *PoolBaseOffset,
-      LS,          OutTrampolines,    Sleds,        Elf,      Liveness,
-      KernelStats, OutScratchPatches, *ControlFlow, Profile};
+  PatchContext Ctx{Config,         Decoded,         Text,
+                   TextSize,       *PoolBaseOffset, LS,
+                   OutTrampolines, Sleds,           Elf,
+                   Liveness,       KernelStats,     OutScratchPatches,
+                   *ControlFlow,   Profile,         DeclaredEntries->Entries};
 
   const HotswapPatchVTable &VT = getHotswapPatchVTable();
 

amd/comgr/src/comgr-hotswap-internal.h

@@ -1133,6 +1133,16 @@ struct DirectControlFlowInfo {
   bool HasUnresolvedTargets = false;
 };
 
+// Per-instruction persistent gfx1250 VGPR-MSB mode (packed src0/src1/src2/dst,
+// two bits each, values 0-255) recovered by the WMMA split pass's
+// whole-function CFG fixed point. The sentinels distinguish "not analyzed",
+// "validated unreachable", and "reachable but ambiguous" so a required WMMA
+// split can fail closed when the incoming mode cannot be proven. See
+// comgr-hotswap-patch-wmma-split.cpp.
+inline constexpr int8_t VgprMsbUnanalyzed = -3;
+inline constexpr int8_t VgprMsbUnreachable = -2;
+inline constexpr int8_t VgprMsbUnknown = -1;
+
 /// Mutable per-run context threaded through all patch passes. Bundles the
 /// input config, decoded instruction stream, raw .text bytes, MC state,
 /// output streams (trampolines / scratch info), and the shared ELF view +
@@ -1158,10 +1168,21 @@ struct PatchContext {
   // Per-rewrite profiling session (inert unless AMD_COMGR_TIME_STATISTICS is
   // set). Deep patch sites record into its lock-free local array.
   HotswapProfile &Profile;
+  // Text-relative declared entry offsets (function symbol starts and kernel
+  // descriptor entries) from collectDeclaredTextEntries(). The WMMA split
+  // pass's VGPR-MSB analysis seeds each in-range entry as an ABI entry point so
+  // an interior kernel-descriptor entry is analyzed from the real entry mode
+  // instead of being misclassified as unreachable.
+  llvm::ArrayRef<uint64_t> DeclaredEntries;
   // Required patches are transformations whose unpatched original code is
   // unsafe to return when the selected rewrite policy needs the patch.
   bool RequiredPatchFailed = false;
   bool RequiredPatchApplied = false;
+  // Packed per-instruction gfx1250 VGPR-MSB mode, lazily populated by the WMMA
+  // split pass (empty until then). Indexed by position in Decoded; each entry
+  // is a VgprMsb* sentinel or a 0-255 mode. See
+  // comgr-hotswap-patch-wmma-split.cpp.
+  std::vector<int16_t> VgprMsbModeBefore;
   // Sum of the bytes already queued in OutTrampolines. Keeping this in the
   // per-rewrite context makes each new pool-position calculation constant
   // time even for code objects with many thousands of patch sites.
@@ -1277,6 +1298,27 @@ std::optional<uint64_t>
 evaluateDirectControlFlowTarget(const InternalDecodedInst &DI,
                                 const LLVMState &LS);
 
+/// A canonical get-PC/add materialized address feeding a PC-materialized
+/// transfer (s_get_pc_i64 / s_add_nc_u64 / s_{swap,set}_pc_i64). \p Target is
+/// the absolute in-.text virtual address; \p SequenceStart is the .text offset
+/// of the s_get_pc_i64 that begins the sequence.
+struct MaterializedPcSequence {
+  uint64_t Target = 0;
+  uint64_t SequenceStart = 0;
+};
+
+/// Resolve the target of a PC-materialized transfer whose target register is
+/// \p TargetReg and whose transfer instruction (an s_swap_pc_i64 call or an
+/// s_set_pc_i64 jump) is \p Decoded[TransferIndex]. Scans backward for the
+/// single s_add_nc_u64 then s_get_pc_i64 that define \p TargetReg, stopping at
+/// the first control-flow boundary or unexpected clobber so any variation
+/// stays unresolved (nullopt) for fail-closed callers. \p TextAddr is the
+/// .text base virtual address.
+std::optional<MaterializedPcSequence>
+resolveMaterializedPcTarget(llvm::ArrayRef<InternalDecodedInst> Decoded,
+                            size_t TransferIndex, llvm::MCRegister TargetReg,
+                            const LLVMState &LS, uint64_t TextAddr);
+
 /// Collect branch and call targets used to protect interior entry points from
 /// trampoline coalescing. Absolute addresses in TextAddr .. TextAddr +
 /// TextSize are converted to text-relative offsets. Canonical PC-materialized