252 lines · plain
1// RUN: mlir-opt -normalize-memrefs %s -split-input-file| FileCheck %s2 3// For all these cases, we test if MemRefs Normalization works with the test4// operations. These are test cases for MemRefs with dynamic dimension5// and tiled-layout map.6// * test.op_norm: this operation has the MemRefsNormalizable attribute. The tests7// that include this operation are constructed so that the normalization should8// happen.9 10#map_tiled = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 floordiv 32, d3 floordiv 64, d2 mod 32, d3 mod 64)>11 12// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d1)>13// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d2 ceildiv 32)>14// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (32)>15 16// Test with op_norm and maps in arguments and in the operations in the function.17// Memref has two dynamic dimensions.18 19// CHECK-LABEL: test_norm_dynamic1220// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x1x?x64xf32>) {21func.func @test_norm_dynamic12(%arg0 : memref<1x?x?x14xf32, #map_tiled>) -> () {22 %c1 = arith.constant 1 : index23 %c2 = arith.constant 2 : index24 %0 = memref.dim %arg0, %c1 :memref<1x?x?x14xf32, #map_tiled>25 %1 = memref.dim %arg0, %c2 :memref<1x?x?x14xf32, #map_tiled>26 %2 = memref.alloc(%0, %1) : memref<1x?x?x14xf32, #map_tiled>27 "test.op_norm"(%arg0, %2) : (memref<1x?x?x14xf32, #map_tiled>, memref<1x?x?x14xf32, #map_tiled>) -> ()28 memref.dealloc %2 : memref<1x?x?x14xf32, #map_tiled>29 return30 // CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index31 // CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index32 // CHECK-NOT: separator of consecutive DAGs33 // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x1x?x64xf32>34 // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x1x?x64xf32>35 // CHECK-DAG: [[CST_1_1_:%.+]] = arith.constant 1 : index36 // CHECK-DAG: [[CST_14_:%.+]] = arith.constant 14 : index37 // CHECK-NOT: separator of consecutive DAGs38 // CHECK-DAG: [[VAR_2_:%.+]] = affine.apply #[[$MAP0]]([[CST_1_1_]], [[DIM_0_]], [[DIM_1_]], [[CST_14_]])39 // CHECK-DAG: [[VAR_3_:%.+]] = affine.apply #[[$MAP1]]([[CST_1_1_]], [[DIM_0_]], [[DIM_1_]], [[CST_14_]])40 // CHECK-DAG: [[VAR_4_:%.+]] = affine.apply #[[$MAP2]]([[CST_1_1_]], [[DIM_0_]], [[DIM_1_]], [[CST_14_]])41 // CHECK: [[RES_:%.+]] = memref.alloc([[VAR_2_]], [[VAR_3_]], [[VAR_4_]]) : memref<1x?x?x1x?x64xf32>42 // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x1x?x64xf32>, memref<1x?x?x1x?x64xf32>) -> ()43 // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x1x?x64xf32>44 // CHECK: return45}46 47// -----48 49// Test with op_norm and maps in arguments and in the operations in the function.50// All of dimensions are dynamic.51 52#map_tiled1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, (d2 floordiv 4) floordiv 32, (d3 mod 8) floordiv 64, (d2 floordiv 4) mod 32, (d3 mod 8) mod 64)>53 54// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d1)>55// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> ((d2 floordiv 4) ceildiv 32)>56// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (32)>57// CHECK-DAG: #[[$MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0)>58// CHECK-DAG: #[[$MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> ((d3 mod 8) ceildiv 64)>59// CHECK-DAG: #[[$MAP5:.+]] = affine_map<(d0, d1, d2, d3) -> (64)>60 61// CHECK-LABEL: test_norm_dynamic123462// CHECK-SAME: ([[ARG_0_:%.+]]: memref<?x?x?x?x?x?xf32>) {63func.func @test_norm_dynamic1234(%arg0 : memref<?x?x?x?xf32, #map_tiled1>) -> () {64 %c0 = arith.constant 0 : index65 %c1 = arith.constant 1 : index66 %c2 = arith.constant 2 : index67 %c3 = arith.constant 3 : index68 %0 = memref.dim %arg0, %c0 :memref<?x?x?x?xf32, #map_tiled1>69 %1 = memref.dim %arg0, %c1 :memref<?x?x?x?xf32, #map_tiled1>70 %2 = memref.dim %arg0, %c2 :memref<?x?x?x?xf32, #map_tiled1>71 %3 = memref.dim %arg0, %c3 :memref<?x?x?x?xf32, #map_tiled1>72 %4 = memref.alloc(%0, %1, %2, %3) : memref<?x?x?x?xf32, #map_tiled1>73 "test.op_norm"(%arg0, %4) : (memref<?x?x?x?xf32, #map_tiled1>, memref<?x?x?x?xf32, #map_tiled1>) -> ()74 memref.dealloc %4 : memref<?x?x?x?xf32, #map_tiled1>75 return76 // CHECK-DAG: [[CST_0_:%.+]] = arith.constant 0 : index77 // CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index78 // CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index79 // CHECK-DAG: [[CST_3_:%.+]] = arith.constant 3 : index80 // CHECK-NOT: separator of consecutive DAGs81 // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_0_]] : memref<?x?x?x?x?x?xf32>82 // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<?x?x?x?x?x?xf32>83 // CHECK-DAG: [[DIM_2_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<?x?x?x?x?x?xf32>84 // CHECK-DAG: [[DIM_3_:%.+]] = memref.dim [[ARG_0_]], [[CST_3_]] : memref<?x?x?x?x?x?xf32>85 // CHECK-NOT: separator of consecutive DAGs86 // CHECK-DAG: [[VAR_4_:%.+]] = affine.apply #[[$MAP3]]([[DIM_0_]], [[DIM_1_]], [[DIM_2_]], [[DIM_3_]])87 // CHECK-DAG: [[VAR_5_:%.+]] = affine.apply #[[$MAP0]]([[DIM_0_]], [[DIM_1_]], [[DIM_2_]], [[DIM_3_]])88 // CHECK-DAG: [[VAR_6_:%.+]] = affine.apply #[[$MAP1]]([[DIM_0_]], [[DIM_1_]], [[DIM_2_]], [[DIM_3_]])89 // CHECK-DAG: [[VAR_7_:%.+]] = affine.apply #[[$MAP4]]([[DIM_0_]], [[DIM_1_]], [[DIM_2_]], [[DIM_3_]])90 // CHECK-DAG: [[VAR_8_:%.+]] = affine.apply #[[$MAP2]]([[DIM_0_]], [[DIM_1_]], [[DIM_2_]], [[DIM_3_]])91 // CHECK-DAG: [[VAR_9_:%.+]] = affine.apply #[[$MAP5]]([[DIM_0_]], [[DIM_1_]], [[DIM_2_]], [[DIM_3_]])92 // CHECK: [[RES_:%.+]] = memref.alloc([[VAR_4_]], [[VAR_5_]], [[VAR_6_]], [[VAR_7_]], [[VAR_8_]], [[VAR_9_]]) : memref<?x?x?x?x?x?xf32>93 // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<?x?x?x?x?x?xf32>, memref<?x?x?x?x?x?xf32>) -> ()94 // CHECK: memref.dealloc [[RES_]] : memref<?x?x?x?x?x?xf32>95 // CHECK: return96}97 98// -----99 100// Same test with maps that are not tiled layout maps in the arguments and the operations in the function.101// This is not normalized since this is not tiled-layout map. No mod and floordiv.102 103#map_not_tiled0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2)>104 105// CHECK-DAG: #[[$MAP6:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2)>106 107// CHECK-LABEL: func @test_norm_dynamic_not_tiled0108// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP6]]>) {109func.func @test_norm_dynamic_not_tiled0(%arg0 : memref<1x?x?x14xf32, #map_not_tiled0>) -> () {110 %c1 = arith.constant 1 : index111 %c2 = arith.constant 2 : index112 %0 = memref.dim %arg0, %c1 :memref<1x?x?x14xf32, #map_not_tiled0>113 %1 = memref.dim %arg0, %c2 :memref<1x?x?x14xf32, #map_not_tiled0>114 %2 = memref.alloc(%0, %1) : memref<1x?x?x14xf32, #map_not_tiled0>115 "test.op_norm"(%arg0, %2) : (memref<1x?x?x14xf32, #map_not_tiled0>, memref<1x?x?x14xf32, #map_not_tiled0>) -> ()116 memref.dealloc %2 : memref<1x?x?x14xf32, #map_not_tiled0>117 return118 // CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index119 // CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index120 // CHECK-NOT: separator of consecutive DAGs121 // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP6]]>122 // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP6]]>123 // CHECK: [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP6]]>124 // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP6]]>, memref<1x?x?x14xf32, #[[$MAP6]]>) -> ()125 // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP6]]>126 // CHECK: return127}128 129// -----130 131// Same test with maps that are not tiled layout maps in the arguments and the operations in the function.132// This is not normalized since this is not tiled-layout map. No floordiv.133 134#map_not_tiled1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2, d2 mod 32, d3 mod 64)>135 136// CHECK-DAG: #[[$MAP6:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2, d2 mod 32, d3 mod 64)>137 138// CHECK-LABEL: func @test_norm_dynamic_not_tiled1139// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP6]]>) {140func.func @test_norm_dynamic_not_tiled1(%arg0 : memref<1x?x?x14xf32, #map_not_tiled1>) -> () {141 %c1 = arith.constant 1 : index142 %c2 = arith.constant 2 : index143 %0 = memref.dim %arg0, %c1 :memref<1x?x?x14xf32, #map_not_tiled1>144 %1 = memref.dim %arg0, %c2 :memref<1x?x?x14xf32, #map_not_tiled1>145 %2 = memref.alloc(%0, %1) : memref<1x?x?x14xf32, #map_not_tiled1>146 "test.op_norm"(%arg0, %2) : (memref<1x?x?x14xf32, #map_not_tiled1>, memref<1x?x?x14xf32, #map_not_tiled1>) -> ()147 memref.dealloc %2 : memref<1x?x?x14xf32, #map_not_tiled1>148 return149 // CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index150 // CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index151 // CHECK-NOT: separator of consecutive DAGs152 // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP6]]>153 // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP6]]>154 // CHECK: [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP6]]>155 // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP6]]>, memref<1x?x?x14xf32, #[[$MAP6]]>) -> ()156 // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP6]]>157 // CHECK: return158}159 160// -----161 162// Same test with maps that are not tiled layout maps in the arguments and the operations in the function.163// This is not normalized since this is not tiled-layout map. RHS of floordiv is different from that of mod.164 165#map_not_tiled2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 floordiv 64, d2 mod 32, d3 mod 32)>166 167// CHECK-DAG: #[[$MAP7:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 floordiv 64, d2 mod 32, d3 mod 32)>168 169// CHECK-LABEL: func @test_norm_dynamic_not_tiled2170// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP7]]>) {171func.func @test_norm_dynamic_not_tiled2(%arg0 : memref<1x?x?x14xf32, #map_not_tiled2>) -> () {172 %c1 = arith.constant 1 : index173 %c2 = arith.constant 2 : index174 %0 = memref.dim %arg0, %c1 :memref<1x?x?x14xf32, #map_not_tiled2>175 %1 = memref.dim %arg0, %c2 :memref<1x?x?x14xf32, #map_not_tiled2>176 %2 = memref.alloc(%0, %1) : memref<1x?x?x14xf32, #map_not_tiled2>177 "test.op_norm"(%arg0, %2) : (memref<1x?x?x14xf32, #map_not_tiled2>, memref<1x?x?x14xf32, #map_not_tiled2>) -> ()178 memref.dealloc %2 : memref<1x?x?x14xf32, #map_not_tiled2>179 return180 // CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index181 // CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index182 // CHECK-NOT: separator of consecutive DAGs183 // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP7]]>184 // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP7]]>185 // CHECK: [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP7]]>186 // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP7]]>, memref<1x?x?x14xf32, #[[$MAP7]]>) -> ()187 // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP7]]>188 // CHECK: return189}190 191// -----192 193// Same test with maps that are not tiled layout maps in the arguments and the operations in the function.194// This is not normalized since this is not tiled-layout map. Multiple mod with the same LHS and RHS.195 196#map_not_tiled3 = affine_map<(d0, d1, d2, d3) -> (d0, d1 floordiv 32, d2, d3, d1 mod 32, d1 mod 32)>197 198// CHECK-DAG: #[[$MAP8:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 floordiv 32, d2, d3, d1 mod 32, d1 mod 32)>199 200// CHECK-LABEL: func @test_norm_dynamic_not_tiled3201// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP8]]>) {202func.func @test_norm_dynamic_not_tiled3(%arg0 : memref<1x?x?x14xf32, #map_not_tiled3>) -> () {203 %c1 = arith.constant 1 : index204 %c2 = arith.constant 2 : index205 %0 = memref.dim %arg0, %c1 :memref<1x?x?x14xf32, #map_not_tiled3>206 %1 = memref.dim %arg0, %c2 :memref<1x?x?x14xf32, #map_not_tiled3>207 %2 = memref.alloc(%0, %1) : memref<1x?x?x14xf32, #map_not_tiled3>208 "test.op_norm"(%arg0, %2) : (memref<1x?x?x14xf32, #map_not_tiled3>, memref<1x?x?x14xf32, #map_not_tiled3>) -> ()209 memref.dealloc %2 : memref<1x?x?x14xf32, #map_not_tiled3>210 return211 // CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index212 // CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index213 // CHECK-NOT: separator of consecutive DAGs214 // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP8]]>215 // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP8]]>216 // CHECK: [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP8]]>217 // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP8]]>, memref<1x?x?x14xf32, #[[$MAP8]]>) -> ()218 // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP8]]>219 // CHECK: return220}221 222// -----223 224// Same test with maps that are not tiled layout maps in the arguments and the operations in the function.225// This is not normalized since this is not tiled-layout map. floordiv and mod with the same LHS and RHS(d0 floordiv 32 and d0 mod 32), but, unrelaed d0 exists in other position.226 227#map_not_tiled4 = affine_map<(d0, d1, d2, d3) -> (d0 floordiv 32, d1 floordiv 32, d0, d3, d0 mod 32, d1 mod 32)>228 229// CHECK-DAG: #[[$MAP9:.+]] = affine_map<(d0, d1, d2, d3) -> (d0 floordiv 32, d1 floordiv 32, d0, d3, d0 mod 32, d1 mod 32)>230 231// CHECK-LABEL: func @test_norm_dynamic_not_tiled4232// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP9]]>) {233func.func @test_norm_dynamic_not_tiled4(%arg0 : memref<1x?x?x14xf32, #map_not_tiled4>) -> () {234 %c1 = arith.constant 1 : index235 %c2 = arith.constant 2 : index236 %0 = memref.dim %arg0, %c1 :memref<1x?x?x14xf32, #map_not_tiled4>237 %1 = memref.dim %arg0, %c2 :memref<1x?x?x14xf32, #map_not_tiled4>238 %2 = memref.alloc(%0, %1) : memref<1x?x?x14xf32, #map_not_tiled4>239 "test.op_norm"(%arg0, %2) : (memref<1x?x?x14xf32, #map_not_tiled4>, memref<1x?x?x14xf32, #map_not_tiled4>) -> ()240 memref.dealloc %2 : memref<1x?x?x14xf32, #map_not_tiled4>241 return242 // CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index243 // CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index244 // CHECK-NOT: separator of consecutive DAGs245 // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP9]]>246 // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP9]]>247 // CHECK: [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP9]]>248 // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP9]]>, memref<1x?x?x14xf32, #[[$MAP9]]>) -> ()249 // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP9]]>250 // CHECK: return251}252