1// RUN: mlir-opt -normalize-memrefs %s -split-input-file| FileCheck %s 2 3// For all these cases, we test if MemRefs Normalization works with the test 4// operations. These are test cases for MemRefs with dynamic dimension 5// and tiled-layout map. 6// * test.op_norm: this operation has the MemRefsNormalizable attribute. The tests 7// that include this operation are constructed so that the normalization should 8// 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_dynamic12 20// 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 : index 23 %c2 = arith.constant 2 : index 24 %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 return 30 // CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index 31 // CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index 32 // CHECK-NOT: separator of consecutive DAGs 33 // 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 : index 36 // CHECK-DAG: [[CST_14_:%.+]] = arith.constant 14 : index 37 // CHECK-NOT: separator of consecutive DAGs 38 // 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: return 45} 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_dynamic1234 62// 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 : index 65 %c1 = arith.constant 1 : index 66 %c2 = arith.constant 2 : index 67 %c3 = arith.constant 3 : index 68 %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 return 76 // CHECK-DAG: [[CST_0_:%.+]] = arith.constant 0 : index 77 // CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index 78 // CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index 79 // CHECK-DAG: [[CST_3_:%.+]] = arith.constant 3 : index 80 // CHECK-NOT: separator of consecutive DAGs 81 // 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 DAGs 86 // 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: return 96} 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_tiled0 108// 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 : index 111 %c2 = arith.constant 2 : index 112 %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 return 118 // CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index 119 // CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index 120 // CHECK-NOT: separator of consecutive DAGs 121 // 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: return 127} 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_tiled1 139// 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 : index 142 %c2 = arith.constant 2 : index 143 %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 return 149 // CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index 150 // CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index 151 // CHECK-NOT: separator of consecutive DAGs 152 // 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: return 158} 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_tiled2 170// 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 : index 173 %c2 = arith.constant 2 : index 174 %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 return 180 // CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index 181 // CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index 182 // CHECK-NOT: separator of consecutive DAGs 183 // 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: return 189} 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_tiled3 201// 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 : index 204 %c2 = arith.constant 2 : index 205 %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 return 211 // CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index 212 // CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index 213 // CHECK-NOT: separator of consecutive DAGs 214 // 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: return 220} 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_tiled4 232// 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 : index 235 %c2 = arith.constant 2 : index 236 %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 return 242 // CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index 243 // CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index 244 // CHECK-NOT: separator of consecutive DAGs 245 // 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: return 251} 252