140 lines · plain
1// RUN: transform-opt-ch4 %s --transform-interpreter --verify-diagnostics2//3// RUN: transform-opt-ch4 %s \4// RUN: --transform-interpreter='entry-point=__transform_main_v2' \5// RUN: --verify-diagnostics6 7// ****************************** IMPORTANT NOTE ******************************8//9// If you are changing this file, you may also need to change10// mlir/docs/Tutorials/Transform accordingly.11//12// ****************************************************************************13 14// Original function to optimize.15func.func @fc_relu(%lhs: tensor<512x512xf32>, %rhs: tensor<512x512xf32>,16 %bias: tensor<512x512xf32>, %output: tensor<512x512xf32>)17 -> tensor<512x512xf32> {18 // Matrix-matrix multiplication.19 // expected-remark @below {{matmul}}20 %matmul = linalg.matmul ins(%lhs, %rhs: tensor<512x512xf32>, tensor<512x512xf32>)21 outs(%output: tensor<512x512xf32>) -> tensor<512x512xf32>22 23 // Elementwise addition.24 // expected-remark @below {{elementwise binary}}25 %biased = linalg.elementwise kind=#linalg.elementwise_kind<add>26 ins(%matmul, %bias : tensor<512x512xf32>, tensor<512x512xf32>)27 outs(%output : tensor<512x512xf32>) -> tensor<512x512xf32>28 29 // Elementwise max with 0 (ReLU).30 %c0f = arith.constant dense<0.0> : tensor<512x512xf32>31 // expected-remark @below {{elementwise binary}}32 %relued = linalg.elementwise kind=#linalg.elementwise_kind<max_signed>33 ins(%biased, %c0f : tensor<512x512xf32>, tensor<512x512xf32>)34 outs(%output : tensor<512x512xf32>) -> tensor<512x512xf32>35 func.return %relued : tensor<512x512xf32>36}37 38// The module containing named sequences must have an attribute allowing them39// to enable verification.40module @transforms attributes { transform.with_named_sequence } {41 // Entry point. This takes as the only argument the root operation (typically42 // pass root) given to the transform interpreter.43 transform.named_sequence @__transform_main(44 %root: !transform.any_op {transform.readonly}) {45 // Collect operations that match the criteria specified in the named46 // sequence. If the named sequence fails with a silenceable failure,47 // silences it (the message is forwarded to the debug stream). If the named48 // sequence succeeds, appends its results to the results of this operation.49 %elemwise = transform.collect_matching @match_elemwise in %root50 : (!transform.any_op) -> !transform.any_op51 %matmul = transform.collect_matching @match_matmul in %root52 : (!transform.any_op) -> !transform.any_op53 54 transform.include @print_elemwise failures(propagate) (%elemwise)55 : (!transform.any_op) -> ()56 transform.include @print_matmul failures(propagate) (%matmul)57 : (!transform.any_op) -> ()58 59 transform.yield60 }61 62 // Alternative entry point.63 transform.named_sequence @__transform_main_v2(64 %root: !transform.any_op {transform.readonly}) {65 // Collect groups of operations that match the criteria specified in the66 // named sequence.67 %matmul, %el1, %el2 = transform.collect_matching @match_matmul_elemwise in %root 68 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)69 %elemwise = transform.merge_handles %el1, %el2 : !transform.any_op70 71 transform.include @print_elemwise failures(propagate) (%elemwise)72 : (!transform.any_op) -> ()73 transform.include @print_matmul failures(propagate) (%matmul)74 : (!transform.any_op) -> ()75 76 transform.yield77 }78 79 // This is a matcher sequence. It is given an operation to match and the80 // match is considered successful unless any nested operation produces a81 // failure. The values yielded by this operation will be forwarded to the82 // rewriter sequence on success.83 transform.named_sequence @match_elemwise(84 %entry: !transform.any_op {transform.readonly}) -> !transform.any_op {85 transform.match.operation_name %entry ["linalg.elementwise"] 86 : !transform.any_op87 transform.yield %entry : !transform.any_op88 }89 transform.named_sequence @match_matmul(90 %entry: !transform.any_op {transform.readonly}) -> !transform.any_op {91 transform.match.operation_name %entry ["linalg.matmul"] : !transform.any_op92 transform.yield %entry : !transform.any_op93 }94 95 // This is an action sequence.96 transform.named_sequence @print_elemwise(97 %elemwise_binary: !transform.any_op {transform.readonly}) {98 transform.debug.emit_remark_at99 %elemwise_binary, "elementwise binary" : !transform.any_op100 transform.yield101 }102 transform.named_sequence @print_matmul(103 %matmul: !transform.any_op {transform.readonly}) {104 transform.debug.emit_remark_at %matmul, "matmul" : !transform.any_op105 transform.yield106 }107 108 // This is also a matcher sequence. It is similarly given an operation to109 // match and nested operations must succeed in order for a match to be deemed110 // successful. It starts matching from the last operation in the use-def chain111 // and goes back because each operand (use) has exactly one definition.112 transform.named_sequence @match_matmul_elemwise(113 %last: !transform.any_op {transform.readonly}) 114 -> (!transform.any_op, !transform.any_op, !transform.any_op) {115 // The last operation must be an elementwise binary.116 transform.match.operation_name %last ["linalg.elementwise"]117 : !transform.any_op118 // Its first operand must be defined by another operation, to which we119 // will get a handle here. We are guaranteed that the first operand exists120 // because we know the operation is binary, but even in absence of such a121 // guarantee, this operation would have produced a silenceable failure when122 // `%last` does not have enough operands.123 %middle = transform.get_producer_of_operand %last[0]124 : (!transform.any_op) -> !transform.any_op125 // The defining operation must itself be an elementwise binary.126 transform.match.operation_name %middle ["linalg.elementwise"]127 : !transform.any_op128 // And the first operand of that operation must be defined by yet another129 // operation.130 %matmul = transform.get_producer_of_operand %middle[0]131 : (!transform.any_op) -> !transform.any_op132 // And that operation is a matmul.133 transform.match.operation_name %matmul ["linalg.matmul"] : !transform.any_op134 // We will yield the handles to the matmul and the two elementwise135 // operations separately. 136 transform.yield %matmul, %middle, %last137 : !transform.any_op, !transform.any_op, !transform.any_op138 }139}140