[CUDA] Fix CUDA transform tests for generalized mmt (#16729) Fixes the test failures reported in https://github.com/openxla/iree/pull/16726#issuecomment-1988681370. ci-extra: test_a100
diff --git a/tests/transform_dialect/cuda/mma_elemwise_layout_analysis_codegen_spec.mlir b/tests/transform_dialect/cuda/mma_elemwise_layout_analysis_codegen_spec.mlir index ecb3410..d1362d0 100644 --- a/tests/transform_dialect/cuda/mma_elemwise_layout_analysis_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/mma_elemwise_layout_analysis_codegen_spec.mlir
@@ -6,7 +6,8 @@ // Step 1. Find the fill, matmul and generic ops // =========================================================================== %fill = transform.structured.match ops{["linalg.fill"]} in %variant_op : (!transform.any_op) -> !transform.any_op - %matmul = transform.structured.match ops{["linalg.matmul_transpose_b"]} + %matmul = transform.structured.match ops{["linalg.generic"]} + attributes{iterator_types = [#linalg.iterator_type<parallel>, #linalg.iterator_type<parallel>, #linalg.iterator_type<reduction>]} in %variant_op : (!transform.any_op) -> !transform.any_op %generic = transform.structured.match ops{["linalg.generic"]} attributes{iterator_types = [#linalg.iterator_type<parallel>, #linalg.iterator_type<parallel>]}
diff --git a/tests/transform_dialect/cuda/mma_reduction_layout_analysis_codegen_spec.mlir b/tests/transform_dialect/cuda/mma_reduction_layout_analysis_codegen_spec.mlir index d246a9e..a7c32a4 100644 --- a/tests/transform_dialect/cuda/mma_reduction_layout_analysis_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/mma_reduction_layout_analysis_codegen_spec.mlir
@@ -7,7 +7,8 @@ // Step 1. Find the fill, matmul and generic ops // =========================================================================== %fill = transform.structured.match ops{["linalg.fill"]} in %variant_op : (!transform.any_op) -> !transform.any_op - %matmul = transform.structured.match ops{["linalg.matmul_transpose_b"]} + %matmul = transform.structured.match ops{["linalg.generic"]} + attributes{iterator_types = [#linalg.iterator_type<parallel>, #linalg.iterator_type<parallel>, #linalg.iterator_type<reduction>]} in %variant_op : (!transform.any_op) -> !transform.any_op %reduce = transform.structured.match ops{["linalg.generic"]} attributes{iterator_types = [#linalg.iterator_type<parallel>, #linalg.iterator_type<reduction>]}