feat(chs): port pinned memory host-device transfer benchmark to chs

Dockerfile

@@ -2,7 +2,7 @@
 FROM python:3.12-slim
 
 # Install system dependencies
-RUN apt-get update && apt-get install -y git curl gnupg && \
+RUN apt-get update && apt-get install -y git curl gnupg numactl && \
     echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
     curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg && \
     apt-get update && apt-get install -y google-cloud-cli && \

Ironwood/configs/host_device/comprehensive_8dev_experiments.yaml

@@ -0,0 +1,10 @@
+benchmarks:
+- benchmark_name: host_device
+  num_runs: 20
+  benchmark_sweep_params:
+  - transfer_type_list: ["pinned_memory", "simple", "pipelined"]
+    num_devices_list: [8]
+    data_size_mib_list: [1, 16, 128, 256, 512, 1024, 2048, 4096, 8192, 16384]
+    input_type_list: ["numpy", "jax"]
+    dtype_list: ["float32"]
+  csv_path: "/tmp/microbenchmarks/host_device_numactl"

Ironwood/configs/host_device/comprehensive_experiments.yaml

@@ -0,0 +1,10 @@
+benchmarks:
+- benchmark_name: host_device
+  num_runs: 20
+  benchmark_sweep_params:
+  - transfer_type_list: ["pinned_memory", "simple", "pipelined"]
+    num_devices_list: [1, 2, 8]
+    data_size_mib_list: [1, 16, 128, 256, 512, 1024, 2048, 4096, 8192, 16384]
+    input_type_list: ["numpy", "jax"]
+    dtype_list: ["float32"]
+  csv_path: "/tmp/microbenchmarks/final_experiments"

Ironwood/configs/host_device/host_device.yaml

@@ -2,8 +2,11 @@ benchmarks:
 - benchmark_name: host_device
   num_runs: 20
   benchmark_sweep_params:
-  - {
-      data_size_mib_list: [1, 16, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768]
-    }
+  - transfer_type: "simple"
+    data_size_mib_list: [1, 16, 128, 256, 512, 1024, 2048, 4096, 8192, 16384]
+  - transfer_type: "pipelined"
+    data_size_mib_list: [1, 16, 128, 256, 512, 1024, 2048, 4096, 8192, 16384]
+  - transfer_type: "pinned_memory"
+    data_size_mib_list: [1, 16, 128, 256, 512, 1024, 2048, 4096, 8192, 16384]
   csv_path: "../microbenchmarks/host_device"
   trace_dir: "../microbenchmarks/host_device/trace"

Ironwood/guides/host_device_benchmark/README.md

@@ -0,0 +1,39 @@
+# Host-Device Benchmark Guide
+
+This directory contains instructions and a script to run comprehensive Host-to-Device (H2D) and Device-to-Host (D2H) microbenchmarks on Cloud TPUs.
+
+## Overview
+
+The benchmarks measure the transfer bandwidth for various configurations:
+- **Transfer Strategies**: `simple`, `pipelined`, `pinned_memory`
+- **Data Sizes**: Ranging from 1 MiB to 16,384 MiB.
+- **Input Types**: Replicating inputs using `numpy` arrays or pre-allocated `jax` arrays.
+- **Device Counts**: Scaling across 1, 2, and 8 TPU devices.
+- **NUMA Settings**: Testing the impact of `--interleave=all` with `numactl`.
+
+## How to Run
+
+A convenience script `run_comprehensive.sh` is provided. It executes two suites of tests sequentially from the `Ironwood/` directory:
+
+1. **Comprehensive Suite**: Sweeps through all transfer strategies, data sizes, input types, and device configurations (1, 2, and 8 devices) without any specific NUMA configuration.
+2. **8-Device NUMA Suite**: Executes an 8-device specific sweep while enforcing `numactl --interleave=all` at the process level to balance memory allocations across NUMA nodes, heavily improving the pipelined D2H transfer bottleneck on multi-chip architectures.
+
+### Execution
+
+Simply execute the script on your TPU VM:
+
+```bash
+bash run_comprehensive.sh
+```
+
+## Configuration Files
+
+The executed configurations are located at:
+- `Ironwood/configs/host_device/comprehensive_experiments.yaml`
+- `Ironwood/configs/host_device/comprehensive_8dev_experiments.yaml`
+
+Refer to these files to adjust the tested parameters or trace directories.
+
+## Analyzing Results
+
+Resulting logs and TSV files will be exported to the directory specified within the configurations (default typically output to the console and/or a timestamped TSV). You can use simple Pandas scripts to analyze and extract the max P50 bandwidths.

Ironwood/guides/host_device_benchmark/run_comprehensive.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+set -e
+
+# Change to the root of the Ironwood directory assuming this script is run from anywhere
+DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+IRONWOOD_DIR="$(dirname $(dirname "$DIR"))"
+cd "$IRONWOOD_DIR"
+
+echo "Running comprehensive benchmarks across 1, 2, and 8 devices..."
+python3 src/run_benchmark.py --config configs/host_device/comprehensive_experiments.yaml
+
+echo "Running 8-device benchmark with numactl interleaving..."
+numactl --interleave=all python3 src/run_benchmark.py --config configs/host_device/comprehensive_8dev_experiments.yaml
+
+echo "Benchmarks completed successfully."

Ironwood/scripts/run_host_device_benchmark.sh

@@ -4,13 +4,15 @@
 CONFIG_DIR="Ironwood/configs/host_device"
 SPECIFIC_CONFIG=""
 INTERLEAVED=false
+NUMACTL_BIND=false
 
 # Helper function for usage
 usage() {
     echo "Usage: $0 [OPTIONS]"
     echo "Options:"
     echo "  --config <path>       Path to specific config file (optional)"
     echo "  --interleaved         Run with numactl --interleave=all"
+    echo "  --numactl_bind        Run with numactl --cpunodebind=0 --membind=0"
     echo "  --help                Show this help message"
     exit 1
 }
@@ -20,6 +22,7 @@ while [[ "$#" -gt 0 ]]; do
     case $1 in
         --config) SPECIFIC_CONFIG="$2"; shift ;;
         --interleaved) INTERLEAVED=true ;;
+        --numactl_bind) NUMACTL_BIND=true ;;
         --help) usage ;;
         *) echo "Unknown parameter passed: $1"; usage ;;
     esac
@@ -33,6 +36,7 @@ echo "********************************************************"
 echo ""
 echo "Configuration:"
 echo "    Interleaved: $INTERLEAVED"
+echo "    numactl Bound: $NUMACTL_BIND"
 echo ""
 
 if [ -n "$SPECIFIC_CONFIG" ]; then
@@ -44,6 +48,11 @@ else
     shopt -u nullglob
 fi
 
+if [[ "$INTERLEAVED" = true && "$NUMACTL_BIND" = true ]]; then
+    echo "Only one of --interleaved and --numactl_bind is allowed to be set at once."
+    exit 1
+fi
+
 if [ ${#CONFIGS[@]} -eq 0 ]; then
     echo "No configuration files found!"
     exit 1
@@ -61,6 +70,14 @@ for CONFIG_FILE in "${CONFIGS[@]}"; do
             echo "Warning: numactl not found. Running without interleaving."
             $CMD
         fi
+    elif [ "$NUMACTL_BIND" = true ]; then
+        if command -v numactl &> /dev/null; then
+            echo "Running with numactl --cpunodebind=0 --membind=0"
+            numactl --cpunodebind=0 --membind=0 $CMD
+        else
+            echo "Warning: numactl not found. Running without binding."
+            $CMD
+        fi
     else
         $CMD
     fi