# Code Modification Readme — `asc26-plan-a`

**Baseline branch:** `baseline`
**Target branch:** `asc26-plan-a`
**Date:** 2026-05-19

---

## Overview

This branch delivers two major performance overhauls to the AMSS-NCKU numerical relativity codebase:

1. **TwoPunctureABE Multithreading** — OpenMP parallelization of the TwoPunctures initial-data solver, combined with a BLAS-driven spectral derivative engine, MKL/LAPACK integration, and C/C++ rewrites of hot Fortran kernel subroutines.

2. **ABE GPU Rewrite** — Complete replacement of the legacy `bssn_gpu_class` abstraction layer with direct, monolithic CUDA kernels for BSSN, Z4C, and Shell-Patch evolution, plus GPU-resident state management and CUDA-aware MPI.

**Total diff:** 84 files changed, +57,919 / −33,795 lines.

---

## Part 1 — TwoPunctureABE Multithreading

### 1.1 Spectral Derivative Engine: BLAS Matrix-Multiplication Rewrite

**Files:** `AMSS_NCKU_source/TwoPunctures.C`, `AMSS_NCKU_source/TwoPunctures.h`

The original `Derivatives_AB3` computed spectral derivatives (Chebyshev in A/B, Fourier in phi) with nested scalar loops over every grid point. The new `Derivatives_AB3_MatMul` expresses all derivatives as matrix-matrix products over pencil-shaped data slices, dispatched to Intel MKL `cblas_dgemm`.

- **Precomputed derivative matrices** — `precompute_derivative_matrices()` builds `D1_A`, `D2_A`, `D1_B`, `D2_B` (Chebyshev collocation derivative matrices) and `DF1_phi`, `DF2_phi` (Fourier derivative matrices) once at construction time.
- **Pencil-based GEMM** — data is gathered into 2D arrays where one dimension is the spectral direction and the other enumerates all remaining degrees of freedom (variables × orthogonal grid indices). Each derivative direction becomes a single `cblas_dgemm` call. The pure derivatives (d/dA, d/dB, d/dphi) and all mixed derivatives (d²/dAdB, d²/dAdphi, d²/dBdphi) are computed this way.
- **`build_cheb_deriv_matrices` / `build_fourier_deriv_matrices`** — construct the standard Chebyshev and Fourier collocation derivative matrices.

### 1.2 OpenMP Parallelization of TwoPunctures

**Files:** `AMSS_NCKU_source/TwoPunctures.C`, `AMSS_NCKU_source/TwoPunctures.h`

Three critical regions are parallelized:

| Region | Directive | Strategy |
|--------|-----------|----------|
| `F_of_v` residual evaluation | `#pragma omp parallel for collapse(3) schedule(dynamic,1)` | Each (i,j,k) thread stack-allocates its own `l_U` (derivs struct) and `l_values[]` to eliminate heap contention and data races |
| `relax_omp` line relaxation | `#pragma omp parallel for schedule(static)` over k-slices | Alternating be/al sweeps, each thread uses pre-allocated per-thread Thomas-algorithm workspace (`ws_*_be[tid]`, `ws_*_al[tid]`) |
| `LineRelax_be_omp` / `LineRelax_al_omp` | Called from `relax_omp` with explicit `tid` | Thread-safe tridiagonal solves using the thread's private scratch arrays |

**Per-thread workspace** — `allocate_workspace()` allocates independent Thomas-algorithm buffers (`diag`, `e`, `f`, `b`, `x`, `l`, `u`, `d`, `y`) for each OpenMP thread in both be and al directions, avoiding lock contention in the inner Newton iteration.

### 1.3 MKL BLAS / LAPACK Integration

**Files:** `AMSS_NCKU_source/TwoPunctures.C`, `AMSS_NCKU_source/gaussj.C`

