Merge branch 'grigorik/development' of https://github.com/EngineeringSoftware/pkdb into grigorik/development

Author: IvanGrigorik

benchmarks/run_boltzmann_debuggers.py

@@ -7,6 +7,7 @@
 """
 
 import argparse
+import re
 import subprocess
 import sys
 import time
@@ -21,10 +22,21 @@
     system_and_suffix,
     write_checkpoint,
 )
-from run_boltzmann import parse_total_time_s
 
 DEFAULT_SPACES = ["DebugOpenMP", "DebugCuda"]
 
+
+def parse_total_time_s(stdout: str) -> float | None:
+    """Last "total time = <float>" line (warmup + RESET produce an earlier one)."""
+    matches = list(
+        re.finditer(
+            r"^total time = ([0-9]+(?:\.[0-9]+)?(?:[eE][+-]?[0-9]+)?)\s*$",
+            stdout,
+            re.MULTILINE,
+        )
+    )
+    return float(matches[-1].group(1)) if matches else None
+
 # Step counts passed to boltzmann/main.py via -s. Every run here pays Python-level tracing on top
 # of the kernels, so the wall clock is far higher than run_boltzmann.py's own timing.
 STEP_COUNTS = [

benchmarks/run_ewald_debuggers.py

@@ -7,6 +7,7 @@
 """
 
 import argparse
+import re
 import subprocess
 import sys
 import time
@@ -21,10 +22,61 @@
     system_and_suffix,
     write_checkpoint,
 )
-from run_ewald import parse_walltimes
 
 DEFAULT_SPACES = ["DebugOpenMP", "DebugCuda"]
 
+KERNEL_KEYS = ["P2P", "P2G", "FFT", "CNV", "IFFT", "G2P"]
+
+
+def parse_walltimes(stdout: str) -> dict[str, float | None] | None:
+    """Last "Ewald Walltimes" section: TOTAL plus each KERNEL_KEYS component's "tot" value."""
+    lines = [re.sub(r"\x1b\[[0-9;]*m", "", line).strip() for line in stdout.splitlines()]
+
+    start = None
+    for i in range(len(lines) - 1, -1, -1):
+        if "Ewald Walltimes" in lines[i]:
+            start = i
+            break
+    if start is None:
+        return None
+
+    end = len(lines)
+    for i in range(start + 1, len(lines)):
+        if "Ewald Cost Model" in lines[i]:
+            end = i
+            break
+
+    result: dict[str, float | None] = {
+        "TOTAL": None,
+        **{key: None for key in KERNEL_KEYS},
+    }
+    component = None
+    for line in lines[start + 1 : end]:
+        header = re.match(r"^(P2P|P2G|FFT|CNV|IFFT|G2P)\s*:", line)
+        if header:
+            component = header.group(1)
+            continue
+
+        m_total = re.match(r"^TOTAL:\s*([0-9eE+\-.]+)", line)
+        if m_total:
+            try:
+                result["TOTAL"] = float(m_total.group(1))
+            except ValueError:
+                pass
+            component = None
+            continue
+
+        if component is not None:
+            m_tot = re.search(r"['\"]?tot['\"]?\s*:\s*([0-9eE+\-.]+)", line)
+            if m_tot:
+                try:
+                    result[component] = float(m_tot.group(1))
+                except ValueError:
+                    pass
+                component = None
+
+    return result
+
 # Atom counts passed to Ewald via --atoms. Every run here pays Python-level tracing on top of the
 # kernels, so the wall clock is far higher than run_ewald.py's own timing.
 ATOM_SIZES = [

benchmarks/run_examinimd_debuggers.py

@@ -21,10 +21,32 @@
     system_and_suffix,
     write_checkpoint,
 )
-from run_examinimd import parse_total_seconds
 
 DEFAULT_SPACES = ["DebugOpenMP", "DebugCuda"]
 
