// RUN: mlir-opt -allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation \
// RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
// RUN: mlir-opt -allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null
// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file --verify-diagnostics > /dev/null
// Test Case: Nested regions - This test defines a BufferBasedOp inside the
// region of a RegionBufferBasedOp.
// BufferDeallocation expected behavior: The AllocOp for the BufferBasedOp
// should remain inside the region of the RegionBufferBasedOp and it should insert
// the missing DeallocOp in the same region. The missing DeallocOp should be
// inserted after CopyOp.
func.func @nested_regions_and_cond_branch(
%arg0: i1,
%arg1: memref<2xf32>,
%arg2: memref<2xf32>) {
cf.cond_br %arg0, ^bb1, ^bb2
^bb1:
cf.br ^bb3(%arg1 : memref<2xf32>)
^bb2:
%0 = memref.alloc() : memref<2xf32>
test.region_buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>) {
^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
%1 = memref.alloc() : memref<2xf32>
test.buffer_based in(%arg1: memref<2xf32>) out(%1: memref<2xf32>)
%tmp1 = math.exp %gen1_arg0 : f32
test.region_yield %tmp1 : f32
}
cf.br ^bb3(%0 : memref<2xf32>)
^bb3(%1: memref<2xf32>):
test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
return
}
// CHECK-LABEL: func @nested_regions_and_cond_branch
// CHECK-SAME: ([[ARG0:%.+]]: i1, [[ARG1:%.+]]: memref<2xf32>, [[ARG2:%.+]]: memref<2xf32>)
// CHECK: ^bb1:
// CHECK-NOT: bufferization.clone
// CHECK-NOT: bufferization.dealloc
// CHECK: cf.br ^bb3([[ARG1]], %false
// CHECK: ^bb2:
// CHECK: [[ALLOC0:%.+]] = memref.alloc()
// CHECK: test.region_buffer_based
// CHECK: [[ALLOC1:%.+]] = memref.alloc()
// CHECK: test.buffer_based
// CHECK: bufferization.dealloc ([[ALLOC1]] : memref<2xf32>) if (%true
// CHECK-NEXT: test.region_yield
// CHECK-NOT: bufferization.clone
// CHECK-NOT: bufferization.dealloc
// CHECK: cf.br ^bb3([[ALLOC0]], %true
// CHECK: ^bb3([[A0:%.+]]: memref<2xf32>, [[COND0:%.+]]: i1):
// CHECK: test.copy
// CHECK-NEXT: [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[A0]]
// CHECK-NEXT: bufferization.dealloc ([[BASE]] : {{.*}}) if ([[COND0]])
// CHECK: return
// -----
// Test Case: nested region control flow
// The alloc %1 flows through both if branches until it is finally returned.
// Hence, it does not require a specific dealloc operation. However, %3
// requires a dealloc.
func.func @nested_region_control_flow(
%arg0 : index,
%arg1 : index) -> memref<?x?xf32> {
%0 = arith.cmpi eq, %arg0, %arg1 : index
%1 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
%2 = scf.if %0 -> (memref<?x?xf32>) {
scf.yield %1 : memref<?x?xf32>
} else {
%3 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
"test.read_buffer"(%3) : (memref<?x?xf32>) -> ()
scf.yield %1 : memref<?x?xf32>
}
return %2 : memref<?x?xf32>
}
// CHECK-LABEL: func @nested_region_control_flow
// CHECK: [[ALLOC:%.+]] = memref.alloc(
// CHECK: [[V0:%.+]]:2 = scf.if
// CHECK: scf.yield [[ALLOC]], %false
// CHECK: [[ALLOC1:%.+]] = memref.alloc(
// CHECK: bufferization.dealloc ([[ALLOC1]] :{{.*}}) if (%true{{[0-9_]*}})
// CHECK-NOT: retain
// CHECK: scf.yield [[ALLOC]], %false
// CHECK: [[V1:%.+]] = scf.if [[V0]]#1
// CHECK: scf.yield [[V0]]#0
// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0
// CHECK: scf.yield [[CLONE]]
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] : {{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] :
// CHECK: return [[V1]]
// -----
// Test Case: nested region control flow with a nested buffer allocation in a
// divergent branch.
// Buffer deallocation places a copy for both %1 and %3, since they are
// returned in the end.
func.func @nested_region_control_flow_div(
%arg0 : index,
%arg1 : index) -> memref<?x?xf32> {
%0 = arith.cmpi eq, %arg0, %arg1 : index
%1 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
%2 = scf.if %0 -> (memref<?x?xf32>) {
scf.yield %1 : memref<?x?xf32>
} else {
%3 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
scf.yield %3 : memref<?x?xf32>
}
return %2 : memref<?x?xf32>
}
// CHECK-LABEL: func @nested_region_control_flow_div
// CHECK: [[ALLOC:%.+]] = memref.alloc(
// CHECK: [[V0:%.+]]:2 = scf.if
// CHECK: scf.yield [[ALLOC]], %false
// CHECK: [[ALLOC1:%.+]] = memref.alloc(
// CHECK: scf.yield [[ALLOC1]], %true
// CHECK: [[V1:%.+]] = scf.if [[V0]]#1
// CHECK: scf.yield [[V0]]#0
// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0
// CHECK: scf.yield [[CLONE]]
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] :
// CHECK: return [[V1]]
// -----
// Test Case: nested region control flow within a region interface.
// No copies are required in this case since the allocation finally escapes
// the method.
func.func @inner_region_control_flow(%arg0 : index) -> memref<?x?xf32> {
%0 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
%1 = test.region_if %0 : memref<?x?xf32> -> (memref<?x?xf32>) then {
^bb0(%arg1 : memref<?x?xf32>):
test.region_if_yield %arg1 : memref<?x?xf32>
} else {
^bb0(%arg1 : memref<?x?xf32>):
test.region_if_yield %arg1 : memref<?x?xf32>
} join {
^bb0(%arg1 : memref<?x?xf32>):
test.region_if_yield %arg1 : memref<?x?xf32>
}
return %1 : memref<?x?xf32>
}
// CHECK-LABEL: func.func @inner_region_control_flow
// CHECK: [[ALLOC:%.+]] = memref.alloc(
// CHECK: [[V0:%.+]]:2 = test.region_if [[ALLOC]], %false
// CHECK: ^bb0([[ARG1:%.+]]: memref<?x?xf32>, [[ARG2:%.+]]: i1):
// CHECK: test.region_if_yield [[ARG1]], [[ARG2]]
// CHECK: ^bb0([[ARG1:%.+]]: memref<?x?xf32>, [[ARG2:%.+]]: i1):
// CHECK: test.region_if_yield [[ARG1]], [[ARG2]]
// CHECK: ^bb0([[ARG1:%.+]]: memref<?x?xf32>, [[ARG2:%.+]]: i1):
// CHECK: test.region_if_yield [[ARG1]], [[ARG2]]
// CHECK: [[V1:%.+]] = scf.if [[V0]]#1
// CHECK: scf.yield [[V0]]#0
// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0
// CHECK: scf.yield [[CLONE]]
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] :
// CHECK: return [[V1]]
// -----
func.func @nestedRegionsAndCondBranchAlloca(
%arg0: i1,
%arg1: memref<2xf32>,
%arg2: memref<2xf32>) {
cf.cond_br %arg0, ^bb1, ^bb2
^bb1:
cf.br ^bb3(%arg1 : memref<2xf32>)
^bb2:
%0 = memref.alloc() : memref<2xf32>
test.region_buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>) {
^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
%1 = memref.alloca() : memref<2xf32>
test.buffer_based in(%arg1: memref<2xf32>) out(%1: memref<2xf32>)
%tmp1 = math.exp %gen1_arg0 : f32
test.region_yield %tmp1 : f32
}
cf.br ^bb3(%0 : memref<2xf32>)
^bb3(%1: memref<2xf32>):
test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
return
}
// CHECK-LABEL: func @nestedRegionsAndCondBranchAlloca
// CHECK-SAME: ([[ARG0:%.+]]: i1, [[ARG1:%.+]]: memref<2xf32>, [[ARG2:%.+]]: memref<2xf32>)
// CHECK: ^bb1:
// CHECK: cf.br ^bb3([[ARG1]], %false
// CHECK: ^bb2:
// CHECK: [[ALLOC:%.+]] = memref.alloc()
// CHECK: test.region_buffer_based
// CHECK: memref.alloca()
// CHECK: test.buffer_based
// CHECK-NOT: bufferization.dealloc
// CHECK-NOT: bufferization.clone
// CHECK: test.region_yield
// CHECK: }
// CHECK: cf.br ^bb3([[ALLOC]], %true
// CHECK: ^bb3([[A0:%.+]]: memref<2xf32>, [[COND:%.+]]: i1):
// CHECK: test.copy
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[A0]]
// CHECK: bufferization.dealloc ([[BASE]] :{{.*}}) if ([[COND]])
// -----
func.func @nestedRegionControlFlowAlloca(
%arg0 : index, %arg1 : index, %arg2: f32) -> memref<?x?xf32> {
%0 = arith.cmpi eq, %arg0, %arg1 : index
%1 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
%2 = scf.if %0 -> (memref<?x?xf32>) {
scf.yield %1 : memref<?x?xf32>
} else {
%3 = memref.alloca(%arg0, %arg1) : memref<?x?xf32>
%c0 = arith.constant 0 : index
memref.store %arg2, %3[%c0, %c0] : memref<?x?xf32>
scf.yield %1 : memref<?x?xf32>
}
return %2 : memref<?x?xf32>
}
// CHECK-LABEL: func @nestedRegionControlFlowAlloca
// CHECK: [[ALLOC:%.+]] = memref.alloc(
// CHECK: [[V0:%.+]]:2 = scf.if
// CHECK: scf.yield [[ALLOC]], %false
// CHECK: memref.alloca(
// CHECK: scf.yield [[ALLOC]], %false
// CHECK: [[V1:%.+]] = scf.if [[V0]]#1
// CHECK: scf.yield [[V0]]#0
// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0
// CHECK: scf.yield [[CLONE]]
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] :
// CHECK: return [[V1]]
// -----
// Test Case: structured control-flow loop using a nested alloc.
// The iteration argument %iterBuf has to be freed before yielding %3 to avoid
// memory leaks.
func.func @loop_alloc(
%lb: index,
%ub: index,
%step: index,
%buf: memref<2xf32>,
%res: memref<2xf32>) {
%0 = memref.alloc() : memref<2xf32>
"test.read_buffer"(%0) : (memref<2xf32>) -> ()
%1 = scf.for %i = %lb to %ub step %step
iter_args(%iterBuf = %buf) -> memref<2xf32> {
%2 = arith.cmpi eq, %i, %ub : index
%3 = memref.alloc() : memref<2xf32>
scf.yield %3 : memref<2xf32>
}
test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
return
}
// CHECK-LABEL: func @loop_alloc
// CHECK-SAME: ([[ARG0:%.+]]: index, [[ARG1:%.+]]: index, [[ARG2:%.+]]: index, [[ARG3:%.+]]: memref<2xf32>, [[ARG4:%.+]]: memref<2xf32>)
// CHECK: [[ALLOC:%.+]] = memref.alloc()
// CHECK: [[V0:%.+]]:2 = scf.for {{.*}} iter_args([[ARG6:%.+]] = [[ARG3]], [[ARG7:%.+]] = %false
// CHECK: [[ALLOC1:%.+]] = memref.alloc()
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG6]]
// CHECK: bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG7]])
// CHECK-NOT: retain
// CHECK: scf.yield [[ALLOC1]], %true
// CHECK: test.copy
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK: bufferization.dealloc ([[ALLOC]] :{{.*}}) if (%true
// CHECK-NOT: retain
// CHECK: bufferization.dealloc ([[BASE]] :{{.*}}) if ([[V0]]#1)
// CHECK-NOT: retain
// -----
// Test Case: structured control-flow loop with a nested if operation.
// The loop yields buffers that have been defined outside of the loop and the
// backedges only use the iteration arguments (or one of its aliases).
// Therefore, we do not have to (and are not allowed to) free any buffers
// that are passed via the backedges.
func.func @loop_nested_if_no_alloc(
%lb: index,
%ub: index,
%step: index,
%buf: memref<2xf32>,
%res: memref<2xf32>) {
%0 = memref.alloc() : memref<2xf32>
%1 = scf.for %i = %lb to %ub step %step
iter_args(%iterBuf = %buf) -> memref<2xf32> {
%2 = arith.cmpi eq, %i, %ub : index
%3 = scf.if %2 -> (memref<2xf32>) {
scf.yield %0 : memref<2xf32>
} else {
scf.yield %iterBuf : memref<2xf32>
}
scf.yield %3 : memref<2xf32>
}
test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
return
}
// CHECK-LABEL: func @loop_nested_if_no_alloc
// CHECK-SAME: ({{.*}}, [[ARG3:%.+]]: memref<2xf32>, [[ARG4:%.+]]: memref<2xf32>)
// CHECK: [[ALLOC:%.+]] = memref.alloc()
// CHECK: [[V0:%.+]]:2 = scf.for {{.*}} iter_args([[ARG6:%.+]] = [[ARG3]], [[ARG7:%.+]] = %false
// CHECK: [[V1:%.+]]:2 = scf.if
// CHECK: scf.yield [[ALLOC]], %false
// CHECK: scf.yield [[ARG6]], %false
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG6]]
// CHECK: [[OWN:%.+]] = bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG7]]) retain ([[V1]]#0 :
// CHECK: [[OWN_AGG:%.+]] = arith.ori [[OWN]], [[V1]]#1
// CHECK: scf.yield [[V1]]#0, [[OWN_AGG]]
// CHECK: test.copy
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1)
// TODO: we know statically that the inner dealloc will never deallocate
// anything, i.e., we can optimize it away
// -----
// Test Case: structured control-flow loop with a nested if operation using
// a deeply nested buffer allocation.
func.func @loop_nested_if_alloc(
%lb: index,
%ub: index,
%step: index,
%buf: memref<2xf32>) -> memref<2xf32> {
%0 = memref.alloc() : memref<2xf32>
%1 = scf.for %i = %lb to %ub step %step
iter_args(%iterBuf = %buf) -> memref<2xf32> {
%2 = arith.cmpi eq, %i, %ub : index
%3 = scf.if %2 -> (memref<2xf32>) {
%4 = memref.alloc() : memref<2xf32>
scf.yield %4 : memref<2xf32>
} else {
scf.yield %0 : memref<2xf32>
}
scf.yield %3 : memref<2xf32>
}
return %1 : memref<2xf32>
}
// CHECK-LABEL: func @loop_nested_if_alloc
// CHECK-SAME: ({{.*}}, [[ARG3:%.+]]: memref<2xf32>)
// CHECK: [[ALLOC:%.+]] = memref.alloc()
// CHECK: [[V0:%.+]]:2 = scf.for {{.*}} iter_args([[ARG5:%.+]] = [[ARG3]], [[ARG6:%.+]] = %false
// CHECK: [[V1:%.+]]:2 = scf.if
// CHECK: [[ALLOC1:%.+]] = memref.alloc()
// CHECK: scf.yield [[ALLOC1]], %true
// CHECK: scf.yield [[ALLOC]], %false
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG5]]
// CHECK: [[OWN:%.+]] = bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG6]]) retain ([[V1]]#0 :
// CHECK: [[OWN_AGG:%.+]] = arith.ori [[OWN]], [[V1]]#1
// CHECK: scf.yield [[V1]]#0, [[OWN_AGG]]
// CHECK: }
// CHECK: [[V2:%.+]] = scf.if [[V0]]#1
// CHECK: scf.yield [[V0]]#0
// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0
// CHECK: scf.yield [[CLONE]]
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V2]] :
// CHECK: return [[V2]]
// -----
// Test Case: several nested structured control-flow loops with a deeply nested
// buffer allocation inside an if operation.
func.func @loop_nested_alloc(
%lb: index,
%ub: index,
%step: index,
%buf: memref<2xf32>,
%res: memref<2xf32>) {
%0 = memref.alloc() : memref<2xf32>
"test.read_buffer"(%0) : (memref<2xf32>) -> ()
%1 = scf.for %i = %lb to %ub step %step
iter_args(%iterBuf = %buf) -> memref<2xf32> {
%2 = scf.for %i2 = %lb to %ub step %step
iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {
%3 = scf.for %i3 = %lb to %ub step %step
iter_args(%iterBuf3 = %iterBuf2) -> memref<2xf32> {
%4 = memref.alloc() : memref<2xf32>
"test.read_buffer"(%4) : (memref<2xf32>) -> ()
%5 = arith.cmpi eq, %i, %ub : index
%6 = scf.if %5 -> (memref<2xf32>) {
%7 = memref.alloc() : memref<2xf32>
scf.yield %7 : memref<2xf32>
} else {
scf.yield %iterBuf3 : memref<2xf32>
}
scf.yield %6 : memref<2xf32>
}
scf.yield %3 : memref<2xf32>
}
scf.yield %2 : memref<2xf32>
}
test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
return
}
// CHECK-LABEL: func @loop_nested_alloc
// CHECK: ({{.*}}, [[ARG3:%.+]]: memref<2xf32>, {{.*}}: memref<2xf32>)
// CHECK: [[ALLOC:%.+]] = memref.alloc()
// CHECK: [[V0:%.+]]:2 = scf.for {{.*}} iter_args([[ARG6:%.+]] = [[ARG3]], [[ARG7:%.+]] = %false
// CHECK: [[V1:%.+]]:2 = scf.for {{.*}} iter_args([[ARG9:%.+]] = [[ARG6]], [[ARG10:%.+]] = %false
// CHECK: [[V2:%.+]]:2 = scf.for {{.*}} iter_args([[ARG12:%.+]] = [[ARG9]], [[ARG13:%.+]] = %false
// CHECK: [[ALLOC1:%.+]] = memref.alloc()
// CHECK: [[V3:%.+]]:2 = scf.if
// CHECK: [[ALLOC2:%.+]] = memref.alloc()
// CHECK: scf.yield [[ALLOC2]], %true
// CHECK: } else {
// CHECK: scf.yield [[ARG12]], %false
// CHECK: }
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG12]]
// CHECK: [[OWN:%.+]] = bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG13]]) retain ([[V3]]#0 :
// CHECK: bufferization.dealloc ([[ALLOC1]] :{{.*}}) if (%true{{[0-9_]*}})
// CHECK-NOT: retain
// CHECK: [[OWN_AGG:%.+]] = arith.ori [[OWN]], [[V3]]#1
// CHECK: scf.yield [[V3]]#0, [[OWN_AGG]]
// CHECK: }
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG9]]
// CHECK: [[OWN:%.+]] = bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG10]]) retain ([[V2]]#0 :
// CHECK: [[OWN_AGG:%.+]] = arith.ori [[OWN]], [[V2]]#1
// CHECK: scf.yield [[V2]]#0, [[OWN_AGG]]
// CHECK: }
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG6]]
// CHECK: [[OWN:%.+]] = bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG7]]) retain ([[V1]]#0 :
// CHECK: [[OWN_AGG:%.+]] = arith.ori [[OWN]], [[V1]]#1
// CHECK: scf.yield [[V1]]#0, [[OWN_AGG]]
// CHECK: }
// CHECK: test.copy
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK: bufferization.dealloc ([[ALLOC]] :{{.*}}) if (%true
// CHECK: bufferization.dealloc ([[BASE]] :{{.*}}) if ([[V0]]#1)
// TODO: all the retain operands could be removed by doing some more thorough analysis
// -----
func.func @affine_loop() -> f32 {
%buffer = memref.alloc() : memref<1024xf32>
%sum_init_0 = arith.constant 0.0 : f32
%res = affine.for %i = 0 to 10 step 2 iter_args(%sum_iter = %sum_init_0) -> f32 {
%t = affine.load %buffer[%i] : memref<1024xf32>
%sum_next = arith.addf %sum_iter, %t : f32
affine.yield %sum_next : f32
}
return %res : f32
}
// CHECK-LABEL: func @affine_loop
// CHECK: [[ALLOC:%.+]] = memref.alloc()
// CHECK: affine.for {{.*}} iter_args(%arg1 = %cst)
// CHECK: affine.yield
// CHECK: bufferization.dealloc ([[ALLOC]] :{{.*}}) if (%true
// -----
func.func @assumingOp(
%arg0: !shape.witness,
%arg2: memref<2xf32>,
%arg3: memref<2xf32>) {
// Confirm the alloc will be dealloc'ed in the block.
%1 = shape.assuming %arg0 -> memref<2xf32> {
%0 = memref.alloc() : memref<2xf32>
"test.read_buffer"(%0) : (memref<2xf32>) -> ()
shape.assuming_yield %arg2 : memref<2xf32>
}
// Confirm the alloc will be returned and dealloc'ed after its use.
%3 = shape.assuming %arg0 -> memref<2xf32> {
%2 = memref.alloc() : memref<2xf32>
shape.assuming_yield %2 : memref<2xf32>
}
test.copy(%3, %arg3) : (memref<2xf32>, memref<2xf32>)
return
}
// CHECK-LABEL: func @assumingOp
// CHECK: ({{.*}}, [[ARG1:%.+]]: memref<2xf32>, {{.*}}: memref<2xf32>)
// CHECK: [[V0:%.+]]:2 = shape.assuming
// CHECK: [[ALLOC:%.+]] = memref.alloc()
// CHECK: bufferization.dealloc ([[ALLOC]] :{{.*}}) if (%true{{[0-9_]*}})
// CHECK-NOT: retain
// CHECK: shape.assuming_yield [[ARG1]], %false
// CHECK: }
// CHECK: [[V1:%.+]]:2 = shape.assuming
// CHECK: [[ALLOC:%.+]] = memref.alloc()
// CHECK: shape.assuming_yield [[ALLOC]], %true
// CHECK: }
// CHECK: test.copy
// CHECK: [[BASE0:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK: [[BASE1:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V1]]#0
// CHECK: bufferization.dealloc ([[BASE1]] :{{.*}}) if ([[V1]]#1)
// CHECK-NOT: retain
// CHECK: bufferization.dealloc ([[BASE0]] :{{.*}}) if ([[V0]]#1)
// CHECK-NOT: retain
// CHECK: return
// -----
// Test Case: The op "test.one_region_with_recursive_memory_effects" does not
// implement the RegionBranchOpInterface. This is allowed during buffer
// deallocation because the operation's region does not deal with any MemRef
// values.
func.func @noRegionBranchOpInterface() {
%0 = "test.one_region_with_recursive_memory_effects"() ({
%1 = "test.one_region_with_recursive_memory_effects"() ({
%2 = memref.alloc() : memref<2xi32>
"test.read_buffer"(%2) : (memref<2xi32>) -> ()
"test.return"() : () -> ()
}) : () -> (i32)
"test.return"() : () -> ()
}) : () -> (i32)
"test.return"() : () -> ()
}
// -----
// Test Case: The second op "test.one_region_with_recursive_memory_effects" does
// not implement the RegionBranchOpInterface but has buffer semantics. This is
// not allowed during buffer deallocation.
func.func @noRegionBranchOpInterface() {
%0 = "test.one_region_with_recursive_memory_effects"() ({
// expected-error@+1 {{All operations with attached regions need to implement the RegionBranchOpInterface.}}
%1 = "test.one_region_with_recursive_memory_effects"() ({
%2 = memref.alloc() : memref<2xi32>
"test.read_buffer"(%2) : (memref<2xi32>) -> ()
"test.return"(%2) : (memref<2xi32>) -> ()
}) : () -> (memref<2xi32>)
"test.return"() : () -> ()
}) : () -> (i32)
"test.return"() : () -> ()
}
// -----
func.func @while_two_arg(%arg0: index) {
%a = memref.alloc(%arg0) : memref<?xf32>
scf.while (%arg1 = %a, %arg2 = %a) : (memref<?xf32>, memref<?xf32>) -> (memref<?xf32>, memref<?xf32>) {
// This op has a side effect, but it's not an allocate/free side effect.
%0 = "test.side_effect_op"() {effects = [{effect="read"}]} : () -> i1
scf.condition(%0) %arg1, %arg2 : memref<?xf32>, memref<?xf32>
} do {
^bb0(%arg1: memref<?xf32>, %arg2: memref<?xf32>):
%b = memref.alloc(%arg0) : memref<?xf32>
scf.yield %arg1, %b : memref<?xf32>, memref<?xf32>
}
return
}
// CHECK-LABEL: func @while_two_arg
// CHECK: [[ALLOC:%.+]] = memref.alloc(
// CHECK: [[V0:%.+]]:4 = scf.while ({{.*}} = [[ALLOC]], {{.*}} = [[ALLOC]], {{.*}} = %false{{[0-9_]*}}, {{.*}} = %false{{[0-9_]*}})
// CHECK: scf.condition
// CHECK: ^bb0([[ARG1:%.+]]: memref<?xf32>, [[ARG2:%.+]]: memref<?xf32>, [[ARG3:%.+]]: i1, [[ARG4:%.+]]: i1):
// CHECK: [[ALLOC1:%.+]] = memref.alloc(
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG2]]
// CHECK: [[OWN:%.+]] = bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG4]]) retain ([[ARG1]] :
// CHECK: [[OWN_AGG:%.+]] = arith.ori [[OWN]], [[ARG3]]
// CHECK: scf.yield [[ARG1]], [[ALLOC1]], [[OWN_AGG]], %true
// CHECK: [[BASE0:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK: [[BASE1:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#1
// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE0]], [[BASE1]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#2, [[V0]]#3)
// -----
func.func @while_three_arg(%arg0: index) {
%a = memref.alloc(%arg0) : memref<?xf32>
scf.while (%arg1 = %a, %arg2 = %a, %arg3 = %a) : (memref<?xf32>, memref<?xf32>, memref<?xf32>) -> (memref<?xf32>, memref<?xf32>, memref<?xf32>) {
// This op has a side effect, but it's not an allocate/free side effect.
%0 = "test.side_effect_op"() {effects = [{effect="read"}]} : () -> i1
scf.condition(%0) %arg1, %arg2, %arg3 : memref<?xf32>, memref<?xf32>, memref<?xf32>
} do {
^bb0(%arg1: memref<?xf32>, %arg2: memref<?xf32>, %arg3: memref<?xf32>):
%b = memref.alloc(%arg0) : memref<?xf32>
%q = memref.alloc(%arg0) : memref<?xf32>
scf.yield %q, %b, %arg2: memref<?xf32>, memref<?xf32>, memref<?xf32>
}
return
}
// CHECK-LABEL: func @while_three_arg
// CHECK: [[ALLOC:%.+]] = memref.alloc(
// CHECK: [[V0:%.+]]:6 = scf.while ({{.*}} = [[ALLOC]], {{.*}} = [[ALLOC]], {{.*}} = [[ALLOC]], {{.*}} = %false{{[0-9_]*}}, {{.*}} = %false{{[0-9_]*}}, {{.*}} = %false
// CHECK: scf.condition
// CHECK: ^bb0([[ARG1:%.+]]: memref<?xf32>, [[ARG2:%.+]]: memref<?xf32>, [[ARG3:%.+]]: memref<?xf32>, [[ARG4:%.+]]: i1, [[ARG5:%.+]]: i1, [[ARG6:%.+]]: i1):
// CHECK: [[ALLOC1:%.+]] = memref.alloc(
// CHECK: [[ALLOC2:%.+]] = memref.alloc(
// CHECK: [[BASE0:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG1]]
// CHECK: [[BASE2:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG3]]
// CHECK: [[OWN:%.+]] = bufferization.dealloc ([[BASE0]], [[BASE2]] :{{.*}}) if ([[ARG4]], [[ARG6]]) retain ([[ARG2]] :
// CHECK: [[OWN_AGG:%.+]] = arith.ori [[OWN]], [[ARG5]]
// CHECK: scf.yield [[ALLOC2]], [[ALLOC1]], [[ARG2]], %true{{[0-9_]*}}, %true{{[0-9_]*}}, [[OWN_AGG]] :
// CHECK: }
// CHECK: [[BASE0:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK: [[BASE1:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#1
// CHECK: [[BASE2:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#2
// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE0]], [[BASE1]], [[BASE2]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#3, [[V0]]#4, [[V0]]#5)
// TODO: better alias analysis could simplify the dealloc inside the body further
// -----
// Memref allocated in `then` region and passed back to the parent if op.
#set = affine_set<() : (0 >= 0)>
func.func @test_affine_if_1(%arg0: memref<10xf32>) -> memref<10xf32> {
%0 = affine.if #set() -> memref<10xf32> {
%alloc = memref.alloc() : memref<10xf32>
affine.yield %alloc : memref<10xf32>
} else {
affine.yield %arg0 : memref<10xf32>
}
return %0 : memref<10xf32>
}
// CHECK-LABEL: func @test_affine_if_1
// CHECK-SAME: ([[ARG0:%.*]]: memref<10xf32>)
// CHECK: [[V0:%.+]]:2 = affine.if
// CHECK: [[ALLOC:%.+]] = memref.alloc()
// CHECK: affine.yield [[ALLOC]], %true
// CHECK: affine.yield [[ARG0]], %false
// CHECK: [[V1:%.+]] = scf.if [[V0]]#1
// CHECK: scf.yield [[V0]]#0
// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0
// CHECK: scf.yield [[CLONE]]
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK: bufferization.dealloc ([[BASE]] :{{.*}}) if ([[V0]]#1) retain ([[V1]] :
// CHECK: return [[V1]]
// TODO: the dealloc could be optimized away since the memref to be deallocated
// either aliases with V1 or the condition is false
// -----
// Memref allocated before parent IfOp and used in `then` region.
// Expected result: deallocation should happen after affine.if op.
#set = affine_set<() : (0 >= 0)>
func.func @test_affine_if_2() -> memref<10xf32> {
%alloc0 = memref.alloc() : memref<10xf32>
%0 = affine.if #set() -> memref<10xf32> {
affine.yield %alloc0 : memref<10xf32>
} else {
%alloc = memref.alloc() : memref<10xf32>
affine.yield %alloc : memref<10xf32>
}
return %0 : memref<10xf32>
}
// CHECK-LABEL: func @test_affine_if_2
// CHECK: [[ALLOC:%.+]] = memref.alloc()
// CHECK: [[V0:%.+]]:2 = affine.if
// CHECK: affine.yield [[ALLOC]], %false
// CHECK: [[ALLOC1:%.+]] = memref.alloc()
// CHECK: affine.yield [[ALLOC1]], %true
// CHECK: [[V1:%.+]] = scf.if [[V0]]#1
// CHECK: scf.yield [[V0]]#0
// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0
// CHECK: scf.yield [[CLONE]]
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] :
// CHECK: return [[V1]]
// -----
// Memref allocated before parent IfOp and used in `else` region.
// Expected result: deallocation should happen after affine.if op.
#set = affine_set<() : (0 >= 0)>
func.func @test_affine_if_3() -> memref<10xf32> {
%alloc0 = memref.alloc() : memref<10xf32>
%0 = affine.if #set() -> memref<10xf32> {
%alloc = memref.alloc() : memref<10xf32>
affine.yield %alloc : memref<10xf32>
} else {
affine.yield %alloc0 : memref<10xf32>
}
return %0 : memref<10xf32>
}
// CHECK-LABEL: func @test_affine_if_3
// CHECK: [[ALLOC:%.+]] = memref.alloc()
// CHECK: [[V0:%.+]]:2 = affine.if
// CHECK: [[ALLOC1:%.+]] = memref.alloc()
// CHECK: affine.yield [[ALLOC1]], %true
// CHECK: affine.yield [[ALLOC]], %false
// CHECK: [[V1:%.+]] = scf.if [[V0]]#1
// CHECK: scf.yield [[V0]]#0
// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0
// CHECK: scf.yield [[CLONE]]
// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]]
// CHECK: return [[V1]]
Codebase Browser
llvm