llvm-project/mlir/test/Dialect/Standard/func-bufferize.mlir
River Riddle 015192c634 [mlir:DialectConversion] Restructure how argument/target materializations get invoked
The current implementation invokes materializations
whenever an input operand does not have a mapping for the
desired type, i.e. it requires materialization at the earliest possible
point. This conflicts with goal of dialect conversion (and also the
current documentation) which states that a materialization is only
required if the materialization is supposed to persist after the
conversion process has finished.

This revision refactors this such that whenever a target
materialization "might" be necessary, we insert an
unrealized_conversion_cast to act as a temporary materialization.
This allows for deferring the invocation of the user
materialization hooks until the end of the conversion process,
where we actually have a better sense if it's actually
necessary. This has several benefits:

* In some cases a target materialization hook is no longer
   necessary
When performing a full conversion, there are some situations
where a temporary materialization is necessary. Moving forward,
these users won't need to provide any target materializations,
as the temporary materializations do not require the user to
provide materialization hooks.

* getRemappedValue can now handle values that haven't been
   converted yet
Before this commit, it wasn't well supported to get the remapped
value of a value that hadn't been converted yet (making it
difficult/impossible to convert multiple operations in many
situations). This commit updates getRemappedValue to properly
handle this case by inserting temporary materializations when
necessary.

Another code-health related benefit is that with this change we
can move a majority of the complexity related to materializations
to the end of the conversion process, instead of handling adhoc
while conversion is happening.

Differential Revision: https://reviews.llvm.org/D111620
2021-10-27 02:09:04 +00:00

84 lines
3.0 KiB
MLIR

// RUN: mlir-opt %s -func-bufferize -split-input-file -verify-diagnostics | FileCheck %s
// CHECK-LABEL: func @identity(
// CHECK-SAME: %[[ARG:.*]]: memref<f32>) -> memref<f32> {
// CHECK: return %[[ARG]] : memref<f32>
func @identity(%arg0: tensor<f32>) -> tensor<f32> {
return %arg0 : tensor<f32>
}
// CHECK-LABEL: func @block_arguments(
// CHECK-SAME: %[[ARG:.*]]: memref<f32>) -> memref<f32> {
// CHECK: br ^bb1(%[[ARG]] : memref<f32>)
// CHECK: ^bb1(%[[BBARG:.*]]: memref<f32>):
// CHECK: return %[[BBARG]] : memref<f32>
func @block_arguments(%arg0: tensor<f32>) -> tensor<f32> {
br ^bb1(%arg0: tensor<f32>)
^bb1(%bbarg: tensor<f32>):
return %bbarg : tensor<f32>
}
// CHECK-LABEL: func private @source() -> memref<f32>
// CHECK-LABEL: func @call_source() -> memref<f32> {
// CHECK: %[[RET:.*]] = call @source() : () -> memref<f32>
// CHECK: return %[[RET]] : memref<f32>
func private @source() -> tensor<f32>
func @call_source() -> tensor<f32> {
%0 = call @source() : () -> tensor<f32>
return %0 : tensor<f32>
}
// CHECK-LABEL: func @call_sink(
// CHECK-SAME: %[[ARG:.*]]: memref<f32>) {
// CHECK: call @sink(%[[ARG]]) : (memref<f32>) -> ()
// CHECK: return
func private @sink(tensor<f32>)
func @call_sink(%arg0: tensor<f32>) {
call @sink(%arg0) : (tensor<f32>) -> ()
return
}
// CHECK-LABEL: func @unconverted_op_in_body() -> memref<f32> {
// CHECK: %[[TENSOR:.*]] = "test.source"() : () -> tensor<f32>
// CHECK: %[[MEMREF:.*]] = memref.buffer_cast %[[TENSOR]] : memref<f32>
// CHECK: return %[[MEMREF]] : memref<f32>
func @unconverted_op_in_body() -> tensor<f32> {
%0 = "test.source"() : () -> tensor<f32>
return %0 : tensor<f32>
}
// -----
// Because this pass updates block arguments, it needs to also atomically
// update all terminators and issue an error if that is not possible.
func @unable_to_update_terminator(%arg0: tensor<f32>) -> tensor<f32> {
%0 = arith.constant true
cond_br %0, ^bb1(%arg0: tensor<f32>), ^bb2(%arg0: tensor<f32>)
^bb1(%bbarg0: tensor<f32>):
// expected-error @+1 {{failed to legalize operation 'test.terminator'}}
"test.terminator"() : () -> ()
^bb2(%bbarg1: tensor<f32>):
return %bbarg1 : tensor<f32>
}
// -----
// There was a bug in func-bufferize pass which caused terminators without
// ReturnLike and BranchOpInterface traits (e.g. scf.condition) to always
// fail to legalize even if bufferization doesn't needed.
// Check the pass succedeed.
// CHECK: bufferize_while
// CHECK: scf.while
// CHECK: scf.condition
func @bufferize_while(%arg0: i64, %arg1: i64) -> i64 {
%c2_i64 = arith.constant 2 : i64
%0:2 = scf.while (%arg2 = %arg0) : (i64) -> (i64, i64) {
%1 = arith.cmpi slt, %arg2, %arg1 : i64
scf.condition(%1) %arg2, %arg2 : i64, i64
} do {
^bb0(%arg2: i64, %arg3: i64):
%1 = arith.muli %arg3, %c2_i64 : i64
scf.yield %1 : i64
}
return %0#1 : i64
}