[comgr][hotswap] Emit address-carrying half last when splitting ds_load_2addr_* with dst/addr overlap

amd/comgr/src/comgr-hotswap-patch-trampoline.cpp

@@ -248,19 +248,76 @@ splitDstPair(MCRegister CompoundReg, bool IsB64, const MCRegisterInfo &MRI) {
 }
 
 // Expand a DS 2-address load into two single-address loads (dst, addr).
+//
+// Register-overlap hazard: the source address (Regs[1]) may alias ANY register
+// in the destination window (Regs[0]) -- e.g. any of v2..v5 in
+// `ds_load_2addr_b64 v[2:5], vX`. The original single 2-addr instruction is
+// safe because it reads the address once before writing the destination. The
+// naive split into two loads is not: whichever half executes first writes its
+// destination, and if that write lands on the address register, the second load
+// then computes its LDS location from a corrupted address.
+//
+// Fix: scan the whole destination window, determine which half carries the
+// address, and emit that half LAST. A single DS load reads its address before
+// writing its own destination, so the address-carrying half is safe as the
+// final use; the other (disjoint) half, emitted first, never touches the
+// address. This covers every position of the address within the window (low or
+// high half) and the no-overlap case. Only loads have this hazard (only loads
+// write the destination); the (dst-half, offset) pairing is preserved
+// regardless of emission order.
 std::vector<std::string> expandDs2AddrLoad(const DsOperands &Ops,
                                            StringRef ToMnem) {
   if (Ops.Regs.size() < 2)
     return {};
+  const MCRegisterInfo &MRI = *Ops.MRI;
   std::pair<std::string, std::string> Dst =
-      splitDstPair(Ops.Regs[0], Ops.IsB64, *Ops.MRI);
+      splitDstPair(Ops.Regs[0], Ops.IsB64, MRI);
   if (Dst.first.empty())
     return {};
-  std::string Addr = toAsmRegName(*Ops.MRI, Ops.Regs[1]);
-  return {
-      ToMnem.str() + " " + Dst.first + ", " + Addr + fmtOffset(Ops.Off0),
-      ToMnem.str() + " " + Dst.second + ", " + Addr + fmtOffset(Ops.Off1),
-  };
+  std::string Addr = toAsmRegName(MRI, Ops.Regs[1]);
+  std::string LoadLow =
+      ToMnem.str() + " " + Dst.first + ", " + Addr + fmtOffset(Ops.Off0);
+  std::string LoadHigh =
+      ToMnem.str() + " " + Dst.second + ", " + Addr + fmtOffset(Ops.Off1);
+
+  // Scan the ENTIRE destination window and classify which half the address
+  // aliases. The 32-bit fragments are ordered low->high; the first half
+  // (Subs[0 .. Half-1]) backs Dst.first and the rest back Dst.second, matching
+  // splitDstPair. b64: Half == 2 (window v2..v5); b32: Half == 1 (window
+  // v2,v3).
+  SmallVector<MCRegister, 4> Subs = getDirectSubRegs(Ops.Regs[0], MRI);
+  unsigned Half = Subs.size() / 2;
+  bool AddrInLow = false, AddrInHigh = false;
+  for (unsigned I = 0; I < Subs.size(); ++I) {
+    if (!MRI.regsOverlap(Ops.Regs[1], Subs[I]))
+      continue;
+    if (I < Half)
+      AddrInLow = true;
+    else
+      AddrInHigh = true;
+  }
+
+  if (AddrInLow && AddrInHigh) {
+    // Unreachable for real DS loads: the address is a single 32-bit VGPR, so it
+    // can alias at most one half. If it somehow spanned both, no reordering is
+    // safe (either half's load clobbers the address the other still needs), so
+    // skip the patch and leave the atomic original in place rather than emit
+    // code that corrupts the address.
+    log() << "hotswap: error: ds_2addr load: address " << Addr
+          << " overlaps both destination halves; cannot reorder safely, "
+             "leaving original instruction in place\n";
+    return {};
+  }
+  if (AddrInLow) {
+    // Address in the low half: emit the (disjoint) high half first so the
+    // address survives until the low-half load reads-then-overwrites it.
+    log() << "hotswap: ds_2addr load: address " << Addr
+          << " aliases low destination half; emitting high half first\n";
+    return {LoadHigh, LoadLow};
+  }
+  // Address in the high half, or no overlap at all: the default low-then-high
+  // order already places the address-carrying half (if any) last.
+  return {LoadLow, LoadHigh};
 }
 
 // Expand a DS 2-address store into two single-address stores (addr, data).

