461dafd2a3
This commit allows for One-Shot Bufferize to be used through the transform dialect. No op handle is currently returned for the bufferized IR. Differential Revision: https://reviews.llvm.org/D125098
125 lines
3.8 KiB
MLIR
125 lines
3.8 KiB
MLIR
// RUN: mlir-opt --test-transform-dialect-interpreter %s -split-input-file -verify-diagnostics | FileCheck %s
|
|
|
|
// Test One-Shot Bufferize.
|
|
|
|
transform.with_pdl_patterns {
|
|
^bb0(%arg0: !pdl.operation):
|
|
sequence %arg0 {
|
|
^bb0(%arg1: !pdl.operation):
|
|
%0 = pdl_match @pdl_target in %arg1
|
|
transform.bufferization.one_shot_bufferize %0
|
|
{target_is_module = false}
|
|
}
|
|
|
|
pdl.pattern @pdl_target : benefit(1) {
|
|
%0 = operation "func.func"
|
|
rewrite %0 with "transform.dialect"
|
|
}
|
|
}
|
|
|
|
// 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: memref.dealloc %[[alloc]]
|
|
// CHECK: return %[[res_tensor]]
|
|
return %0 : tensor<?xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test analysis of One-Shot Bufferize only.
|
|
|
|
transform.with_pdl_patterns {
|
|
^bb0(%arg0: !pdl.operation):
|
|
sequence %arg0 {
|
|
^bb0(%arg1: !pdl.operation):
|
|
%0 = pdl_match @pdl_target in %arg1
|
|
transform.bufferization.one_shot_bufferize %0
|
|
{target_is_module = false, test_analysis_only = true}
|
|
}
|
|
|
|
pdl.pattern @pdl_target : benefit(1) {
|
|
%0 = operation "func.func"
|
|
rewrite %0 with "transform.dialect"
|
|
}
|
|
}
|
|
|
|
// 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`.
|
|
|
|
transform.with_pdl_patterns {
|
|
^bb0(%arg0: !pdl.operation):
|
|
sequence %arg0 {
|
|
^bb0(%arg1: !pdl.operation):
|
|
%0 = pdl_match @pdl_target in %arg1
|
|
// expected-error @+1 {{bufferization failed}}
|
|
transform.bufferization.one_shot_bufferize %0 {target_is_module = false}
|
|
}
|
|
|
|
pdl.pattern @pdl_target : benefit(1) {
|
|
%0 = operation "func.func"
|
|
rewrite %0 with "transform.dialect"
|
|
}
|
|
}
|
|
|
|
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>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test One-Shot Bufferize transform failure with a module op.
|
|
|
|
transform.with_pdl_patterns {
|
|
^bb0(%arg0: !pdl.operation):
|
|
sequence %arg0 {
|
|
^bb0(%arg1: !pdl.operation):
|
|
// %arg1 is the module
|
|
transform.bufferization.one_shot_bufferize %arg1
|
|
}
|
|
}
|
|
|
|
module {
|
|
// 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: memref.dealloc %[[alloc]]
|
|
// CHECK: return %[[res_tensor]]
|
|
return %0 : tensor<?xf32>
|
|
}
|
|
} |