diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..33473dd --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,52 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Project Overview + +SASS FMA Bypass Tool — a SASS-level binary translator for NVIDIA GPUs (Volta sm_70 / Ampere sm_80) that intercepts CUDA module loading via `LD_PRELOAD` and rewrites defective `FFMA` instructions into `FMUL` + `FADD` sequences. Targets hardware where the fused multiply-add unit is physically defective. + +## Build and Test + +```bash +# Build everything (interceptor.so + all test executables + cubin/sass artifacts) +./build.sh + +# Run all test suites (basic, comprehensive, stress, runtime JIT) +./run_test.sh + +# Use with any CUDA application +CUDA_FORCE_PTX_JIT=1 LD_PRELOAD=./interceptor.so ./your_cuda_app +``` + +Requires: CUDA Toolkit (`nvcc`, `nvdisasm`), Python 3, g++. The Makefile only builds `interceptor.so`; `build.sh` also compiles all CUDA test binaries with `nvcc -arch=sm_70`. + +## Architecture + +Two-layer design: + +1. **`interceptor.cpp`** — C++ shared library loaded via `LD_PRELOAD`. Hooks `cuModuleLoadDataEx`, `cuModuleLoadData`, `cuModuleLoad`, and `__cudaRegisterFatBinary`. On each cubin load, writes the binary to a temp file, spawns `python3 rewriter.py`, and loads the patched result. Thread-safe via mutex. + +2. **`rewriter.py`** — Python engine that performs the actual transformation: + - Disassembles cubin via `nvdisasm -hex`, auto-detects sm_70 vs sm_80 + - Parses SASS into per-kernel sections (`.text.`) + - Builds CFG and runs iterative backward liveness analysis to find dead registers for use as temporaries + - Splits each `FFMA` into `FMUL` + `FADD`, preserving predicates and scoreboard barriers + - Manually expands the ELF `.text` sections and shifts all section/program headers + - Patches `.nv.info` metadata (REGCOUNT) and jump table offsets in `.nv.constant*` sections + - Updates branch targets (BRA/SSY/PBK) using a `pc_map` of old-to-new PC values + +## Key Invariants + +- Instructions are 128-bit (16 bytes each). All PC values are multiples of 16. +- FFMA splitting doubles instruction count for each FMA — the ELF must physically grow. +- When `r_dest == r_srcC` in an FFMA, a spare register is required (found via liveness or by incrementing `max_reg`). +- Jump table entries in `.nv.constant*` are only patched if they match known PC values (verified via BRX/LDC tracing to avoid corrupting float literals). +- The symtab (SHT_SYMTAB, type 11) function symbol sizes are updated to reflect expanded `.text` sections. + +## Test Suites + +- `test_kernel.cu` — Basic FMA test +- `comprehensive.cu` / `comprehensive_test.cu` — Accumulator aliasing, predicated FMAs, loops with backward branches, switch/jump tables +- `stress_test.cu` — 1024-FMA unrolled expansion, register saturation (250+ regs) +- `runtime_test.cu` — Runtime JIT path via `CUDA_FORCE_PTX_JIT=1` diff --git a/comprehensive_test.cu b/comprehensive_test.cu index d1e8f7d..36aa377 100644 --- a/comprehensive_test.cu +++ b/comprehensive_test.cu @@ -52,6 +52,10 @@ int main() { // Load the cubin (Interceptor will hook this) std::ifstream ifs("comprehensive.cubin", std::ios::binary | std::ios::ate); + if (!ifs.is_open()) { + std::cerr << "Error: comprehensive.cubin not found. Run ./build.sh first." << std::endl; + return 1; + } std::streamsize size = ifs.tellg(); ifs.seekg(0, std::ios::beg); std::vector buffer(size); diff --git a/interceptor.cpp b/interceptor.cpp index 1a74314..a297b0d 100644 --- a/interceptor.cpp +++ b/interceptor.cpp @@ -17,7 +17,7 @@ typedef void** (*cudaRegisterFatBinary_t)(void*); static std::mutex g_mutex; -// --- FatBinary Structures (Simplied) --- +// --- FatBinary Structures (Simplified) --- struct __fatBinC_Header { uint32_t magic; uint16_t version; @@ -43,13 +43,23 @@ extern "C" void** __cudaRegisterFatBinary(void* fatb) { size_t get_elf_size(const void* data) { const Elf64_Ehdr* ehdr = (const Elf64_Ehdr*)data; if (memcmp(ehdr->e_ident, ELFMAG, SELFMAG) != 0) return 0; - size_t sh_end = ehdr->e_shoff + (ehdr->e_shentsize * ehdr->e_shnum); - size_t ph_end = ehdr->e_phoff + (ehdr->e_phentsize * ehdr->e_phnum); - return ((sh_end > ph_end) ? sh_end : ph_end) + 8192; + const uint8_t* base = (const uint8_t*)data; + size_t max_end = ehdr->e_shoff + (size_t)ehdr->e_shentsize * ehdr->e_shnum; + size_t ph_end = ehdr->e_phoff + (size_t)ehdr->e_phentsize * ehdr->e_phnum; + if (ph_end > max_end) max_end = ph_end; + for (int i = 0; i < ehdr->e_shnum; i++) { + const Elf64_Shdr* shdr = (const Elf64_Shdr*)(base + ehdr->e_shoff + (size_t)i * ehdr->e_shentsize); + if (shdr->sh_type != SHT_NOBITS) { + size_t sec_end = shdr->sh_offset + shdr->sh_size; + if (sec_end > max_end) max_end = sec_end; + } + } + return max_end; } // Internal helper to run the rewriter on a memory buffer CUresult patch_and_load(cuModuleLoadDataEx_t original_func, CUmodule* module, const void* image, unsigned int numOptions, CUjit_option* options, void** optionValues) { + std::lock_guard lock(g_mutex); size_t size = get_elf_size(image); if (size == 0) return original_func(module, image, numOptions, options, optionValues); @@ -59,8 +69,13 @@ CUresult patch_and_load(cuModuleLoadDataEx_t original_func, CUmodule* module, co int fd_out = mkstemps(tmp_out, 6); close(fd_out); - write(fd_in, image, size); + ssize_t written = write(fd_in, image, size); close(fd_in); + if (written < 0 || (size_t)written != size) { + unlink(tmp_in); + unlink(tmp_out); + return original_func(module, image, numOptions, options, optionValues); + } std::string cmd = "python3 rewriter.py " + std::string(tmp_in) + " " + std::string(tmp_out); int status = system(cmd.c_str()); diff --git a/requirements.txt b/requirements.txt index cd699e9..e69de29 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +0,0 @@ -lief -pyelftools diff --git a/rewriter.py b/rewriter.py index 75379fc..ac67983 100644 --- a/rewriter.py +++ b/rewriter.py @@ -4,22 +4,18 @@ import re import struct -# Unified Constants for sm_70/sm_80 FMUL_BASE_H1 = 0x7220 FADD_BASE_H1 = 0x7221 -BRA_BASE_H1 = 0x7947 -BRA_BASE_H2 = 0x000fc0000383ffff class SassRewriter: def __init__(self, input_path, output_path): self.input_path = input_path self.output_path = output_path self.arch = "sm_70" - self.sections = {} + self.sections = {} self.pc_map = {} # (sec_name, old_pc) -> new_pc - self.verified_jump_tables = set() - self.used_registers = set() - self.max_reg = 0 + self.verified_jump_tables = set() + self.section_max_reg = {} # per-kernel max register number def disassemble(self): """Runs nvdisasm and detects architecture.""" @@ -29,15 +25,19 @@ def disassemble(self): if "SM80" in raw: self.arch = "sm_80" print("[Rewriter] Detected SM80 (Ampere) architecture.") + print("[Rewriter] WARNING: SM80 instruction encoding is not fully validated.") else: self.arch = "sm_70" print("[Rewriter] Detected SM70 (Volta) architecture.") return raw - except: return "" + except Exception: + return "" def get_regs(self, op, args): - regs = [int(r) for r in re.findall(r"R(\d+)", args)] - if not regs: return set(), set() + raw = [int(r) for r in re.findall(r"R(\d+)", args)] + regs = [r for r in raw if r != 255] # R255 is the zero register + if not regs: + return set(), set() if any(x in op for x in ["ST", "BRA", "EXIT", "RET", "JMP", "CAL", "BAR", "MEMBAR"]): return set(), set(regs) return {regs[0]}, set(regs[1:]) @@ -46,12 +46,13 @@ def solve_liveness(self, instrs): pc_to_idx = {ins['pc']: i for i, ins in enumerate(instrs)} for i, ins in enumerate(instrs): succs = [] - if "EXIT" not in ins['op'] and "RET" not in ins['op'] and i+1 < len(instrs): - succs.append(i+1) + if "EXIT" not in ins['op'] and "RET" not in ins['op'] and i + 1 < len(instrs): + succs.append(i + 1) if "BRA" in ins['op']: off = struct.unpack('> 32) & 0xFFFFFFFF))[0] target = ins['pc'] + 16 + off - if target in pc_to_idx: succs.append(pc_to_idx[target]) + if target in pc_to_idx: + succs.append(pc_to_idx[target]) ins['succs'] = succs ins['live_in'], ins['live_out'] = set(), set() @@ -61,18 +62,18 @@ def solve_liveness(self, instrs): for i in reversed(range(len(instrs))): ins = instrs[i] new_out = set() - for s_idx in ins['succs']: new_out.update(instrs[s_idx]['live_in']) + for s_idx in ins['succs']: + new_out.update(instrs[s_idx]['live_in']) new_in = ins['uses'].union(new_out - ins['defs']) if new_in != ins['live_in'] or new_out != ins['live_out']: ins['live_in'], ins['live_out'] = new_in, new_out changed = True def parse_sass(self, raw_text): - for r in re.findall(r"R(\d+)", raw_text): self.max_reg = max(self.max_reg, int(r)) sec_re = re.compile(r"//-+\s+(?P\.text\.\S+)\s+-+") ins_re = re.compile(r"/\*(?P[0-9a-f]+)\*/\s+(?P[\w\.]+)(?P[^;]*);\s+/\*\s+(?P

