"""A [ck_prefetch] marker in the merged .s ↔ objdump table."""

idx: int # index in merged list

direction: str # DIR_FORWARD or DIR_BACKWARD

"""Write to both stdout and a log file simultaneously."""

def __init__(self, log_path: Path):

self._file = log_path.open("w", encoding="utf-8")

self._stdout = sys.stdout

def write(self, data: str) -> int:

self._stdout.write(data)

self._file.write(data)

return len(data)

def flush(self) -> None:

self._stdout.flush()

self._file.flush()

def close(self) -> None:

self._file.close()

"""Return True if the line is an instruction (not a comment/label/directive/blank)."""

s = line.strip()

if not s:

return False

if s[0] in (';', '/', '.', '#'):

return False

if s.split()[0].endswith(':'):

return False

print(f"[run] {' '.join(cmd)}", flush=True)

result = subprocess.run(cmd, cwd=cwd, text=True, capture_output=True)

if result.stdout:

print(result.stdout, end="")

if result.stderr:

print(result.stderr, end="", file=sys.stderr)

if result.returncode != 0:

sys.exit(f"Command failed with exit code {result.returncode}")

"""Find the GPU .s file produced by --save-temps."""

all_candidates = sorted(

search_dir.rglob(f"{cpp_stem}*.s"),

key=lambda p: p.stat().st_mtime, reverse=True,

)

if not all_candidates:

sys.exit(

f"No .s file matching '{cpp_stem}*.s' found under {search_dir}.\n"

"Make sure --save-temps is in the target's compile options."

)

def is_gpu(p: Path) -> bool:

n = p.name

if "-host-" in n:

return False

if "-hip-" in n:

return True

if gpu_arch and gpu_arch in n:

return True

return False

gpu = [p for p in all_candidates if is_gpu(p)]

chosen = gpu[0] if gpu else all_candidates[0]

if not gpu:

print(f"[warn] No GPU .s found; falling back to {chosen.name}")

"""Find the most recent .o for the given CMake target.

candidates = sorted(

build_dir.rglob(f"{target}.dir/**/*.o"),

key=lambda p: p.stat().st_mtime, reverse=True,

)

if not candidates:

sys.exit(

f"No .o file found under '*/{target}.dir/' in {build_dir}.\n"

"Check that the target was built before running the patch script."

)

cmd = [objdump_path, f"--mcpu={mcpu}", "-d", str(obj_path)]

print(f"[run] {' '.join(cmd)}", flush=True)

result = subprocess.run(cmd, text=True, capture_output=True)

if result.returncode != 0:

sys.exit(f"objdump failed:\n{result.stderr}")

"""Extract the GPU architecture from .amdgcn_target directive in the .s file.

m = re.search(r'\.amdgcn_target\s+"[^"]*--(gfx[0-9a-zA-Z]+)', asm_text)

"""Return unique label names from [ck_prefetch] comments in the .s file."""

label_re = re.compile(r';\s*\[ck_prefetch\].*\bname\s*=\s*(\w+)')

seen: dict[str, None] = {}

for line in asm_text.splitlines():

m = label_re.search(line)

if m:

seen.setdefault(m.group(1), None)

"""Assemble GPU .s → temp .o via llvm-mc. Returns path (caller deletes)."""

llvm_mc = str(Path(objdump_path).parent / "llvm-mc")

out_obj = asm_file.with_suffix(".ck_tmp_patching.o")

run([llvm_mc, f"--mcpu={mcpu}", "--triple=amdgcn-amd-amdhsa",

"--filetype=obj", "-o", str(out_obj), str(asm_file)], asm_file.parent)

"""Parse llvm-objdump -d output into per-function (addr, instr_text) lists."""

instr_re = re.compile(r'^\t(.+?)//\s*([0-9a-fA-F]+):\s+[0-9a-fA-F]')

result: dict[str, list[tuple[int, str]]] = {}

cur_name: str | None = None

cur_entries: list[tuple[int, str]] = []

for line in objdump_text.splitlines():

m = OBJDUMP_FUNC_RE.match(line)

if m:

if cur_name is not None:

result[cur_name] = cur_entries

cur_name = m.group(1)

cur_entries = []

elif cur_name is not None:

m2 = instr_re.match(line)

if m2:

cur_entries.append((int(m2.group(2), 16), m2.group(1).strip()))

