70 lines · plain
1// RUN: mlir-opt %s -transform-interpreter | FileCheck %s2 3func.func @mmt4d_to_fma(%A: tensor<16x16x8x1xf32>, %B: tensor<16x16x8x1xf32>, %C_in: tensor<16x16x8x8xf32>) -> tensor<16x16x8x8xf32> {4 %res = linalg.mmt4d5 ins(%A, %B: tensor<16x16x8x1xf32>, tensor<16x16x8x1xf32>)6 outs(%C_in: tensor<16x16x8x8xf32>)7 -> tensor<16x16x8x8xf32>8 return %res : tensor<16x16x8x8xf32>9}10 11 12// CHECK-LABEL: @mmt4d_to_fma13// CHECK-COUNT-8: vector.fma14 15module attributes {transform.with_named_sequence} {16 transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {17 %func = transform.structured.match ops{["func.func"]} in %module : (!transform.any_op) -> !transform.op<"func.func">18 19 %mmt4d = transform.structured.match ops{["linalg.mmt4d"]} in %func : (!transform.op<"func.func">) -> !transform.any_op20 21 // Step 1: Tile22 // Tile parallel dims23 %tiled_linalg_op_p, %loops:4 = transform.structured.tile_using_for %mmt4d tile_sizes [1, 1, 0, 8, 8, 0]24 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)25 // Tile reduction dims26 %tiled_linalg_op_r, %loops2:2 = transform.structured.tile_using_for %tiled_linalg_op_p tile_sizes [0, 0, 1, 0, 0, 1]27 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)28 29 // Step 2: Vectorize30 transform.structured.vectorize %tiled_linalg_op_r : !transform.any_op31 32 // Step 3: Simplify33 // vector.multi_reduction --> vector.contract34 // Generates a 6-dim vector.contract with the dim matching the original MMT4D Op35 // and with the following split into parallel and reduction dims:36 // * parallel, parallel, reduction, parallel, parallel, reduction37 transform.apply_patterns to %func {38 transform.apply_patterns.vector.reduction_to_contract39 // Reduce the rank of xfer ops. This transforms vector.contract to be40 // more matmul-like and to enable the lowering to outer product Ops.41 transform.apply_patterns.vector.transfer_permutation_patterns42 } : !transform.op<"func.func">43 44 // Hoisting and LICM - not strictly required45 %func_h = transform.structured.hoist_redundant_vector_transfers %func46 : (!transform.op<"func.func">) -> !transform.op<"func.func">47 %all_loops = transform.structured.match interface{LoopLikeInterface} in %func_h48 : (!transform.op<"func.func">) -> !transform.any_op49 transform.apply_licm to %all_loops : !transform.any_op50 transform.loop.hoist_loop_invariant_subsets %all_loops : !transform.any_op51 52 // Simplify the 6-dim vector.contract into a 3-dim matmul-like53 // vector.contract with the following split into parallel and reduction54 // dims:55 // * parallel, parallel, reduction56 transform.apply_patterns to %func_h {57 transform.apply_patterns.vector.reduction_to_contract58 transform.apply_patterns.vector.cast_away_vector_leading_one_dim59 transform.apply_patterns.canonicalization60 } : !transform.op<"func.func">61 62 // Step 4: Lower vector.contract to vector.fma via vector.outerproduct63 transform.apply_patterns to %func_h {64 transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"65 transform.apply_patterns.vector.lower_outerproduct66 } : !transform.op<"func.func">67 transform.yield68 }69}70