[GPU] Fix e2e matmul generator to extract the input element type. (#1…

…8140) It is a follow-up for #18134 (comment) ci-extra: build_packages,test_amd_mi250,test_amd_mi300,test_amd_w7900,test_nvidia_t4 --------- Signed-off-by: hanhanW <[email protected]>
iree-org · Aug 7, 2024 · 31bfc93 · 31bfc93
1 parent cc5566c
commit 31bfc93
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/tests/e2e/matmul/generate_e2e_matmul_tests.py b/tests/e2e/matmul/generate_e2e_matmul_tests.py
@@ -309,10 +309,10 @@ def get_rocm_test_compilation_infos(
     for schedule in schedules:
         # Skip schedules with an intrinsic which element type does not
         # match the requested one.
-        # Extracts the input type from strings containing either 'MFMA' or 'WMMA'
-        # followed by an underscore.
-        extract_input_type = lambda s: re.search(r"(?:MFMA|WMMA)_([^_]+)_", s).group(1)
-        if lhs_rhs_type.value.upper() != extract_input_type(schedule.intrinsic):
+        # Extracts the input type from strings. The naming convention is
+        # [output_type]_MxNxK_[input_type].
+        input_type = schedule.intrinsic.split("_")[-1]
+        if lhs_rhs_type.value.upper() != input_type:
             continue
 
         if schedule.intrinsic == "MFMA_F32_16x16x4_F32":