if cur_name is not None:

result[cur_name] = cur_entries

"""Split the .s file into per-function blocks."""

blocks: list[tuple[str, list[str]]] = []

current_name = "<top>"

current_lines: list[str] = []

for line in asm_text.splitlines():

m = FUNC_LABEL_RE.match(line)

if m:

if current_lines:

blocks.append((current_name, current_lines))

current_name = m.group(1)

current_lines = [line]

else:

current_lines.append(line)

if current_lines:

blocks.append((current_name, current_lines))

obj_entries: list[tuple[int, str]]) -> list[tuple[int | None, str]]:

"""Pair each .s instruction with its objdump entry by mnemonic matching.

MAX_LOOKAHEAD = 32

p2align_re = re.compile(r'\.p2align\s+(\d+)')

merged: list[tuple[int | None, str]] = []

obj_idx = 0

for line in s_lines:

m = p2align_re.search(line)

if m:

if obj_idx < len(obj_entries):

align = 1 << int(m.group(1))

while obj_idx < len(obj_entries) and (obj_entries[obj_idx][0] % align) != 0:

obj_idx += 1

merged.append((None, line))

continue

if is_asm_instruction(line):

if obj_idx < len(obj_entries):

s_mnem = line.strip().split()[0].lower()

for scan in range(obj_idx, min(obj_idx + MAX_LOOKAHEAD, len(obj_entries))):

if obj_entries[scan][1].split()[0].lower() == s_mnem:

obj_idx = scan

break

addr = obj_entries[obj_idx][0]

obj_idx += 1

else:

addr = None

merged.append((addr, line))

else:

merged.append((None, line))

dump_dir: Path | None = None

) -> dict[str, list[tuple[int | None, str]]]:

"""Merge .s ↔ objdump once per function. Returns {funcname: merged_list}.

s_blocks = split_functions(asm_text)

obj_funcs = parse_objdump_functions(objdump_text)

merged_funcs: dict[str, list[tuple[int | None, str]]] = {}

for name, s_lines in s_blocks:

if name not in obj_funcs:

continue

merged = _merge_s_and_objdump(s_lines, obj_funcs[name])

merged_funcs[name] = merged

if dump_dir is not None:

safe = re.sub(r'[^A-Za-z0-9_]', '_', name)[:80]

dump_path = dump_dir / f"merged_{safe}.txt"

with dump_path.open('w', encoding='utf-8') as fh:

for idx, (addr, line) in enumerate(merged):

addr_str = f'0x{addr:08x}' if addr is not None else ' '

fh.write(f'[{idx:5d}] {addr_str} {line.rstrip()}\n')

print(f"[dump] Merged table written to {dump_path}")

merged: list[tuple[int | None, str]],

tgt_idx: int,

tgt_mode: int,

name: str,

) -> int | None:

"""Resolve the target address for a prefetch's INST_PREFETCH_TARGET marker.

if tgt_mode == PLACE_MODE_BLOCK_ENTRY:

block_idx: int | None = None

for k in range(tgt_idx - 1, -1, -1):

if BLOCK_LABEL_RE.match(merged[k][1].strip()):

block_idx = k

break

scan_from = block_idx if block_idx is not None else tgt_idx

target: int | None = None

for k in range(scan_from + 1, len(merged)):

if merged[k][0] is not None:

target = merged[k][0]

break

if block_idx is not None and target is not None:

orig_target: int | None = None

for k in range(tgt_idx + 1, len(merged)):

if merged[k][0] is not None:

orig_target = merged[k][0]

break

if orig_target is not None and target != orig_target:

print(f"[adjust] {name[:60]!r}: BLOCK_ENTRY mode — "

f"block label at merged[{block_idx}] "

f"→ target 0x{target:x} (was 0x{orig_target:x}, "

f"saved {orig_target - target}B)")

return target

# DEFAULT mode (mode=0): first instruction after [ck_label].

for k in range(tgt_idx + 1, len(merged)):

if merged[k][0] is not None:

return merged[k][0]

name: str,

pair_idx: int,

pc_next: int,

target: int,

orig_klength: int,

direction: str,

offset_cl: int,

func_end: int,

) -> tuple[int, int] | None:

"""Compute (koffset, klength) for one prefetch pair with OOB clamping.

