1// RUN: mlir-opt --test-transform-dialect-interpreter --split-input-file %s | FileCheck %s
2
3transform.with_pdl_patterns {
4^bb0(%arg0: !pdl.operation):
5  sequence %arg0 {
6    ^bb0(%arg1: !pdl.operation):
7      %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
8      %1, %loops:3 = transform.structured.tile %0 [4, 4, 4]
9  }
10}
11
12// CHECK-LABEL: func @tile_linalg_matmul(
13// CHECK-SAME:    %[[TA:[0-9a-z]+]]: tensor<128x128xf32>
14// CHECK-SAME:    %[[TB:[0-9a-z]+]]: tensor<128x128xf32>
15// CHECK-SAME:    %[[TC:[0-9a-z]+]]: tensor<128x128xf32>
16// CHECK-SAME:  -> tensor<128x128xf32> {
17func.func @tile_linalg_matmul(
18  %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>)
19    -> tensor<128x128xf32> {
20//      CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<128x128xf32>) {
21//      CHECK:   %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<128x128xf32>) {
22//      CHECK:     %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<128x128xf32>) {
23//      CHECK:       %[[sTA:.*]] = tensor.extract_slice %[[TA]][{{.*}}] : tensor<128x128xf32> to tensor<4x4xf32>
24//      CHECK:       %[[sTB:.*]] = tensor.extract_slice %[[TB]][{{.*}}] : tensor<128x128xf32> to tensor<4x4xf32>
25//      CHECK:       %[[sTC:.*]] = tensor.extract_slice %[[TC2]][{{.*}}] : tensor<128x128xf32> to tensor<4x4xf32>
26//      CHECK:       %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor<4x4xf32>, tensor<4x4xf32>)
27// CHECK-SAME:                                   outs(%[[sTC]] : tensor<4x4xf32>)  -> tensor<4x4xf32>
28//      CHECK:       %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}]  : tensor<4x4xf32> into tensor<128x128xf32>
29//      CHECK:       scf.yield %[[TD]] : tensor<128x128xf32>
30//      CHECK:     scf.yield %[[TD2]] : tensor<128x128xf32>
31//      CHECK:   scf.yield %[[TD1]] : tensor<128x128xf32>
32  %0 = linalg.matmul  ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
33                     outs(%arg2: tensor<128x128xf32>)
34    -> tensor<128x128xf32>
35
36//      CHECK: return %[[TD0]] : tensor<128x128xf32>
37  return %0 : tensor<128x128xf32>
38}
39
40// -----
41
42transform.with_pdl_patterns {
43^bb0(%arg0: !pdl.operation):
44  sequence %arg0 {
45    ^bb0(%arg1: !pdl.operation):
46      %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
47      %1 = transform.structured.match ops{["func.call"]} in %arg1
48      %2, %loops:3 = transform.structured.tile %0 [%1, %1, 4]
49  }
50}
51
52func.func private @get_dynamic_tile_size() -> index
53
54// CHECK-LABEL: func @tile_linalg_matmul_dynamic(
55// CHECK-SAME:    %[[TA:[0-9a-z]+]]: tensor<128x128xf32>
56// CHECK-SAME:    %[[TB:[0-9a-z]+]]: tensor<128x128xf32>
57// CHECK-SAME:    %[[TC:[0-9a-z]+]]: tensor<128x128xf32>
58// CHECK-SAME:  -> tensor<128x128xf32> {
59func.func @tile_linalg_matmul_dynamic(
60  %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>)
61    -> tensor<128x128xf32> {
62//      CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<128x128xf32>) {
63//      CHECK:   %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<128x128xf32>) {
64//      CHECK:     %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<128x128xf32>) {
65//      CHECK:       %[[sTA:.*]] = tensor.extract_slice %[[TA]][{{.*}}] : tensor<128x128xf32> to tensor<?x4xf32>
66//      CHECK:       %[[sTB:.*]] = tensor.extract_slice %[[TB]][{{.*}}] : tensor<128x128xf32> to tensor<4x?xf32>
67//      CHECK:       %[[sTC:.*]] = tensor.extract_slice %[[TC2]][{{.*}}] : tensor<128x128xf32> to tensor<?x?xf32>
68//      CHECK:       %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor<?x4xf32>, tensor<4x?xf32>)
69// CHECK-SAME:                                   outs(%[[sTC]] : tensor<?x?xf32>)  -> tensor<?x?xf32>
70//      CHECK:       %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}]  : tensor<?x?xf32> into tensor<128x128xf32>
71//      CHECK:       scf.yield %[[TD]] : tensor<128x128xf32>
72//      CHECK:     scf.yield %[[TD2]] : tensor<128x128xf32>
73//      CHECK:   scf.yield %[[TD1]] : tensor<128x128xf32>
74  %sz = func.call @get_dynamic_tile_size() : () -> index
75  %0 = linalg.matmul  ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
76                     outs(%arg2: tensor<128x128xf32>)
77    -> tensor<128x128xf32>
78
79//      CHECK: return %[[TD0]] : tensor<128x128xf32>
80  return %0 : tensor<128x128xf32>
81}
82