Add batch size 2 tests

nod-ai · Sep 5, 2024 · 232bc19 · 232bc19
1 parent 8795d9d
commit 232bc19
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/build_tools/ci/cpu_comparison/run_test.py b/build_tools/ci/cpu_comparison/run_test.py
@@ -664,12 +664,15 @@ def run(self, config):
         )
 
         # Test(s) of the form batch_matmul(A,B) where A:BxMxK, B:BxKxN
-        test_name = output_dir / "test_from_template_bmm.mlir"
         template_name = matmul_template_dir / "batch_matmul_BxMxK_BxKxN.mlir"
         for (lhs_type, acc_type) in zip(["i32", "bf16"], ["i32", "f32"]):
+            test_name = output_dir / f"test_from_template_bmm_1_{lhs_type}_{acc_type}.mlir"
             generate_matmul_test(test_name, template_name, 128, 128, 256, lhs_type, acc_type, b=1)
             aie_vs_llvm_cpu(config, test_name, tile_pipeline="pack-peel", lower_to_aie_pipeline="objectFifo")
 
+            test_name = output_dir / f"test_from_template_bmm_2_{lhs_type}_{acc_type}.mlir"
+            generate_matmul_test(test_name, template_name, 64, 64, 64, lhs_type, acc_type, b=2)
+            aie_vs_llvm_cpu(config, test_name, tile_pipeline="pack-peel", lower_to_aie_pipeline="objectFifo")
 
 class SmokeSet(TestSet):
     def __init__(self):