target_aligned = target & ~(CACHELINE_SIZE - 1)

offset_bytes = offset_cl * CACHELINE_SIZE

klength = orig_klength

if direction == DIR_BACKWARD:

region_end = target_aligned + CACHELINE_SIZE + offset_bytes

region_start = region_end - (klength + 1) * CACHELINE_SIZE

min_base = (pc_next & ~(CACHELINE_SIZE - 1)) + CACHELINE_SIZE

if region_start < min_base:

region_start = min_base

usable = (region_end - region_start) // CACHELINE_SIZE

if usable <= 0:

klength = NOP_KLENGTH_SENTINEL

print(f"[nop] {name[:60]!r}: backward prefetch fully OOB "

f"(min_base 0x{min_base:x} >= region_end 0x{region_end:x}), "

f"replacing with 2× s_nop")

else:

klength = usable - 1

print(f"[clamp] {name[:60]!r}: backward start clamped "

f"(first cacheline after pc_next: 0x{min_base:x}), "

f"klength {orig_klength} → {klength}")

if klength == NOP_KLENGTH_SENTINEL:

print(f"[debug] func={name[:60]!r} pair {pair_idx}: "

f"pc_next=0x{pc_next:x} dir=backward → NOP (0 cachelines in bounds)")

return (0, NOP_KLENGTH_SENTINEL)

prefetch_base = region_start

koffset = prefetch_base - pc_next

print(f"[debug] func={name[:60]!r} pair {pair_idx}: "

f"pc_next=0x{pc_next:x} target=0x{target:x} dir=backward "

f"offset={offset_cl}cl prefetch_base=0x{prefetch_base:x} "

f"koffset=0x{koffset:x} ({koffset}B) klength={klength} "

f"region=[0x{region_start:x}, 0x{region_end:x}) "

f"({(region_end - region_start)}B = {klength + 1} cachelines)")

return (koffset, klength)

# ── Forward direction ────────────────────────────────────────────────

prefetch_base = target_aligned + offset_bytes

koffset = prefetch_base - pc_next

if koffset < 0:

print(f"[warn] {name[:60]!r}: negative koffset — target before prefetch, skipping")

return None

region_end = prefetch_base + (klength + 1) * CACHELINE_SIZE

if region_end > func_end:

needed = max(0, (func_end - prefetch_base + CACHELINE_SIZE - 1) // CACHELINE_SIZE)

if needed == 0:

klength = NOP_KLENGTH_SENTINEL

print(f"[nop] {name[:60]!r}: forward prefetch fully OOB "

f"(prefetch_base 0x{prefetch_base:x} >= func_end 0x{func_end:x}), "

f"replacing with 2× s_nop")

else:

klength = needed - 1

region_end = prefetch_base + (klength + 1) * CACHELINE_SIZE

print(f"[clamp] {name[:60]!r}: forward end clamped "

f"(func_end 0x{func_end:x}), klength {orig_klength} → {klength}")

if klength == NOP_KLENGTH_SENTINEL:

print(f"[debug] func={name[:60]!r} pair {pair_idx}: "

f"pc_next=0x{pc_next:x} dir=forward → NOP (0 cachelines in bounds)")

else:

region_start = prefetch_base

print(f"[debug] func={name[:60]!r} pair {pair_idx}: "

f"pc_next=0x{pc_next:x} target=0x{target:x} dir=forward "

f"offset={offset_cl}cl "

f"koffset=0x{koffset:x} ({koffset}B) klength={klength} "

f"region=[0x{region_start:x}, 0x{region_end:x}) "

f"({(region_end - region_start)}B = {klength + 1} cachelines)")

label: str) -> dict[str, list[tuple[int, int]]]:

