From e6239e14f7ca338cf95c2eab48b39c07ea2df23a Mon Sep 17 00:00:00 2001
From: "PoYen, Chen" <PoYen.Chen@amd.com>
Date: Fri, 16 Aug 2024 12:46:16 +0000
Subject: [PATCH] Re-organize bash functions

---
 .../ck_tile/01_fmha/script/smoke_test_fwd.sh  | 35 +++++++++----------
 1 file changed, 16 insertions(+), 19 deletions(-)

diff --git a/example/ck_tile/01_fmha/script/smoke_test_fwd.sh b/example/ck_tile/01_fmha/script/smoke_test_fwd.sh
index a1b902cf32..3d7ed70998 100755
--- a/example/ck_tile/01_fmha/script/smoke_test_fwd.sh
+++ b/example/ck_tile/01_fmha/script/smoke_test_fwd.sh
@@ -7,19 +7,15 @@ KNAME=1
 export CK_WARMUP=0
 export CK_REPEAT=1
 
-P_DROP=(0.0)
-NUM_SPLITS=(1)
-PAGE_BLOCK_SIZE=(0)
-
 COMMON_ARGS='-v=1 -warmup=0 -repeat=1'
 # mode=0
 # export HIP_VISIBLE_DEVICES=4
 
 TEST_SPLITKV=0
 TEST_APPENDKV=0
-# usage:
+# options:
 #    -s: run splitkv tests
-#    -a: run appendkv tests 
+#    -a: run appendkv tests
 while getopts ":sa" opt; do
     case "${opt}" in
         s)
@@ -33,14 +29,15 @@ while getopts ":sa" opt; do
     esac
 done
 
-if [[ ($TEST_SPLITKV -eq 1) || ($TEST_APPENDKV -eq 1)]] ; then
-    NUM_SPLITS+=(2 3)
-    PAGE_BLOCK_SIZE+=(128)
-else
-    P_DROP+=(0.2)
-fi
-
 run_fp16_bf16_tests() {
+    local NUM_SPLITS=(1)
+    local PAGE_BLOCK_SIZE=(0)
+
+    if [ $TEST_SPLITKV -eq 1 ] ; then
+        NUM_SPLITS+=(2 3)
+        PAGE_BLOCK_SIZE+=(128)
+    fi
+
     for prec in "fp16" "bf16" ; do
     for mode in 1 0 ; do
     for perm in 0 1 ; do
@@ -48,7 +45,7 @@ run_fp16_bf16_tests() {
     for hdim in 32 64 128 256 ; do
     for lse in 0 1 ; do
     for bias in "n" "e" "a" ; do
-    for p_drop in "${P_DROP[@]}" ; do
+    for p_drop in 0.0 0.2 ; do
     for num_splits in "${NUM_SPLITS[@]}" ; do
     for page_block_size in "${PAGE_BLOCK_SIZE[@]}" ; do
 
@@ -61,7 +58,7 @@ run_fp16_bf16_tests() {
     $EXE -prec=$prec -mode=$mode -b=3 -h=2 -h_k=1 -d=$hdim -s=200 -s_k=520 -bias=$bias -p_drop=$p_drop -lse=$lse -iperm=$perm -operm=$perm -mask=t:128,30 -vlayout=$vlayout -num_splits=$num_splits -page_block_size=$page_block_size -kname=$KNAME $COMMON_ARGS  
     $EXE -prec=$prec -mode=$mode -b=2 -h=1 -d=$hdim -s=99 -s_k=32 -bias=$bias -p_drop=$p_drop -lse=$lse -iperm=$perm -operm=$perm -mask=b:4,35 -vlayout=$vlayout -num_splits=$num_splits -page_block_size=$page_block_size -kname=$KNAME $COMMON_ARGS  
     $EXE -prec=$prec -mode=$mode -b=1 -h=2 -h_k=1 -d=$hdim -s=33 -s_k=0 -bias=$bias -p_drop=$p_drop -lse=$lse -iperm=$perm -operm=$perm -mask=2 -vlayout=$vlayout -num_splits=$num_splits -page_block_size=$page_block_size -kname=$KNAME $COMMON_ARGS  
-    $EXE -prec=$prec -mode=$mode -b=1 -h=2 -h_k=1 -d=$hdim -s=1 -s_k=10 -s_kpad=32 -bias=$bias -lse=$lse -iperm=$perm -operm=$perm -mask=2 -vlayout=$vlayout -num_splits=$num_splits -page_block_size=$page_block_size -kname=$KNAME $COMMON_ARGS  
+    $EXE -prec=$prec -mode=$mode -b=1 -h=2 -h_k=1 -d=$hdim -s=1 -s_k=10 -s_kpad=32 -bias=$bias -p_drop=$p_drop -lse=$lse -iperm=$perm -operm=$perm -mask=2 -vlayout=$vlayout -num_splits=$num_splits -page_block_size=$page_block_size -kname=$KNAME $COMMON_ARGS  
 
     done ; done ; done ; done ; done
     done ; done ; done ; done ; done
@@ -78,16 +75,16 @@ run_fp8_tests() {
     done ; done ; done ; done
 }
 
-run_appendkv_tests() {
+run_fp16_appendkv_tests() {
     for s in $(seq 63 1 65) ; do
     for s_k in 65 129 ; do
     for s_knew in 64 $s_k ; do
     for hdim in 32 64 128 256 ; do
     for ri in 0 1 ; do
     for rdim in 0 16 32 $hdim ; do
-    for page_block_size in "${PAGE_BLOCK_SIZE[@]}" ; do
+    for page_block_size in 0 128 ; do
 
-    $EXE -prec=fp16 -b=3 -h=3 -d=$hdim -s=$s -s_k=$s_k -s_knew=$s_knew -rotary_dim=$rdim -rotary_interleaved=$ri -iperm=1 -operm=1 -kname=1 $COMMON_ARGS
+    $EXE -prec=fp16 -b=3 -h=3 -d=$hdim -s=$s -s_k=$s_k -s_knew=$s_knew -rotary_dim=$rdim -rotary_interleaved=$ri -page_block_size=$page_block_size -iperm=1 -operm=1 -kname=1 $COMMON_ARGS
 
     done ; done ; done ; done ; done 
     done ; done
@@ -99,7 +96,7 @@ run_fp16_bf16_tests
 run_fp8_tests
 
 if [ $TEST_APPENDKV -eq 1 ] ; then
-    run_appendkv_tests
+    run_fp16_appendkv_tests
 fi
 
 set +x
\ No newline at end of file