+
+def parse_total_seconds(stdout: str) -> float | None:
+    """Sum the Total(s) column of the kernel profile table ExaMiniMD prints."""
+    total = 0.0
+    found = False
+    in_table = False
+    for line in stdout.splitlines():
+        stripped = line.strip()
+        if not in_table:
+            in_table = stripped.startswith("Kernel") and "Total(s)" in stripped and "Avg(ms)" in stripped
+            continue
+        if not stripped:
+            break
+        parts = stripped.split()
+        # Data rows look like: <name> <Count> <Total(s)> <Avg(ms)> <Percent%>
+        if len(parts) >= 5 and parts[-1].endswith("%") and parts[-4].isdigit():
+            try:
+                total += float(parts[-3])
+                found = True
+            except ValueError:
+                pass
+    return total if found else None
+
 # Atom counts passed to ExaMiniMD via --atoms. Small next to run_examinimd.py's list: every
 # run here pays Python-level tracing on top of the kernels, so the wall clock is far higher.
 ATOM_SIZES = [

sbatch-scripts/amd_boltzmann.sh sbatch-scripts/AMD/amd_boltzmann.shsbatch-scripts/amd_boltzmann.sh renamed to sbatch-scripts/AMD/amd_boltzmann.sh

@@ -0,0 +1,34 @@
+# TACC sbatch scripts and apptainer file
+
+First, build container with apptainer:
+
+```bash
+module load tacc-apptainer
+apptainer build pkdb_container.sif pkdb_container.def
+```
+
+Do `pkdb` installation of `PyKokkos` and `pkdb`:
+
+```bash
+apptainer shell --nv --fakeroot conda_git.sif
+Apptainer> cd [pkdb_dir]/pykokkos
+Apptainer> conda create -n pkdb python=3.13 -y
+Apptainer> conda env update -n pkdb -f base/environment.yml
+Apptainer> conda activate pkdb
+Apptainer> python install_base.py install --verbose -- -DENABLE_LAYOUTS=ON -DENABLE_MEMORY_TRAITS=OFF -DENABLE_VIEW_RANKS=4 -DENABLE_THREADS=OFF -DENABLE_OPENMP=ON -DENABLE_CUDA=ON
+Apptainer> conda install -c conda-forge pybind11 patchelf -y
+Apptainer> pip install -e .
+Apptainer> cd [pkdb_dir]/pkdb
+Apptainer> pip install -r requirements.txt
+Apptainer> pip install -e .
+```
+
+Now, when our environment is ready we can run sbatch scripts. Execute scripts
+and pass apptainer directory (`./`, if you are doing it from current directory)
+and `pkdb` directory (`../../` if you are doing it from current directory). It's
+not required to pass absolute path.  
+For example:
+
+```bash
+sbatch tacc_examinimd.sh ./ ../../
+```

sbatch-scripts/amd_ewald.sh sbatch-scripts/AMD/amd_ewald.shsbatch-scripts/amd_ewald.sh renamed to sbatch-scripts/AMD/amd_ewald.sh

@@ -0,0 +1,44 @@
+Bootstrap: docker
+From: ubuntu:22.04
+
+%environment
+    export PATH="/opt/conda/bin:$PATH"
+
+%post
+    set -e
+
+    # Set apt downloads to root user env
+    echo 'APT::Sandbox::User "root";' > /etc/apt/apt.conf.d/99no-sandbox
+    apt-get update
+
+    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+        wget \
+        bzip2 \
+        ca-certificates \
+        git \
+        gdb \
+        build-essential
+
+    # Install miniconda
+    ARCH=$(uname -m)
+    case "${ARCH}" in
+        x86_64)  MINICONDA_ARCH="x86_64" ;;
+        aarch64|arm64) MINICONDA_ARCH="aarch64" ;;
+        *) echo "Unsupported architecture: ${ARCH}" >&2; exit 1 ;;
+    esac
+    MINICONDA_INSTALLER=Miniconda3-latest-Linux-${MINICONDA_ARCH}.sh
+    wget -q https://repo.anaconda.com/miniconda/${MINICONDA_INSTALLER} -O /tmp/${MINICONDA_INSTALLER}
+    bash /tmp/${MINICONDA_INSTALLER} -b -p /opt/conda
+    rm /tmp/${MINICONDA_INSTALLER}
+    export PATH="/opt/conda/bin:$PATH"
+
+    # conda TOS and prerequisites
+    conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main
+    conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r
+    conda install -y -c nvidia cuda-nvcc cuda-gdb
+    conda clean -afy
+
+    # Pass existing tool references to container
+    ln -sf /opt/conda/bin/nvcc /usr/local/bin/nvcc
+    ln -sf /opt/conda/bin/conda /usr/local/bin/conda
+    ln -sf /opt/conda/bin/cuda-gdb /usr/local/bin/cuda-gdb