"""Compute per-function (koffset, klength) for INST_PREFETCH/INST_PREFETCH_TARGET label pairs.

# [ck_prefetch] marks INST_PREFETCH sites, [ck_label] marks INST_PREFETCH_TARGET targets.

prefetch_re = re.compile(rf";\s*\[ck_prefetch\].*\bname\s*=\s*{re.escape(label)}\b")

target_re = re.compile(rf";\s*\[ck_label\].*\bname\s*=\s*{re.escape(label)}\b")

either_re = re.compile(rf";\s*(?:\[ck_label\]|\[ck_prefetch\]).*\bname\s*=\s*{re.escape(label)}\b")

klength_re = re.compile(r's_prefetch_inst_pc_rel\s+\S+\s*,\s*\S+\s*,\s*(\d+)')

mode_re = re.compile(r'\bmode\s*=\s*(\d+)')

dir_re = re.compile(r'\bdir\s*=\s*(\w+)')

offset_re = re.compile(r'\boffset\s*=\s*(-?\d+)')

results: dict[str, list[tuple[int, int]]] = {}

for name, merged in merged_funcs.items():

if not any(either_re.search(line) for _, line in merged):

continue

# Determine function end address (for OOB clamping).

func_end: int = 0

for addr, _ in reversed(merged):

if addr is not None:

func_end = addr + 8 # conservative: largest instruction is 8 bytes

break

# Classify markers from .s comments.

prefetch_sites: list[PrefetchSite] = []

target_sites: list[TargetSite] = []

for idx, (_addr, line) in enumerate(merged):

if prefetch_re.search(line):

m_dir = dir_re.search(line)

m_off = offset_re.search(line)

prefetch_sites.append(PrefetchSite(

idx=idx,

direction=m_dir.group(1) if m_dir else DIR_FORWARD,

offset_cl=int(m_off.group(1)) if m_off else 0,

))

elif target_re.search(line):

m_mode = mode_re.search(line)

target_sites.append(TargetSite(

idx=idx,

mode=int(m_mode.group(1)) if m_mode else PLACE_MODE_DEFAULT,

))

pairs: list[tuple[int, int]] = []

for pf in prefetch_sites:

# Find the s_prefetch_inst_pc_rel instruction and parse its klength.

pf_instr_idx: int | None = None

orig_klength = 3 # default

j = pf.idx + 1

while j < len(merged) and 's_prefetch_inst_pc_rel' not in merged[j][1]:

j += 1

if j < len(merged):

pf_instr_idx = j

m_kl = klength_re.search(merged[j][1])

if m_kl:

orig_klength = int(m_kl.group(1))

# pc_next = address of the instruction after s_prefetch_inst_pc_rel.

pc_next: int | None = None

if pf_instr_idx is not None:

for k in range(pf_instr_idx + 1, len(merged)):

if merged[k][0] is not None:

pc_next = merged[k][0]

break

if pc_next is None:

print(f"[warn] {name[:60]!r}: no pc_next for prefetch at merged[{pf.idx}], skipping")

continue

# Find the nearest INST_PREFETCH_TARGET after this INST_PREFETCH.

tgt: TargetSite | None = None

for t in target_sites:

if t.idx > pf.idx:

tgt = t

break

if tgt is None:

# Unpaired prefetch — treat as forward with target=pc_next.

print(f"[warn] {name[:60]!r}: unpaired prefetch label at merged[{pf.idx}] — "

f"using koffset=0, clamping klength")

result = _clamp_prefetch_region(

name, len(pairs), pc_next, pc_next, orig_klength,

DIR_FORWARD, 0, func_end)

if result is not None:

pairs.append(result)

continue

target = _resolve_target_address(merged, tgt.idx, tgt.mode, name)

if target is None:

print(f"[warn] {name[:60]!r}: no target address for label at merged[{tgt.idx}]")

continue

result = _clamp_prefetch_region(

name, len(pairs), pc_next, target, orig_klength,

pf.direction, pf.offset_cl, func_end)

if result is not None:

pairs.append(result)

if pairs:

results[name] = pairs

if not results:

print(f"[skip] Label '{label}' not found in any matching function block.")

else:

total = sum(len(v) for v in results.values())

print(f"[offsets] {len(results)} function(s), {total} pair(s) with koffset for '{label}'.")

"""Patch s_prefetch_inst_pc_rel koffset and klength operands in the .s file.

prefetch_re = re.compile(

r'(s_prefetch_inst_pc_rel\s+)(?:0x[0-9a-fA-F]+|0|-?\d+)'

r'(\s*,\s*null\s*,\s*)(?:\d+)')

# Full-line regex to capture indentation and the entire prefetch instruction

# (used for NOP replacement).

prefetch_full_re = re.compile(

r'^(\s*)s_prefetch_inst_pc_rel\s+\S+\s*,\s*\S+\s*,\s*\d+(.*)$')

text = asm_file.read_text(encoding="utf-8", errors="replace")