| Function | Old | New | Benefit |
|----------|-----|-----|---------|
| `norm2` | scalar `sqrt(sum(v[i]²))` loop | `cblas_dnrm2` | BLAS Level 1, SIMD-optimized |
| `scalarproduct` | scalar `sum(v[i]*w[i])` loop | `cblas_ddot` | BLAS Level 1, SIMD-optimized |
| `gaussj` | hand-written Gauss-Jordan elimination (~100 lines) | `LAPACKE_dgesv` + `LAPACKE_dgetrf` + `LAPACKE_dgetri` | LAPACK LU with partial pivoting, asymptotically faster for the `n~50` matrix sizes used in spectral elliptic solves |

### 1.4 C/C++ Rewrite of Hot Fortran Kernels

**Files (new):**
- `AMSS_NCKU_source/fderivs_c.C` (167 lines) — first derivatives, 2nd/4th order
- `AMSS_NCKU_source/fdderivs_c.C` (332 lines) — second derivatives, 2nd/4th order
- `AMSS_NCKU_source/kodiss_c.C` (117 lines) — Kreiss-Oliger dissipation
- `AMSS_NCKU_source/lopsided_c.C` (255 lines) — lopsided advection
- `AMSS_NCKU_source/lopsided_kodis_c.C` (248 lines) — fused advection + dissipation
- `AMSS_NCKU_source/rungekutta4_rout_c.C` (212 lines) — RK4 time-stepper
- `AMSS_NCKU_source/bssn_rhs_c.C` (1,287 lines) — full BSSN RHS kernel
- `AMSS_NCKU_source/z4c_rhs_c.C` (725 lines) — full Z4C RHS kernel

Every C rewrite follows a consistent optimization pattern:
- **64-byte aligned allocation** (`aligned_alloc(64, ...)`) for AVX-512 compatibility.
- **Static buffer caching** — scratch arrays (e.g., the padded `fh` ghost-zone buffer) persist across calls via a `static` pointer + capacity check, avoiding repeated `malloc`/`free`.
- **Two-pass strategy** — 2nd-order finite differences are computed on the full domain first, then the interior sub-volume is overwritten with 4th-order stencils. This eliminates the per-point `if/elseif` branching of the original Fortran.
- **Non-overlapping shell pass** — in `fdderivs_c.C`, the 2nd-order pass skips points that will be overwritten by the 4th-order pass, avoiding redundant computation.

### 1.5 Fortran Kernel Fusion: lopsided_kodis

**File:** `AMSS_NCKU_source/lopsidediff.f90`

A new `lopsided_kodis` subroutine fuses the advection (lopsided) and Kreiss-Oliger dissipation (kodis) operators into a single pass over the grid. Both operators previously called `symmetry_bd` independently to fill ghost zones — the fused version calls it once and shares the padded `fh` array, halving ghost-zone fill overhead for this hot path.

### 1.6 Build System for TwoPunctures

**Files:** `AMSS_NCKU_source/makefile`, `AMSS_NCKU_source/makefile.inc`

- **`TP_OPTFLAGS`** — TwoPunctures and TwoPunctureABE are compiled with a dedicated, more aggressive optimization flag set (`-O3 -march=znver5 -fp-model fast=2 -fma -ipo`) separate from the main code.
- **`USE_CXX_KERNELS`** — selects between the C rewrites and the original Fortran kernels (`bssn_rhs.f90`, etc.) for the CPU path.
- **`USE_CXX_RK4`** — independently selects between the C and Fortran RK4 stepper.
- **Intel oneTBB allocator** (`libtbbmalloc.so`) — replaces the system `malloc` with a scalable thread-safe allocator, critical for multi-threaded TwoPunctures performance.
- **PGO support** — `PGO_MODE=opt|instrument` for profile-guided optimization (currently disabled after testing showed negative benefit).
- **Toolchains** — Intel oneAPI (`TOOLCHAIN=intel`, default) and NVIDIA HPC SDK (`TOOLCHAIN=nvhpc`).

---

## Part 2 — ABE GPU Rewrite

### 2.1 Architecture: From Class Wrapper to Direct CUDA Kernels

