423 lines · plain
1// RUN: mlir-opt %s \2// RUN: -canonicalize="test-convergence" \3// RUN: --split-input-file -allow-unregistered-dialect | \4// RUN: FileCheck %s5 6// Basic folding of to_tensor(to_buffer(t)) -> t7// CHECK-LABEL: func @tensor_load_of_buffer_cast(8func.func @tensor_load_of_buffer_cast(%arg0: tensor<?xf32>) -> tensor<?xf32> {9 %0 = bufferization.to_buffer %arg0 : tensor<?xf32> to memref<?xf32>10 %1 = bufferization.to_tensor %0 : memref<?xf32> to tensor<?xf32>11 return %1 : tensor<?xf32>12}13// CHECK-SAME: %[[TENSOR:.*]]: tensor<?xf32>) -> tensor<?xf32> {14// CHECK: return %[[TENSOR]]15 16// -----17 18// Basic folding of to_buffer(to_tensor(m)) -> m19// CHECK-LABEL: func @buffer_cast_of_tensor_load(20func.func @buffer_cast_of_tensor_load(%arg0: memref<?xf32>) -> memref<?xf32> {21 %0 = bufferization.to_tensor %arg0 : memref<?xf32> to tensor<?xf32>22 %1 = bufferization.to_buffer %0 : tensor<?xf32> to memref<?xf32>23 return %1 : memref<?xf32>24}25// CHECK-SAME: %[[MEMREF:.*]]: memref<?xf32>) -> memref<?xf32> {26// CHECK: return %[[MEMREF]]27 28// -----29 30// If the memrefs are not the same type, don't fold them.31// If the memrefs are not cast-compatible but one can be copied into the other32// (e.g. different address space), canonicalize them to add + copy.33// CHECK-LABEL: func @canonicalize_buffer_cast_of_tensor_load_different_address_space(34// CHECK-SAME: %[[MEMREF_ADDRSPACE2:.*]]: memref<?xf32, 2>)35// CHECK-SAME: -> memref<?xf32, 7> {36// CHECK-NOT: bufferization.to_tensor37// CHECK-NOT: bufferization.to_buffer38// CHECK: %[[C0:.*]] = arith.constant 0 : index39// CHECK: %[[DIM:.*]] = memref.dim %[[MEMREF_ADDRSPACE2]], %[[C0]] : memref<?xf32, 2>40// CHECK: %[[MEMREF_ADDRSPACE7:.*]] = memref.alloc(%[[DIM]]) : memref<?xf32, 7>41// CHECK: memref.copy %[[MEMREF_ADDRSPACE2]], %[[MEMREF_ADDRSPACE7]]42// CHECK-SAME: memref<?xf32, 2> to memref<?xf32, 7>43// CHECK: return %[[MEMREF_ADDRSPACE7]]44func.func @canonicalize_buffer_cast_of_tensor_load_different_address_space(%arg0: memref<?xf32, 2>)45 -> memref<?xf32, 7> {46 %0 = bufferization.to_tensor %arg0 : memref<?xf32, 2> to tensor<?xf32, 7>47 %1 = bufferization.to_buffer %0 : tensor<?xf32, 7> to memref<?xf32, 7>48 return %1 : memref<?xf32, 7>49}50 51// -----52 53// If the memrefs are definitely cast-compatible, canonicalize to54// cast.55// CHECK-LABEL: func @canonicalize_buffer_cast_of_tensor_load(56// CHECK-SAME: %[[M:.*]]: memref<?xf32, strided<[1], offset: 3>>)57// CHECK-SAME: -> memref<?xf32, strided<[1], offset: ?>> {58// CHECK-NOT: bufferization.to_tensor59// CHECK-NOT: bufferization.to_buffer60// CHECK: %[[R:.*]] = memref.cast %[[M]]61// CHECK-SAME: memref<?xf32, strided<[1], offset: 3>> to memref<?xf32, strided<[1], offset: ?>>62// CHECK: return %[[R]]63func.func @canonicalize_buffer_cast_of_tensor_load(64 %arg0: memref<?xf32, strided<[1], offset: 3>>)65 -> memref<?xf32, strided<[1], offset: ?>>66{67 %0 = bufferization.to_tensor %arg0 : memref<?xf32, strided<[1], offset: 3>> to tensor<?xf32>68 %1 = bufferization.to_buffer %0 : tensor<?xf32> to memref<?xf32, strided<[1], offset: ?>>69 return %1 : memref<?xf32, strided<[1], offset: ?>>70}71 72// -----73 74// If the memrefs are potentially cast-compatible, canonicalize to75// copy.76// CHECK-LABEL: func @canonicalize_buffer_cast_of_tensor_load_to_copy(77func.func @canonicalize_buffer_cast_of_tensor_load_to_copy(78 %arg0: memref<?xf32, strided<[1], offset: ?>>)79 -> memref<?xf32, strided<[1], offset: 3>> {80 %0 = bufferization.to_tensor %arg0 : memref<?xf32, strided<[1], offset: ?>> to tensor<?xf32>81 %1 = bufferization.to_buffer %0 : tensor<?xf32> to memref<?xf32, strided<[1], offset: 3>>82 return %1 : memref<?xf32, strided<[1], offset: 3>>83}84// CHECK-SAME: %[[M:.*]]: memref<?xf32, strided<[1], offset: ?>>)85// CHECK-SAME: -> memref<?xf32, strided<[1], offset: 3>> {86// CHECK-NOT: bufferization.to_tensor87// CHECK-NOT: bufferization.to_buffer88// CHECK: %[[C0:.*]] = arith.constant 0 : index89// CHECK: %[[DIM:.*]] = memref.dim %[[M]], %[[C0]] : memref<?xf32, strided<[1], offset: ?>>90// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM]]) : memref<?xf32, strided<[1], offset: 3>>91// CHECK: memref.copy %[[M]], %[[ALLOC]]92// CHECK-SAME: memref<?xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: 3>>93// CHECK: return %[[ALLOC]]94 95// -----96 97 98// Basic folding of tensor.dim(to_tensor(m)) -> memref.dim(m).99// CHECK-LABEL: func @dim_of_tensor_load(100// CHECK-SAME: %[[MEMREF:[0-9a-z]*]]: memref<?xf32>101// CHECK: %[[C0:.*]] = arith.constant 0102// CHECK: %[[D:.*]] = memref.dim %[[MEMREF]], %[[C0]]103// CHECK: return %[[D]] : index104func.func @dim_of_tensor_load(%arg0: memref<?xf32>) -> index {105 %c0 = arith.constant 0 : index106 %0 = bufferization.to_tensor %arg0 : memref<?xf32> to tensor<?xf32>107 %1 = tensor.dim %0, %c0 : tensor<?xf32>108 return %1 : index109}110 111// -----112 113// CHECK-LABEL: @clone_before_dealloc114func.func @clone_before_dealloc(%arg0: memref<?xf32>) -> memref<?xf32> {115 %0 = bufferization.clone %arg0 : memref<?xf32> to memref<?xf32>116 memref.dealloc %arg0 : memref<?xf32>117 return %0 : memref<?xf32>118}119// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>120// CHECK-NEXT: return %[[ARG]]121 122// -----123 124// CHECK-LABEL: @clone_before_dealloc125func.func @clone_before_dealloc(%arg0: memref<?xf32>) -> memref<?xf32> {126 %0 = bufferization.clone %arg0 : memref<?xf32> to memref<?xf32>127 "use"(%0) : (memref<?xf32>) -> ()128 memref.dealloc %0 : memref<?xf32>129 return %arg0 : memref<?xf32>130}131// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>132// CHECK-NEXT: "use"(%arg0)133// CHECK-NEXT: return %[[ARG]]134 135// -----136 137// CHECK-LABEL: @clone_after_cast138func.func @clone_after_cast(%arg0: memref<?xf32>) -> memref<32xf32> {139 %0 = memref.cast %arg0 : memref<?xf32> to memref<32xf32>140 %1 = bufferization.clone %0 : memref<32xf32> to memref<32xf32>141 return %1 : memref<32xf32>142}143// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>144// CHECK-NEXT: bufferization.clone %[[ARG]] : memref<?xf32> to memref<32xf32>145// CHECK-NOT: memref.cast146 147// -----148 149// CHECK-LABEL: @clone_and_cast150func.func @clone_and_cast(%arg0: memref<?xf32>) -> memref<32xf32> {151 %0 = bufferization.clone %arg0 : memref<?xf32> to memref<32xf32>152 memref.dealloc %arg0 : memref<?xf32>153 return %0 : memref<32xf32>154}155// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>156// CHECK-NEXT: %[[RES:.*]] = memref.cast %[[ARG]]157// CHECK-SAME: memref<?xf32> to memref<32xf32>158// CHECK-NEXT: return %[[RES]]159 160// -----161 162// CHECK-LABEL: @clone_incompatible163func.func @clone_incompatible(%arg0: memref<32xf32, strided<[2]>>) -> memref<32xf32> {164 %0 = bufferization.clone %arg0 : memref<32xf32, strided<[2]>> to memref<32xf32>165 memref.dealloc %arg0 : memref<32xf32, strided<[2]>>166 return %0 : memref<32xf32>167}168// CHECK-SAME: %[[ARG:.*]]: memref<32xf32, strided<[2]>>169// CHECK-NEXT: bufferization.clone %[[ARG]] : memref<32xf32, strided<[2]>> to memref<32xf32>170// CHECK-NOT: memref.cast171 172// -----173 174// CHECK-LABEL: @alias_is_freed175func.func @alias_is_freed(%arg0 : memref<?xf32>) {176 %0 = memref.cast %arg0 : memref<?xf32> to memref<32xf32>177 %1 = bufferization.clone %0 : memref<32xf32> to memref<32xf32>178 memref.dealloc %arg0 : memref<?xf32>179 "use"(%1) : (memref<32xf32>) -> ()180 memref.dealloc %1 : memref<32xf32>181 return182}183// CHECK: bufferization.clone184// CHECK: memref.dealloc185// CHECK: memref.dealloc186 187// -----188 189// Verify SimplifyClones skips clones with multiple deallocations.190// CHECK-LABEL: @clone_multiple_dealloc_of_source191func.func @clone_multiple_dealloc_of_source(%arg0: memref<?xf32>) -> memref<?xf32> {192 %0 = bufferization.clone %arg0 : memref<?xf32> to memref<?xf32>193 "if_else"() ({194 memref.dealloc %arg0 : memref<?xf32>195 }, {196 memref.dealloc %arg0 : memref<?xf32>197 }) : () -> ()198 return %0 : memref<?xf32>199}200// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>201// CHECK-NEXT: %[[RES:.*]] = bufferization.clone %[[ARG]]202// CHECK: memref.dealloc %[[ARG]]203// CHECK: memref.dealloc %[[ARG]]204// CHECK: return %[[RES]]205 206// -----207 208// CHECK-LABEL: @clone_multiple_dealloc_of_clone209// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>210func.func @clone_multiple_dealloc_of_clone(%arg0: memref<?xf32>) -> memref<?xf32> {211 // CHECK-NEXT: %[[CLONE:.*]] = bufferization.clone %[[ARG]]212 // CHECK: memref.dealloc %[[CLONE]]213 // CHECK: memref.dealloc %[[CLONE]]214 // CHECK: return %[[ARG]]215 %0 = bufferization.clone %arg0 : memref<?xf32> to memref<?xf32>216 "use"(%0) : (memref<?xf32>) -> ()217 "if_else"() ({218 memref.dealloc %0 : memref<?xf32>219 }, {220 memref.dealloc %0 : memref<?xf32>221 }) : () -> ()222 return %arg0 : memref<?xf32>223}224 225// -----226 227// Verify SimplifyClones skips clones followed by realloc.228// CHECK-LABEL: @clone_and_realloc229func.func @clone_and_realloc(%arg0: memref<?xf32>) {230 %0 = bufferization.clone %arg0 : memref<?xf32> to memref<32xf32>231 "use"(%0) : (memref<32xf32>) -> ()232 %1 = memref.realloc %0 : memref<32xf32> to memref<64xf32>233 memref.dealloc %1 : memref<64xf32>234 return235}236// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>237// CHECK-NOT: %cast = memref.cast %[[ARG]]238 239// -----240 241// Verify SimplifyClones skips clones with preceding deallocation.242// CHECK-LABEL: @clone_and_preceding_dealloc243func.func @clone_and_preceding_dealloc(%arg0: memref<?xf32>) -> memref<32xf32> {244 memref.dealloc %arg0 : memref<?xf32>245 %0 = bufferization.clone %arg0 : memref<?xf32> to memref<32xf32>246 return %0 : memref<32xf32>247}248// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>249// CHECK-NOT: %cast = memref.cast %[[ARG]]250 251// -----252 253// CHECK-LABEL: func @tensor_cast_to_buffer254// CHECK-SAME: %[[ARG0:.+]]: tensor<4x6x16x32xi8>255func.func @tensor_cast_to_buffer(%arg0 : tensor<4x6x16x32xi8>) ->256 memref<?x?x16x32xi8> {257 %0 = tensor.cast %arg0 : tensor<4x6x16x32xi8> to tensor<?x?x16x32xi8>258 %1 = bufferization.to_buffer %0 read_only : tensor<?x?x16x32xi8> to memref<?x?x16x32xi8>259 return %1 : memref<?x?x16x32xi8>260}261// CHECK: %[[M:.+]] = bufferization.to_buffer %[[ARG0]] read_only : tensor<4x6x16x32xi8>262// CHECK: %[[M1:.+]] = memref.cast %[[M]]263// CHECK-SAME: memref<4x6x16x32xi8> to memref<?x?x16x32xi8>264// CHECK: return %[[M1]] : memref<?x?x16x32xi8>265 266// CHECK-LABEL: func @tensor_cast_to_unranked_buffer267// CHECK-SAME: %[[ARG0:.+]]: tensor<4x6x16x32xi8>268func.func @tensor_cast_to_unranked_buffer(%arg0 : tensor<4x6x16x32xi8>) ->269 memref<*xi8> {270 %0 = tensor.cast %arg0 : tensor<4x6x16x32xi8> to tensor<*xi8>271 %1 = bufferization.to_buffer %0 read_only : tensor<*xi8> to memref<*xi8>272 return %1 : memref<*xi8>273}274// CHECK: %[[M:.+]] = bufferization.to_buffer %[[ARG0]] read_only : tensor<4x6x16x32xi8>275// CHECK: %[[M1:.+]] = memref.cast %[[M]]276// CHECK-SAME: memref<4x6x16x32xi8> to memref<*xi8>277// CHECK: return %[[M1]] : memref<*xi8>278 279// -----280 281// CHECK-LABEL: func @tensor_cast_to_buffer282// CHECK-SAME: %[[ARG0:.+]]: tensor<4x6x16x32xi8>283func.func @tensor_cast_to_buffer_layout_and_memspace(%arg0 : tensor<4x6x16x32xi8>) ->284 memref<?x?x16x32xi8, strided<[?, ?, ?, 1], offset: ?>, 1> {285 %0 = tensor.cast %arg0 : tensor<4x6x16x32xi8> to tensor<?x?x16x32xi8>286 %1 = bufferization.to_buffer %0 : tensor<?x?x16x32xi8> to memref<?x?x16x32xi8, strided<[?, ?, ?, 1], offset: ?>, 1>287 return %1 : memref<?x?x16x32xi8, strided<[?, ?, ?, 1], offset: ?>, 1>288}289// CHECK: %[[M:.+]] = bufferization.to_buffer %[[ARG0]] : tensor<4x6x16x32xi8>290// CHECK: %[[M1:.+]] = memref.cast %[[M]]291// CHECK-SAME: memref<4x6x16x32xi8, strided<[?, ?, ?, 1], offset: ?>, 1>292// CHECK-SAME: to memref<?x?x16x32xi8, strided<[?, ?, ?, 1], offset: ?>, 1>293// CHECK: return %[[M1]] : memref<?x?x16x32xi8, strided<[?, ?, ?, 1], offset: ?>, 1>294 295// -----296 297// Folding of memref.load(to_buffer(%v, %idxs)) -> tensor.extract(%v, %idx)298// CHECK-LABEL: func @load_from_buffer_cast(299func.func @load_from_buffer_cast(%arg0: index, %arg1: index,300 %arg2: tensor<?x?xf32>) -> f32 {301 %0 = bufferization.to_buffer %arg2 : tensor<?x?xf32> to memref<?x?xf32>302 %1 = memref.load %0[%arg0, %arg1] : memref<?x?xf32>303 return %1 : f32304}305// CHECK-SAME: %[[IDX0:[0-9a-z]+]]: index, %[[IDX1:[0-9a-z]+]]: index306// CHECK-SAME: %[[TENSOR:[0-9a-z]+]]: tensor<?x?xf32>307// CHECK: %[[RES:.*]] = tensor.extract %[[TENSOR]][%[[IDX0]], %[[IDX1]]]308// CHECK-NOT: memref.load309// CHECK: return %[[RES]] : f32310 311 312// -----313 314func.func @alloc_tensor_canonicalize() -> (tensor<4x5x?xf32>) {315 %c6 = arith.constant 6 : index316 %0 = bufferization.alloc_tensor(%c6) : tensor<4x5x?xf32>317 return %0 : tensor<4x5x?xf32>318}319// CHECK: func @alloc_tensor_canonicalize320// CHECK: %[[T0:.+]] = bufferization.alloc_tensor() : tensor<4x5x6xf32>321// CHECK: %[[T1:.+]] = tensor.cast %[[T0]] : tensor<4x5x6xf32> to tensor<4x5x?xf32>322// CHECK: return %[[T1]]323 324// -----325 326func.func @dealloc_canonicalize_clone_removal(%arg0: memref<?xindex>) -> memref<*xf32> {327 %c1 = arith.constant 1 : index328 %0 = memref.alloc(%c1) : memref<?xf32>329 %1 = memref.reshape %0(%arg0) : (memref<?xf32>, memref<?xindex>) -> memref<*xf32>330 %2 = bufferization.clone %1 : memref<*xf32> to memref<*xf32>331 memref.dealloc %0 : memref<?xf32>332 return %2 : memref<*xf32>333}334// CHECK-LABEL: @dealloc_canonicalize_clone_removal335// CHECK-NOT: bufferization.clone336// CHECK-NOT: memref.dealloc337// CHECK: return {{.*}}338 339// -----340 341func.func @dealloc_canonicalize_duplicates(%arg0: memref<2xi32>, %arg1: i1, %arg2: i1, %arg3: memref<2xi32>, %arg4: memref<2xi32>, %arg5: memref<2xi32>) -> (i1, i1, i1) {342 %0:3 = bufferization.dealloc (%arg4, %arg0, %arg0 : memref<2xi32>, memref<2xi32>, memref<2xi32>) if (%arg1, %arg1, %arg1) retain (%arg3, %arg5, %arg3 : memref<2xi32>, memref<2xi32>, memref<2xi32>)343 bufferization.dealloc (%arg0, %arg0 : memref<2xi32>, memref<2xi32>) if (%arg1, %arg2)344 return %0#0, %0#1, %0#2 : i1, i1, i1345}346 347// CHECK-LABEL: func @dealloc_canonicalize_duplicates348// CHECK-SAME: ([[ARG0:%.+]]: memref<2xi32>, [[ARG1:%.+]]: i1, [[ARG2:%.+]]: i1, [[ARG3:%.+]]: memref<2xi32>, [[ARG4:%.+]]: memref<2xi32>, [[ARG5:%.+]]: memref<2xi32>)349// CHECK-NEXT: [[V0:%.+]]:2 = bufferization.dealloc ([[ARG4]], [[ARG0]] : memref<2xi32>, memref<2xi32>) if ([[ARG1]], [[ARG1]]) retain ([[ARG3]], [[ARG5]] : memref<2xi32>, memref<2xi32>)350// CHECK-NEXT: [[NEW_COND:%.+]] = arith.ori [[ARG1]], [[ARG2]] : i1351// CHECK-NEXT: bufferization.dealloc ([[ARG0]] : memref<2xi32>) if ([[NEW_COND]])352// CHECK-NEXT: return [[V0]]#0, [[V0]]#1, [[V0]]#0 :353 354// -----355 356func.func @dealloc_erase_empty(%arg0: memref<2xi32>, %arg1: i1, %arg2: memref<2xi32>) -> i1 {357 bufferization.dealloc358 %0 = bufferization.dealloc retain (%arg0 : memref<2xi32>)359 return %0 : i1360}361 362// CHECK-LABEL: func @dealloc_erase_empty363// CHECK-SAME: ([[ARG0:%.+]]: memref<2xi32>, [[ARG1:%.+]]: i1, [[ARG2:%.+]]: memref<2xi32>)364// CHECK-NEXT: [[FALSE:%.+]] = arith.constant false365// CHECK-NEXT: return [[FALSE]] :366 367// -----368 369func.func @dealloc_always_false_condition(%arg0: memref<2xi32>, %arg1: memref<2xi32>, %arg2: i1) {370 %false = arith.constant false371 bufferization.dealloc (%arg0, %arg1 : memref<2xi32>, memref<2xi32>) if (%false, %arg2)372 return373}374 375// CHECK-LABEL: func @dealloc_always_false_condition376// CHECK-SAME: ([[ARG0:%.+]]: memref<2xi32>, [[ARG1:%.+]]: memref<2xi32>, [[ARG2:%.+]]: i1)377// CHECK-NEXT: bufferization.dealloc ([[ARG1]] : {{.*}}) if ([[ARG2]])378// CHECK-NEXT: return379 380// -----381 382func.func @dealloc_base_memref_extract_of_alloc(%arg0: memref<2xi32>, %arg1: i1, %arg2: i1, %arg3: memref<2xi32>) -> memref<2xi32> {383 %alloc = memref.alloc() : memref<2xi32>384 %base0, %size0, %stride0, %offset0 = memref.extract_strided_metadata %alloc : memref<2xi32> -> memref<i32>, index, index, index385 %base1, %size1, %stride1, %offset1 = memref.extract_strided_metadata %arg3 : memref<2xi32> -> memref<i32>, index, index, index386 bufferization.dealloc (%base0, %arg0, %base1 : memref<i32>, memref<2xi32>, memref<i32>) if (%arg1, %arg2, %arg2)387 return %alloc : memref<2xi32>388}389 390// CHECK-LABEL: func @dealloc_base_memref_extract_of_alloc391// CHECK-SAME: ([[ARG0:%.+]]: memref<2xi32>, [[ARG1:%.+]]: i1, [[ARG2:%.+]]: i1, [[ARG3:%.+]]: memref<2xi32>)392// CHECK-NEXT: [[ALLOC:%.+]] = memref.alloc() : memref<2xi32>393// CHECK-NEXT: [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[ARG3]] :394// CHECK-NEXT: bufferization.dealloc ([[ALLOC]], [[ARG0]], [[BASE]] : memref<2xi32>, memref<2xi32>, memref<i32>) if ([[ARG1]], [[ARG2]], [[ARG2]])395// CHECK-NEXT: return396 397// -----398 399func.func @dealloc_base_memref_extract_of_alloc(%arg0: memref<2xi32>) {400 %true = arith.constant true401 %alloc = memref.alloc() : memref<2xi32>402 bufferization.dealloc (%alloc, %arg0 : memref<2xi32>, memref<2xi32>) if (%true, %true)403 return404}405 406// CHECK-LABEL: func @dealloc_base_memref_extract_of_alloc407// CHECK-SAME:([[ARG0:%.+]]: memref<2xi32>)408// CHECK-NOT: memref.alloc(409// CHECK: bufferization.dealloc ([[ARG0]] : memref<2xi32>) if (%true410 411// -----412 413// CHECK-LABEL: func @negative_input414func.func @negative_input() -> tensor<?x?x?xf16> {415 %idx27 = index.constant 27416 %idx-3 = index.constant -3 // negative integer?417 %c10 = arith.constant 10 : index418// CHECK: bufferization.alloc_tensor419// CHECK-SAME: tensor<10x?x27xf16>420 %11 = bufferization.alloc_tensor(%c10, %idx-3, %idx27) : tensor<?x?x?xf16>421 return %11 : tensor<?x?x?xf16>422}423