out_lines: list[str] = []

current_func = "<top>"

func_pf_idx: dict[str, int] = {}

n_patched = 0

n_nopped = 0

for line in text.splitlines(keepends=True):

m = FUNC_LABEL_RE.match(line.rstrip())

if m:

current_func = m.group(1)

if current_func in func_koffsets:

pair_list = func_koffsets[current_func]

idx = func_pf_idx.get(current_func, 0)

if idx < len(pair_list):

koffset, klength = pair_list[idx]

if klength == NOP_KLENGTH_SENTINEL:

# Replace 8-byte prefetch with 2× 4-byte s_nop 0

m_full = prefetch_full_re.match(line.rstrip('\n\r'))

if m_full:

indent = m_full.group(1)

trailing = m_full.group(2) # e.g. comment

eol = line[len(line.rstrip('\n\r')):] # preserve \n

nop_lines = (f"{indent}s_nop 0{trailing}{eol}"

f"{indent}s_nop 0{eol}")

print(f"[patch-s] {current_func[:60]}: prefetch[{idx}] → "

f"2× s_nop 0 (OOB)")

func_pf_idx[current_func] = idx + 1

n_nopped += 1

n_patched += 1

out_lines.append(nop_lines)

continue

else:

koffset_str = hex(koffset)

new_line, n = prefetch_re.subn(

rf'\g<1>{koffset_str}\g<2>{klength}', line)

if n:

print(f"[patch-s] {current_func[:60]}: prefetch[{idx}] → "

f"koffset={koffset_str} klength={klength}")

func_pf_idx[current_func] = idx + 1

n_patched += n

out_lines.append(new_line)

continue

out_lines.append(line)

new_text = ''.join(out_lines)

print(f"[patch-s] {n_patched} operand(s) patched ({n_nopped} replaced with NOPs).")

if new_text == text:

print("[patch-s] No change.")

return False

asm_file.write_text(new_text, encoding="utf-8")

print(f"[patch-s] Written: {asm_file}")

"""Find .text section in an ELF image starting at data[base:].

import struct as _s

d = data[base:]

if len(d) < 64 or d[:4] != b'\x7fELF':

return None

ei_class, ei_data = d[4], d[5]

endian = '<' if ei_data == 1 else '>'

try:

if ei_class == 2:

e_shoff, = _s.unpack_from(f'{endian}Q', d, 40)

e_shentsize, = _s.unpack_from(f'{endian}H', d, 58)

e_shnum, = _s.unpack_from(f'{endian}H', d, 60)

e_shstrndx, = _s.unpack_from(f'{endian}H', d, 62)

addr_in_shdr, off_in_shdr, sz_in_shdr = 16, 24, 32

addr_fmt, off_fmt, sz_fmt = f'{endian}Q', f'{endian}Q', f'{endian}Q'

else:

e_shoff, = _s.unpack_from(f'{endian}I', d, 32)

e_shentsize, = _s.unpack_from(f'{endian}H', d, 46)

e_shnum, = _s.unpack_from(f'{endian}H', d, 48)

e_shstrndx, = _s.unpack_from(f'{endian}H', d, 50)

addr_in_shdr, off_in_shdr, sz_in_shdr = 12, 16, 20

addr_fmt, off_fmt, sz_fmt = f'{endian}I', f'{endian}I', f'{endian}I'

shstr_sh = e_shoff + e_shstrndx * e_shentsize

shstr_off, = _s.unpack_from(off_fmt, d, shstr_sh + off_in_shdr)

for i in range(e_shnum):

sh = e_shoff + i * e_shentsize

name_idx, = _s.unpack_from(f'{endian}I', d, sh)

ns = shstr_off + name_idx

ne = d.index(b'\x00', ns)

if d[ns:ne] == b'.text':

sec_addr, = _s.unpack_from(addr_fmt, d, sh + addr_in_shdr)

sec_off, = _s.unpack_from(off_fmt, d, sh + off_in_shdr)

sec_sz, = _s.unpack_from(sz_fmt, d, sh + sz_in_shdr)

return (sec_off, sec_sz, sec_addr)

except (_s.error, ValueError):

pass

) -> tuple[int, int, int, str] | None:

"""Locate the GPU bundle in a fat .o / fatbin.

import struct as _s

MAGIC = b'__CLANG_OFFLOAD_BUNDLE__'

