|
2 | 2 | // DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule |\
|
3 | 3 | // DEFINE: mlir-opt --test-linalg-transform-patterns="test-decompose-tensor-pack"\
|
4 | 4 | // DEFINE: --test-transform-dialect-erase-schedule \
|
5 |
| -// DEFINE: -one-shot-bufferize="bufferize-function-boundaries" \ |
6 |
| -// DEFINE: -buffer-deallocation-pipeline="private-function-dynamic-ownership" \ |
7 |
| -// DEFINE: -cse -canonicalize -test-lower-to-llvm -o %t |
| 5 | +// DEFINE: -test-lower-to-llvm -o %t |
8 | 6 | // DEFINE: %{entry_point} = main
|
9 | 7 | // DEFINE: %{run} = mlir-cpu-runner %t -e %{entry_point} -entry-point-result=void \
|
10 | 8 | // DEFINE: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils
|
@@ -84,12 +82,30 @@ func.func private @pack(%A: tensor<7x16xi32>) {
|
84 | 82 | }
|
85 | 83 |
|
86 | 84 | module @transforms attributes { transform.with_named_sequence } {
|
87 |
| - transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) { |
| 85 | + transform.named_sequence @__transform_main(%module: !transform.any_op {transform.consume}) { |
88 | 86 | %pack = transform.structured.match ops{["tensor.pack"]} in %module : (!transform.any_op) -> !transform.any_op
|
89 | 87 |
|
90 |
| - %tiled_linalg_op_p, %loops:2 = transform.structured.tile_using_for %pack tile_sizes [1, 1] |
| 88 | + // 1. Tile so that we can decompose tensor.pack into tensor.pad, |
| 89 | + // linalg.transpose, etc (see step 2). |
| 90 | + %tiled_pack_op_p, %loops:2 = transform.structured.tile_using_for %pack tile_sizes [1, 1] |
91 | 91 | : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
|
92 | 92 |
|
| 93 | + // 2. Decompose the tiled Op into tensor.pad etc |
| 94 | + %func_1 = transform.get_parent_op %tiled_pack_op_p {isolated_from_above} : (!transform.any_op) -> !transform.any_op |
| 95 | + transform.apply_patterns to %func_1 { |
| 96 | + transform.apply_patterns.linalg.decompose_pack_unpack |
| 97 | + } : !transform.any_op |
| 98 | + |
| 99 | + // 3. Bufferize before lowering to LLVM |
| 100 | + %bufferize = transform.bufferization.one_shot_bufferize %module |
| 101 | + {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op |
| 102 | + |
| 103 | + // 4. Canonicalize |
| 104 | + %func_2 = transform.structured.match ops{["func.func"]} in %bufferize : (!transform.any_op) -> !transform.op<"func.func"> |
| 105 | + transform.apply_patterns to %func_2 { |
| 106 | + transform.apply_patterns.canonicalization |
| 107 | + } : !transform.op<"func.func"> |
| 108 | + |
93 | 109 | transform.yield
|
94 | 110 | }
|
95 | 111 | }
|
|
0 commit comments