sbatch-scripts/amd_examinimd.sh sbatch-scripts/AMD/amd_examinimd.shsbatch-scripts/amd_examinimd.sh renamed to sbatch-scripts/AMD/amd_examinimd.sh

@@ -0,0 +1,38 @@
+#!/bin/bash
+
+#SBATCH -A <Your project>
+#SBATCH -J pkdb_boltzmann_evaluation
+#SBATCH -o boltzmann_job.out
+#SBATCH -p gh
+#SBATCH -N 1
+#SBATCH -n 1
+#SBATCH -t 06:00:00
+
+# ------------------
+# Main code section
+# ------------------
+
+if [ -z "$1" ] || [ -z "$2" ]; then
+    printf "Script requires two paths: \n\t(1) apptainer .sif image path\n\t(2) pkdb root directory\n"
+    exit 1
+fi
+
+APPTAINER_PATH="$(realpath "$1")"
+PKDB_ROOT="$(realpath "$2")"
+
+module load tacc-apptainer
+apptainer exec --nv --fakeroot "$APPTAINER_PATH" bash -c '
+    set -e
+    export CUDACXX=/opt/conda/bin/nvcc
+    export CXX=/usr/bin/g++
+    export CC=/usr/bin/gcc
+    eval "$(conda shell.bash hook)"
+    conda activate pkdb
+
+    export OMP_NUM_THREADS=72
+    BENCH_ROOT="$(realpath "$1/benchmarks")"
+    cd "$BENCH_ROOT"
+    python run_boltzmann_debuggers.py --spaces "DebugCuda, DebugOpenMP"
+' _ "$PKDB_ROOT"
+
+echo "Done."

sbatch-scripts/TACC/README.md

@@ -0,0 +1,38 @@
+#!/bin/bash
+
+#SBATCH -A <Your project>
+#SBATCH -J pkdb_ewald_evaluation
+#SBATCH -o ewald_job.out
+#SBATCH -p gh
+#SBATCH -N 1
+#SBATCH -n 1
+#SBATCH -t 06:00:00
+
+# ------------------
+# Main code section
+# ------------------
+
+if [ -z "$1" ] || [ -z "$2" ]; then
+    printf "Script requires two paths: \n\t(1) apptainer .sif image path\n\t(2) pkdb root directory\n"
+    exit 1
+fi
+
+APPTAINER_PATH="$(realpath "$1")"
+PKDB_ROOT="$(realpath "$2")"
+
+module load tacc-apptainer
+apptainer exec --nv --fakeroot "$APPTAINER_PATH" bash -c '
+    set -e
+    export CUDACXX=/opt/conda/bin/nvcc
+    export CXX=/usr/bin/g++
+    export CC=/usr/bin/gcc
+    eval "$(conda shell.bash hook)"
+    conda activate pkdb
+
+    export OMP_NUM_THREADS=72
+    BENCH_ROOT="$(realpath "$1/benchmarks")"
+    cd "$BENCH_ROOT"
+    python run_ewald_debuggers.py --spaces "DebugCuda, DebugOpenMP"
+' _ "$PKDB_ROOT"
+
+echo "Done."