63 changes: 63 additions & 0 deletions mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,4 +68,67 @@ func @empty_func() -> () {
return
}

// -----

// CHECK-LABEL: func @read_after_write_conflict(
func @read_after_write_conflict(%cst : f32, %idx : index, %idx2 : index)
-> (f32, f32) {
// CHECK-DAG: %[[alloc:.*]] = memref.alloc
// CHECK-DAG: %[[dummy:.*]] = "test.dummy_op"
// CHECK-DAG: %[[dummy_m:.*]] = bufferization.to_memref %[[dummy]]
%t = "test.dummy_op"() : () -> (tensor<10xf32>)

// CHECK: memref.copy %[[dummy_m]], %[[alloc]]
// CHECK: memref.store %{{.*}}, %[[alloc]]
%write = tensor.insert %cst into %t[%idx2] : tensor<10xf32>

// CHECK: %[[read:.*]] = "test.some_use"(%[[dummy]])
%read = "test.some_use"(%t) : (tensor<10xf32>) -> (f32)
// CHECK: %[[read2:.*]] = memref.load %[[alloc]]
%read2 = tensor.extract %write[%idx] : tensor<10xf32>

// CHECK: memref.dealloc %[[alloc]]
// CHECK: return %[[read]], %[[read2]]
return %read, %read2 : f32, f32
}

// -----

// CHECK-LABEL: func @copy_deallocated(
func @copy_deallocated() -> tensor<10xf32> {
// CHECK: %[[alloc:.*]] = memref.alloc()
%0 = linalg.init_tensor[10] : tensor<10xf32>
// CHECK: %[[alloc_tensor:.*]] = bufferization.to_tensor %[[alloc]]
// CHECK: memref.dealloc %[[alloc]]
// CHECK: return %[[alloc_tensor]]
return %0 : tensor<10xf32>
}

// -----

// CHECK-LABEL: func @buffer_not_deallocated(
// CHECK-SAME: %[[t:.*]]: tensor<?xf32>
func @buffer_not_deallocated(%t : tensor<?xf32>, %c : i1) -> tensor<?xf32> {
// CHECK: %[[r:.*]] = scf.if %{{.*}} {
%r = scf.if %c -> tensor<?xf32> {
// CHECK: %[[some_op:.*]] = "test.some_op"
// CHECK: %[[alloc:.*]] = memref.alloc(%[[some_op]])
// CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
// CHECK-NOT: dealloc
// CHECK: scf.yield %[[casted]]
%sz = "test.some_op"() : () -> (index)
%0 = linalg.init_tensor[%sz] : tensor<?xf32>
scf.yield %0 : tensor<?xf32>
} else {
// CHECK: } else {
// CHECK: %[[m:.*]] = bufferization.to_memref %[[t]]
// CHECK: scf.yield %[[m]]
scf.yield %t : tensor<?xf32>
}
// CHECK: }
// CHECK-NOT: dealloc
// CHECK: %[[r_tensor:.*]] = bufferization.to_tensor %[[r]]
// CHECK: return %[[r_tensor]]
return %r : tensor<?xf32>
}

4 changes: 0 additions & 4 deletions mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,6 @@ func @not_inplace(
// CHECK: linalg.fill ins(%[[F0]] : f32) outs(%[[ALLOC]] : memref<?xf32>)
%r = linalg.fill ins(%f0 : f32) outs(%A : tensor<?xf32>) -> tensor<?xf32>

// CHECK: dealloc %[[ALLOC]] : memref<?xf32>
// CHECK: return %[[ALLOC]] : memref<?xf32>
return %r: tensor<?xf32>
}
Expand Down Expand Up @@ -292,7 +291,6 @@ func @insert_slice_fun_not_inplace(
// CHECK: memref.copy %[[A]], %[[ALLOC]] : memref<?xf32{{.*}} to memref<?xf32>
// CHECK: %[[SV:.*]] = memref.subview %[[ALLOC]][0] [4] [1] : memref<?xf32> to memref<4xf32>
// CHECK: memref.copy %[[t]], %[[SV]] : memref<4xf32, #map> to memref<4xf32>
// CHECK: memref.dealloc %[[ALLOC]] : memref<?xf32>
%r0 = tensor.insert_slice %t into %A[0][4][1] : tensor<4xf32> into tensor<?xf32>

// CHECK: return %{{.*}} : memref<?xf32>
Expand Down Expand Up @@ -329,7 +327,6 @@ func @scf_for_yield_only(%A : tensor<?xf32> {linalg.inplaceable = false},
scf.yield %t : tensor<?xf32>
}

// CHECK: memref.dealloc %[[ALLOC_FOR_A]] : memref<?xf32>
// CHECK: return %[[CASTED]] : memref<?xf32, #[[$map_1d_dyn]]>
return %r0, %r1: tensor<?xf32>, tensor<?xf32>
}
Expand Down Expand Up @@ -395,7 +392,6 @@ func @scf_for_with_tensor.insert_slice(
scf.yield %ttA, %ttB : tensor<?xf32>, tensor<?xf32>
}

// CHECK: memref.dealloc %[[ALLOC_FOR_A]] : memref<?xf32>
// CHECK: return %[[CASTED]] : memref<?xf32, #[[$map_1d_dyn]]>
return %r0#0, %r0#1: tensor<?xf32>, tensor<?xf32>
}
Expand Down