commit 3baa82b53d397dbe0fac03f412441653a856143a
author Jakub Kuderski <jakub@nod-labs.com> Mon Mar 11 13:26:56 2024 -0400
committer GitHub <noreply@github.com> Mon Mar 11 17:26:56 2024 +0000
tree 3e6df98743561804248498dda4ec381529d1f72f
parent 27b837c177e82f6b85a5a12f6bbcdf337c62b8f8

[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

tests/transform_dialect/cuda/mma_elemwise_layout_analysis_codegen_spec.mlir

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>]}

tests/transform_dialect/cuda/mma_reduction_layout_analysis_codegen_spec.mlir

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>]}