[GPU] Add transpose to set of generalized named ops

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
iree-org · Jan 29, 2025 · d7fafa4 · d7fafa4
1 parent 60c3c40
commit d7fafa4
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUGeneralizeNamedOps.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUGeneralizeNamedOps.cpp
@@ -55,7 +55,7 @@ struct GPUGeneralizeNamedOpsPass final
     SmallVector<linalg::LinalgOp> namedOpCandidates;
     funcOp.walk([&](linalg::LinalgOp linalgOp) {
       if (isa<linalg::BatchMatmulTransposeBOp, linalg::MatmulTransposeBOp,
-              linalg::VecmatOp, linalg::MatvecOp>(linalgOp))
+              linalg::VecmatOp, linalg::MatvecOp, linalg::TransposeOp>(linalgOp))
         namedOpCandidates.push_back(linalgOp);
     });
 

diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_generalize_named_ops.mlir b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_generalize_named_ops.mlir
@@ -111,3 +111,17 @@ func.func @lowering_config(%arg0: tensor<512x128xf16>, %arg1: tensor<512x128xf16
 // CHECK-LABEL: func.func @lowering_config
 //       CHECK:   linalg.generic
 //  CHECK-SAME:     lowering_config = #[[$CONFIG]]
+
+// -----
+
+func.func @transpose_op(%arg0: tensor<16x32xf16>) -> tensor<32x16xf16> {
+  %empty = tensor.empty() : tensor<32x16xf16>
+  %transpose = linalg.transpose
+      ins(%arg0 : tensor<16x32xf16>)
+      outs(%empty : tensor<32x16xf16>)
+      permutation = [1, 0]
+  return %transpose : tensor<32x16xf16>
+}
+
+// CHECK-LABEL: func.func @transpose_op
+//       CHECK:   linalg.generic