From 0829b2ce891b4972ec62b109cf067809729936d0 Mon Sep 17 00:00:00 2001
From: yzhang93 <zhyuhang88@gmail.com>
Date: Thu, 12 Dec 2024 15:12:28 -0800
Subject: [PATCH] [Tests] Add missing tests for matmul-transpose-b and 2x2
 cores

---
 build_tools/ci/cpu_comparison/run.py | 16 +++++++++++++++
 build_tools/ci/run_matmul_test.sh    | 29 +++++++++++++++++++++++++---
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/build_tools/ci/cpu_comparison/run.py b/build_tools/ci/cpu_comparison/run.py
index aafa2777a..7d0854368 100755
--- a/build_tools/ci/cpu_comparison/run.py
+++ b/build_tools/ci/cpu_comparison/run.py
@@ -1446,6 +1446,22 @@ def __init__(self):
             )
         )
 
+        # Matmul test on 2(rows)x2(cols) cores
+        self.register(
+                Matmul(
+                    32,
+                    32,
+                    32,
+                    "bf16",
+                    "f32",
+                    aie_compilation_flags=[
+                        "--iree-amdaie-num-rows=2", 
+                        "--iree-amdaie-num-cols=2"
+                        ],
+                    name_suffix="2rows_2cols",
+                )
+            )
+
         performance_tests = [
             {
                 "M": 512,
diff --git a/build_tools/ci/run_matmul_test.sh b/build_tools/ci/run_matmul_test.sh
index 81bdb99bc..3298f0335 100755
--- a/build_tools/ci/run_matmul_test.sh
+++ b/build_tools/ci/run_matmul_test.sh
@@ -629,8 +629,6 @@ run_matmul_test_on_shapes ${i32_shapes_medium[@]} \
     --acc_type "i32" \
     --num_repeat_runs "2"
 
-# bf16 Matmul tests.
-
 bf16_i8_shapes_small=(
   '64x64x64'
   '128x256x128'
@@ -643,7 +641,7 @@ bf16_i8_shapes_medium=(
   '4096x2048x4096'
 )
 
-
+# bf16 Matmul tests.
 run_matmul_test_on_shapes ${bf16_i8_shapes_small[@]} \
     --name_prefix "small_bf16" \
     --lower_to_aie_pipeline "objectFifo" \
@@ -677,6 +675,31 @@ run_matmul_test_on_shapes ${bf16_i8_shapes_medium[@]} \
     --acc_type "i32" \
     --num_repeat_runs "2"
 
+# matmul_transpose_b tests.
+transpose_shapes=(
+  '64x64x64'
+  '128x32x256'
+  '512x128x256'
+  '1536x2048x1536'
+)
+
+run_matmul_test_on_shapes ${transpose_shapes[@]} \
+    --name_prefix "transpose_bf16" \
+    --lower_to_aie_pipeline "objectFifo" \
+    --tile_pipeline "pack-peel" \
+    --lhs_rhs_type "bf16" \
+    --acc_type "f32" \
+    --num_repeat_runs "2" \
+    --do_transpose_rhs "1"
+
+run_matmul_test_on_shapes ${transpose_shapes[@]} \
+    --name_prefix "transpose_i8" \
+    --lower_to_aie_pipeline "objectFifo" \
+    --tile_pipeline "pack-peel" \
+    --lhs_rhs_type "i8" \
+    --acc_type "i32" \
+    --num_repeat_runs "2" \
+    --do_transpose_rhs "1"
 
 
 # note this will not actually show any devices because --xrt_lite_n_core_rows --xrt_lite_n_core_cols are not passed