441b672bbd
With this PR I am trying to address: https://github.com/llvm/llvm-project/issues/63230. What changed: - While merging identical blocks, don't add a block argument if it is "identical" to another block argument. I.e., if the two block arguments refer to the same `Value`. The operations operands in the block will point to the argument we already inserted. This needs to happen to all the arguments we pass to the different successors of the parent block - After merged the blocks, get rid of "unnecessary" arguments. I.e., if all the predecessors pass the same block argument, there is no need to pass it as an argument. - This last simplification clashed with `BufferDeallocationSimplification`. The reason, I think, is that the two simplifications are clashing. I.e., `BufferDeallocationSimplification` contains an analysis based on the block structure. If we simplify the block structure (by merging and/or dropping block arguments) the analysis is invalid . The solution I found is to do a more prudent simplification when running that pass. **Note-1**: I ran all the integration tests (`-DMLIR_INCLUDE_INTEGRATION_TESTS=ON`) and they passed. **Note-2**: I fixed a bug found by @Dinistro in #97697 . The issue was that, when looking for redundant arguments, I was not considering that the block might have already some arguments. So the index (in the block args list) of the i-th `newArgument` is `i+numOfOldArguments`.
178 lines
5.3 KiB
MLIR
178 lines
5.3 KiB
MLIR
// RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -pass-pipeline="builtin.module(func.func(linalg-detensorize))" | FileCheck %s
|
|
|
|
#map0 = affine_map<() -> ()>
|
|
|
|
#attrs = {
|
|
indexing_maps = [#map0, #map0, #map0],
|
|
iterator_types = []
|
|
}
|
|
|
|
func.func @main() -> (tensor<i32>) attributes {} {
|
|
%c0 = arith.constant 0 : i32
|
|
%0 = tensor.from_elements %c0 : tensor<i32>
|
|
%c10 = arith.constant 10 : i32
|
|
%1 = tensor.from_elements %c10 : tensor<i32>
|
|
cf.br ^bb1(%0 : tensor<i32>)
|
|
|
|
^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2
|
|
%3 = tensor.empty() : tensor<i1>
|
|
%4 = linalg.generic #attrs
|
|
ins(%2, %1 : tensor<i32>, tensor<i32>)
|
|
outs(%3 : tensor<i1>) {
|
|
^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
|
|
%8 = arith.cmpi slt, %arg0, %arg1 : i32
|
|
linalg.yield %8 : i1
|
|
} -> tensor<i1>
|
|
%5 = tensor.extract %4[] : tensor<i1>
|
|
cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>)
|
|
|
|
^bb2(%6: tensor<i32>): // pred: ^bb1
|
|
%7 = tensor.empty() : tensor<i32>
|
|
%8 = linalg.generic #attrs
|
|
ins(%6, %6 : tensor<i32>, tensor<i32>)
|
|
outs(%7 : tensor<i32>) {
|
|
^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
|
|
%9 = arith.addi %arg0, %arg1 : i32
|
|
linalg.yield %9 : i32
|
|
} -> tensor<i32>
|
|
cf.br ^bb3(%8 : tensor<i32>)
|
|
|
|
^bb3(%10: tensor<i32>): // pred: ^bb1
|
|
return %10 : tensor<i32>
|
|
}
|
|
|
|
// CHECK-LABEL: func @main()
|
|
// CHECK-DAG: %[[cst:.*]] = arith.constant dense<0>
|
|
// CHECK-DAG: arith.constant true
|
|
// CHECK: cf.br
|
|
// CHECK-NEXT: ^[[bb1:.*]]:
|
|
// CHECK-NEXT: cf.cond_br %{{.*}}, ^[[bb2:.*]], ^bb3
|
|
// CHECK-NEXT: ^[[bb2]]
|
|
// CHECK-NEXT: cf.br ^[[bb3:.*]]
|
|
// CHECK-NEXT: ^[[bb3]]
|
|
// CHECK-NEXT: return %[[cst]]
|
|
// CHECK-NEXT: }
|
|
|
|
// -----
|
|
|
|
// Similar to the above test with one change: one of the block after the
|
|
// if-condition passes/forwards its tensor argument to another block.
|
|
|
|
#map0 = affine_map<() -> ()>
|
|
|
|
#attrs = {
|
|
indexing_maps = [#map0, #map0, #map0],
|
|
iterator_types = []
|
|
}
|
|
|
|
func.func @main() -> (tensor<i32>) attributes {} {
|
|
%c0 = arith.constant 0 : i32
|
|
%0 = tensor.from_elements %c0 : tensor<i32>
|
|
%c10 = arith.constant 10 : i32
|
|
%1 = tensor.from_elements %c10 : tensor<i32>
|
|
cf.br ^bb1(%0 : tensor<i32>)
|
|
|
|
^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2
|
|
%3 = tensor.empty() : tensor<i1>
|
|
%4 = linalg.generic #attrs
|
|
ins(%2, %1 : tensor<i32>, tensor<i32>)
|
|
outs(%3 : tensor<i1>) {
|
|
^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
|
|
%8 = arith.cmpi slt, %arg0, %arg1 : i32
|
|
linalg.yield %8 : i1
|
|
} -> tensor<i1>
|
|
%5 = tensor.extract %4[] : tensor<i1>
|
|
cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>)
|
|
|
|
^bb2(%6: tensor<i32>): // pred: ^bb1
|
|
%7 = tensor.empty() : tensor<i32>
|
|
%8 = linalg.generic #attrs
|
|
ins(%6, %6 : tensor<i32>, tensor<i32>)
|
|
outs(%7 : tensor<i32>) {
|
|
^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
|
|
%9 = arith.addi %arg0, %arg1 : i32
|
|
linalg.yield %9 : i32
|
|
} -> tensor<i32>
|
|
cf.br ^bb3(%8 : tensor<i32>)
|
|
|
|
^bb3(%10: tensor<i32>): // pred: ^bb1
|
|
cf.br ^bb4(%10 : tensor<i32>)
|
|
|
|
^bb4(%11: tensor<i32>): // pred: ^bb1
|
|
return %11 : tensor<i32>
|
|
}
|
|
|
|
// CHECK-LABEL: func @main()
|
|
// CHECK-DAG: %[[cst:.*]] = arith.constant dense<0>
|
|
// CHECK-DAG: arith.constant true
|
|
// CHECK: cf.br ^[[bb1:.*]]
|
|
// CHECK-NEXT: ^[[bb1:.*]]:
|
|
// CHECK-NEXT: cf.cond_br %{{.*}}, ^[[bb2:.*]], ^bb3
|
|
// CHECK-NEXT: ^[[bb2]]:
|
|
// CHECK-NEXT: cf.br ^[[bb3:.*]]
|
|
// CHECK-NEXT: ^[[bb3]]:
|
|
// CHECK-NEXT: cf.br ^[[bb4:.*]]
|
|
// CHECK-NEXT: ^[[bb4]]:
|
|
// CHECK-NEXT: return %[[cst]]
|
|
// CHECK-NEXT: }
|
|
|
|
// -----
|
|
|
|
#map0 = affine_map<() -> ()>
|
|
|
|
#attrs = {
|
|
indexing_maps = [#map0, #map0, #map0],
|
|
iterator_types = []
|
|
}
|
|
|
|
func.func @main() -> (tensor<i32>) attributes {} {
|
|
%c0 = arith.constant 0 : i32
|
|
%0 = tensor.from_elements %c0 : tensor<i32>
|
|
%c10 = arith.constant 10 : i32
|
|
%1 = tensor.from_elements %c10 : tensor<i32>
|
|
cf.br ^bb1(%0 : tensor<i32>)
|
|
|
|
^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2
|
|
%3 = tensor.empty() : tensor<i1>
|
|
%4 = linalg.generic #attrs
|
|
ins(%2, %1 : tensor<i32>, tensor<i32>)
|
|
outs(%3 : tensor<i1>) {
|
|
^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
|
|
%8 = arith.cmpi slt, %arg0, %arg1 : i32
|
|
linalg.yield %8 : i1
|
|
} -> tensor<i1>
|
|
%5 = tensor.extract %4[] : tensor<i1>
|
|
// This cf.cond_br intentionally has bb2 as it's target for both branches. This
|
|
// is to make sure that the "forward phase" of the cost-model correctly adds
|
|
// the users of a block argument (in this case bb2's argument) to the work
|
|
// list.
|
|
cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb2(%2 : tensor<i32>)
|
|
|
|
^bb2(%6: tensor<i32>): // pred: ^bb1
|
|
%12 = tensor.from_elements %c10 : tensor<i32>
|
|
%7 = tensor.empty() : tensor<i32>
|
|
%8 = linalg.generic #attrs
|
|
ins(%6, %12 : tensor<i32>, tensor<i32>)
|
|
outs(%7 : tensor<i32>) {
|
|
^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
|
|
%9 = arith.addi %arg0, %arg1 : i32
|
|
linalg.yield %9 : i32
|
|
} -> tensor<i32>
|
|
cf.br ^bb3(%8 : tensor<i32>)
|
|
|
|
^bb3(%10: tensor<i32>): // pred: ^bb1
|
|
return %10 : tensor<i32>
|
|
}
|
|
|
|
// CHECK-LABEL: func @main()
|
|
// CHECK-DAG: %[[cst:.*]] = arith.constant dense<10>
|
|
// CHECK-DAG: arith.constant true
|
|
// CHECK: cf.br ^[[bb1:.*]]
|
|
// CHECK-NEXT: ^[[bb1]]:
|
|
// CHECK-NEXT: cf.cond_br %{{.*}}, ^[[bb2:.*]], ^bb2
|
|
// CHECK-NEXT: ^[[bb2]]
|
|
// CHECK-NEXT: cf.br ^[[bb3:.*]]
|
|
// CHECK-NEXT: ^[[bb3]]
|
|
// CHECK-NEXT: return %[[cst]]
|
|
// CHECK-NEXT: }
|