28d6aa90b0
bufferization.clone does not currently support lowering to memref for
unranked memrefs. This interferes with bufferizing unranked tensors at
boundaries where a clone operation is needed.
```
func.func @foo(%input: memref<*xf32>, %shape: memref<?xindex>) -> memref<*xf32>
{
%reshape = memref.reshape %input(%shape) : (memref<*xf32>, memref<?xindex>) -> memref<*xf32>
%copy = bufferization.clone %reshape : memref<*xf32> to memref<*xf32>
return %copy : memref<*xf32>
}
```
Patterns such as that are possibly when bufferizing functions with input
and output unranked tensors. The clone operation currently fails to
legalize during the bufferization-to-memref conversion with unranked
memrefs.
This change modifies the conversion of bufferization.clone to memref to
generate the runtime calculations and allocation to allow for cloning an
unranked memref.
99 lines
3.9 KiB
MLIR
99 lines
3.9 KiB
MLIR
// RUN: mlir-opt -verify-diagnostics -convert-bufferization-to-memref -split-input-file %s | FileCheck %s
|
|
|
|
// CHECK-LABEL: @conversion_static
|
|
func.func @conversion_static(%arg0 : memref<2xf32>) -> memref<2xf32> {
|
|
%0 = bufferization.clone %arg0 : memref<2xf32> to memref<2xf32>
|
|
memref.dealloc %arg0 : memref<2xf32>
|
|
return %0 : memref<2xf32>
|
|
}
|
|
|
|
// CHECK: %[[ALLOC:.*]] = memref.alloc
|
|
// CHECK-NEXT: memref.copy %[[ARG:.*]], %[[ALLOC]]
|
|
// CHECK-NEXT: memref.dealloc %[[ARG]]
|
|
// CHECK-NEXT: return %[[ALLOC]]
|
|
|
|
// -----
|
|
|
|
// CHECK-LABEL: @conversion_dynamic
|
|
func.func @conversion_dynamic(%arg0 : memref<?xf32>) -> memref<?xf32> {
|
|
%1 = bufferization.clone %arg0 : memref<?xf32> to memref<?xf32>
|
|
memref.dealloc %arg0 : memref<?xf32>
|
|
return %1 : memref<?xf32>
|
|
}
|
|
|
|
// CHECK: %[[CONST:.*]] = arith.constant
|
|
// CHECK: %[[DIM:.*]] = memref.dim %[[ARG:.*]], %[[CONST]]
|
|
// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc(%[[DIM]])
|
|
// CHECK-NEXT: memref.copy %[[ARG]], %[[ALLOC]]
|
|
// CHECK-NEXT: memref.dealloc %[[ARG]]
|
|
// CHECK-NEXT: return %[[ALLOC]]
|
|
|
|
// -----
|
|
|
|
// CHECK-LABEL: @conversion_unknown
|
|
func.func @conversion_unknown(%arg0 : memref<*xf32>) -> memref<*xf32> {
|
|
%1 = bufferization.clone %arg0 : memref<*xf32> to memref<*xf32>
|
|
memref.dealloc %arg0 : memref<*xf32>
|
|
return %1 : memref<*xf32>
|
|
}
|
|
|
|
// CHECK: %[[RANK:.*]] = memref.rank %[[ARG:.*]]
|
|
// CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca(%[[RANK]])
|
|
// CHECK-NEXT: %[[FOR:.*]] = scf.for
|
|
// CHECK-NEXT: %[[DIM:.*]] = memref.dim %[[ARG:.*]] %[[ARG:.*]]
|
|
// CHECK-NEXT: memref.store %[[DIM:.*]], %[[ALLOCA:.*]][%[[ARG:.*]]]
|
|
// CHECK-NEXT: %[[MUL:.*]] = arith.muli %[[ARG:.*]], %[[DIM:.*]]
|
|
// CHECK-NEXT: scf.yield %[[MUL:.*]]
|
|
// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[FOR:.*]])
|
|
// CHECK-NEXT: %[[RESHAPE:.*]] = memref.reshape %[[ALLOC:.*]]
|
|
// CHECK-NEXT: memref.copy
|
|
// CHECK-NEXT: memref.dealloc
|
|
// CHECK-NEXT: return %[[RESHAPE:.*]]
|
|
|
|
// -----
|
|
|
|
// CHECK-LABEL: func @conversion_with_layout_map(
|
|
// CHECK-SAME: %[[ARG:.*]]: memref<?xf32, strided<[?], offset: ?>>
|
|
// CHECK: %[[C0:.*]] = arith.constant 0 : index
|
|
// CHECK: %[[DIM:.*]] = memref.dim %[[ARG]], %[[C0]]
|
|
// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM]]) : memref<?xf32>
|
|
// CHECK: %[[CASTED:.*]] = memref.cast %[[ALLOC]] : memref<?xf32> to memref<?xf32, strided<[?], offset: ?>>
|
|
// CHECK: memref.copy
|
|
// CHECK: memref.dealloc
|
|
// CHECK: return %[[CASTED]]
|
|
func.func @conversion_with_layout_map(%arg0 : memref<?xf32, strided<[?], offset: ?>>) -> memref<?xf32, strided<[?], offset: ?>> {
|
|
%1 = bufferization.clone %arg0 : memref<?xf32, strided<[?], offset: ?>> to memref<?xf32, strided<[?], offset: ?>>
|
|
memref.dealloc %arg0 : memref<?xf32, strided<[?], offset: ?>>
|
|
return %1 : memref<?xf32, strided<[?], offset: ?>>
|
|
}
|
|
|
|
// -----
|
|
|
|
// This bufferization.clone cannot be lowered because a buffer with this layout
|
|
// map cannot be allocated (or casted to).
|
|
|
|
func.func @conversion_with_invalid_layout_map(%arg0 : memref<?xf32, strided<[10], offset: ?>>)
|
|
-> memref<?xf32, strided<[10], offset: ?>> {
|
|
// expected-error@+1 {{failed to legalize operation 'bufferization.clone' that was explicitly marked illegal}}
|
|
%1 = bufferization.clone %arg0 : memref<?xf32, strided<[10], offset: ?>> to memref<?xf32, strided<[10], offset: ?>>
|
|
memref.dealloc %arg0 : memref<?xf32, strided<[10], offset: ?>>
|
|
return %1 : memref<?xf32, strided<[10], offset: ?>>
|
|
}
|
|
|
|
// -----
|
|
// Test: check that the dealloc lowering pattern is registered.
|
|
|
|
// CHECK-NOT: func @deallocHelper
|
|
// CHECK-LABEL: func @conversion_dealloc_simple
|
|
// CHECK-SAME: [[ARG0:%.+]]: memref<2xf32>
|
|
// CHECK-SAME: [[ARG1:%.+]]: i1
|
|
func.func @conversion_dealloc_simple(%arg0: memref<2xf32>, %arg1: i1) {
|
|
bufferization.dealloc (%arg0 : memref<2xf32>) if (%arg1)
|
|
return
|
|
}
|
|
|
|
// CHECK: scf.if [[ARG1]] {
|
|
// CHECK-NEXT: memref.dealloc [[ARG0]] : memref<2xf32>
|
|
// CHECK-NEXT: }
|
|
// CHECK-NEXT: return
|