107 lines · plain
1module @transforms attributes { transform.with_named_sequence } {2 3 //===----------------------------------------------------------------------===//4 // TD sequence _without_ vectorization5 //===----------------------------------------------------------------------===//6 transform.named_sequence @__transform_main_basic(%module: !transform.any_op {transform.consume}) {7 %pack = transform.structured.match ops{["linalg.pack"]} in %module : (!transform.any_op) -> !transform.any_op8 %unpack = transform.structured.match ops{["linalg.unpack"]} in %module : (!transform.any_op) -> !transform.any_op9 10 // 1.1 Tile the linalg.pack Op so that we can decompose it into e.g. tensor.pad11 // and other lower-level Ops (see step 2.1)12 %tiled_pack_op_p, %loops_pack:2 = transform.structured.tile_using_for %pack tile_sizes [1, 1]13 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)14 15 // 1.2 Tile the linalg.unpack Op so that we can decompose it into e.g. tensor.pad16 // and other lower-level Ops (see step 2.2)17 %tiled_unpack_op_p, %loops_unpack:2 = transform.structured.tile_using_for %unpack tile_sizes [4, 1]18 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)19 20 // 2.1. Decompose tiled PackOp into lower-level Ops21 %func_op_pack = transform.get_parent_op %tiled_pack_op_p {isolated_from_above} : (!transform.any_op) -> !transform.op<"func.func">22 transform.apply_patterns to %func_op_pack {23 transform.apply_patterns.linalg.decompose_pack_unpack24 transform.apply_patterns.linalg.decompose_pad25 } : !transform.op<"func.func">26 27 transform.apply_patterns to %func_op_pack {28 transform.apply_patterns.tensor.fold_tensor_subset_ops29 transform.apply_patterns.canonicalization30 } : !transform.op<"func.func">31 32 // 2.2. Decompose tiled UnpackOp into lower-level Ops33 %func_op_unpack = transform.get_parent_op %tiled_unpack_op_p {isolated_from_above} : (!transform.any_op) -> !transform.op<"func.func">34 transform.apply_patterns to %func_op_unpack {35 transform.apply_patterns.linalg.decompose_pack_unpack36 } : !transform.op<"func.func">37 38 transform.apply_patterns to %func_op_unpack {39 transform.apply_patterns.tensor.fold_tensor_subset_ops40 transform.apply_patterns.canonicalization41 } : !transform.op<"func.func">42 43 // 3. Bufferize before lowering to LLVM44 %bufferize = transform.bufferization.one_shot_bufferize %module45 {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op46 47 // 4. Canonicalize48 %func_op_bufferized = transform.structured.match ops{["func.func"]} in %bufferize : (!transform.any_op) -> !transform.op<"func.func">49 transform.apply_patterns to %func_op_bufferized {50 transform.apply_patterns.canonicalization51 } : !transform.op<"func.func">52 53 transform.yield54 }55 56 //===----------------------------------------------------------------------===//57 // TD sequence _with_ vectorization58 //===----------------------------------------------------------------------===//59 transform.named_sequence @__transform_main_vectorized(%module: !transform.any_op {transform.consume}) {60 %pack = transform.structured.match ops{["linalg.pack"]} in %module : (!transform.any_op) -> !transform.any_op61 %unpack = transform.structured.match ops{["linalg.unpack"]} in %module : (!transform.any_op) -> !transform.any_op62 63 // 1.1 Tile the linalg.pack Op so that we can decompose it into e.g. tensor.pad64 // and other lower-level Ops (see step 2.1)65 %tiled_pack_op_p, %loops_pack:2 = transform.structured.tile_using_for %pack tile_sizes [1, 1]66 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)67 68 // 1.2 Tile the linalg.unpack Op 69 %tiled_unpack_op_p, %loops_unpack:2 = transform.structured.tile_using_for %unpack tile_sizes [1, 1]70 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)71 72 // 2.1. Decompose tiled PackOp into lower-level Ops73 %func_op_pack = transform.get_parent_op %tiled_pack_op_p {isolated_from_above} : (!transform.any_op) -> !transform.op<"func.func">74 transform.apply_patterns to %func_op_pack {75 transform.apply_patterns.linalg.decompose_pack_unpack76 transform.apply_patterns.linalg.decompose_pad77 } : !transform.op<"func.func">78 79 transform.apply_patterns to %func_op_pack {80 transform.apply_patterns.tensor.fold_tensor_subset_ops81 transform.apply_patterns.canonicalization82 } : !transform.op<"func.func">83 84 // 2.2. Vectorize tiled UnpackOp into lower-level Ops85 %func_op_unpack = transform.get_parent_op %tiled_unpack_op_p {isolated_from_above} : (!transform.any_op) -> !transform.op<"func.func">86 transform.structured.vectorize %tiled_unpack_op_p vector_sizes [1, 1, 4, [4]] {assume_dynamic_dims_match_vec_sizes} : !transform.any_op87 88 transform.apply_patterns to %func_op_unpack {89 transform.apply_patterns.vector.transfer_permutation_patterns90 transform.apply_patterns.vector.lower_masked_transfers91 transform.apply_patterns.vector.sink_ops92 } : !transform.op<"func.func">93 94 // 3. Bufferize95 %bufferize = transform.bufferization.one_shot_bufferize %module96 {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op97 98 // 4. Canonicalize99 %func_op_bufferized = transform.structured.match ops{["func.func"]} in %bufferize : (!transform.any_op) -> !transform.op<"func.func">100 transform.apply_patterns to %func_op_bufferized {101 transform.apply_patterns.canonicalization102 } : !transform.op<"func.func">103 104 transform.yield105 }106}107