magic_idx = data.find(MAGIC)

if magic_idx < 0:

print(f"[{tag}] __CLANG_OFFLOAD_BUNDLE__ magic not found")

return None

hdr = magic_idx + len(MAGIC)

if hdr + 8 > len(data):

print(f"[{tag}] Truncated fatbin header")

return None

num_bundles, = _s.unpack_from('<Q', data, hdr)

cur = hdr + 8

gpu_off = gpu_sz = 0

gpu_triple = ""

for _ in range(num_bundles):

if cur + 24 > len(data):

break

off, sz, triple_sz = _s.unpack_from('<QQQ', data, cur)

cur += 24

if triple_sz == 0 or triple_sz > 512 or cur + triple_sz > len(data):

break

triple = data[cur:cur + triple_sz].decode('utf-8', errors='replace')

cur += triple_sz

if 'amdgcn' in triple or (triple.startswith('hip') and 'host' not in triple):

gpu_off, gpu_sz, gpu_triple = off, sz, triple

if not gpu_triple:

print(f"[{tag}] No GPU entry found in fatbin header")

return None

mcpu: str, objdump_path: str, tmp_path: Path,

tag: str = "fatbin") -> str | None:

"""Extract the GPU ELF from *data*, run objdump -d, return the text or None."""

try:

tmp_path.write_bytes(bytes(data[abs_gpu_start:abs_gpu_start + gpu_sz]))

result = subprocess.run(

[objdump_path, f"--mcpu={mcpu}", "-d", str(tmp_path)],

text=True, capture_output=True,

)

if result.returncode != 0:

print(f"[{tag}] objdump on GPU ELF failed: {result.stderr[:200]}")

return None

return result.stdout

finally:

idx: int, new_koffset: int, new_klength: int) -> None:

"""Patch a single s_prefetch_inst_pc_rel at *instr_pos* in *fat_data*.

import struct as _struct

old_dw0 = _struct.unpack_from('<I', fat_data, instr_pos)[0]

old_dw1 = _struct.unpack_from('<I', fat_data, instr_pos + 4)[0]

raw_before = fat_data[instr_pos:instr_pos + 8]

print(f"[patch-obj] VA 0x{instr_va:x}: BEFORE "

f"dw0=0x{old_dw0:08x} dw1=0x{old_dw1:08x} "

f"raw={raw_before.hex()} "

f"klength_bits12_6={(old_dw0 >> KLENGTH_SHIFT) & 0x7F}")

if new_klength == NOP_KLENGTH_SENTINEL:

_struct.pack_into('<I', fat_data, instr_pos, S_NOP_ENCODING)

_struct.pack_into('<I', fat_data, instr_pos + 4, S_NOP_ENCODING)

raw_after = fat_data[instr_pos:instr_pos + 8]

print(f"[patch-obj] VA 0x{instr_va:x}: AFTER "

f"2× s_nop 0 (0x{S_NOP_ENCODING:08x}) "

f"raw={raw_after.hex()}")

print(f"[patch-obj] VA 0x{instr_va:x}: prefetch[{idx}] → "

f"2× s_nop 0 (OOB)")

else:

first_dword = _struct.unpack_from('<I', fat_data, instr_pos)[0]

first_dword = (first_dword & ~KLENGTH_MASK) | ((new_klength & 0x7F) << KLENGTH_SHIFT)

_struct.pack_into('<I', fat_data, instr_pos, first_dword)

second_dword = _struct.unpack_from('<I', fat_data, instr_pos + 4)[0]

second_dword = (second_dword & ~KOFFSET_MASK) | (new_koffset & KOFFSET_MASK)

_struct.pack_into('<I', fat_data, instr_pos + 4, second_dword)

raw_after = fat_data[instr_pos:instr_pos + 8]

new_dw0 = _struct.unpack_from('<I', fat_data, instr_pos)[0]

new_dw1 = _struct.unpack_from('<I', fat_data, instr_pos + 4)[0]

print(f"[patch-obj] VA 0x{instr_va:x}: AFTER "

f"dw0=0x{new_dw0:08x} dw1=0x{new_dw1:08x} "

f"raw={raw_after.hex()} "

f"klength_bits12_6={(new_dw0 >> KLENGTH_SHIFT) & 0x7F}")

print(f"[patch-obj] VA 0x{instr_va:x}: prefetch[{idx}] → "

func_koffsets: dict[str, list[tuple[int, int]]]) -> bool:

"""Patch s_prefetch_inst_pc_rel koffsets and klengths directly in the GPU ELF

