xref: /llvm-project/mlir/test/Integration/Dialect/Linalg/CPU/test-conv-2d-nhwc-hwcf-call.mlir (revision eb206e9ea84eff0a0596fed2de8316d924f946d1)
1// RUN: mlir-opt %s -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf  -expand-strided-metadata -lower-affine -convert-arith-to-llvm -convert-scf-to-cf --finalize-memref-to-llvm -convert-func-to-llvm -convert-cf-to-llvm -reconcile-unrealized-casts | \
2// RUN: mlir-runner -e main -entry-point-result=void \
3// RUN:   -shared-libs=%mlir_runner_utils \
4// RUN: | FileCheck %s
5
6// RUN: mlir-opt %s -transform-interpreter -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf \
7// RUN:    -expand-strided-metadata -lower-affine -convert-arith-to-llvm -convert-scf-to-cf --finalize-memref-to-llvm -convert-func-to-llvm -convert-cf-to-llvm -reconcile-unrealized-casts | \
8// RUN: mlir-runner -e main -entry-point-result=void \
9// RUN:   -shared-libs=%mlir_runner_utils \
10// RUN: | FileCheck %s
11
12func.func private @printMemrefF32(memref<*xf32>)
13
14// Creates and returns 4-D buffer of size (%s1, %s2, %s3, %s4) filled with the value %f
15func.func @alloc_4d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %s4 : index, %f : f32) -> memref<?x?x?x?xf32> {
16  %buf = memref.alloc(%s1, %s2, %s3, %s4) : memref<?x?x?x?xf32>
17  linalg.fill ins(%f : f32) outs(%buf : memref<?x?x?x?xf32>)
18  return %buf : memref<?x?x?x?xf32>
19}
20
21func.func @conv_2d_nhwc_hwcf(%arg0: memref<?x?x?x?xf32>, %arg1: memref<?x?x?x?xf32>, %arg2: memref<?x?x?x?xf32>) {
22  linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>,
23                          strides = dense<1> : tensor<2xi64>}
24     ins (%arg0, %arg1: memref<?x?x?x?xf32>, memref<?x?x?x?xf32>)
25    outs (%arg2: memref<?x?x?x?xf32>)
26  return
27}
28
29module attributes {transform.with_named_sequence} {
30  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
31    %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op
32    %1, %loops:4 = transform.structured.tile_using_for %0 tile_sizes [2, 3, 3, 2] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
33    transform.yield
34  }
35}
36
37func.func @main() {
38  %c0 = arith.constant 0 : index
39  %c1 = arith.constant 1 : index
40  %c3 = arith.constant 3 : index
41  %c6 = arith.constant 6 : index
42  %c8 = arith.constant 8 : index
43  %f10 = arith.constant 10.00000e+00 : f32
44  %val = arith.constant 2.00000e+00 : f32
45  %zero = arith.constant 0.00000e+00 : f32
46
47  %filter2D_nhwc = call @alloc_4d_filled_f32(%c3, %c3, %c3, %c1, %val) :(index, index, index, index, f32) -> (memref<?x?x?x?xf32>)
48  %in2D_nhwc = call @alloc_4d_filled_f32(%c3, %c8, %c8, %c3, %val) : (index, index, index, index, f32) -> (memref<?x?x?x?xf32>)
49  %out2D_nhwc = call @alloc_4d_filled_f32(%c3, %c6, %c6, %c1, %zero) : (index, index, index, index, f32) -> (memref<?x?x?x?xf32>)
50
51  memref.store %f10, %in2D_nhwc[%c0, %c0, %c3, %c0] : memref<?x?x?x?xf32>
52  call @conv_2d_nhwc_hwcf(%in2D_nhwc, %filter2D_nhwc, %out2D_nhwc) : (memref<?x?x?x?xf32>, memref<?x?x?x?xf32>, memref<?x?x?x?xf32>) -> ()
53  %out2D_nhwc_ = memref.cast %out2D_nhwc : memref<?x?x?x?xf32> to memref<*xf32>
54  call @printMemrefF32(%out2D_nhwc_): (memref<*xf32>) -> ()
55
56  memref.dealloc %filter2D_nhwc : memref<?x?x?x?xf32>
57  memref.dealloc %in2D_nhwc : memref<?x?x?x?xf32>
58  memref.dealloc %out2D_nhwc : memref<?x?x?x?xf32>
59  return
60}
61
62// CHECK:       Unranked Memref {{.*}}
63// CHECK-NEXT:  [
64// CHECK-SAME:   [
65// CHECK-SAME:    [
66// CHECK-SAME:     [108],
67// CHECK-COUNT-3:  [124],
68// CHECK-COUNT-2:  [108]
69// CHECK-SAME:    ],
70// CHECK-NEXT:    [
71// CHECK-COUNT-6:  [108]
72// CHECK-SAME:    ],
73// CHECK-NEXT:    [
74// CHECK-COUNT-6:  [108]
75// CHECK-SAME:    ],
76// CHECK-NEXT:    [
77// CHECK-COUNT-6:  [108]
78// CHECK-SAME:    ],
79// CHECK-NEXT:    [
80// CHECK-COUNT-6:  [108]
81// CHECK-SAME:    ],
82// CHECK-NEXT:    [
83// CHECK-COUNT-6:  [108]
84// CHECK-SAME:    ]
85// CHECK-SAME:   ],
86// CHECK-NEXT:   [
87// CHECK-SAME:    [
88// CHECK-COUNT-6:  [108]
89// CHECK-SAME:    ],
90// CHECK-NEXT:    [
91// CHECK-COUNT-6:  [108]
92// CHECK-SAME:    ],
93// CHECK-NEXT:    [
94// CHECK-COUNT-6:  [108]
95// CHECK-SAME:    ],
96// CHECK-NEXT:    [
97// CHECK-COUNT-6:  [108]
98// CHECK-SAME:    ],
99// CHECK-NEXT:    [
100// CHECK-COUNT-6:  [108]
101// CHECK-SAME:    ],
102// CHECK-NEXT:    [
103// CHECK-COUNT-6:  [108]
104// CHECK-SAME:    ]
105// CHECK-SAME:   ],
106// CHECK-NEXT:   [
107// CHECK-SAME:    [
108// CHECK-COUNT-6:  [108]
109// CHECK-SAME:    ],
110// CHECK-NEXT:    [
111// CHECK-COUNT-6:  [108]
112// CHECK-SAME:    ],
113// CHECK-NEXT:    [
114// CHECK-COUNT-6:  [108]
115// CHECK-SAME:    ],
116// CHECK-NEXT:    [
117// CHECK-COUNT-6:  [108]
118// CHECK-SAME:    ],
119// CHECK-NEXT:    [
120// CHECK-COUNT-6:  [108]
121// CHECK-SAME:    ],
122// CHECK-NEXT:    [
123// CHECK-COUNT-6:  [108]
124// CHECK-SAME:    ]
125// CHECK-SAME:   ]
126// CHECK-SAME:  ]
127