405 lines · plain
1// RUN: mlir-opt %s -transform-interpreter --split-input-file | FileCheck %s2 3!A_mk = tensor<1023x255xf32>4!B_kn = tensor<255x127xf32>5!C_mn = tensor<1023x127xf32>6 7// Normalized dims are: ( k, m, n)(kk, mm, nn)8// CHECK-DAG: #[[$mk_kkmm:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d0, d3, d4)>9// CHECK-DAG: #[[$kn_kknn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d2, d3, d5)>10// CHECK-DAG: #[[$mn_mmnn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d2, d4, d5)>11 12// CHECK-LABEL: @matmul_mk_kn_mn(13func.func @matmul_mk_kn_mn(%A : !A_mk, %B : !B_kn, %C : !C_mn) -> !C_mn {14 // CHECK: linalg.generic15 // CHECK-SAME: indexing_maps = [#[[$mk_kkmm]], #[[$kn_kknn]], #[[$mn_mmnn]]]16 // CHECK-SAME: ["reduction", "parallel", "parallel", "reduction", "parallel", "parallel"]}17 // CHECK-SAME: ins(%{{.*}} : tensor<128x8x32x8xf32>, tensor<8x8x32x16xf32>)18 // CHECK-SAME: outs(%{{.*}} : tensor<128x8x8x16xf32>)19 %0 = linalg.matmul ins(%A, %B : !A_mk, !B_kn) outs(%C : !C_mn) -> !C_mn20 return %0 : !C_mn21}22 23module attributes {transform.with_named_sequence} {24 transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {25 %matmul = transform.structured.match ops{["linalg.matmul"]} in %module_op26 : (!transform.any_op) -> !transform.op<"linalg.matmul">27 transform.structured.pack_greedily %matmul28 matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]29 : (!transform.op<"linalg.matmul">) -> !transform.op<"linalg.generic">30 transform.yield31 }32}33 34// -----35 36!A_mk = tensor<1023x255xf32>37!B_nk = tensor<127x255xf32>38!C_nm = tensor<127x1023xf32>39 40#mkn_accesses = [41 affine_map<(m, n, k) -> (m, k)>,42 affine_map<(m, n, k) -> (n, k)>,43 affine_map<(m, n, k) -> (n, m)>44]45#mkn_trait = {46 indexing_maps = #mkn_accesses,47 iterator_types = ["parallel", "parallel", "reduction"]48}49 50// Normalized dims are: ( k, m, n)(kk, mm, nn)51// CHECK-DAG: #[[$km_kkmm:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d0, d3, d4)>52// CHECK-DAG: #[[$kn_kknn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d0, d3, d5)>53// CHECK-DAG: #[[$mn_mmnn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d1, d4, d5)>54 55// CHECK-LABEL: @matmul_mk_nk_nm(56func.func @matmul_mk_nk_nm(%A : !A_mk, %B : !B_nk, %C : !C_nm) -> !C_nm {57 // CHECK: linalg.generic58 // CHECK-SAME: indexing_maps = [#[[$mk_kkmm]], #[[$kn_kknn]], #[[$mn_mmnn]]]59 // CHECK-SAME: ["reduction", "parallel", "parallel", "reduction", "parallel", "parallel"]}60 // CHECK-SAME: ins(%{{.*}} : tensor<128x8x32x8xf32>, tensor<8x8x32x16xf32>)61 // CHECK-SAME: outs(%{{.*}} : tensor<8x128x8x16xf32>)62 %0 = linalg.generic #mkn_trait ins(%A, %B : !A_mk, !B_nk) outs(%C : !C_nm) {63 ^bb0(%a: f32, %b: f32, %c: f32):64 %d = arith.mulf %a, %b : f3265 %e = arith.addf %c, %d : f3266 linalg.yield %e : f3267 } -> !C_nm68 return %0 : !C_nm69}70 71module attributes {transform.with_named_sequence} {72 transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {73 %generic = transform.structured.match ops{["linalg.generic"]} in %module_op : (!transform.any_op) -> !transform.op<"linalg.generic">74 transform.structured.pack_greedily %generic75 matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]76 : (!transform.op<"linalg.generic">) -> !transform.op<"linalg.generic">77 transform.yield78 }79}80 81// -----82 83!A_mk = tensor<1023x255xf32>84!B_nk = tensor<127x255xf32>85!C_nm = tensor<127x1023xf32>86 87#mkn_accesses = [88 affine_map<(k, m, n) -> (m, k)>,89 affine_map<(k, m, n) -> (n, k)>,90 affine_map<(k, m, n) -> (n, m)>91]92#mkn_trait = {93 indexing_maps = #mkn_accesses,94 iterator_types = ["reduction", "parallel", "parallel"]95}96 97// Normalized dims are: ( k, m, n)(kk, mm, nn)98// CHECK-DAG: #[[$mk_kkmm:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d0, d3, d4)>99// CHECK-DAG: #[[$kn_kknn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d0, d3, d5)>100// CHECK-DAG: #[[$mn_mmnn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d1, d4, d5)>101 102// CHECK-LABEL: @matmul_mk_nk_nm_transposed(103func.func @matmul_mk_nk_nm_transposed(%A : !A_mk, %B : !B_nk, %C : !C_nm) -> !C_nm {104 // CHECK: linalg.generic105 // CHECK-SAME: indexing_maps = [#[[$mk_kkmm]], #[[$kn_kknn]], #[[$mn_mmnn]]]106 // CHECK-SAME: ["reduction", "parallel", "parallel", "reduction", "parallel", "parallel"]}107 // CHECK-SAME: ins(%{{.*}} : tensor<128x8x32x8xf32>, tensor<8x8x32x16xf32>)108 // CHECK-SAME: outs(%{{.*}} : tensor<8x128x8x16xf32>)109 %0 = linalg.generic #mkn_trait ins(%A, %B : !A_mk, !B_nk) outs(%C : !C_nm) {110 ^bb0(%a: f32, %b: f32, %c: f32):111 %d = arith.mulf %a, %b : f32112 %e = arith.addf %c, %d : f32113 linalg.yield %e : f32114 } -> !C_nm115 return %0 : !C_nm116}117 118module attributes {transform.with_named_sequence} {119 transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {120 %generic = transform.structured.match ops{["linalg.generic"]} in %module_op : (!transform.any_op) -> !transform.op<"linalg.generic">121 transform.structured.pack_greedily %generic122 matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]123 : (!transform.op<"linalg.generic">) -> !transform.op<"linalg.generic">124 transform.yield125 }126}127 128// -----129 130!A_bmkm2 = tensor<42x1023x255x33xf32>131!B_nkb = tensor<127x255x42xf32>132!C_nbm = tensor<127x42x1023xf32>133 134#mkn_accesses = [135 affine_map<(k, m, n, b, m2) -> (b, m, k, m2)>,136 affine_map<(k, m, n, b, m2) -> (n, k, b)>,137 affine_map<(k, m, n, b, m2) -> (n, b, m)>138]139#mkn_trait = {140 indexing_maps = #mkn_accesses,141 iterator_types = ["reduction", "parallel", "parallel", "parallel", "parallel"]142}143 144// Normalized dims are: ( ?, ?, k, m, n)(kk, mm, nn)145// CHECK-DAG: #[[$bmkm2_kkmm:.*]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d3, d2, d1, d5, d6)>146// CHECK-DAG: #[[$nkb_kknn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d4, d2, d0, d5, d7)>147// CHECK-DAG: #[[$nbm_mmnn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d4, d0, d3, d6, d7)>148 149// CHECK-LABEL: @contraction_bmkm2_nkb_nbm(150func.func @contraction_bmkm2_nkb_nbm(%A : !A_bmkm2, %B : !B_nkb, %C : !C_nbm) -> !C_nbm {151 // CHECK: linalg.generic152 // CHECK-SAME: indexing_maps = [#[[$bmkm2_kkmm]], #[[$nkb_kknn]], #[[$nbm_mmnn]]]153 // CHECK-SAME: ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction", "parallel", "parallel"]}154 // CHECK-SAME: ins(%{{.*}} : tensor<42x128x8x33x32x8xf32>, tensor<8x8x42x32x16xf32>)155 // CHECK-SAME: outs(%{{.*}} : tensor<8x42x128x8x16xf32>)156 %0 = linalg.generic #mkn_trait ins(%A, %B : !A_bmkm2, !B_nkb) outs(%C : !C_nbm) {157 ^bb0(%a: f32, %b: f32, %c: f32):158 %d = arith.mulf %a, %b : f32159 %e = arith.addf %c, %d : f32160 linalg.yield %e : f32161 } -> !C_nbm162 return %0 : !C_nbm163}164 165module attributes {transform.with_named_sequence} {166 transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {167 %generic = transform.structured.match ops{["linalg.generic"]} in %module_op : (!transform.any_op) -> !transform.op<"linalg.generic">168 transform.structured.pack_greedily %generic169 matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]170 : (!transform.op<"linalg.generic">) -> !transform.op<"linalg.generic">171 transform.yield172 }173}174 175// -----176 177// Conv linguo: h w kh kw c n f cc nn ff178// Normalized dims are: ( ?, ?, ?, ?, k, m, n)(kk, mm, nn)179// n c h + kh w + kw cc nn180// CHECK-DAG: #[[$M1:.*]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8, d9) -> (d5, d4, d0 + d2, d1 + d3, d7, d8)>181// f c kh kw cc ff182// CHECK-DAG: #[[$M2:.*]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8, d9) -> (d6, d4, d2, d3, d7, d9)>183// n f h w nn ff184// CHECK-DAG: #[[$M3:.*]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8, d9) -> (d5, d6, d0, d1, d8, d9)>185 186// CHECK-LABEL: @conv_2d_nchw_fchw187func.func @conv_2d_nchw_fchw(%arg0: tensor<?x47x16x16xf32>, %arg2: tensor<?x16x14x14xf32>) -> tensor<?x16x14x14xf32> {188 %c0 = arith.constant dense<0.1> : tensor<16x47x3x3xf32>189 // CHECK: linalg.generic190 // CHECK-SAME: indexing_maps = [#[[$M1]], #[[$M2]], #[[$M3]]]191 // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction", "reduction", "reduction", "parallel", "parallel", "reduction", "parallel", "parallel"]192 // CHECK-SAME: ins(%{{.*}} : tensor<?x2x16x16x32x8xf32>, tensor<1x2x3x3x32x16xf32>)193 // CHECK-SAME: outs(%{{.*}} : tensor<?x1x14x14x8x16xf32>)194 %0 = linalg.conv_2d_nchw_fchw195 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }196 ins(%arg0, %c0: tensor<?x47x16x16xf32>, tensor<16x47x3x3xf32>)197 outs(%arg2: tensor<?x16x14x14xf32>) -> tensor<?x16x14x14xf32>198 return %0 : tensor<?x16x14x14xf32>199}200 201module attributes {transform.with_named_sequence} {202 transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {203 %conv = transform.structured.match ops{["linalg.conv_2d_nchw_fchw"]} in %module_op204 : (!transform.any_op) -> !transform.op<"linalg.conv_2d_nchw_fchw">205 transform.structured.pack_greedily %conv206 matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]207 : (!transform.op<"linalg.conv_2d_nchw_fchw">) -> !transform.op<"linalg.generic">208 transform.yield209 }210}211 212 213// -----214 215// These should fail to pack for now as they don't contain a contraction.216// CHECK-LABEL: @reduce_and_map217func.func @reduce_and_map(%arg0: tensor<10x100xf32>,218 %arg1: tensor<10x100xf32>, %output: tensor<10xf32>) -> tensor<10xf32> {219 %map_init = tensor.empty() : tensor<10x100xf32>220 // CHECK: linalg.map221 %mapped = linalg.map { arith.addf }222 ins(%arg0, %arg1 : tensor<10x100xf32>, tensor<10x100xf32>)223 outs(%map_init : tensor<10x100xf32>)224 // CHECK: linalg.reduce225 %res = linalg.reduce { arith.addf }226 ins(%mapped: tensor<10x100xf32>)227 outs(%output: tensor<10xf32>)228 dimensions = [1]229 return %res : tensor<10xf32>230}231 232module attributes {transform.with_named_sequence} {233 transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {234 %generic = transform.structured.match ops{["linalg.generic"]} in %module_op : (!transform.any_op) -> !transform.op<"linalg.generic">235 transform.structured.pack_greedily %generic236 matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]237 : (!transform.op<"linalg.generic">) -> !transform.op<"linalg.generic">238 transform.yield239 }240}241 242// -----243 244!A_mk = tensor<1023x255xf32>245!B_nk = tensor<127x255xf32>246!C_nm = tensor<127x1023xf32>247 248#mkn_accesses = [249 affine_map<(m, n, k) -> (m, k)>,250 affine_map<(m, n, k) -> (n, k)>,251 affine_map<(m, n, k) -> (n, m)>252]253#mkn_trait = {254 indexing_maps = #mkn_accesses,255 iterator_types = ["parallel", "parallel", "reduction"]256}257 258// Normalized dims are: ( k, m, n)(kk, mm, nn)259// CHECK-DAG: #[[$km_kkmm:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d0, d3, d4)>260// CHECK-DAG: #[[$kn_kknn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d0, d3, d5)>261// CHECK-DAG: #[[$mn_mmnn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d1, d4, d5)>262 263// CHECK-LABEL: @matmul_mk_nk_nm(264func.func @matmul_mk_nk_nm(%A : !A_mk, %B : !B_nk, %C : !C_nm) -> !C_nm {265 // CHECK: linalg.generic266 // CHECK-SAME: indexing_maps = [#[[$mk_kkmm]], #[[$kn_kknn]], #[[$mn_mmnn]]]267 // CHECK-SAME: ["reduction", "parallel", "parallel", "reduction", "parallel", "parallel"]}268 // CHECK-SAME: ins(%{{.*}} : tensor<128x8x32x8xf32>, tensor<1x8x32x130xf32>)269 // CHECK-SAME: outs(%{{.*}} : tensor<1x128x8x130xf32>)270 %0 = linalg.generic #mkn_trait ins(%A, %B : !A_mk, !B_nk) outs(%C : !C_nm) {271 ^bb0(%a: f32, %b: f32, %c: f32):272 %d = arith.mulf %a, %b : f32273 %e = arith.addf %c, %d : f32274 linalg.yield %e : f32275 } -> !C_nm276 return %0 : !C_nm277}278 279module attributes {transform.with_named_sequence} {280 transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {281 %generic = transform.structured.match ops{["linalg.generic"]} in %module_op : (!transform.any_op) -> !transform.op<"linalg.generic">282 transform.structured.pack_greedily %generic283 // In this spec, the "k" dimension is not packed but rather padded to the284 // next multiple of 10 (i.e. 130).285 matmul_packed_sizes = [8, 0, 32]286 matmul_padded_sizes_next_multiple_of = [0, 10, 0]287 matmul_inner_dims_order = [1, 2, 0]288 : (!transform.op<"linalg.generic">) -> !transform.op<"linalg.generic">289 transform.yield290 }291}292 293 294// -----295 296!A_mk = tensor<1023x255xf32>297!B_nk = tensor<127x255xf32>298!C_nm = tensor<127x1023xf32>299 300#mkn_accesses = [301 affine_map<(m, n, k) -> (m, k)>,302 affine_map<(m, n, k) -> (n, k)>,303 affine_map<(m, n, k) -> (n, m)>304]305#mkn_trait = {306 indexing_maps = #mkn_accesses,307 iterator_types = ["parallel", "parallel", "reduction"]308}309 310// Normalized dims are: ( k, m, n)(kk, mm)311// CHECK-DAG: #[[$km_kkmm:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d0, d3)>312// CHECK-DAG: #[[$kn_kknn:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d0, d3, d4)>313// CHECK-DAG: #[[$mn_mmnn:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d1, d4)>314 315// CHECK-LABEL: @matmul_mk_nk_nm(316func.func @matmul_mk_nk_nm(%A : !A_mk, %B : !B_nk, %C : !C_nm) -> !C_nm {317 // CHECK: linalg.generic318 // CHECK-SAME: indexing_maps = [#[[$mk_kkmm]], #[[$kn_kknn]], #[[$mn_mmnn]]]319 // CHECK-SAME: ["reduction", "parallel", "parallel", "reduction", "parallel"]}320 // CHECK-SAME: ins(%{{.*}} : tensor<1023x8x32xf32>, tensor<1x8x32x130xf32>)321 // CHECK-SAME: outs(%{{.*}} : tensor<1x1023x130xf32>)322 %0 = linalg.generic #mkn_trait ins(%A, %B : !A_mk, !B_nk) outs(%C : !C_nm) {323 ^bb0(%a: f32, %b: f32, %c: f32):324 %d = arith.mulf %a, %b : f32325 %e = arith.addf %c, %d : f32326 linalg.yield %e : f32327 } -> !C_nm328 return %0 : !C_nm329}330 331module attributes {transform.with_named_sequence} {332 transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {333 %generic = transform.structured.match ops{["linalg.generic"]} in %module_op : (!transform.any_op) -> !transform.op<"linalg.generic">334 transform.structured.pack_greedily %generic335 // In this spec, the "n" dimension is neither packed not unpacked.336 // We don't end up with an innermost matmul after packing but only with an337 // innermost matvec.338 matmul_packed_sizes = [0, 0, 32]339 matmul_padded_sizes_next_multiple_of = [0, 10, 0]340 matmul_inner_dims_order = [1, 2, 0]341 : (!transform.op<"linalg.generic">) -> !transform.op<"linalg.generic">342 transform.yield343 }344}345 346// -----347 348!A = tensor<1023x255xf32>349!X = tensor<255xf32>350!Y = tensor<1023xf32>351 352// CHECK-LABEL: @matvec_fail(353func.func @matvec_fail(%A : !A, %x : !X, %y : !Y) -> !Y {354 // CHECK: linalg.matvec355 %0 = linalg.matvec ins(%A, %x : !A, !X) outs(%y : !Y) -> !Y356 return %0 : !Y357}358 359module attributes {transform.with_named_sequence} {360 transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {361 %matmul = transform.structured.match ops{["linalg.matvec"]} in %module_op362 : (!transform.any_op) -> !transform.op<"linalg.matvec">363 transform.structured.pack_greedily %matmul364 matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]365 : (!transform.op<"linalg.matvec">) -> !transform.any_op366 transform.yield367 }368}369 370// -----371 372func.func @no_padding_on_packs(%A: tensor<32x32xf32>, %B: tensor<32x32xf32>, %C: tensor<32x32xf32>)373 -> tensor<32x32xf32> {374 %0 = linalg.matmul ins(%A, %B: tensor<32x32xf32>, tensor<32x32xf32>)375 outs(%C: tensor<32x32xf32>)376 -> tensor<32x32xf32>377 return %0 : tensor<32x32xf32>378}379 380// CHECK-LABEL: no_padding_on_packs381// CHECK: linalg.pack %{{.+}} inner_dims_pos = [0, 1] inner_tiles = [8, 4]382// CHECK-SAME: into %{{.+}} : tensor<32x32xf32> -> tensor<4x8x8x4xf32>383// CHECK: linalg.pack %{{.+}} outer_dims_perm = [1, 0]384// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [4, 16] into %{{.+}} : tensor<32x32xf32> -> tensor<2x8x4x16xf32>385// CHECK: linalg.pack %{{.+}} inner_dims_pos = [0, 1] inner_tiles = [8, 16]386// CHECK-SAME: into %{{.+}} : tensor<32x32xf32> -> tensor<4x2x8x16xf32>387 388module attributes {transform.with_named_sequence} {389 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {390 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1391 : (!transform.any_op) -> !transform.op<"linalg.matmul">392 %1 = transform.structured.pack_greedily %0393 matmul_packed_sizes = [8, 16, 4] matmul_inner_dims_order = [0, 1, 2]394 : (!transform.op<"linalg.matmul">) -> !transform.op<"linalg.generic">395 %pack = transform.get_producer_of_operand %1[1]396 : (!transform.op<"linalg.generic">) -> (!transform.op<"linalg.pack">)397 %2, %pack_2, %empty_unpack_2 =398 transform.structured.pack_transpose %pack with_compute_op(%1)399 outer_perm = [1, 0] inner_perm = [1, 0]400 : (!transform.op<"linalg.pack">, !transform.op<"linalg.generic">)401 -> (!transform.op<"linalg.generic">, !transform.op<"linalg.pack">, !transform.any_op)402 transform.yield403 }404}405