[Mlir-commits] [mlir] [MLIR][Linalg] Improve morphism semantics (PR #222959)

[Adam Siemieniuk]

================
@@ -0,0 +1,377 @@
+// RUN: mlir-opt %s -split-input-file -linalg-categorize-ops | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=named-to-category | FileCheck %s -check-prefix=NAMED
+// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=generic-to-category | FileCheck %s -check-prefix=GENERIC
+
+func.func @categorize_matmul_buffer(%A : memref<16x8xf32>, %B: memref<8x32xf32>, %C: memref<16x32xf32>) {
+  linalg.matmul ins(%A, %B: memref<16x8xf32>, memref<8x32xf32>)
+               outs(%C: memref<16x32xf32>)
+  return
+}
+
+// CHECK-DAG: #[[A_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[B_MAP:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK-DAG: #[[C_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK: func @categorize_matmul_buffer
+// CHECK-SAME: %[[A:.+]]: memref<16x8xf32>
+// CHECK-SAME: %[[B:.+]]: memref<8x32xf32>
+// CHECK-SAME: %[[C:.+]]: memref<16x32xf32>
+
+// CHECK-NOT: linalg.matmul
+// CHECK: linalg.contract
+// CHECK-SAME: indexing_maps = [#[[A_MAP]], #[[B_MAP]], #[[C_MAP]]]
+// CHECK-SAME: ins(%[[A]], %[[B]] : memref<16x8xf32>, memref<8x32xf32>)
+// CHECK-SAME: outs(%[[C]] : memref<16x32xf32>)
+
+// NAMED-LABEL: func @categorize_matmul_buffer
+// NAMED: linalg.contract
+
+// GENERIC-LABEL: func @categorize_matmul_buffer
+// GENERIC-NOT: linalg.contract
+
+// -----
+
+func.func @categorize_matmul_tensor(%A : tensor<16x8xf32>, %B: tensor<8x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
+  %0 = linalg.matmul ins(%A, %B: tensor<16x8xf32>, tensor<8x32xf32>)
+                    outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
+  return %0: tensor<16x32xf32>
+}
+
+// CHECK-DAG: #[[A_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[B_MAP:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK-DAG: #[[C_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK: func @categorize_matmul_tensor
+
+// CHECK-NOT: linalg.matmul
+// CHECK: linalg.contract
+// CHECK-SAME: indexing_maps = [#[[A_MAP]], #[[B_MAP]], #[[C_MAP]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<16x8xf32>, tensor<8x32xf32>)
+// CHECK-SAME: outs(%{{.+}} : tensor<16x32xf32>)
+
+// NAMED-LABEL: func @categorize_matmul_tensor
+// NAMED: linalg.contract
+
+// GENERIC-LABEL: func @categorize_matmul_tensor
+// GENERIC-NOT: linalg.contract
+
+// -----
+
+// Unsigned cast is preserved on the category op.
+func.func @categorize_matmul_unsigned_cast(%A: tensor<16x8xi16>, %B: tensor<8x32xi16>,
+                                           %C: tensor<16x32xi32>) -> tensor<16x32xi32> {
+  %0 = linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
----------------
adam-smnk wrote:

nit: Maybe one test with custom indexing maps for completness?

https://github.com/llvm/llvm-project/pull/222959