The old GPU path (`baseline`) was organized as:

```
bssn_gpu_class.C/h    — C++ class managing GPU state and kernel launches
bssn_step_gpu.C       — RK4 stepper with per-substep GPU/CPU synchronisation
bssn_gpu.cu           — CUDA kernel implementations called through the class
```

The new GPU path (`asc26-plan-a`) replaces all of the above with:

```
bssn_rhs_cuda.cu/h    — 10,381-line monolithic CUDA BSSN RHS kernel
z4c_rhs_cuda.cu/h     —  7,909-line monolithic CUDA Z4C RHS kernel
fd_cuda_helpers.cuh   —    412-line shared finite-difference device functions
bssn_gpu_rhs_ss.cu    —  (retained, lightly modified) Shell-Patch GPU RHS
```

**Key architectural differences:**
- The old `bssn_gpu_class` managed GPU memory through a C++ class with explicit allocate/free/sync methods scattered across the time-stepping logic. The new kernels operate directly on raw device pointers with a clear resident/transient memory model.
- The old code launched many small kernels (one per derivative or algebraic term). The new code is a **single monolithic kernel per formulation** — all 24 BSSN evolution variables are computed in one launch with on-the-fly finite differences, eliminating kernel-launch latency and intermediate global-memory round-trips.
- The old `bssn_step_gpu.C` performed per-substep GPU→CPU downloads for boundary conditions and analysis. The new model supports **GPU-resident state** — variables stay on device across timesteps unless explicitly requested.

### 2.2 GPU-Resident State Model

A central theme across ~20 commits is the "resident-sync" optimization:

| Commit | What it does |
|--------|-------------|
| `22c1e71` | Optimize BSSN CUDA resident state and CUDA-aware MPI |
| `090d865` | Optimize BSSN CUDA state transfers |
| `68eab03` | Add opt-in BSSN CUDA resident AMR path |
| `1ee229a` | Add keyed BSSN CUDA resident banks |
| `18e9c9c` | Optimize BSSN CUDA resident AMR prolong |
| `8486532` | Add resident BSSN GPU point interpolation |
| `b1974ef` | Stabilize device AMR restrict across regrid |
| `ae64a22` | Complete BSSN-EScalar CUDA resident transfers |
| `83afaf1` | Skip zero EM resident downloads |
| `35b6cef` | Broaden cached CUDA sync paths |

The resident model works as follows:
- BSSN grid functions are allocated once on the GPU and persist across timesteps.
- Inter-processor ghost-zone exchanges use **CUDA-aware MPI** — MPI directly reads/writes device memory without staging through host buffers.
- AMR prolongation and restriction operate directly on device memory.
- Boundary conditions and analysis routines download only the specific slices/points they need, not the full grid.
- When EM fields are zero (pure-gravity runs), EM downloads are skipped entirely.

### 2.3 Z4C and Shell-Patch GPU Acceleration

**Files:** `AMSS_NCKU_source/z4c_rhs_cuda.cu`, `AMSS_NCKU_source/bssn_gpu_rhs_ss.cu`

- The Z4C constraint-damped formulation gets its own 7,909-line monolithic CUDA kernel (`z4c_rhs_cuda.cu`), matching the BSSN kernel's architecture.
- **Shell-Patch GPU acceleration** — the spherical shell boundary patches now compute on GPU with dedicated kernels in `bssn_gpu_rhs_ss.cu`.
- Z4C + Shell-Patch can coexist on GPU (Phase 3 commits).
- A CPU-side wrapper (`z4c_rhs_c.C`) handles the trKd + TZ_rhs contribution that remains on CPU, minimizing GPU/CPU traffic.

### 2.4 Finite-Difference Order Flexibility

**File:** `AMSS_NCKU_source/fd_cuda_helpers.cuh`

Shared device functions for finite-difference stencils support **2nd, 4th, 6th, and 8th order** at compile time via preprocessor switches. This enables:
- Per-run selection of convergence order without recompilation of the full kernel.
- 8th-order AMR transfers (`1064a68`) for BSSN-EM.
- 6th-order optimized AMR stencils (`0076b3c`).