amd/comgr/test-lit/hotswap-trampoline-ds-load-overlap.s

@@ -0,0 +1,287 @@
+// COM: Test the register-overlap reorder in expandDs2AddrLoad: when a
+// COM: ds_load_2addr_* address register aliases the destination window, the
+// COM: split must emit the half that carries the address LAST. A single DS
+// COM: load reads its address before writing its destination, so the
+// COM: address-carrying half is safe as the final use; the other (disjoint)
+// COM: half, emitted first, never clobbers the address.
+// COM:
+// COM: Motivating real case (gfx1250 rocThrust, hotswap OFF disasm):
+// COM:   ds_load_2addr_b64 v[2:5], v2 offset1:1
+// COM: address v2 aliases the low destination pair v[2:3]; the naive low-first
+// COM: split clobbers v2 before the high load reads it (intermittent torn read).
+// COM:
+// COM: Covers the whole destination window for b64 (address = v2,v3 in the low
+// COM: pair / v4,v5 in the high pair) and b32 (address = v2 low / v3 high).
+// COM: Only loads are affected (only loads write the destination); the
+// COM: store/exchange split paths are unchanged and covered by the other
+// COM: hotswap-trampoline-ds*.s tests, which also cover the no-overlap order.
+
+// RUN: %clang -target amdgcn-amd-amdhsa -mcpu=gfx1250 -nostdlib %s -o %t.elf
+
+// RUN: hotswap-rewrite %t.elf \
+// RUN:   amdgcn-amd-amdhsa--gfx1250 amdgcn-amd-amdhsa--gfx1250 \
+// RUN:   --output %t.out.elf \
+// RUN:   | %FileCheck --check-prefix=API %s
+// API: RESULT: SUCCESS
+
+// RUN: %llvm-objdump -d %t.out.elf | %FileCheck --check-prefix=DISASM %s
+
+.amdgcn_target "amdgcn-amd-amdhsa--gfx1250"
+.text
+
+// ---- Kernel 1: b64, address in LOW half (v2) -- the real rocThrust case -----
+// COM: offset0 defaults to 0 (omitted in disasm); offset1:1 -> byte 8. v2
+// COM: aliases the low pair v[2:3] -> reorder: high v[4:5] first, low v[2:3]
+// COM: last so v2 stays live until the low load reads-then-overwrites it.
+// DISASM-LABEL: <test_ds_load_b64_addr_low_v2>:
+// DISASM-NOT: ds_load_2addr_b64
+// DISASM: s_branch
+// DISASM: s_wait_dscnt 0x0
+// DISASM: ds_load_b64 v[4:5], v2 offset:8
+// DISASM-NEXT: ds_load_b64 v[2:3], v2
+// DISASM: s_branch
+.globl test_ds_load_b64_addr_low_v2
+.p2align 8
+.type test_ds_load_b64_addr_low_v2,@function
+test_ds_load_b64_addr_low_v2:
+  ds_load_2addr_b64 v[2:5], v2 offset1:1
+  s_wait_dscnt 0x0
+  s_endpgm
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+.Ltest_ds_load_b64_addr_low_v2_end:
+.size test_ds_load_b64_addr_low_v2, .Ltest_ds_load_b64_addr_low_v2_end-test_ds_load_b64_addr_low_v2
+
+// ---- Kernel 2: b64, address in LOW half (v3, upper dword of the low pair) ---
+// COM: offset0:1 offset1:2 -> bytes 8 and 16. v3 is still in low pair v[2:3],
+// COM: so this is a low-half overlap -> reorder (high v[4:5] first).
+// DISASM-LABEL: <test_ds_load_b64_addr_low_v3>:
+// DISASM-NOT: ds_load_2addr_b64
+// DISASM: s_branch
+// DISASM: s_wait_dscnt 0x0
+// DISASM: ds_load_b64 v[4:5], v3 offset:16
+// DISASM-NEXT: ds_load_b64 v[2:3], v3 offset:8
+// DISASM: s_branch
+.globl test_ds_load_b64_addr_low_v3
+.p2align 8
+.type test_ds_load_b64_addr_low_v3,@function
+test_ds_load_b64_addr_low_v3:
+  ds_load_2addr_b64 v[2:5], v3 offset0:1 offset1:2
+  s_wait_dscnt 0x0
+  s_endpgm
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+.Ltest_ds_load_b64_addr_low_v3_end:
+.size test_ds_load_b64_addr_low_v3, .Ltest_ds_load_b64_addr_low_v3_end-test_ds_load_b64_addr_low_v3
+
+// ---- Kernel 3: b64, address in HIGH half (v4) -- no reorder (default) -------
+// COM: offset0:1 offset1:2 -> bytes 8 and 16. The low load v[2:3] is disjoint
+// COM: from v4, so the default low-first order is already safe; the high load
+// COM: (which overwrites v4) is the address's last use.
+// DISASM-LABEL: <test_ds_load_b64_addr_high_v4>:
+// DISASM-NOT: ds_load_2addr_b64
+// DISASM: s_branch
+// DISASM: s_wait_dscnt 0x0
+// DISASM: ds_load_b64 v[2:3], v4 offset:8
+// DISASM-NEXT: ds_load_b64 v[4:5], v4 offset:16
+// DISASM: s_branch
+.globl test_ds_load_b64_addr_high_v4
+.p2align 8
+.type test_ds_load_b64_addr_high_v4,@function
+test_ds_load_b64_addr_high_v4:
+  ds_load_2addr_b64 v[2:5], v4 offset0:1 offset1:2
+  s_wait_dscnt 0x0
+  s_endpgm
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+.Ltest_ds_load_b64_addr_high_v4_end:
+.size test_ds_load_b64_addr_high_v4, .Ltest_ds_load_b64_addr_high_v4_end-test_ds_load_b64_addr_high_v4
+
+// ---- Kernel 4: b64, address in HIGH half (v5) -- no reorder -----------------
+// COM: v5 is the upper dword of the high pair v[4:5]; still a high-half
+// COM: overlap, so default low-first order is retained.
+// DISASM-LABEL: <test_ds_load_b64_addr_high_v5>:
+// DISASM-NOT: ds_load_2addr_b64
+// DISASM: s_branch
+// DISASM: s_wait_dscnt 0x0
+// DISASM: ds_load_b64 v[2:3], v5 offset:8
+// DISASM-NEXT: ds_load_b64 v[4:5], v5 offset:16
+// DISASM: s_branch
+.globl test_ds_load_b64_addr_high_v5
+.p2align 8
+.type test_ds_load_b64_addr_high_v5,@function
+test_ds_load_b64_addr_high_v5:
+  ds_load_2addr_b64 v[2:5], v5 offset0:1 offset1:2
+  s_wait_dscnt 0x0
+  s_endpgm
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+.Ltest_ds_load_b64_addr_high_v5_end:
+.size test_ds_load_b64_addr_high_v5, .Ltest_ds_load_b64_addr_high_v5_end-test_ds_load_b64_addr_high_v5
+
+// ---- Kernel 5: b32, address in LOW half (v2) -- reorder ---------------------
+// COM: b32 split: low half = v2, high half = v3 (Half == 1). offset0:1
+// COM: offset1:2 -> bytes 4 and 8. v2 in low -> reorder: high v3 first.
+// DISASM-LABEL: <test_ds_load_b32_addr_low_v2>:
+// DISASM-NOT: ds_load_2addr_b32
+// DISASM: s_branch
+// DISASM: s_wait_dscnt 0x0
+// DISASM: ds_load_b32 v3, v2 offset:8
+// DISASM-NEXT: ds_load_b32 v2, v2 offset:4
+// DISASM: s_branch
+.globl test_ds_load_b32_addr_low_v2
+.p2align 8
+.type test_ds_load_b32_addr_low_v2,@function
+test_ds_load_b32_addr_low_v2:
+  ds_load_2addr_b32 v[2:3], v2 offset0:1 offset1:2
+  s_wait_dscnt 0x0
+  s_endpgm
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+.Ltest_ds_load_b32_addr_low_v2_end:
+.size test_ds_load_b32_addr_low_v2, .Ltest_ds_load_b32_addr_low_v2_end-test_ds_load_b32_addr_low_v2
+
+// ---- Kernel 6: b32, address in HIGH half (v3) -- no reorder -----------------
+// COM: v3 is the high half -> default low-first order (low v2 first).
+// DISASM-LABEL: <test_ds_load_b32_addr_high_v3>:
+// DISASM-NOT: ds_load_2addr_b32
+// DISASM: s_branch
+// DISASM: s_wait_dscnt 0x0
+// DISASM: ds_load_b32 v2, v3 offset:4
+// DISASM-NEXT: ds_load_b32 v3, v3 offset:8
+// DISASM: s_branch
+.globl test_ds_load_b32_addr_high_v3
+.p2align 8
+.type test_ds_load_b32_addr_high_v3,@function
+test_ds_load_b32_addr_high_v3:
+  ds_load_2addr_b32 v[2:3], v3 offset0:1 offset1:2
+  s_wait_dscnt 0x0
+  s_endpgm
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+  s_nop 0
+.Ltest_ds_load_b32_addr_high_v3_end:
+.size test_ds_load_b32_addr_high_v3, .Ltest_ds_load_b32_addr_high_v3_end-test_ds_load_b32_addr_high_v3
+
+// COM: Idempotency: rewriting the output again is a no-op (no DS2 mnemonic
+// COM: remains, second pass produces identical bytes).
+// RUN: hotswap-rewrite %t.out.elf \
+// RUN:   amdgcn-amd-amdhsa--gfx1250 amdgcn-amd-amdhsa--gfx1250 \
+// RUN:   --check-idempotent \
+// RUN:   | %FileCheck --check-prefix=IDEM %s
+// IDEM: IDEMPOTENT: YES
+
+.rodata
+.p2align 8
+.amdhsa_kernel test_ds_load_b64_addr_low_v2
+  .amdhsa_next_free_vgpr 6
+  .amdhsa_next_free_sgpr 1
+.end_amdhsa_kernel
+
+.amdhsa_kernel test_ds_load_b64_addr_low_v3
+  .amdhsa_next_free_vgpr 6
+  .amdhsa_next_free_sgpr 1
+.end_amdhsa_kernel
+
+.amdhsa_kernel test_ds_load_b64_addr_high_v4
+  .amdhsa_next_free_vgpr 6
+  .amdhsa_next_free_sgpr 1
+.end_amdhsa_kernel
+
+.amdhsa_kernel test_ds_load_b64_addr_high_v5
+  .amdhsa_next_free_vgpr 6
+  .amdhsa_next_free_sgpr 1
+.end_amdhsa_kernel
+
+.amdhsa_kernel test_ds_load_b32_addr_low_v2
+  .amdhsa_next_free_vgpr 4
+  .amdhsa_next_free_sgpr 1
+.end_amdhsa_kernel
+
+.amdhsa_kernel test_ds_load_b32_addr_high_v3
+  .amdhsa_next_free_vgpr 4
+  .amdhsa_next_free_sgpr 1
+.end_amdhsa_kernel