[Automation] Add BMM into automation script

Author: hylin2002

Ironwood/guides/automation/aggregator.py

@@ -29,6 +29,13 @@
         "tflops_per_sec_per_device_avg", "tflops_per_sec_per_device_min",
         "tflops_per_sec_per_device_max",
     ],
+    "bmm": [
+        "b", "m", "n", "k", "dtype", "step_time_ms_num_runs",
+        "tflops_per_sec_per_device_p50", "tflops_per_sec_per_device_p90",
+        "tflops_per_sec_per_device_p95", "tflops_per_sec_per_device_p99",
+        "tflops_per_sec_per_device_avg", "tflops_per_sec_per_device_min",
+        "tflops_per_sec_per_device_max",
+    ],
 }
 
 def download_from_gcs(bucket_path: str, local_dir: str):
@@ -86,15 +93,27 @@ def aggregate_gemm(directories: list[str], picked_columns: list[str]) -> pd.Data
             aggregated_df = pd.concat([aggregated_df, df[picked_columns].rename(columns={"step_time_ms_num_runs": "num_runs"})], ignore_index=True)
     return aggregated_df
 
+def aggregate_bmm(directories: list[str], picked_columns: list[str]) -> pd.DataFrame:
+    if len(directories) == 0:
+        return None
+    aggregated_df = pd.DataFrame()
+    for directory in directories:
+        files = glob.glob(f"{directory}/*.tsv")
+        for file in files:
+            df = pd.read_csv(file, sep='\t')
+            aggregated_df = pd.concat([aggregated_df, df[picked_columns].rename(columns={"step_time_ms_num_runs": "num_runs"})], ignore_index=True)
+    return aggregated_df
+
 aggregate_function = {
     "collectives": aggregate_collectives,
     "hbm": aggregate_hbm,
     "host_device": aggregate_host_device,
     "gemm": aggregate_gemm,
+    "bmm": aggregate_bmm,
 }
 
 def aggregate_results(bucket_path: str, local_dir: str):
-    categories = ["collectives", "hbm", "host_device", "gemm"]
+    categories = ["collectives", "hbm", "host_device", "gemm", "bmm"]
     directories = {}
     results = {}
     for category in categories:

Ironwood/guides/automation/automation_launch.sh

@@ -15,6 +15,7 @@ yaml_names=(
     "tpu7x-2x2x1-hbm.yaml"
     "tpu7x-2x2x1-host_device.yaml"
     "tpu7x-2x2x1-gemm.yaml"
+    "tpu7x-2x2x1-bmm.yaml"
     "tpu7x-2x2x1-collectives.yaml"
     "tpu7x-2x2x2-collectives.yaml"
     "tpu7x-2x2x4-collectives.yaml"

Ironwood/guides/automation/tpu7x-2x2x1-bmm.yaml

@@ -0,0 +1,61 @@
+apiVersion: v1
+kind: Service
+metadata:
+  name: headless-svc-${JOB_NAME}
+spec:
+  clusterIP: None
+  selector:
+    job-name: ${JOB_NAME}
+---
+apiVersion: batch/v1
+kind: Job
+metadata:
+  name: ${JOB_NAME}
+  labels:
+    kueue.x-k8s.io/queue-name: user-queue-2x2x1
+spec:
+  completionMode: Indexed
+  suspend: true
+  parallelism: 1
+  completions: 1
+  backoffLimit: 0
+  template:
+    spec:
+      subdomain: headless-svc-${JOB_NAME}
+      serviceAccountName: ${GCS_SA_NAME}
+      restartPolicy: Never
+      nodeSelector:
+        cloud.google.com/gke-tpu-accelerator: tpu7x
+        cloud.google.com/gke-tpu-topology: 2x2x1
+      containers:
+      - name: jax-tpu
+        image: python:3.12
+        securityContext:
+          privileged: false
+        env:
+        - name: JAX_PLATFORMS
+          value: "tpu,cpu"
+        - name: TPU_VMODULE
+          value: "singleton_tpu_system_manager=10,tpu_version_flag=10,device_util=10,device_scanner=10,mesh_builder=10,master=10"
+        - name: XLA_IR_DEBUG
+          value: "1"
+        - name: XLA_HLO_DEBUG
+          value: "1"
+        command:
+        - bash
+        - -c
+        - |
+          set -ex
+
+          git clone https://github.com/AI-Hypercomputer/accelerator-microbenchmarks.git
+          cd accelerator-microbenchmarks
+          git checkout tpu7x-auto
+          pip install -r requirements.txt
+
+          GCS_BUCKET_DIR=${GCS_PATH}
+          python Ironwood/src/run_benchmark.py --config="Ironwood/configs/bmm/single_device_bmm.yaml" --gcs-bucket-csv-dir=${GCS_BUCKET_DIR}
+        resources:
+          requests:
+            google.com/tpu: 4
+          limits:
+            google.com/tpu: 4

Ironwood/src/benchmark_bmm.py

@@ -131,4 +131,6 @@ def single_device_bmm_calculate_metrics(
         total_flops,
         total_flops_all_devices,
         PEAK_FLOPS_PER_DEVICE,
+        dtype=dtype.dtype.name,
+        b=b,
     )

Ironwood/src/benchmark_utils.py

@@ -1111,6 +1111,7 @@ def unified_flops_metrics(
     total_flops_all_devices: int,
     peak_TFLOPS_per_device: float,
     dtype: str = None,
+    b: int = None,
 ) -> Dict[str, Any]:
     """Calculates the metrics for the naive matmul benchmark."""
     # Build dictionary of all the parameters in the function