### 2.5 GPU Diagnostics and Quality Assurance

**File:** `AMSS_NCKU_GPUCheck.py` (559 lines, new)

A Python-based GPU correctness verification tool that compares GPU and CPU evolution outputs. The GPU build pipeline includes optional kernel profiling switches (`7683459`) for performance debugging.

**GPU-specific bug fixes:**
- `f226498` — Fix CUDA AMR symmetry drift (incorrect ghost-zone handling under symmetry boundary conditions)
- `2317e4a` — Fix BSSN GPU resident AMR sync default
- `fea2dcc` — Fix BSSN-EM runtime crash
- `dd0e20d` — Fix BSSN-EScalar CUDA boundary and scalar KO
- `5eb4994` — Fix AHF crash under CUDA resident-sync mode

### 2.6 Build Integration

**Makefile switches:**
- `USE_CUDA_BSSN=0/1` — route BSSN RHS through GPU or CPU
- `USE_CUDA_Z4C=0/1` — route Z4C RHS through GPU or CPU
- `CUDA_ARCH=sm_80` — target NVIDIA Ampere (A100)
- `NVHPC_ROOT` — path to NVIDIA HPC SDK for the `nvcc` compiler wrapper
- CUDA compilation flags: `-O3 --ptxas-options=-v -arch=$(CUDA_ARCH)`

---

## Part 3 — Shared Infrastructure

### 3.1 Interp_Points Load-Balance Profiler

**Files:** `AMSS_NCKU_source/interp_lb_profile.C`, `interp_lb_profile.h`, `interp_lb_profile_data.h`, `generate_interp_lb_header.py`

A two-pass instrumentation system for load-balancing the `Interp_Points` parallel interpolation routine:
- **Pass 1** (`INTERP_LB_MODE=profile`): instrument each MPI rank's interpolation calls with timing, write a binary profile.
- **Pass 2** (`INTERP_LB_MODE=optimize`): read the profile and rebalance work across MPI ranks.

### 3.2 Helper Headers

**Files:** `AMSS_NCKU_source/tool.h` (33 lines), `AMSS_NCKU_source/share_func.h` (246 lines)

- `tool.h` — shared indexing macros (`idx_ex`, `idx_fh_F_ord2`) and the `symmetry_bd` declaration used by all C kernel rewrites.
- `share_func.h` — common utility functions shared across the C++ source files.

### 3.3 Plot-Only Restart Script

**File:** `parallel_plot_helper.py` (29 lines)

A lightweight restart script that skips recomputation when plotting was interrupted — reads existing checkpoint data and replots without re-running the simulation.

---

## Performance Summary

| Component | Optimization | Expected Impact |
|-----------|-------------|-----------------|
| TwoPunctures `Derivatives_AB3` | Scalar loops → MKL GEMM | 5-20× speedup for spectral derivative computation |
| TwoPunctures `F_of_v` | OpenMP collapse(3) + stack-local variables | Near-linear scaling with core count for residual evaluation |
| TwoPunctures `gaussj` | Hand-written Gauss-Jordan → LAPACK LU | 2-5× speedup for N~50 matrix inversion |
| BSSN RHS (GPU) | Many small kernels → one monolithic kernel | Eliminates kernel-launch overhead; 2-5× GPU throughput improvement |
| GPU state transfers | Per-step download → resident model | Eliminates ~80% of GPU↔CPU PCIe traffic |
| `lopsided_kodis` fusion | Two `symmetry_bd` calls → one shared call | ~30% reduction in ghost-zone fill cost for this operator pair |
| Memory allocator | System malloc → Intel TBB malloc | Significant reduction in malloc contention under OpenMP |
| C kernel rewrites | Fortran → C with aligned alloc + static buffers | Enables Intel compiler IPO across C/C++/Fortran boundaries; better SIMD codegen |

---