Fix critical correctness bugs in SASS rewriter and interceptor

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.<kernel_name>`)
+   - 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`

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<char> buffer(size);

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<std::mutex> 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());

requirements.txt

@@ -1,2 +0,0 @@
-lief
-pyelftools