xref: /llvm-project/mlir/test/Dialect/Linalg/tile-conv.mlir (revision 2c1c67674cb3beb4e091a9f446de5858631cf8ae)
1// RUN: mlir-opt %s -transform-interpreter -canonicalize | FileCheck %s
2
3//  CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0] -> (-d0 + s0, 2)>
4//  CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0] -> (-d0 + s0, 3)>
5//  CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0)[s0] -> (d0 + s0 - 1)>
6
7func.func @conv(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?xf32>, %arg2 : memref<?x?xf32>) {
8  linalg.conv_2d ins(%arg0, %arg1 : memref<?x?xf32>, memref<?x?xf32>) outs(%arg2 : memref<?x?xf32>)
9  return
10}
11
12module attributes {transform.with_named_sequence} {
13  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
14    %0 = transform.structured.match ops{["linalg.conv_2d"]} in %arg1 : (!transform.any_op) -> !transform.any_op
15    %1, %loop:2 = transform.structured.tile_using_for %0 tile_sizes [2, 3] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
16    transform.yield
17  }
18}
19
20//       CHECK: func @conv
21//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]*]]: memref<?x?xf32>
22//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]*]]: memref<?x?xf32>
23//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]*]]: memref<?x?xf32>
24//   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
25//   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
26//   CHECK-DAG:   %[[C2:.*]] = arith.constant 2 : index
27//   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
28//   CHECK-DAG:   %[[KH:.*]] = memref.dim %[[ARG1]], %[[C0]]
29//   CHECK-DAG:   %[[KW:.*]] = memref.dim %[[ARG1]], %[[C1]]
30//   CHECK-DAG:   %[[H:.*]] = memref.dim %[[ARG2]], %[[C0]]
31//   CHECK-DAG:   %[[W:.*]] = memref.dim %[[ARG2]], %[[C1]]
32//       CHECK:   scf.for %[[I:.*]] = %[[C0]] to %[[H]] step %[[C2]]
33//       CHECK:     scf.for %[[J:.*]] = %[[C0]] to %[[W]] step %[[C3]]
34//   CHECK-DAG:     %[[T4:.*]] = affine.min #[[MAP0]](%[[I]])[%[[H]]]
35//   CHECK-DAG:       %[[T5:.*]] = affine.min #[[MAP1]](%[[J]])[%[[W]]]
36//   CHECK-DAG:       %[[T6:.*]] = affine.apply #[[MAP2]](%[[T4]])[%[[KH]]]
37//   CHECK-DAG:       %[[T7:.*]] = affine.apply #[[MAP2]](%[[T5]])[%[[KW]]]
38//   CHECK-DAG:       %[[SVIN:.*]] = memref.subview %[[ARG0]][%[[I]], %[[J]]] [%[[T6]], %[[T7]]]
39//   CHECK-DAG:       %[[SVKER:.*]] = memref.subview %[[ARG1]][0, 0] [%[[KH]], %[[KW]]]
40//   CHECK-DAG:       %[[SVOUT:.*]] = memref.subview %[[ARG2]][%[[I]], %[[J]]] [%[[T4]], %[[T5]]]
41//       CHECK:       linalg.conv_2d
42//  CHECK-SAME:         ins(%[[SVIN]], %[[SVKER]]
43//  CHECK-SAME:         outs(%[[SVOUT]]
44