a36348c586
AllocTensorElimination does currently not support chains where the type is changing. AllocTensorElimination used to generate invalid IR for such inputs. With this commit, AllocTensorElimination does no longer apply to such inputs. (It can be extended to support such IR if needed.) Differential Revision: https://reviews.llvm.org/D131880
139 lines
5.8 KiB
MLIR
139 lines
5.8 KiB
MLIR
// RUN: mlir-opt %s -eliminate-alloc-tensors -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs" -canonicalize -split-input-file | FileCheck %s
|
|
|
|
// CHECK: func @buffer_forwarding_conflict(
|
|
// CHECK-SAME: %[[FUNC_ARG:[0-9a-zA-Z]*]]: memref<?xf32>
|
|
// CHECK-SAME: %[[sz:[0-9a-zA-Z]*]]: index
|
|
func.func @buffer_forwarding_conflict(
|
|
%t: tensor<?xf32> {bufferization.buffer_layout = affine_map<(d0) -> (d0)>, bufferization.writable = true},
|
|
%sz: index)
|
|
-> (tensor<?xf32>, tensor<?xf32>)
|
|
{
|
|
%f0 = arith.constant 0.0: f32
|
|
|
|
// CHECK: %[[EXTRACT_SLICE_ALLOC:.*]] = memref.alloc(%[[sz]])
|
|
// CHECK: linalg.fill ins({{.*}} : f32) outs(%[[EXTRACT_SLICE_ALLOC]] : memref<?xf32>)
|
|
// Alloc is needed for the **first** insert_slice (due to backward traversal during analysis).
|
|
// CHECK: %[[DIM:.*]] = memref.dim %[[FUNC_ARG]]
|
|
// This allocs the whole dim to allow for a full clone of t.
|
|
// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM]])
|
|
// alloc_tensor itself does not alloc but forwards to the **second**
|
|
// insert_slice. AllocTensorOp replaces the alloc_tensor with an out-of-place
|
|
// extract_slice.
|
|
%a = bufferization.alloc_tensor(%sz) : tensor<?xf32>
|
|
%f = linalg.fill ins(%f0 : f32) outs(%a : tensor<?xf32>) -> tensor<?xf32>
|
|
|
|
// CHECK: memref.copy %[[FUNC_ARG]], %[[ALLOC]] : memref<?xf32> to memref<?xf32>
|
|
// CHECK: %[[SV0_ALLOC:.*]] = memref.subview %[[ALLOC]][0] [%[[sz]]] [1] : memref<?xf32> to memref<?xf32>
|
|
// CHECK: memref.copy %[[EXTRACT_SLICE_ALLOC]], %[[SV0_ALLOC]] : memref<?xf32> to memref<?xf32>
|
|
%r0 = tensor.insert_slice %f into %t[0][%sz][1]: tensor<?xf32> into tensor<?xf32>
|
|
|
|
// CHECK: %[[T_SUBVIEW:.*]] = memref.subview %[[FUNC_ARG]][42] [%[[sz]]] [1]
|
|
// CHECK: memref.copy %[[EXTRACT_SLICE_ALLOC]], %[[T_SUBVIEW]]
|
|
%r1 = tensor.insert_slice %f into %t[42][%sz][1]: tensor<?xf32> into tensor<?xf32>
|
|
|
|
return %r0, %r1: tensor<?xf32>, tensor<?xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// CHECK: func @buffer_forwarding_no_conflict(
|
|
// CHECK-SAME: %[[FUNC_ARG:[0-9a-zA-Z]*]]: memref<?xf32>
|
|
// CHECK-SAME: %[[sz:[0-9a-zA-Z]*]]: index
|
|
func.func @buffer_forwarding_no_conflict(
|
|
%t: tensor<?xf32> {bufferization.buffer_layout = affine_map<(d0) -> (d0)>, bufferization.writable = true},
|
|
%sz: index)
|
|
-> (tensor<?xf32>)
|
|
{
|
|
%f0 = arith.constant 0.0: f32
|
|
|
|
// alloc_tensor itself does not alloc but forwards to the insert_slice.
|
|
// InitTensorOp replaces the alloc_tensor with an inplace extract_slice.
|
|
// CHECK: %[[T_SUBVIEW:.*]] = memref.subview %[[FUNC_ARG]][42] [%[[sz]]] [1]
|
|
%a = bufferization.alloc_tensor(%sz) : tensor<?xf32>
|
|
|
|
// CHECK: linalg.fill ins({{.*}} : f32) outs(%[[T_SUBVIEW]] : memref<?xf32
|
|
%f = linalg.fill ins(%f0 : f32) outs(%a : tensor<?xf32>) -> tensor<?xf32>
|
|
|
|
// Self-copy canonicalizes away later.
|
|
%r1 = tensor.insert_slice %f into %t[42][%sz][1]: tensor<?xf32> into tensor<?xf32>
|
|
|
|
return %r1: tensor<?xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// CHECK: func @insertion_point_inside_loop(
|
|
// CHECK-SAME: %[[t:.*]]: memref<?xf32, #{{.*}}>, %[[sz:.*]]: index)
|
|
func.func @insertion_point_inside_loop(%t : tensor<?xf32>, %sz : index) -> (tensor<?xf32>) {
|
|
%c0 = arith.constant 0 : index
|
|
%c1 = arith.constant 1 : index
|
|
%c5 = arith.constant 5 : index
|
|
|
|
// CHECK-NOT: memref.alloc
|
|
%blank = bufferization.alloc_tensor() : tensor<5xf32>
|
|
|
|
// CHECK: scf.for %[[iv:.*]] = %{{.*}} to %[[sz]] step %{{.*}} {
|
|
%r = scf.for %iv = %c0 to %sz step %c5 iter_args(%bb = %t) -> (tensor<?xf32>) {
|
|
// CHECK: %[[subview:.*]] = memref.subview %[[t]][%[[iv]]] [5] [1]
|
|
%iv_i32 = arith.index_cast %iv : index to i32
|
|
%f = arith.sitofp %iv_i32 : i32 to f32
|
|
|
|
// CHECK: linalg.fill ins(%{{.*}}{{.*}}outs(%[[subview]]
|
|
%filled = linalg.fill ins(%f : f32) outs(%blank : tensor<5xf32>) -> tensor<5xf32>
|
|
|
|
// CHECK-NOT: memref.copy
|
|
%inserted = tensor.insert_slice %filled into %bb[%iv][5][1] : tensor<5xf32> into tensor<?xf32>
|
|
scf.yield %inserted : tensor<?xf32>
|
|
}
|
|
|
|
return %r : tensor<?xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// CHECK: func @insertion_point_outside_loop(
|
|
// CHECK-SAME: %[[t:.*]]: memref<?xf32, #{{.*}}>, %[[sz:.*]]: index, %[[idx:.*]]: index)
|
|
func.func @insertion_point_outside_loop(%t : tensor<?xf32>, %sz : index,
|
|
%idx : index) -> (tensor<?xf32>) {
|
|
%c0 = arith.constant 0 : index
|
|
%c1 = arith.constant 1 : index
|
|
%c5 = arith.constant 5 : index
|
|
|
|
// CHECK-NOT: memref.alloc
|
|
// CHECK: %[[subview:.*]] = memref.subview %[[t]][%[[idx]]] [5] [1]
|
|
%blank = bufferization.alloc_tensor() : tensor<5xf32>
|
|
|
|
// CHECK: scf.for %[[iv:.*]] = %{{.*}} to %[[sz]] step %{{.*}} {
|
|
%r = scf.for %iv = %c0 to %sz step %c5 iter_args(%bb = %t) -> (tensor<?xf32>) {
|
|
%iv_i32 = arith.index_cast %iv : index to i32
|
|
%f = arith.sitofp %iv_i32 : i32 to f32
|
|
|
|
// CHECK: linalg.fill ins(%{{.*}}{{.*}}outs(%[[subview]]
|
|
%filled = linalg.fill ins(%f : f32) outs(%blank : tensor<5xf32>) -> tensor<5xf32>
|
|
|
|
// CHECK-NOT: memref.copy
|
|
%inserted = tensor.insert_slice %filled into %bb[%idx][5][1] : tensor<5xf32> into tensor<?xf32>
|
|
scf.yield %inserted : tensor<?xf32>
|
|
}
|
|
|
|
return %r : tensor<?xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// AllocTensorElimination does currently not apply to chains where the type is
|
|
// changing. This test just ensures that we do not crash or generate IR that
|
|
// does not verify.
|
|
|
|
// CHECK-LABEL: func @shape_mismatch
|
|
func.func @shape_mismatch(%t: tensor<5x6x128xf32>) -> tensor<5x6x128xf32> {
|
|
%cst = arith.constant 8.0 : f32
|
|
%0 = bufferization.alloc_tensor() : tensor<128xf32>
|
|
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<128xf32>) -> tensor<128xf32>
|
|
%2 = tensor.expand_shape %1 [[0, 1, 2]]
|
|
: tensor<128xf32> into tensor<1x1x128xf32>
|
|
%3 = tensor.insert_slice %2 into %t[2, 3, 0][1, 1, 128][1, 1, 1]
|
|
: tensor<1x1x128xf32> into tensor<5x6x128xf32>
|
|
return %3 : tensor<5x6x128xf32>
|
|
}
|