From 2264119e9d7169c725588c5533b8fc1ce14632ef Mon Sep 17 00:00:00 2001
From: "lcy.seso" <lcy.seso@gmail.com>
Date: Thu, 19 Sep 2024 05:03:00 +0000
Subject: [PATCH] Simplify ncu scripts.

---
 artifacts/run_all_ncu_cutlass.sh | 5 +++--
 artifacts/run_all_ncu_flash2.sh  | 3 ++-
 artifacts/run_all_ncu_ft.sh      | 6 ++++--
 artifacts/run_all_ncu_pt.sh      | 3 ++-
 artifacts/run_all_ncu_triton.sh  | 6 ++++--
 5 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/artifacts/run_all_ncu_cutlass.sh b/artifacts/run_all_ncu_cutlass.sh
index e7644b680..36372d1a1 100755
--- a/artifacts/run_all_ncu_cutlass.sh
+++ b/artifacts/run_all_ncu_cutlass.sh
@@ -17,10 +17,11 @@ fi
 # 1. ncu test the mha benchmark
 echo "NCU profiling mha benchmark"
 $ncu_dir/ncu --section "MemoryWorkloadAnalysis" \
-    --csv --set full cutlass_attn --nheads="$nheads" \
+    --metrics "dram__bytes.sum,lts__t_bytes.sum,l1tex__t_bytes.sum" \
+    --csv cutlass_attn --nheads="$nheads" \
         --batch_size=$batch_size \
         --head_size=$head_size \
         --head_size_v=$head_size \
         --seq_length="$length" \
         --seq_length_kv=$length \
-        --causal=false  > $log_dir/cutlass_attention_ncu.csv
\ No newline at end of file
+        --causal=false  > $log_dir/cutlass_attention_ncu.csv
diff --git a/artifacts/run_all_ncu_flash2.sh b/artifacts/run_all_ncu_flash2.sh
index 913dcf599..9e2c6488b 100755
--- a/artifacts/run_all_ncu_flash2.sh
+++ b/artifacts/run_all_ncu_flash2.sh
@@ -9,5 +9,6 @@ mha_dir="$benchmark_dir/multi-head_attention/baseline"
 # 1. ncu test the mha benchmark
 echo "NCU profiling mha benchmark"
 $ncu_dir/ncu --section "MemoryWorkloadAnalysis" \
-    --csv --set full python3 $mha_dir/test_pt_model.py > $log_dir/flash2_attention_ncu.csv
+    --metrics "dram__bytes.sum,lts__t_bytes.sum,l1tex__t_bytes.sum" \
+    --csv python3 $mha_dir/test_pt_model.py > $log_dir/flash2_attention_ncu.csv
 
diff --git a/artifacts/run_all_ncu_ft.sh b/artifacts/run_all_ncu_ft.sh
index fad045c60..82001b344 100755
--- a/artifacts/run_all_ncu_ft.sh
+++ b/artifacts/run_all_ncu_ft.sh
@@ -10,10 +10,12 @@ echo "NCU profiling mha benchmark"
 mha_dir="$benchmark_dir/multi-head_attention/fractaltensor/build"
 mha_exe="$mha_dir/main"
 $ncu_dir/ncu --section "MemoryWorkloadAnalysis" \
-    --csv --set full $mha_exe > $log_dir/ft_attention_ncu.csv
+    --metrics "dram__bytes.sum,lts__t_bytes.sum,l1tex__t_bytes.sum" \
+    --csv $mha_exe > $log_dir/ft_attention_ncu.csv
 
 # 2. ncu test the bigbird benchmark
 bigbird_dir="$benchmark_dir/blocked_sparse_attention/fractaltensor/build"
 echo "NCU profiling BigBird benchmark"
 $ncu_dir/ncu --section "MemoryWorkloadAnalysis" \
-    --csv --set full $bigbird_dir/bigbird > $log_dir/ft_bigbird_ncu.csv
\ No newline at end of file
+    --metrics "dram__bytes.sum,lts__t_bytes.sum,l1tex__t_bytes.sum" \
+    --csv $bigbird_dir/bigbird > $log_dir/ft_bigbird_ncu.csv
\ No newline at end of file
diff --git a/artifacts/run_all_ncu_pt.sh b/artifacts/run_all_ncu_pt.sh
index c2a6d93db..951197d4e 100755
--- a/artifacts/run_all_ncu_pt.sh
+++ b/artifacts/run_all_ncu_pt.sh
@@ -10,4 +10,5 @@ bigbird_dir="$benchmark_dir/blocked_sparse_attention/pytorch"
 # 2. ncu test the bigbird benchmark
 echo "NCU profiling BigBird benchmark"
 $ncu_dir/ncu --section "MemoryWorkloadAnalysis" \
-    --csv --set full python3 $bigbird_dir/main.py > $log_dir/pt_bigbird_ncu.csv
\ No newline at end of file
+    --metrics "dram__bytes.sum,lts__t_bytes.sum,l1tex__t_bytes.sum" \
+    --csv python3 $bigbird_dir/main.py > $log_dir/pt_bigbird_ncu.csv
\ No newline at end of file
diff --git a/artifacts/run_all_ncu_triton.sh b/artifacts/run_all_ncu_triton.sh
index 0498aa506..7682bbb38 100755
--- a/artifacts/run_all_ncu_triton.sh
+++ b/artifacts/run_all_ncu_triton.sh
@@ -9,12 +9,14 @@ mha_dir="$benchmark_dir/multi-head_attention/baseline"
 # 1. ncu test the mha benchmark
 echo "NCU profiling mha benchmark"
 $ncu_dir/ncu --section "MemoryWorkloadAnalysis" \
-    --csv --set full python3 $mha_dir/test_triton_model.py > $log_dir/triton_attention_ncu.csv
+    --metrics "dram__bytes.sum,lts__t_bytes.sum,l1tex__t_bytes.sum" \
+    --csv python3 $mha_dir/test_triton_model.py > $log_dir/triton_attention_ncu.csv
 
 
 bigbird_dir="$benchmark_dir/blocked_sparse_attention/triton"
 # 2. ncu test the bigbird benchmark
 echo "NCU profiling BigBird benchmark"
 $ncu_dir/ncu --section "MemoryWorkloadAnalysis" \
-    --csv --set full python3 $bigbird_dir/main.py \
+    --metrics "dram__bytes.sum,lts__t_bytes.sum,l1tex__t_bytes.sum" \
+    --csv python3 $bigbird_dir/main.py \
     --default_test True > $log_dir/triton_bigbird_ncu.csv
\ No newline at end of file