Add smart_chunking flag and config

Author: linamy85

Ironwood/configs/host_device/host_device_single_chip_smart_chunking.yaml

@@ -0,0 +1,9 @@
+benchmarks:
+- benchmark_name: host_device
+  num_runs: 20
+  benchmark_sweep_params:
+  # Single Chip (1 Chip, 2 Devices)
+  - {mesh_shape: "1x2", data_size_mb_list: [1, 16, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768], smart_chunking: true}
+
+  csv_path: "../microbenchmarks/host_device/single_chip_smart"
+  trace_dir: "../microbenchmarks/host_device/single_chip_smart/trace"

Ironwood/guides/host_device/host_device.md

@@ -16,6 +16,11 @@ To run the microbenchmarks, apply the following Kubernetes configuration:
 kubectl apply -f tpu7x-host-device-benchmark.yaml
 ```
 
+To run with "Smart Chunking" enabled (advanced optimizations):
+```bash
+kubectl apply -f tpu7x-host-device-benchmark-smart-chunking.yaml
+```
+
 To extract the log of the microbenchmark, use `kubectl logs`:
 ```bash
 kubectl logs tpu7x-host-device-benchmark

Ironwood/guides/host_device/tpu7x-host-device-benchmark-smart-chunking.yaml

@@ -0,0 +1,34 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  name: tpu7x-host-device-benchmark-smart-chunking
+spec:
+  restartPolicy: Never
+  nodeSelector:
+    cloud.google.com/gke-tpu-accelerator: tpu7x
+    cloud.google.com/gke-tpu-topology: 2x2x1
+  containers:
+  - name: tpu-job
+    image: python:3.12
+    ports:
+    - containerPort: 8431
+    securityContext:
+      privileged: false
+    command:
+    - bash
+    - -c
+    - |
+      set -ex
+
+      git clone https://github.com/AI-Hypercomputer/accelerator-microbenchmarks.git
+      cd accelerator-microbenchmarks
+      pip install -r requirements.txt
+
+      echo "Running Smart Chunking Benchmark..."
+      bash Ironwood/scripts/run_host_device_benchmark.sh --smart-chunking
+
+    resources:
+      requests:
+        google.com/tpu: 4
+      limits:
+        google.com/tpu: 4

Ironwood/scripts/run_host_device_benchmark.sh

@@ -4,13 +4,15 @@
 CONFIG_DIR="Ironwood/configs/host_device"
 SPECIFIC_CONFIG=""
 INTERLEAVED=false
+SMART_CHUNKING=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 "  --smart-chunking      Run with smart chunking enabled (uses *_smart_chunking.yaml configs)"
     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 ;;
+        --smart-chunking) SMART_CHUNKING=true ;;
         --help) usage ;;
         *) echo "Unknown parameter passed: $1"; usage ;;
     esac
@@ -33,14 +36,26 @@ echo "********************************************************"
 echo ""
 echo "Configuration:"
 echo "    Interleaved: $INTERLEAVED"
+echo "    Smart Chunking: $SMART_CHUNKING"
 echo ""
 
 if [ -n "$SPECIFIC_CONFIG" ]; then
     CONFIGS=("$SPECIFIC_CONFIG")
 else
     # Use nullglob to handle case where no files match (though unlikely here)
     shopt -s nullglob
-    CONFIGS=("$CONFIG_DIR"/*.yaml)
+    if [ "$SMART_CHUNKING" = true ]; then
+        CONFIGS=("$CONFIG_DIR"/*smart_chunking.yaml)
+    else
+        # Exclude smart_chunking configs by default
+        ALL_CONFIGS=("$CONFIG_DIR"/*.yaml)
+        CONFIGS=()
+        for cfg in "${ALL_CONFIGS[@]}"; do
+            if [[ "$cfg" != *"smart_chunking.yaml" ]]; then
+                CONFIGS+=("$cfg")
+            fi
+        done
+    fi
     shopt -u nullglob
 fi