1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
|
// RUN: transform-opt-ch4 %s --transform-interpreter --verify-diagnostics
//
// RUN: transform-opt-ch4 %s \
// RUN: --transform-interpreter='entry-point=__transform_main_v2' \
// RUN: --verify-diagnostics
// ****************************** IMPORTANT NOTE ******************************
//
// If you are changing this file, you may also need to change
// mlir/docs/Tutorials/Transform accordingly.
//
// ****************************************************************************
// Original function to optimize.
func.func @fc_relu(%lhs: tensor<512x512xf32>, %rhs: tensor<512x512xf32>,
%bias: tensor<512x512xf32>, %output: tensor<512x512xf32>)
-> tensor<512x512xf32> {
// Matrix-matrix multiplication.
// expected-remark @below {{matmul}}
%matmul = linalg.matmul ins(%lhs, %rhs: tensor<512x512xf32>, tensor<512x512xf32>)
outs(%output: tensor<512x512xf32>) -> tensor<512x512xf32>
// Elementwise addition.
// expected-remark @below {{elementwise binary}}
%biased = linalg.elemwise_binary { fun = #linalg.binary_fn<add> }
ins(%matmul, %bias : tensor<512x512xf32>, tensor<512x512xf32>)
outs(%output : tensor<512x512xf32>) -> tensor<512x512xf32>
// Elementwise max with 0 (ReLU).
%c0f = arith.constant 0.0 : f32
// expected-remark @below {{elementwise binary}}
%relued = linalg.elemwise_binary { fun = #linalg.binary_fn<max_signed> }
ins(%biased, %c0f : tensor<512x512xf32>, f32)
outs(%output : tensor<512x512xf32>) -> tensor<512x512xf32>
func.return %relued : tensor<512x512xf32>
}
// The module containing named sequences must have an attribute allowing them
// to enable verification.
module @transforms attributes { transform.with_named_sequence } {
// Entry point. This takes as the only argument the root operation (typically
// pass root) given to the transform interpreter.
transform.named_sequence @__transform_main(
%root: !transform.any_op {transform.readonly}) {
// Collect operations that match the criteria specified in the named
// sequence. If the named sequence fails with a silenceable failure,
// silences it (the message is forwarded to the debug stream). If the named
// sequence succeeds, appends its results to the results of this operation.
%elemwise = transform.collect_matching @match_elemwise in %root
: (!transform.any_op) -> !transform.any_op
%matmul = transform.collect_matching @match_matmul in %root
: (!transform.any_op) -> !transform.any_op
transform.include @print_elemwise failures(propagate) (%elemwise)
: (!transform.any_op) -> ()
transform.include @print_matmul failures(propagate) (%matmul)
: (!transform.any_op) -> ()
transform.yield
}
// Alternative entry point.
transform.named_sequence @__transform_main_v2(
%root: !transform.any_op {transform.readonly}) {
// Collect groups of operations that match the criteria specified in the
// named sequence.
%matmul, %el1, %el2 = transform.collect_matching @match_matmul_elemwise in %root
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
%elemwise = transform.merge_handles %el1, %el2 : !transform.any_op
transform.include @print_elemwise failures(propagate) (%elemwise)
: (!transform.any_op) -> ()
transform.include @print_matmul failures(propagate) (%matmul)
: (!transform.any_op) -> ()
transform.yield
}
// This is a matcher sequence. It is given an operation to match and the
// match is considered successful unless any nested operation produces a
// failure. The values yielded by this operation will be forwarded to the
// rewriter sequence on success.
transform.named_sequence @match_elemwise(
%entry: !transform.any_op {transform.readonly}) -> !transform.any_op {
transform.match.operation_name %entry ["linalg.elemwise_binary"]
: !transform.any_op
transform.yield %entry : !transform.any_op
}
transform.named_sequence @match_matmul(
%entry: !transform.any_op {transform.readonly}) -> !transform.any_op {
transform.match.operation_name %entry ["linalg.matmul"] : !transform.any_op
transform.yield %entry : !transform.any_op
}
// This is an action sequence.
transform.named_sequence @print_elemwise(
%elemwise_binary: !transform.any_op {transform.readonly}) {
transform.debug.emit_remark_at
%elemwise_binary, "elementwise binary" : !transform.any_op
transform.yield
}
transform.named_sequence @print_matmul(
%matmul: !transform.any_op {transform.readonly}) {
transform.debug.emit_remark_at %matmul, "matmul" : !transform.any_op
transform.yield
}
// This is also a matcher sequence. It is similarly given an operation to
// match and nested operations must succeed in order for a match to be deemed
// successful. It starts matching from the last operation in the use-def chain
// and goes back because each operand (use) has exactly one definition.
transform.named_sequence @match_matmul_elemwise(
%last: !transform.any_op {transform.readonly})
-> (!transform.any_op, !transform.any_op, !transform.any_op) {
// The last operation must be an elementwise binary.
transform.match.operation_name %last ["linalg.elemwise_binary"]
: !transform.any_op
// Its first operand must be defined by another operation, to which we
// will get a handle here. We are guaranteed that the first operand exists
// because we know the operation is binary, but even in absence of such a
// guarantee, this operation would have produced a silenceable failure when
// `%last` does not have enough operands.
%middle = transform.get_producer_of_operand %last[0]
: (!transform.any_op) -> !transform.any_op
// The defining operation must itself be an elementwise binary.
transform.match.operation_name %middle ["linalg.elemwise_binary"]
: !transform.any_op
// And the first operand of that operation must be defined by yet another
// operation.
%matmul = transform.get_producer_of_operand %middle[0]
: (!transform.any_op) -> !transform.any_op
// And that operation is a matmul.
transform.match.operation_name %matmul ["linalg.matmul"] : !transform.any_op
// We will yield the handles to the matmul and the two elementwise
// operations separately.
transform.yield %matmul, %middle, %last
: !transform.any_op, !transform.any_op, !transform.any_op
}
}
|
