1// RUN: mlir-opt -normalize-memrefs %s | FileCheck %s
2
3// For all these cases, we test if MemRefs Normalization works with the test
4// operations.
5// * test.op_norm: this operation has the MemRefsNormalizable attribute. The tests
6//   that include this operation are constructed so that the normalization should
7//   happen.
8// * test_op_nonnorm: this operation does not have the MemRefsNormalization
9//   attribute. The tests that include this operation are constructed so that the
10//    normalization should not happen.
11
12#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 floordiv 32, d3 floordiv 64, d2 mod 32, d3 mod 64)>
13
14// Test with op_norm and maps in arguments and in the operations in the function.
15
16// CHECK-LABEL: test_norm
17// CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x64xf32>)
18func.func @test_norm(%arg0 : memref<1x16x14x14xf32, #map0>) -> () {
19    %0 = memref.alloc() : memref<1x16x14x14xf32, #map0>
20    "test.op_norm"(%arg0, %0) : (memref<1x16x14x14xf32, #map0>, memref<1x16x14x14xf32, #map0>) -> ()
21    memref.dealloc %0 :  memref<1x16x14x14xf32, #map0>
22
23    // CHECK: %[[v0:.*]] = memref.alloc() : memref<1x16x1x1x32x64xf32>
24    // CHECK: "test.op_norm"(%[[ARG0]], %[[v0]]) : (memref<1x16x1x1x32x64xf32>, memref<1x16x1x1x32x64xf32>) -> ()
25    // CHECK: memref.dealloc %[[v0]] : memref<1x16x1x1x32x64xf32>
26    return
27}
28
29// Same test with op_nonnorm, with maps in the arguments and the operations in the function.
30
31// CHECK-LABEL: test_nonnorm
32// CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x14x14xf32, #map>)
33func.func @test_nonnorm(%arg0 : memref<1x16x14x14xf32, #map0>) -> () {
34    %0 = memref.alloc() : memref<1x16x14x14xf32, #map0>
35    "test.op_nonnorm"(%arg0, %0) : (memref<1x16x14x14xf32, #map0>, memref<1x16x14x14xf32, #map0>) -> ()
36    memref.dealloc %0 :  memref<1x16x14x14xf32, #map0>
37
38    // CHECK: %[[v0:.*]] = memref.alloc() : memref<1x16x14x14xf32, #map>
39    // CHECK: "test.op_nonnorm"(%[[ARG0]], %[[v0]]) : (memref<1x16x14x14xf32, #map>, memref<1x16x14x14xf32, #map>) -> ()
40    // CHECK: memref.dealloc %[[v0]] : memref<1x16x14x14xf32, #map>
41    return
42}
43
44// Test with op_norm, with maps in the operations in the function.
45
46// CHECK-LABEL: test_norm_mix
47// CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x64xf32>
48func.func @test_norm_mix(%arg0 : memref<1x16x1x1x32x64xf32>) -> () {
49    %0 = memref.alloc() : memref<1x16x14x14xf32, #map0>
50    "test.op_norm"(%arg0, %0) : (memref<1x16x1x1x32x64xf32>, memref<1x16x14x14xf32, #map0>) -> ()
51    memref.dealloc %0 :  memref<1x16x14x14xf32, #map0>
52
53    // CHECK: %[[v0:.*]] = memref.alloc() : memref<1x16x1x1x32x64xf32>
54    // CHECK: "test.op_norm"(%[[ARG0]], %[[v0]]) : (memref<1x16x1x1x32x64xf32>, memref<1x16x1x1x32x64xf32>) -> ()
55    // CHECK: memref.dealloc %[[v0]] : memref<1x16x1x1x32x64xf32>
56    return
57}
58
59// Test with maps in load and store ops.
60
61#map_tile = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 floordiv 32, d3 floordiv 32, d2 mod 32, d3 mod 32)>
62
63// CHECK-LABEL: test_load_store
64// CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x14x14xf32>
65func.func @test_load_store(%arg0 : memref<1x16x14x14xf32>) -> () {
66    %0 = memref.alloc() : memref<1x16x14x14xf32, #map_tile>
67    // CHECK: %[[v0:.*]] = memref.alloc() : memref<1x16x1x1x32x32xf32>
68    %1 = memref.alloc() : memref<1x16x14x14xf32>
69    // CHECK: %[[v1:.*]] = memref.alloc() : memref<1x16x14x14xf32>
70    "test.op_norm"(%0, %1) : (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>) -> ()
71    // CHECK: "test.op_norm"(%[[v0]], %[[v1]]) : (memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>) -> ()
72    %cst = arith.constant 3.0 : f32
73    affine.for %i = 0 to 1 {
74      affine.for %j = 0 to 16 {
75        affine.for %k = 0 to 14 {
76          affine.for %l = 0 to 14 {
77            %2 = memref.load %1[%i, %j, %k, %l] : memref<1x16x14x14xf32>
78            // CHECK: memref<1x16x14x14xf32>
79            %3 = arith.addf %2, %cst : f32
80            memref.store %3, %arg0[%i, %j, %k, %l] : memref<1x16x14x14xf32>
81            // CHECK: memref<1x16x14x14xf32>
82          }
83        }
84      }
85    }
86    memref.dealloc %0 :  memref<1x16x14x14xf32, #map_tile>
87    // CHECK: memref.dealloc %[[v0]] : memref<1x16x1x1x32x32xf32>
88    memref.dealloc %1 :  memref<1x16x14x14xf32>
89    // CHECK: memref.dealloc %[[v1]] : memref<1x16x14x14xf32>
90    return
91}
92
93// Test with op_norm_ret, with maps in the results of normalizable operation.
94
95// CHECK-LABEL: test_norm_ret
96// CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x32xf32>) -> (memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>) {
97func.func @test_norm_ret(%arg0: memref<1x16x14x14xf32, #map_tile>) -> (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>) {
98    %0 = memref.alloc() : memref<1x16x14x14xf32, #map_tile>
99    // CHECK-NEXT: %[[v0:.*]] = memref.alloc() : memref<1x16x1x1x32x32xf32>
100    %1, %2 = "test.op_norm_ret"(%arg0) : (memref<1x16x14x14xf32, #map_tile>) -> (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>)
101    // CHECK-NEXT: %[[v1:.*]], %[[v2:.*]] = "test.op_norm_ret"
102    // CHECK-SAME: (memref<1x16x1x1x32x32xf32>) -> (memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>)
103    "test.op_norm"(%1, %0) : (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32, #map_tile>) -> ()
104    // CHECK-NEXT: "test.op_norm"
105    // CHECK-SAME: : (memref<1x16x1x1x32x32xf32>, memref<1x16x1x1x32x32xf32>) -> ()
106    memref.dealloc %0 : memref<1x16x14x14xf32, #map_tile>
107    // CHECK-NEXT: memref.dealloc %[[v0]] : memref<1x16x1x1x32x32xf32>
108    return %1, %2 : memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>
109    // CHECK-NEXT: return %[[v1]], %[[v2]] : memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>
110}
111
112// Test with an arbitrary op that references the function symbol.
113
114"test.op_funcref"() {func = @test_norm_mix} : () -> ()
115
116
117// -----
118
119#map_1d_tile = affine_map<(d0) -> (d0 floordiv 32, d0 mod 32)>
120
121// Test with memref.reinterpret_cast
122
123// CHECK-LABEL: test_norm_reinterpret_cast
124// CHECK-SAME: (%[[ARG0:.*]]: memref<1x32xf32>) -> memref<3x1x1xf32> {
125func.func @test_norm_reinterpret_cast(%arg0 : memref<3xf32, #map_1d_tile>) -> (memref<3x1x1xf32>) {
126    %0 = memref.alloc() : memref<3xf32>
127    "test.op_norm"(%arg0, %0) : (memref<3xf32, #map_1d_tile>, memref<3xf32>) -> ()
128    %1 = memref.reinterpret_cast %0 to offset: [0], sizes: [3, 1, 1], strides: [1, 1, 1] : memref<3xf32> to memref<3x1x1xf32>
129    // CHECK: %[[v0:.*]] = memref.alloc() : memref<3xf32>
130    // CHECK: "test.op_norm"(%[[ARG0]], %[[v0]]) : (memref<1x32xf32>, memref<3xf32>) -> ()
131    // CHECK: memref.reinterpret_cast %[[v0]] to offset: [0], sizes: [3, 1, 1], strides: [1, 1, 1] : memref<3xf32> to memref<3x1x1xf32>
132    return %1 : memref<3x1x1xf32>
133}
134