fat_data = bytearray(fat_obj.read_bytes())

bundle = _find_gpu_bundle(fat_data, tag="patch-obj")

if bundle is None:

return False

magic_idx, gpu_off, gpu_sz, gpu_triple = bundle

abs_gpu_start = magic_idx + gpu_off

print(f"[patch-obj] GPU bundle: '{gpu_triple}' abs=0x{abs_gpu_start:x} size={gpu_sz}")

if fat_data[abs_gpu_start:abs_gpu_start + 4] != b'\x7fELF':

print("[patch-obj] GPU data does not start with ELF magic")

return False

text_info = _find_elf_text_section(fat_data, abs_gpu_start)

if text_info is None:

print("[patch-obj] Cannot find .text in GPU ELF")

return False

text_foff, text_sz, text_va = text_info

print(f"[patch-obj] .text: foff=0x{text_foff:x} size={text_sz} va=0x{text_va:x}")

objdump_text = _objdump_gpu_elf(fat_data, abs_gpu_start, gpu_sz, mcpu, objdump_path,

fat_obj.with_suffix(".ck_gpu_elf_tmp"), tag="patch-obj")

if objdump_text is None:

return False

n_patched = 0

current_func: str | None = None

func_pf_idx: dict[str, int] = {}

for line in objdump_text.splitlines():

m = OBJDUMP_FUNC_RE.match(line)

if m:

current_func = m.group(1)

continue

if not current_func or current_func not in func_koffsets:

continue

if 's_prefetch_inst_pc_rel' not in line:

continue

m2 = OBJDUMP_ADDR_RE.search(line)

if not m2:

continue

idx = func_pf_idx.get(current_func, 0)

pair_list = func_koffsets[current_func]

if idx >= len(pair_list):

continue

instr_va = int(m2.group(1), 16)

new_koffset, new_klength = pair_list[idx]

instr_pos = abs_gpu_start + text_foff + (instr_va - text_va)

_patch_one_prefetch(fat_data, instr_pos, instr_va, idx,

new_koffset, new_klength)

func_pf_idx[current_func] = idx + 1

n_patched += 1

if n_patched == 0:

print("[patch-obj] No s_prefetch_inst_pc_rel found to patch")

return False

fat_obj.write_bytes(bytes(fat_data))

# Sanity check: re-read and verify the write persisted

import hashlib

written_hash = hashlib.md5(fat_obj.read_bytes()).hexdigest()

expected_hash = hashlib.md5(bytes(fat_data)).hexdigest()

if written_hash != expected_hash:

print(f"[patch-obj] WARNING: write verification failed! "

f"expected={expected_hash} written={written_hash}")

print(f"[patch-obj] Patched {n_patched} instruction(s) in {fat_obj.name} "

f"md5={written_hash}")

func_koffsets: dict[str, list[tuple[int, int]]]) -> bool:

"""Verify patched s_prefetch_inst_pc_rel koffsets and klengths. Returns True if all match.

import struct as _struct

data = fat_obj.read_bytes()

bundle = _find_gpu_bundle(data, tag="verify")

if bundle is None:

return False

magic_idx, gpu_off, gpu_sz, _ = bundle

abs_start = magic_idx + gpu_off

# Locate .text for raw byte diagnostics

text_info = _find_elf_text_section(data, abs_start)

text_foff = text_va = 0

if text_info:

text_foff, text_sz, text_va = text_info

objdump_text = _objdump_gpu_elf(data, abs_start, gpu_sz, mcpu, objdump_path,

fat_obj.with_suffix(".ck_verify_tmp"), tag="verify")

if objdump_text is None:

return False

prefetch_re = re.compile(

r's_prefetch_inst_pc_rel\s+(0x[0-9a-fA-F]+|\d+)\s*,\s*\S+\s*,\s*(\d+)')

current_func: str | None = None

func_pf_idx: dict[str, int] = {}

# Track VAs already consumed as part of a NOP pair so the second s_nop

# of a pair (or compiler-emitted s_nops) are not misidentified.

consumed_nop_vas: set[int] = set()

ok = True