From 36b10ac51b15f3a175921b4559763cf85b19c361 Mon Sep 17 00:00:00 2001
From: yangxinyu <yangxinyu.302@bytedance.com>
Date: Wed, 3 Jul 2024 05:11:41 +0000
Subject: [PATCH] fix

---
 tests/numerical_test/testset.py | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/tests/numerical_test/testset.py b/tests/numerical_test/testset.py
index 02585154e..e07d71696 100644
--- a/tests/numerical_test/testset.py
+++ b/tests/numerical_test/testset.py
@@ -9,6 +9,7 @@
 
 CUR_DIR = os.path.dirname(os.path.abspath(__file__))
 
+
 def _get_test_files_from_dir(directory):
     test_files = []
     for filename in os.listdir(directory):
@@ -49,6 +50,10 @@ def _get_test_files_from_dir(directory):
     "transpose1203.mlir",
     "transpose2013.mlir",
     "transpose120.mlir",
+    "gemm_crr_f16f16f32.mlir",
+    "gemm_rrr_f16f16f32.mlir",
+    "bmm_rcr_f16f16f32.mlir",
+    "bmm_rrr_f16f16f32.mlir",
 }
 
 CUDA_ALL_SET = (CUDA_MLIR_TEST_SET | CUDA_TORCH_TEST_SET) - CUDA_XFAIL_SET
@@ -91,3 +96,28 @@ def _get_test_files_from_dir(directory):
 }
 
 CUDA_AIT_ALL_SET = CUDA_AIT_MLIR_TEST_SET | CUDA_AIT_TORCH_TEST_SET
+
+##### CUDA WITH GEMMCODEGEN TEST SET #######
+CUDA_TORCH_MATMUL_TESTS = {test for test in CUDA_TORCH_TEST_SET if "Matmul" in test}
+
+CUDA_GEMMCODEGEN_TESTS = {
+    "gemm_crr_f16f16f32.mlir",
+    "gemm_crr_f32.mlir",
+    "gemm_rrr_f16f16f32.mlir",
+    "bmm_rcr_f16f16f32.mlir",
+    "bmm_rrr_f16f16f32.mlir",
+}
+
+CUDA_WITH_GEMM_CODEGEN_XFAIL_SET = {
+    "MatmulTransposeAF16Module_basic",
+    "MatmulTransposeBF16Module_basic",
+    "BatchMatmulAddF32Module_basic",
+    "MatmulTransposeModule_basic",
+    # TODO: fix bug
+    "gemm_crr_f16f16f32.mlir",
+    "bmm_rcr_f16f16f32.mlir",
+}
+
+CUDA_WITH_GEMM_CODEGEN_SET = (
+    CUDA_TORCH_MATMUL_TESTS | CUDA_GEMMCODEGEN_TESTS
+) - CUDA_WITH_GEMM_CODEGEN_XFAIL_SET