2646c36a86
The existing OneShotModuleBufferize will analyze and bufferize operations which are in nested symbol tables (e.g. nested `builtin.module`, `gpu.module`, or similar operations). This behavior is untested and likely unintentional given other limitations of OneShotModuleBufferize (`func.call` can't call into nested symbol tables). This change reverses the existing behavior so that the operations considered by the analysis and bufferization exclude any operations in nested symbol table scopes. Users who desire to bufferize nested modules can still do so by applying the transformation in a pass pipeline or in a custom pass. This further enables controlling the order in which modules are bufferized as well as allowing use of different options for different kinds of modules.
239 lines
9.2 KiB
MLIR
239 lines
9.2 KiB
MLIR
// RUN: mlir-opt --transform-interpreter %s -split-input-file -verify-diagnostics | FileCheck %s
|
|
|
|
// Test One-Shot Bufferize.
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%1 = transform.bufferization.one_shot_bufferize %0 : (!transform.any_op) -> !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// CHECK-LABEL: func @test_function(
|
|
// CHECK-SAME: %[[A:.*]]: tensor<?xf32>
|
|
func.func @test_function(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
|
|
%c0 = arith.constant 0 : index
|
|
|
|
// CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]]
|
|
// CHECK: %[[dim:.*]] = memref.dim %[[A_memref]]
|
|
// CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])
|
|
// CHECK: memref.copy %[[A_memref]], %[[alloc]]
|
|
// CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
|
|
// CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[alloc]]
|
|
%0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>
|
|
|
|
// CHECK: return %[[res_tensor]]
|
|
return %0 : tensor<?xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Emit linalg.copy instead of memref.copy.
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%1 = transform.bufferization.one_shot_bufferize %0 {memcpy_op = "linalg.copy"} : (!transform.any_op) -> !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// CHECK-LABEL: func @test_function(
|
|
// CHECK-SAME: %[[A:.*]]: tensor<?xf32>
|
|
// CHECK-NOT: memref.copy
|
|
func.func @test_function(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
|
|
%c0 = arith.constant 0 : index
|
|
|
|
// CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]]
|
|
// CHECK: %[[dim:.*]] = memref.dim %[[A_memref]]
|
|
// CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])
|
|
// CHECK: linalg.copy ins(%[[A_memref]] : memref<{{.*}}>) outs(%[[alloc]]
|
|
// CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
|
|
// CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[alloc]]
|
|
%0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>
|
|
|
|
// CHECK: return %[[res_tensor]]
|
|
return %0 : tensor<?xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test analysis of One-Shot Bufferize only.
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%1 = transform.bufferization.one_shot_bufferize %0
|
|
{test_analysis_only = true} : (!transform.any_op) -> !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// CHECK-LABEL: func @test_function_analysis(
|
|
// CHECK-SAME: %[[A:.*]]: tensor<?xf32>
|
|
func.func @test_function_analysis(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
|
|
%c0 = arith.constant 0 : index
|
|
// CHECK: vector.transfer_write
|
|
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "false", "none"]}
|
|
// CHECK-SAME: tensor<?xf32>
|
|
%0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>
|
|
return %0 : tensor<?xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test One-Shot Bufferize transform failure with an unknown op. This would be
|
|
// allowed with `allow_unknown_ops`.
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
// expected-error @+1 {{bufferization failed}}
|
|
%1 = transform.bufferization.one_shot_bufferize %0 : (!transform.any_op) -> !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
func.func @test_unknown_op_failure() -> (tensor<?xf32>) {
|
|
// expected-error @+1 {{op was not bufferized}}
|
|
%0 = "test.dummy_op"() : () -> (tensor<?xf32>)
|
|
return %0 : tensor<?xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.consumed}) {
|
|
// %arg1 is the module
|
|
%0 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// CHECK-LABEL: func @test_function(
|
|
// CHECK-SAME: %[[A:.*]]: tensor<?xf32>
|
|
func.func @test_function(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
|
|
%c0 = arith.constant 0 : index
|
|
|
|
// CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]]
|
|
// CHECK: %[[dim:.*]] = memref.dim %[[A_memref]]
|
|
// CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])
|
|
// CHECK: memref.copy %[[A_memref]], %[[alloc]]
|
|
// CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
|
|
// CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[alloc]]
|
|
%0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>
|
|
|
|
// CHECK: return %[[res_tensor]]
|
|
return %0 : tensor<?xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test we use identity layout at function boundaries.
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.consumed}) {
|
|
%0 = transform.bufferization.one_shot_bufferize layout{IdentityLayoutMap} %arg1
|
|
{ bufferize_function_boundaries = true } : (!transform.any_op) -> !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// CHECK: func.func @matmul(
|
|
// CHECK-SAME: %[[A:.*]]: memref<12x9xf32>,
|
|
// CHECK-SAME: %[[B:.*]]: memref<9x6xf32>,
|
|
// CHECK-SAME: %[[C:.*]]: memref<12x6xf32>) -> memref<12x6xf32> {
|
|
func.func @matmul(%A: tensor<12x9xf32>, %B: tensor<9x6xf32>, %C: tensor<12x6xf32>) -> tensor<12x6xf32> {
|
|
// CHECK: linalg.matmul ins(%[[A]], %[[B]] : memref<12x9xf32>, memref<9x6xf32>) outs(%[[C]] : memref<12x6xf32>)
|
|
%D = linalg.matmul ins(%A, %B: tensor<12x9xf32>, tensor<9x6xf32>) outs(%C: tensor<12x6xf32>) -> tensor<12x6xf32>
|
|
// CHECK: return %[[C]] : memref<12x6xf32>
|
|
return %D : tensor<12x6xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["tensor.empty"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%1 = transform.cast %0 : !transform.any_op to !transform.op<"tensor.empty">
|
|
transform.bufferization.empty_tensor_to_alloc_tensor %1 : (!transform.op<"tensor.empty">) -> !transform.op<"bufferization.alloc_tensor">
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// Expect `bufferization.empty_tensor_to_alloc_tensor` to replace the tensor.empty.
|
|
func.func @empty_to_tensor_alloc() -> tensor<2x2xf32> {
|
|
// CHECK: bufferization.alloc_tensor
|
|
%0 = tensor.empty() : tensor<2x2xf32>
|
|
return %0 : tensor<2x2xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.bufferization.eliminate_empty_tensors %0 : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// CHECK-LABEL: func @empty_tensor_elimination(
|
|
// CHECK: tensor.extract_slice
|
|
// CHECK: linalg.fill
|
|
// CHECK: tensor.insert_slice
|
|
func.func @empty_tensor_elimination(
|
|
%t: tensor<10xf32>, %f: f32) -> tensor<10xf32> {
|
|
%0 = tensor.empty() : tensor<5xf32>
|
|
%1 = linalg.fill ins(%f : f32) outs(%0 : tensor<5xf32>) -> tensor<5xf32>
|
|
%2 = tensor.insert_slice %1 into %t [1][5][1]
|
|
: tensor<5xf32> into tensor<10xf32>
|
|
return %2 : tensor<10xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.bufferization.buffer_loop_hoisting %0 : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// CHECK-LABEL: func @buffer_loop_hoisting(
|
|
// CHECK: memref.alloca
|
|
// CHECK: scf.for
|
|
// CHECK: memref.store
|
|
func.func @buffer_loop_hoisting(%lb: index, %ub: index, %step: index, %f: f32, %pos: index) {
|
|
scf.for %iv = %lb to %ub step %step {
|
|
%0 = memref.alloca() : memref<5xf32>
|
|
memref.store %f, %0[%pos] : memref<5xf32>
|
|
}
|
|
return
|
|
}
|
|
|
|
// -----
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%alloc_tensor = transform.structured.match ops{["bufferization.alloc_tensor"]} in %arg1
|
|
: (!transform.any_op) -> !transform.op<"bufferization.alloc_tensor">
|
|
%2, %new = transform.structured.bufferize_to_allocation %alloc_tensor
|
|
{alloc_op = "memref.alloca"}
|
|
: !transform.op<"bufferization.alloc_tensor">
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// Expect `bufferization.bufferize_to_allocation` to create an alloc.
|
|
// CHECK-LABEL: func.func @empty_to_tensor_alloc()
|
|
func.func @empty_to_tensor_alloc() -> tensor<2x2xf32> {
|
|
// CHECK-NEXT: %[[alloca:.*]] = memref.alloca() : memref<2x2xf32>
|
|
// CHECK-NEXT: %[[tensor:.*]] = bufferization.to_tensor %[[alloca]] restrict writable : memref<2x2xf32>
|
|
// CHECK-NEXT: return %[[tensor]] : tensor<2x2xf32>
|
|
%0 = bufferization.alloc_tensor() : tensor<2x2xf32>
|
|
return %0 : tensor<2x2xf32>
|
|
}
|