241 lines · plain
1// RUN: mlir-opt %s --transform-interpreter --split-input-file -canonicalize | FileCheck %s2 3// For pack op, we use lowerPadLikeWithInsertSlice = false to ensure no insert_slice is generated.4// This allows linalg.transpose to be fused as a producer operation. In below testcase, linalg.transpose5// as a producer operation is fused into the scf.forall loop.6 7module {8 // CHECK-label: func @fuse_pack_as_producer9 // CHECK: scf.forall {{.*}} {10 // CHECK: %[[PRODUCER:.*]] = linalg.transpose11 // CHECK: linalg.generic {{.*}} ins(%[[PRODUCER]]12 // CHECK: scf.forall.in_parallel13 // CHECK: }14 func.func @fuse_pack_as_producer(%src: tensor<128x256xf32>, %other: tensor<4x4x128x256xf32>)15 -> tensor<4x4x128x256xf32> {16 %dest = tensor.empty() : tensor<1x1x128x256xf32>17 %pack = linalg.pack %src inner_dims_pos = [0, 1] inner_tiles = [128, 256]18 into %dest : tensor<128x256xf32> -> tensor<1x1x128x256xf32>19 20 %out = tensor.empty() : tensor<4x4x128x256xf32>21 %res = linalg.generic22 {indexing_maps = [affine_map<(i, j, k, l) -> (0, 0, k, l)>,23 affine_map<(i, j, k, l) -> (i, j, k, l)>,24 affine_map<(i, j, k, l) -> (i, j, k, l)>],25 iterator_types = ["parallel", "parallel", "parallel", "parallel"]}26 ins(%pack, %other: tensor<1x1x128x256xf32>, tensor<4x4x128x256xf32>)27 outs(%out: tensor<4x4x128x256xf32>) {28 ^bb0(%pack_elem: f32, %other_elem: f32, %out_elem: f32):29 %r = arith.addf %pack_elem, %other_elem : f3230 linalg.yield %r : f3231 } -> tensor<4x4x128x256xf32>32 33 return %res : tensor<4x4x128x256xf32>34 }35 36 module attributes {transform.with_named_sequence} {37 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {38 // Find and lower pack operation.39 %pack = transform.structured.match ops{["linalg.pack"]} in %arg140 : (!transform.any_op) -> !transform.op<"linalg.pack">41 %paded, %expanded, %transpose = transform.structured.lower_pack %pack {lowerPadLikeWithInsertSlice = false}42 : (!transform.op<"linalg.pack">)43 -> (!transform.op<"tensor.pad">,44 !transform.op<"tensor.expand_shape">,45 !transform.op<"linalg.transpose">)46 47 %root = transform.structured.match ops{["linalg.generic"]} in %arg148 : (!transform.any_op) -> !transform.any_op49 // Tile the lialg operation with parallel forall loop tiling [4, 4].50 %tiled_op, %forall_op = transform.structured.tile_using_forall %root num_threads [4, 4]51 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)52 53 // Fuse the transpose operation into the tiled loop.54 transform.structured.fuse_into_containing_op %transpose into %forall_op55 : (!transform.op<"linalg.transpose">, !transform.any_op) -> (!transform.any_op, !transform.any_op)56 transform.yield57 }58 }59}60 61// -----62// For pack op, by default lowerPadLikeWithInsertSlice = true, which generates insert_slice and blocks fusion.63// In below testcase, tensor.insert_slice as a producer operation cannot be fused into the scf.forall loop.64 65module {66 // CHECK-label: func @fuse_pack_as_producer_blocked_by_insert_slice67 // CHECK: %[[PRODUCER:.*]] = tensor.insert_slice68 // CHECK: scf.forall {{.*}} {69 // CHECK: linalg.generic {{.*}} ins(%[[PRODUCER]]70 // CHECK: scf.forall.in_parallel71 // CHECK: }72 func.func @fuse_pack_as_producer_blocked_by_insert_slice(%src: tensor<128x256xf32>, %other: tensor<4x4x128x256xf32>)73 -> tensor<4x4x128x256xf32> {74 %dest = tensor.empty() : tensor<1x1x128x256xf32>75 %pack = linalg.pack %src inner_dims_pos = [0, 1] inner_tiles = [128, 256]76 into %dest : tensor<128x256xf32> -> tensor<1x1x128x256xf32>77 78 %out = tensor.empty() : tensor<4x4x128x256xf32>79 %res = linalg.generic80 {indexing_maps = [affine_map<(i, j, k, l) -> (0, 0, k, l)>,81 affine_map<(i, j, k, l) -> (i, j, k, l)>,82 affine_map<(i, j, k, l) -> (i, j, k, l)>],83 iterator_types = ["parallel", "parallel", "parallel", "parallel"]}84 ins(%pack, %other: tensor<1x1x128x256xf32>, tensor<4x4x128x256xf32>)85 outs(%out: tensor<4x4x128x256xf32>) {86 ^bb0(%pack_elem: f32, %other_elem: f32, %out_elem: f32):87 %r = arith.addf %pack_elem, %other_elem : f3288 linalg.yield %r : f3289 } -> tensor<4x4x128x256xf32>90 91 return %res : tensor<4x4x128x256xf32>92 }93 94 module attributes {transform.with_named_sequence} {95 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {96 // Find and lower pack operation.97 %pack = transform.structured.match ops{["linalg.pack"]} in %arg198 : (!transform.any_op) -> !transform.op<"linalg.pack">99 %paded, %expanded, %transpose = transform.structured.lower_pack %pack100 : (!transform.op<"linalg.pack">)101 -> (!transform.op<"tensor.pad">,102 !transform.op<"tensor.expand_shape">,103 !transform.op<"linalg.transpose">)104 105 %root = transform.structured.match ops{["linalg.generic"]} in %arg1106 : (!transform.any_op) -> !transform.any_op107 // Tile the lialg operation with parallel forall loop tiling [4, 4].108 %tiled_op, %forall_op = transform.structured.tile_using_forall %root num_threads [4, 4]109 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)110 111 // Fuse the transpose operation into the tiled loop.112 transform.structured.fuse_into_containing_op %transpose into %forall_op113 : (!transform.op<"linalg.transpose">, !transform.any_op) -> (!transform.any_op, !transform.any_op)114 transform.yield115 }116 }117}118 119// -----120// For unpack op, we use lowerUnpadLikeWithExtractSlice = false to ensure no extract_slice is generated.121// This allows linalg.transpose to be fused as a consumer operation. In below testcase, linalg.transpose122// as a consumer operation is fused into the scf.forall loop.123module {124 // CHECK-label: func @fuse_unpack_as_consumer125 // CHECK: scf.forall {{.*}} {126 // CHECK: %[[CONSUMER:.*]] = linalg.generic127 // CHECK: linalg.transpose ins(%[[CONSUMER]]128 // CHECK: scf.forall.in_parallel129 // CHECK: }130 func.func @fuse_unpack_as_consumer(%src: tensor<4x4x128x256xf32>, %other: tensor<4x4x128x256xf32>)131 -> tensor<128x256xf32> {132 %out = tensor.empty() : tensor<1x1x128x256xf32>133 %res = linalg.generic134 {indexing_maps = [affine_map<(i, j, k, l) -> (i, j, k, l)>,135 affine_map<(i, j, k, l) -> (i, j, k, l)>,136 affine_map<(i, j, k, l) -> (0, 0, k, l)>],137 iterator_types = ["parallel", "parallel", "parallel", "parallel"]}138 ins(%src, %other: tensor<4x4x128x256xf32>, tensor<4x4x128x256xf32>)139 outs(%out: tensor<1x1x128x256xf32>) {140 ^bb0(%unpack_elem: f32, %other_elem: f32, %out_elem: f32):141 %r = arith.addf %unpack_elem, %other_elem : f32142 linalg.yield %r : f32143 } -> tensor<1x1x128x256xf32>144 145 %dest = tensor.empty() : tensor<128x256xf32>146 %unpack = linalg.unpack %res inner_dims_pos = [0, 1] inner_tiles = [128, 256]147 into %dest : tensor<1x1x128x256xf32> -> tensor<128x256xf32>148 149 return %unpack : tensor<128x256xf32>150 }151 152 module attributes {transform.with_named_sequence} {153 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {154 // Find and lower unpack operation.155 %unpack = transform.structured.match ops{["linalg.unpack"]} in %arg1156 : (!transform.any_op) -> !transform.op<"linalg.unpack">157 transform.structured.lower_unpack %unpack {lowerUnpadLikeWithExtractSlice = false}158 : (!transform.op<"linalg.unpack">)159 -> (!transform.op<"tensor.empty">,160 !transform.op<"linalg.transpose">,161 !transform.op<"tensor.collapse_shape">,162 !transform.op<"tensor.extract_slice">)163 164 %root = transform.structured.match ops{["linalg.generic"]} in %arg1165 : (!transform.any_op) -> !transform.any_op166 // Tile the lialg operation with parallel forall loop tiling [4, 4].167 %tiled_op, %forall_op = transform.structured.tile_using_forall %root num_threads [4, 4]168 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)169 170 // Fuse the consumer operation into the tiled loop.171 %slice_op = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %forall_op172 : (!transform.any_op) -> !transform.op<"tensor.parallel_insert_slice">173 transform.test.fuse_consumer_using_slice %slice_op in (%forall_op)174 : (!transform.op<"tensor.parallel_insert_slice">, !transform.any_op) -> (!transform.any_op, !transform.any_op)175 transform.yield176 }177 }178}179 180// -----181// For unpack op, by default lowerUnpadLikeWithExtractSlice = true, which generates extract_slice and blocks fusion.182// In below testcase, tensor.extract_slice as a consumer operation cannot be fused into the scf.forall loop.183module {184 // CHECK-label: func @fuse_unpack_as_consumer_blocked_by_extract_slice185 // CHECK: %[[CONSUMER:.*]] = scf.forall {{.*}} {186 // CHECK: %[[ADDF:.*]] = linalg.generic187 // CHECK: scf.forall.in_parallel188 // CHECK: tensor.parallel_insert_slice %[[ADDF]]189 // CHECK: }190 // CHECK: tensor.extract_slice %[[CONSUMER]]191 func.func @fuse_unpack_as_consumer_blocked_by_extract_slice(%src: tensor<4x4x128x256xf32>, %other: tensor<4x4x128x256xf32>)192 -> tensor<128x256xf32> {193 %out = tensor.empty() : tensor<1x1x128x256xf32>194 %res = linalg.generic195 {indexing_maps = [affine_map<(i, j, k, l) -> (i, j, k, l)>,196 affine_map<(i, j, k, l) -> (i, j, k, l)>,197 affine_map<(i, j, k, l) -> (0, 0, k, l)>],198 iterator_types = ["parallel", "parallel", "parallel", "parallel"]}199 ins(%src, %other: tensor<4x4x128x256xf32>, tensor<4x4x128x256xf32>)200 outs(%out: tensor<1x1x128x256xf32>) {201 ^bb0(%unpack_elem: f32, %other_elem: f32, %out_elem: f32):202 %r = arith.addf %unpack_elem, %other_elem : f32203 linalg.yield %r : f32204 } -> tensor<1x1x128x256xf32>205 206 %dest = tensor.empty() : tensor<128x256xf32>207 %unpack = linalg.unpack %res inner_dims_pos = [0, 1] inner_tiles = [128, 256]208 into %dest : tensor<1x1x128x256xf32> -> tensor<128x256xf32>209 210 return %unpack : tensor<128x256xf32>211 }212 213 module attributes {transform.with_named_sequence} {214 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {215 // Find and lower unpack operation.216 %unpack = transform.structured.match ops{["linalg.unpack"]} in %arg1217 : (!transform.any_op) -> !transform.op<"linalg.unpack">218 transform.structured.lower_unpack %unpack219 : (!transform.op<"linalg.unpack">)220 -> (!transform.op<"tensor.empty">,221 !transform.op<"linalg.transpose">,222 !transform.op<"tensor.collapse_shape">,223 !transform.op<"tensor.extract_slice">)224 225 %root = transform.structured.match ops{["linalg.generic"]} in %arg1226 : (!transform.any_op) -> !transform.any_op227 // Tile the lialg operation with parallel forall loop tiling [4, 4].228 %tiled_op, %forall_op = transform.structured.tile_using_forall %root num_threads [4, 4]229 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)230 231 // Fuse the consumer operation into the tiled loop.232 %slice_op = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %forall_op233 : (!transform.any_op) -> !transform.op<"tensor.parallel_insert_slice">234 // Note that we cannot apply transform.test.fuse_consumer_using_slice here because the extract_slice235 // is not qualified consumer operation. Forcing this will yeild "could not fetch consumer236 // to fuse" error.237 transform.yield238 }239 }240}241