0x[0-9a-f]+)\s+\*/") const_re = re.compile(r"c\s*\[\s*(?P0x[0-9a-f]+)\s*\]\s*\[\s*(?P0x[0-9a-f]+)\s*\]") - reg_sources = {} + reg_sources = {} lines = raw_text.splitlines() cur_sec, i = ".text", 0 @@ -82,151 +83,229 @@ def parse_sass(self, raw_text): if s_m: cur_sec = s_m.group('name') self.sections[cur_sec] = [] - i += 1; continue + i += 1 + continue m = ins_re.search(line) if m: - if cur_sec not in self.sections: self.sections[cur_sec] = [] - pc, op, args, h1 = int(m.group('pc'), 16), m.group('op'), m.group('args'), int(m.group('h1'), 16) + if cur_sec not in self.sections: + self.sections[cur_sec] = [] + pc = int(m.group('pc'), 16) + op = m.group('op') + args = m.group('args') + h1 = int(m.group('h1'), 16) i += 1 h2 = 0 if i < len(lines): h2_m = re.search(r"/\*\s+(?P

0x[0-9a-f]+)\s+\*/", lines[i]) - if h2_m: h2 = int(h2_m.group('h2'), 16) - + if h2_m: + h2 = int(h2_m.group('h2'), 16) + c_m = const_re.search(args) if c_m: - bank, off = int(c_m.group('bank'), 16), int(c_m.group('off'), 16) + bank = int(c_m.group('bank'), 16) + off = int(c_m.group('off'), 16) sec_name = f".nv.constant{bank}" - if "BRX" in op: self.verified_jump_tables.add((sec_name, off)) + if "BRX" in op: + self.verified_jump_tables.add((sec_name, off)) elif "LDC" in op: r_m = re.search(r"(?PR\d+)", args) - if r_m: reg_sources[r_m.group('reg')] = (sec_name, off) + if r_m: + reg_sources[r_m.group('reg')] = (sec_name, off) if any(x in op for x in ["JMP", "CAL"]): r_m = re.search(r"(?PR\d+)", args) - if r_m and r_m.group('reg') in reg_sources: self.verified_jump_tables.add(reg_sources[r_m.group('reg')]) + if r_m and r_m.group('reg') in reg_sources: + self.verified_jump_tables.add(reg_sources[r_m.group('reg')]) defs, uses = self.get_regs(op, args) - self.sections[cur_sec].append({'pc':pc, 'h1':h1, 'h2':h2, 'op':op, 'defs':defs, 'uses':uses}) - else: i += 1 + self.sections[cur_sec].append({ + 'pc': pc, 'h1': h1, 'h2': h2, 'op': op, + 'defs': defs, 'uses': uses + }) + else: + i += 1 for sn, instrs in self.sections.items(): + # Per-kernel max register tracking + sec_max = max((r for ins in instrs for r in ins['defs'] | ins['uses']), default=0) self.solve_liveness(instrs) new_instrs = [] for ins in instrs: self.pc_map[(sn, ins['pc'])] = len(new_instrs) * 16 - if "FMA" in ins['op']: - r_dest, r_srcA, r_srcB, r_srcC = (ins['h1']>>16)&0xFF, (ins['h1']>>24)&0xFF, (ins['h1']>>32)&0xFF, ins['h2']&0xFF - pred, temp_reg = ins['h2']&0xF, r_dest + if ins['op'].startswith('FFMA'): # Only match FFMA, not DFMA/HFMA2 + r_dest = (ins['h1'] >> 16) & 0xFF + r_srcA = (ins['h1'] >> 24) & 0xFF + r_srcB = (ins['h1'] >> 32) & 0xFF + r_srcC = ins['h2'] & 0xFF + pred = ins['h2'] & 0xF + temp_reg = r_dest if r_dest == r_srcC: spare = next((r for r in range(255) if r not in ins['live_out']), None) - if spare is None: self.max_reg += 1; spare = self.max_reg + if spare is None: + sec_max += 1 + spare = sec_max temp_reg = spare - - # 1. FMUL (Common for 70/80) + + # Preserve original predicate encoding (critical for predicated FFMAs + # where temp_reg == r_dest; unconditional FMUL would clobber r_dest + # when the predicate is false) + orig_pred_byte = ins['h2'] & 0xFF + + # 1. FMUL: temp = srcA * srcB fmul_h1 = FMUL_BASE_H1 | (temp_reg << 16) | (r_srcA << 24) | (r_srcB << 32) - fmul_h2 = (ins['h2'] & ~((0x7<<46)|(0xF<<41)|0xFF|0xF)) | (0x0<<46) | (0x1<<41) | 0xFF | pred - new_instrs.append({'h1':fmul_h1, 'h2':fmul_h2, 'op':'FMUL', 'old_pc':ins['pc']}) - # 2. FADD (Common for 70/80) + fmul_h2 = (ins['h2'] & ~((0x7 << 46) | (0xF << 41) | 0xFF)) | (0x1 << 41) | orig_pred_byte + new_instrs.append({'h1': fmul_h1, 'h2': fmul_h2, 'op': 'FMUL', 'old_pc': ins['pc']}) + + # 2. FADD: dest = srcC + temp fadd_h1 = FADD_BASE_H1 | (r_dest << 16) | (r_srcC << 24) | (temp_reg << 32) orig_wb = (ins['h2'] >> 46) & 0x7 - fadd_h2 = (ins['h2'] & ~((0x3F<<52)|(0x7<<46)|0xF)) | (0x1<<52) | (orig_wb<<46) | pred - new_instrs.append({'h1':fadd_h1, 'h2':fadd_h2, 'op':'FADD', 'old_pc':None}) + fadd_h2 = (ins['h2'] & ~((0x3F << 52) | (0x7 << 46) | 0xF)) | (0x1 << 52) | (orig_wb << 46) | pred + new_instrs.append({'h1': fadd_h1, 'h2': fadd_h2, 'op': 'FADD', 'old_pc': None}) else: - new_instrs.append({'h1':ins['h1'], 'h2':ins['h2'], 'op':ins['op'], 'old_pc':ins['pc']}) + new_instrs.append({'h1': ins['h1'], 'h2': ins['h2'], 'op': ins['op'], 'old_pc': ins['pc']}) + self.section_max_reg[sn] = sec_max self.sections[sn] = new_instrs + def _section_name(self, data, str_tab_off, sh_name_idx): + start = str_tab_off + sh_name_idx + return data[start:start + 64].split(b'\x00', 1)[0].decode('ascii', errors='replace') + def patch_and_rebuild(self): - with open(self.input_path, 'rb') as f: data = bytearray(f.read()) + with open(self.input_path, 'rb') as f: + data = bytearray(f.read()) + if data[:4] != b'\x7fELF': - print(f"Error: {self.input_path} is not a valid ELF file.") + print(f"[Rewriter] Error: {self.input_path} is not a valid ELF file.") sys.exit(1) - e_phoff, e_shoff = struct.unpack_from('> 32) & 0xFFFFFFFF))[0] if ins['old_pc'] is not None: + off = struct.unpack('> 32) & 0xFFFFFFFF))[0] target = ins['old_pc'] + 16 + off if (sn, target) in self.pc_map: - ins['h1'] = (ins['h1'] & 0x00000000FFFFFFFF) | (((self.pc_map[(sn, target)] - (i*16+16)) & 0xFFFFFFFF) << 32) + new_off = self.pc_map[(sn, target)] - (i * 16 + 16) + ins['h1'] = (ins['h1'] & 0x00000000FFFFFFFF) | ((new_off & 0xFFFFFFFF) << 32) - cur_data, cum_growth = data, 0 - sh_table = [list(struct.unpack_from(' old_off: sh_table[j][4] += growth - - for j in range(e_phnum): - if ph_table[j][2] <= old_off and (ph_table[j][2] + ph_table[j][5]) > old_off: - ph_table[j][5] += growth # p_filesz - ph_table[j][6] += growth # p_memsz - elif ph_table[j][2] > old_off: - ph_table[j][2] += growth # p_offset - - if e_phoff > old_off: e_phoff += growth - if e_shoff > old_off: e_shoff += growth - cum_growth += growth + if sh_table[j][4] > old_off: + sh_table[j][4] += growth + # Update program headers: expand segments containing this section, + # shift segments after it + for p in range(e_phnum): + p_off = ph_table[p][2] + p_fsz = ph_table[p][5] + if p_fsz > 0 and p_off <= old_off < p_off + p_fsz: + ph_table[p][5] += growth # p_filesz + ph_table[p][6] += growth # p_memsz + elif p_off > old_off: + ph_table[p][2] += growth # p_offset + if e_phoff > old_off: + e_phoff += growth + if e_shoff > old_off: + e_shoff += growth + # Patch .nv.info and .nv.constant sections for i in range(e_shnum): - name = get_name(sh_table[i][0]) + name = self._section_name(data, str_tab_off, sh_table[i][0]) o, s = sh_table[i][4], sh_table[i][5] target_k = name.replace(".nv.info.", ".text.") if ".nv.info." in name else None - if ".nv.info" in name or ".nv.constant" in name: - content = bytearray(cur_data[o:o+s]) + is_info = ".nv.info" in name + is_constant = ".nv.constant" in name and not is_info + if is_info or is_constant: + # Only patch constant sections that contain verified jump tables + if is_constant: + has_jt = any(sn == name for sn, _ in self.verified_jump_tables) + if not has_jt: + continue + content = bytearray(cur_data[o:o + s]) for j in range(0, len(content) - 3, 4): - v = struct.unpack(' len(content): + break + if (content[j + 4] & 0xf) == 2: # STT_FUNC + tx_idx = struct.unpack_from(' buffer(size);