1// Note: Default is function-boundary-type-conversion=infer-layout-map 2// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs" -split-input-file | FileCheck %s 3 4// Run fuzzer with different seeds. 5// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null 6// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null 7// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null 8 9// Test bufferization using memref types that have no layout map. 10// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP 11 12// Test bufferization using memref types that have fully dynamic layout maps. 13// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs function-boundary-type-conversion=fully-dynamic-layout-map" -split-input-file | FileCheck %s --check-prefix=CHECK-FULLY-DYNAMIC-LAYOUT-MAP 14 15 16// Bufferization of bodiless function with no tensor return value. 17 18// CHECK: #[[$map0:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> 19// CHECK: #[[$map1:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)> 20// CHECK-LABEL: func private @private_func(memref<?xf32, 21// CHECK-SAME: #[[$map0]]>) 22// CHECK-NO-LAYOUT-MAP-LABEL: func private @private_func(memref<?xf32>) 23func.func private @private_func(tensor<?xf32>) -> () 24 25// CHECK-LABEL: func private @private_func_2d(memref<?x?xf32, 26// CHECK-SAME: #[[$map1]]>) 27// CHECK-NO-LAYOUT-MAP-LABEL: func private @private_func_2d(memref<?x?xf32>) 28func.func private @private_func_2d(tensor<?x?xf32>) -> () 29 30// CHECK-LABEL: func @empty_func() { 31// CHECK-NO-LAYOUT-MAP-LABEL: func @empty_func() { 32// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-LABEL: func @empty_func() { 33func.func @empty_func() -> () { 34 return 35} 36 37// ----- 38 39// A bodiless function that returns something that is not a tensor. 40 41// CHECK: func private @external_func_with_return_val(memref<4xi32, #{{.*}}>) -> f32 42// CHECK-FULLY-DYNAMIC-LAYOUT-MAP: #[[$map1:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> 43// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-LABEL: func private @external_func_with_return_val(memref<4xi32, 44// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-SAME: #[[$map1]]> 45func.func private @external_func_with_return_val(tensor<4xi32>) -> f32 46 47// ----- 48 49// A function that returns a non-equivalent tensor with layout map. 50 51// CHECK: #[[$map2:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 10 + s0 + d1)> 52// CHECK-LABEL: func @return_extract_slice(%{{.*}}) -> memref<2x?xf32, 53// CHECK-SAME: #[[$map2]]> { 54// CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<20x10xf32> 55// CHECK: %[[subview:.*]] = memref.subview {{.*}} : memref<20x10xf32> to memref<2x?xf32, #[[$map2]]> 56// CHECK: return %[[subview]] 57 58// CHECK-NO-LAYOUT-MAP: #[[$map2:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 10 + s0 + d1)> 59// CHECK-NO-LAYOUT-MAP-LABEL: func @return_extract_slice(%{{.*}}) -> memref<2x?xf32> 60// CHECK-NO-LAYOUT-MAP: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<20x10xf32> 61// CHECK-NO-LAYOUT-MAP: %[[subview:.*]] = memref.subview {{.*}} : memref<20x10xf32> to memref<2x?xf32, #[[$map2]]> 62// CHECK-NO-LAYOUT-MAP: %[[alloc_no_layout:.*]] = memref.alloc(%{{.*}}) : memref<2x?xf32> 63// CHECK-NO-LAYOUT-MAP: memref.copy %[[subview]], %[[alloc_no_layout]] 64// CHECK-NO-LAYOUT-MAP: memref.dealloc %[[alloc]] 65// CHECK-NO-LAYOUT-MAP: return %[[alloc_no_layout]] 66 67// CHECK-FULLY-DYNAMIC-LAYOUT-MAP: #[[$map2a:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)> 68// CHECK-FULLY-DYNAMIC-LAYOUT-MAP: #[[$map2b:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 10 + s0 + d1)> 69// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-LABEL: func @return_extract_slice(%{{.*}}) -> memref<2x?xf32, 70// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-SAME: #[[$map2a]]> { 71func.func @return_extract_slice(%idx: index, %sz: index) -> (tensor<2x?xf32>) 72{ 73 %t = bufferization.alloc_tensor() : tensor<20x10xf32> 74 %0 = tensor.extract_slice %t[%idx, %idx][2, %sz][1, 1] 75 : tensor<20x10xf32> to tensor<2x?xf32> 76 return %0 : tensor<2x?xf32> 77} 78 79// ----- 80 81// CHECK-LABEL: func private @private_func 82func.func private @private_func(tensor<?xf32>) -> (f32) 83 84// private_func may modify the buffer arg, but that's OK because %t is writable. 85// No alloc/copy should be inserted. 86 87// CHECK-LABEL: func @main( 88// CHECK-SAME: %[[t:.*]]: memref<?xf32 89// CHECK-NOT: alloc 90// CHECK-NOT: copy 91// CHECK: call @private_func(%[[t]]) 92func.func @main(%t: tensor<?xf32> {bufferization.writable = true}) -> (f32) { 93 %0 = call @private_func(%t) : (tensor<?xf32>) -> (f32) 94 return %0 : f32 95} 96 97// ----- 98 99// CHECK-LABEL: func private @private_func 100func.func private @private_func(tensor<?xf32>) -> (f32) 101 102// private_func may modify the buffer arg, %t is not writable. A copy is needed. 103 104// CHECK-LABEL: func @main( 105// CHECK-SAME: %[[t:.*]]: memref<?xf32 106// CHECK: %[[alloc:.*]] = memref.alloc 107// CHECK-DAG: memref.copy %[[t]], %[[alloc]] 108// CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]] 109// CHECK: call @private_func(%[[casted]]) 110// CHECK: memref.dealloc %[[alloc]] 111func.func @main(%t: tensor<?xf32> {bufferization.writable = false}) -> (f32) { 112 %0 = call @private_func(%t) : (tensor<?xf32>) -> (f32) 113 return %0 : f32 114} 115 116// ----- 117 118// Test bufferization of a function without tensor args. 119 120// CHECK-LABEL: func @func_without_tensor_args 121func.func @func_without_tensor_args(%v : vector<10xf32>) -> () { 122 // CHECK: %[[alloc:.*]] = memref.alloc() 123 %0 = bufferization.alloc_tensor() : tensor<10xf32> 124 125 %c0 = arith.constant 0 : index 126 // CHECK: vector.transfer_write %{{.*}}, %[[alloc]] 127 %1 = vector.transfer_write %v, %0[%c0] : vector<10xf32>, tensor<10xf32> 128 129 %cst = arith.constant 0.0 : f32 130 // CHECK: vector.transfer_read %[[alloc]] 131 %r = vector.transfer_read %1[%c0], %cst : tensor<10xf32>, vector<11xf32> 132 133 vector.print %r : vector<11xf32> 134 return 135} 136 137// ----- 138 139// Bufferization of a function that is reading and writing. %t0 is writable, so 140// no copy should be inserted. 141 142// CHECK-LABEL: func @inner_func( 143// CHECK-SAME: %[[arg0:.*]]: memref<?xf32 144func.func @inner_func(%t: tensor<?xf32>) -> (tensor<?xf32>, f32) { 145 // CHECK-NOT: copy 146 %f = arith.constant 1.0 : f32 147 %c0 = arith.constant 0 : index 148 %c1 = arith.constant 1 : index 149 // CHECK: memref.store %{{.*}}, %[[arg0]] 150 %0 = tensor.insert %f into %t[%c0] : tensor<?xf32> 151 // CHECK: %[[load:.*]] = memref.load %[[arg0]] 152 %1 = tensor.extract %0[%c1] : tensor<?xf32> 153 // CHECK: return %[[load]] : f32 154 return %0, %1 : tensor<?xf32>, f32 155} 156 157// CHECK-LABEL: func @call_func_with_non_tensor_return( 158// CHECK-SAME: %[[arg0:.*]]: memref<?xf32 159func.func @call_func_with_non_tensor_return( 160 %t0: tensor<?xf32> {bufferization.writable = true}) -> (f32, tensor<?xf32>) { 161 // CHECK-NOT: alloc 162 // CHECK-NOT: copy 163 // CHECK: %[[call:.*]] = call @inner_func(%[[arg0]]) 164 %0, %1 = call @inner_func(%t0) : (tensor<?xf32>) -> (tensor<?xf32>, f32) 165 // CHECK: return %[[call]] : f32 166 return %1, %0 : f32, tensor<?xf32> 167} 168 169// ----- 170 171// Bufferization of a function that is reading and writing. %t0 is not writable, 172// so a copy is needed. 173 174// CHECK-LABEL: func @inner_func( 175// CHECK-SAME: %[[arg0:.*]]: memref<?xf32 176func.func @inner_func(%t: tensor<?xf32>) -> (tensor<?xf32>, f32) { 177 // CHECK-NOT: copy 178 %f = arith.constant 1.0 : f32 179 %c0 = arith.constant 0 : index 180 %c1 = arith.constant 1 : index 181 // CHECK: memref.store %{{.*}}, %[[arg0]] 182 %0 = tensor.insert %f into %t[%c0] : tensor<?xf32> 183 // CHECK: %[[load:.*]] = memref.load %[[arg0]] 184 %1 = tensor.extract %0[%c1] : tensor<?xf32> 185 // CHECK: return %[[load]] : f32 186 return %0, %1 : tensor<?xf32>, f32 187} 188 189// CHECK-LABEL: func @call_func_with_non_tensor_return( 190// CHECK-SAME: %[[arg0:.*]]: memref<?xf32 191func.func @call_func_with_non_tensor_return( 192 %t0: tensor<?xf32> {bufferization.writable = false}) -> (f32, tensor<?xf32>) { 193 // CHECK: %[[alloc:.*]] = memref.alloc 194 // CHECK-DAG: memref.copy %[[arg0]], %[[alloc]] 195 // CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]] 196 // CHECK: %[[call:.*]] = call @inner_func(%[[casted]]) 197 %0, %1 = call @inner_func(%t0) : (tensor<?xf32>) -> (tensor<?xf32>, f32) 198 199 // Note: The tensor return value has folded away. 200 // CHECK: return %[[call]] : f32 201 return %1, %0 : f32, tensor<?xf32> 202} 203 204// ----- 205 206// A chain of function calls. The last function f0 is potentially writing to the 207// buffer. This becomes a problem when bufferizing main and a copy must be 208// inserted then. (No copies in the other functions.) 209 210// CHECK-LABEL: func private @f0( 211func.func private @f0(tensor<?xf32>) -> (f32) 212 213// CHECK-LABEL: func @f1( 214// CHECK-SAME: %[[t1:.*]]: memref<?xf32 215// CHECK: %[[r1:.*]] = call @f0(%[[t1]]) 216// CHECK: return %[[r1]] 217func.func @f1(%t: tensor<?xf32>) -> (f32) { 218 %0 = call @f0(%t) : (tensor<?xf32>) -> (f32) 219 return %0 : f32 220} 221 222// CHECK-LABEL: func @f2( 223// CHECK-SAME: %[[t2:.*]]: memref<?xf32 224// CHECK: %[[r2:.*]] = call @f1(%[[t2]]) 225// CHECK: return %[[r2]] 226func.func @f2(%t: tensor<?xf32>) -> (f32) { 227 %0 = call @f1(%t) : (tensor<?xf32>) -> (f32) 228 return %0 : f32 229} 230 231// CHECK-LABEL: func @main( 232// CHECK-SAME: %[[t3:.*]]: memref<?xf32 233// CHECK: %[[alloc:.*]] = memref.alloc 234// CHECK-DAG: memref.copy %[[t3]], %[[alloc]] 235// CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]] 236// CHECK: call @f2(%[[casted]]) 237// CHECK: memref.dealloc %[[alloc]] 238func.func @main(%t: tensor<?xf32> {bufferization.writable = false}) -> (f32) { 239 %0 = call @f2(%t) : (tensor<?xf32>) -> (f32) 240 return %0 : f32 241} 242 243// ----- 244 245// This function does not read, just write. We need an alloc, but no copy. 246 247// CHECK-LABEL: func @does_not_read( 248// CHECK-NOT: alloc 249// CHECK-NOT: copy 250func.func @does_not_read(%t: tensor<?xf32>) -> tensor<?xf32> { 251 %f0 = arith.constant 0.0 : f32 252 %r = linalg.fill ins(%f0 : f32) outs(%t : tensor<?xf32>) -> tensor<?xf32> 253 return %r : tensor<?xf32> 254} 255 256// CHECK-LABEL: func @main( 257// CHECK-SAME: %[[t:.*]]: memref<?xf32 258// CHECK: %[[alloc:.*]] = memref.alloc 259// CHECK-NOT: copy 260// CHECK: %[[casted:.*]] = memref.cast %[[alloc]] 261// CHECK-NOT: copy 262// CHECK: call @does_not_read(%[[casted]]) 263// CHECK: %[[r:.*]] = memref.load %[[alloc]] 264// CHECK: memref.dealloc %[[alloc]] 265func.func @main(%t: tensor<?xf32> {bufferization.writable = false}) -> f32 { 266 %0 = call @does_not_read(%t) : (tensor<?xf32>) -> (tensor<?xf32>) 267 %idx = arith.constant 4 : index 268 %r = tensor.extract %0[%idx] : tensor<?xf32> 269 return %r : f32 270} 271 272// ----- 273 274// Alloc and copy must be inserted because the arith.constant is read-only. 275 276// CHECK: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> 277 278// CHECK: memref.global "private" constant @__constant_4xi32 : memref<4xi32> = dense<[1, 2, 3, 4]> 279// CHECK: func private @some_external_func(memref<4xi32, #[[$DYN_1D_MAP]]>) 280func.func private @some_external_func(tensor<4xi32>) 281 282// CHECK: func @main() 283func.func @main() { 284// CHECK-DAG: %[[A:.*]] = memref.get_global @__constant_4xi32 : memref<4xi32> 285 %A = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32> 286 287// CHECK-DAG: %[[alloc:.*]] = memref.alloc 288// CHECK-DAG: %[[B:.*]] = memref.cast %[[alloc]] : memref<4xi32> to memref<4xi32, #[[$DYN_1D_MAP]]> 289// CHECK-DAG: memref.copy %[[A]], %[[alloc]] 290// CHECK: call @some_external_func(%[[B]]) : (memref<4xi32, #[[$DYN_1D_MAP]]>) -> () 291 call @some_external_func(%A) : (tensor<4xi32>) -> () 292 293// CHECK: memref.dealloc %[[alloc]] 294 return 295} 296 297// ----- 298 299// Alloc and copy must be inserted because the arith.constant is read-only. The 300// function call is inside of an scf.execute_region. 301 302// CHECK: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> 303 304// CHECK: memref.global "private" constant @__constant_4xi32 : memref<4xi32> = dense<[1, 2, 3, 4]> 305// CHECK: func private @some_external_func_within_scf_execute(memref<4xi32, #[[$DYN_1D_MAP]]>) 306func.func private @some_external_func_within_scf_execute(tensor<4xi32>) 307 308// CHECK: func @main() 309func.func @main() { 310// CHECK-DAG: %[[A:.*]] = memref.get_global @__constant_4xi32 : memref<4xi32> 311 %A = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32> 312 313// Note: The scf.execute_region canonicalizes away. 314 315// CHECK-DAG: %[[alloc:.*]] = memref.alloc 316// CHECK-DAG: %[[B:.*]] = memref.cast %[[alloc]] : memref<4xi32> to memref<4xi32, #[[$DYN_1D_MAP]]> 317// CHECK-DAG: memref.copy %[[A]], %[[alloc]] 318// CHECK: call @some_external_func_within_scf_execute(%[[B]]) : (memref<4xi32, #[[$DYN_1D_MAP]]>) -> () 319 scf.execute_region { 320 func.call @some_external_func_within_scf_execute(%A) : (tensor<4xi32>) -> () 321 scf.yield 322 } 323 324// CHECK: memref.dealloc %[[alloc]] 325 return 326} 327 328// ----- 329 330// A write inside an scf.execute_region. An equivalent tensor is yielded. 331 332// CHECK-LABEL: func @execute_region_test( 333// CHECK-SAME: %[[m1:.*]]: memref<?xf32 334func.func @execute_region_test(%t1 : tensor<?xf32>) 335 -> (f32, tensor<?xf32>, f32) 336{ 337 %f1 = arith.constant 0.0 : f32 338 %f2 = arith.constant 1.0 : f32 339 %idx = arith.constant 7 : index 340 341 // scf.execute_region is canonicalized away after bufferization. So just the 342 // memref.store is left over. 343 344 // CHECK-NOT: alloc 345 // CHECK-NOT: copy 346 // CHECK: memref.store %{{.*}}, %[[m1]][%{{.*}}] 347 %0, %1, %2 = scf.execute_region -> (f32, tensor<?xf32>, f32) { 348 %t2 = tensor.insert %f2 into %t1[%idx] : tensor<?xf32> 349 scf.yield %f1, %t2, %f2 : f32, tensor<?xf32>, f32 350 } 351 352 // CHECK: return %{{.*}}, %{{.*}} : f32, f32 353 return %0, %1, %2 : f32, tensor<?xf32>, f32 354} 355 356// ----- 357 358// CHECK: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> 359 360// CHECK: func private @some_external_func(memref<?xf32, #[[$DYN_1D_MAP]]>) 361func.func private @some_external_func(tensor<?xf32>) 362 363// CHECK: func @scf_for_with_tensor_insert_slice( 364// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$DYN_1D_MAP]]> 365// CHECK-SAME: %[[B:[a-zA-Z0-9]*]]: memref<?xf32, #[[$DYN_1D_MAP]]> 366// CHECK-SAME: %[[C:[a-zA-Z0-9]*]]: memref<4xf32, #[[$DYN_1D_MAP]]> 367func.func @scf_for_with_tensor_insert_slice( 368 %A : tensor<?xf32>, %B : tensor<?xf32>, %C : tensor<4xf32>, 369 %lb : index, %ub : index, %step : index) 370 -> (tensor<?xf32>, tensor<?xf32>) 371{ 372 // CHECK-NEXT: scf.for 373 %r0:2 = scf.for %i = %lb to %ub step %step iter_args(%tA = %A, %tB = %B) 374 -> (tensor<?xf32>, tensor<?xf32>) 375 { 376 // CHECK-NEXT: %[[SVA:.*]] = memref.subview %[[A]] 377 // CHECK-NEXT: memref.copy %[[C]], %[[SVA]] : memref<4xf32, #[[$DYN_1D_MAP]]> to memref<4xf32, #[[$DYN_1D_MAP]]> 378 %ttA = tensor.insert_slice %C into %tA[%i][4][1] : tensor<4xf32> into tensor<?xf32> 379 380 // CHECK-NEXT: %[[SVB:.*]] = memref.subview %[[B]] 381 // CHECK-NEXT: memref.copy %[[C]], %[[SVB]] : memref<4xf32, #[[$DYN_1D_MAP]]> to memref<4xf32, #[[$DYN_1D_MAP]]> 382 %ttB = tensor.insert_slice %C into %tB[%i][4][1] : tensor<4xf32> into tensor<?xf32> 383 384 // scf.yield is empty and is elided 385 // CHECK-NOT: scf.yield 386 scf.yield %ttA, %ttB : tensor<?xf32>, tensor<?xf32> 387 } 388 389 // Swaparoo requires bufferizing the whole function to figure out who's who. 390 return %r0#1, %r0#0: tensor<?xf32>, tensor<?xf32> 391} 392 393// CHECK: func @bar( 394// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$DYN_1D_MAP]]> 395// CHECK-SAME: %[[B:[a-zA-Z0-9]*]]: memref<?xf32, #[[$DYN_1D_MAP]]> 396// CHECK-SAME: %[[C:[a-zA-Z0-9]*]]: memref<4xf32, #[[$DYN_1D_MAP]]> 397func.func @bar( 398 %A : tensor<?xf32> {bufferization.writable = true}, 399 %B : tensor<?xf32> {bufferization.writable = true}, 400 %C : tensor<4xf32> {bufferization.writable = true}, 401 %lb : index, %ub : index, %step : index) 402 -> (tensor<?xf32>, tensor<?xf32>) 403{ 404// CHECK-DAG: call @scf_for_with_tensor_insert_slice(%[[A]], %[[B]], %[[C]] 405 %r0:2 = call @scf_for_with_tensor_insert_slice(%A, %B, %C, %lb, %ub, %step) : 406 (tensor<?xf32>, tensor<?xf32>, tensor<4xf32>, index, index, index) 407 -> (tensor<?xf32>, tensor<?xf32>) 408 409 // %r0#0 requires a copy because we have no idea what the function is doing. 410// CHECK-DAG: %[[alloc:.*]] = memref.alloc 411// CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]] 412// CHECK-DAG: memref.copy %[[B]], %[[alloc]] 413// CHECK-NEXT: call @some_external_func(%[[casted]]) : (memref<?xf32, #[[$DYN_1D_MAP]]>) -> () 414 call @some_external_func(%r0#0) : (tensor<?xf32>) -> () 415 416// CHECK: return 417 return %r0#0, %r0#1: tensor<?xf32>, tensor<?xf32> 418} 419 420// ----- 421 422// CHECK-DAG: #[[$DYN_0D_MAP:.*]] = affine_map<()[s0] -> (s0)> 423// CHECK-DAG: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> 424 425// CHECK: func @init_and_dot( 426// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<64xf32, #[[$DYN_1D_MAP]]> 427// CHECK-SAME: %[[B:[a-zA-Z0-9]*]]: memref<64xf32, #[[$DYN_1D_MAP]]> 428// CHECK-SAME: %[[C:[a-zA-Z0-9]*]]: memref<f32, #[[$DYN_0D_MAP]]> 429func.func @init_and_dot(%a: tensor<64xf32>, %b: tensor<64xf32>, %c: tensor<f32>) -> tensor<f32> { 430 // CHECK-NEXT: %[[C0:.*]] = arith.constant 0{{.*}} : f32 431 %v0 = arith.constant 0.0 : f32 432 433 // CHECK-NEXT: linalg.fill ins(%[[C0]] : f32) outs(%[[C]] : memref<f32, #[[$DYN_0D_MAP]]>) 434 %d = linalg.fill ins(%v0 : f32) outs(%c : tensor<f32>) -> tensor<f32> 435 436 // CHECK-NEXT: linalg.dot ins(%[[A]], %[[B]] : memref<64xf32, #[[$DYN_1D_MAP]]>, memref<64xf32, #[[$DYN_1D_MAP]]>) outs(%[[C]] : memref<f32, #[[$DYN_0D_MAP]]>) 437 %e = linalg.dot ins(%a, %b : tensor<64xf32>,tensor<64xf32>) 438 outs(%d: tensor<f32>) -> tensor<f32> 439 440 // CHECK-NEXT: return 441 return %e : tensor<f32> 442} 443 444// CHECK: func @main() 445func.func @main() { 446 // CHECK-DAG: %[[C0:.*]] = arith.constant 0{{.*}} : f32 447 // CHECK-DAG: %[[C1:.*]] = arith.constant 1{{.*}} : f32 448 // CHECK-DAG: %[[C2:.*]] = arith.constant 2{{.*}} : f32 449 %v0 = arith.constant 0.0 : f32 450 %v1 = arith.constant 1.0 : f32 451 %v2 = arith.constant 2.0 : f32 452 453 // CHECK-NEXT: %[[A:.*]] = memref.alloc() {alignment = 128 : i64} : memref<64xf32> 454 // CHECK-NEXT: %[[B:.*]] = memref.alloc() {alignment = 128 : i64} : memref<64xf32> 455 // CHECK-NEXT: %[[C:.*]] = memref.alloc() {alignment = 128 : i64} : memref<f32> 456 // CHECK-DAG: %[[cA:.*]] = memref.cast %[[A]] : memref<64xf32> to memref<64xf32, #[[$DYN_1D_MAP]]> 457 // CHECK-DAG: %[[cB:.*]] = memref.cast %[[B]] : memref<64xf32> to memref<64xf32, #[[$DYN_1D_MAP]]> 458 // CHECK-DAG: %[[cC:.*]] = memref.cast %[[C]] : memref<f32> to memref<f32, #[[$DYN_0D_MAP]]> 459 %A = bufferization.alloc_tensor() : tensor<64xf32> 460 %B = bufferization.alloc_tensor() : tensor<64xf32> 461 %C = bufferization.alloc_tensor() : tensor<f32> 462 463 // CHECK-DAG: linalg.fill ins(%[[C1]] : f32) outs(%[[A]] : memref<64xf32>) 464 // CHECK-DAG: linalg.fill ins(%[[C2]] : f32) outs(%[[B]] : memref<64xf32>) 465 // CHECK-DAG: linalg.fill ins(%[[C0]] : f32) outs(%[[C]] : memref<f32>) 466 %AA = linalg.fill ins(%v1 : f32) outs(%A : tensor<64xf32>) -> tensor<64xf32> 467 %BB = linalg.fill ins(%v2 : f32) outs(%B : tensor<64xf32>) -> tensor<64xf32> 468 %CC = linalg.fill ins(%v0 : f32) outs(%C : tensor<f32>) -> tensor<f32> 469 470 // CHECK-NEXT: call @init_and_dot(%[[cA]], %[[cB]], %[[cC]]) 471 %res = call @init_and_dot(%AA, %BB, %CC) : 472 (tensor<64xf32>, tensor<64xf32>, tensor<f32>) -> tensor<f32> 473 474 // CHECK-NEXT: %[[dC:.*]] = memref.cast %[[C]] : memref<f32> to memref<*xf32> 475 %res2 = tensor.cast %res: tensor<f32> to tensor<*xf32> 476 477 // CHECK-NEXT: call @printMemrefF32(%[[dC]]) : (memref<*xf32>) -> () 478 call @printMemrefF32(%res2) : (tensor<*xf32>) -> () 479 480 // CHECK-DAG: memref.dealloc %[[A]] : memref<64xf32> 481 // CHECK-DAG: memref.dealloc %[[B]] : memref<64xf32> 482 // CHECK-DAG: memref.dealloc %[[C]] : memref<f32> 483 // CHECK-NEXT: return 484 return 485} 486 487// CHECK: func private @printMemrefF32(memref<*xf32>) 488func.func private @printMemrefF32(tensor<*xf32>) 489 490// ----- 491 492// CHECK: #[[$DYNAMIC:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> 493 494// CHECK: func private @external_func(memref<?xf32, #[[$DYNAMIC]]>) 495func.func private @external_func(tensor<?xf32>) 496 497// CHECK: func @callee( 498// CHECK-SAME: %[[A:[0-9a-zA-Z]*]]: memref<?xf32> 499// CHECK-SAME: %[[B:[0-9a-zA-Z]*]]: memref<?xf32, #[[$DYNAMIC]]> 500// CHECK-SAME: %[[C:[0-9a-zA-Z]*]]: memref<?xf32, #[[$DYNAMIC]]> 501func.func @callee( 502 %A : tensor<?xf32> {bufferization.buffer_layout = affine_map<(i)[s0, s1] -> (i)>}, 503 %B : tensor<?xf32>, 504 %C : tensor<?xf32>) { 505// CHECK-NEXT: %[[CASTED:.*]] = memref.cast %[[A]] : memref<?xf32> to memref<?xf32, #[[$DYNAMIC]]> 506// CHECK-NEXT: call @external_func(%[[CASTED]]) : (memref<?xf32, #[[$DYNAMIC]]>) -> () 507 call @external_func(%A) : (tensor<?xf32>) -> () 508 509// CHECK-NEXT: call @external_func(%[[B]]) : (memref<?xf32, #[[$DYNAMIC]]>) -> () 510 call @external_func(%B) : (tensor<?xf32>) -> () 511 512// CHECK-NEXT: call @external_func(%[[C]]) : (memref<?xf32, #[[$DYNAMIC]]>) -> () 513 call @external_func(%C) : (tensor<?xf32>) -> () 514 515 return 516} 517 518// CHECK: func @entry( 519// CHECK-SAME: %[[A:[0-9a-zA-Z]*]]: memref<?xf32> 520// CHECK-SAME: %[[B:[0-9a-zA-Z]*]]: memref<?xf32> 521// CHECK-SAME: %[[C:[0-9a-zA-Z]*]]: memref<?xf32, #[[$DYNAMIC]]> 522func.func @entry(%A : tensor<?xf32> {bufferization.buffer_layout = affine_map<(i)[s0, s1] -> (i)>, bufferization.writable = false}, 523 %B : tensor<?xf32> {bufferization.buffer_layout = affine_map<(i)[s0, s1] -> (i)>, bufferization.writable = false}, 524 %C : tensor<?xf32> {bufferization.writable = false}) { 525// Note: `callee` does not write to its bbArg directly, but `external_func` 526// does. Inside `callee`, the writes via `external_func` do not cause a 527// conflict. However, inside `entry`, the writes do cause a conflict because 528// %A, %B and %C are not inplaceable. This test case shows that this kind of 529// conflict detection has a "transitive" nature. 530// CHECK-DAG: %[[ALLOC_C:.*]] = memref.alloc 531// CHECK-DAG: %[[CASTED_C:.*]] = memref.cast %[[ALLOC_C]] 532// CHECK-DAG: %[[ALLOC_B:.*]] = memref.alloc 533// CHECK-DAG: %[[CASTED_B:.*]] = memref.cast %[[ALLOC_B]] 534// CHECK-DAG: %[[ALLOC_A:.*]] = memref.alloc 535// CHECK-DAG: %[[CASTED_A:.*]] = memref.cast %[[ALLOC_A]] 536// CHECK-DAG: memref.copy %[[A]], %[[ALLOC_A]] 537// CHECK-DAG: memref.copy %[[B]], %[[ALLOC_B]] 538// CHECK-DAG: memref.copy %[[C]], %[[ALLOC_C]] 539// CHECK-NEXT: call @callee(%[[CASTED_A]], %[[CASTED_B]], %[[CASTED_C]]) 540 call @callee(%A, %B, %C) : (tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) -> () 541 return 542} 543 544// ----- 545 546// No alloc or copy inside of the loop. 547 548// CHECK-LABEL: func @inner_func( 549// CHECK-SAME: %[[arg0:.*]]: memref<?xf32 550func.func @inner_func(%t: tensor<?xf32>) -> tensor<?xf32> { 551 %f = arith.constant 1.0 : f32 552 %c0 = arith.constant 0 : index 553 // CHECK: memref.store %{{.*}}, %[[arg0]] 554 %0 = tensor.insert %f into %t[%c0] : tensor<?xf32> 555 return %0 : tensor<?xf32> 556} 557 558// CHECK-LABEL: func @equivalent_func_arg( 559// CHECK-SAME: %[[arg0:.*]]: memref<?xf32 560func.func @equivalent_func_arg(%t0: tensor<?xf32> {bufferization.writable = true}, 561 %c0: index, %c10: index, %c1: index) -> tensor<?xf32> { 562 // CHECK-NOT: alloc 563 // CHECK-NOT: copy 564 %1 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%t1 = %t0) -> (tensor<?xf32>) { 565 // CHECK: call @inner_func(%[[arg0]]) 566 %3 = func.call @inner_func(%t1) : (tensor<?xf32>) -> tensor<?xf32> 567 scf.yield %3 : tensor<?xf32> 568 } 569 return %1: tensor<?xf32> 570} 571 572// ----- 573 574// inner_func_2 modifies the bbArg, but the loop yields the original value. A 575// buffer copy must be inserted inside the loop. 576 577// CHECK-LABEL: func @inner_func_2( 578// CHECK-SAME: %[[arg0:.*]]: memref<?xf32 579func.func @inner_func_2(%t: tensor<?xf32>) -> tensor<?xf32> { 580 %f = arith.constant 1.0 : f32 581 %c0 = arith.constant 0 : index 582 // CHECK: memref.store %{{.*}}, %[[arg0]] 583 %0 = tensor.insert %f into %t[%c0] : tensor<?xf32> 584 return %0 : tensor<?xf32> 585} 586 587// CHECK-LABEL: func @equivalent_func_arg_2( 588// CHECK-SAME: %[[arg0:.*]]: memref<?xf32 589func.func @equivalent_func_arg_2(%t0: tensor<?xf32> {bufferization.writable = true}, 590 %c0: index, %c10: index, %c1: index) -> tensor<?xf32> { 591 // CHECK: scf.for {{.*}} { 592 %1 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%t1 = %t0) -> (tensor<?xf32>) { 593 // CHECK: %[[alloc:.*]] = memref.alloc 594 // CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]] 595 // CHECK-DAG: memref.copy %[[arg0]], %[[alloc]] 596 // CHECK: call @inner_func_2(%[[casted]]) 597 // CHECK: memref.dealloc %[[alloc]] 598 // CHECK-NOT: scf.yield 599 %3 = func.call @inner_func_2(%t1) : (tensor<?xf32>) -> tensor<?xf32> 600 scf.yield %t1 : tensor<?xf32> 601 } 602 return %1: tensor<?xf32> 603} 604 605// ----- 606 607// Bufferize without fully dynamic layout maps. 608 609// CHECK-LABEL: func @transfer_read(%{{.*}}: memref<?xf32, #map>) -> vector<4xf32> { 610// CHECK-NO-LAYOUT-MAP-LABEL: func @transfer_read(%{{.*}}: memref<?xf32>) -> vector<4xf32> 611func.func @transfer_read( 612 %A : tensor<?xf32> {bufferization.writable = false}) 613 -> (vector<4xf32>) 614{ 615 %c0 = arith.constant 0 : index 616 %f0 = arith.constant 0.0 : f32 617 618// CHECK: %[[RES:.*]] = vector.transfer_read {{.*}} : memref<?xf32, #{{.*}}>, vector<4xf32> 619 %0 = vector.transfer_read %A[%c0], %f0 : tensor<?xf32>, vector<4xf32> 620 621// CHECK: return %[[RES]] : vector<4xf32> 622 return %0 : vector<4xf32> 623} 624