-
Notifications
You must be signed in to change notification settings - Fork 1
/
add_python_operator.html
355 lines (235 loc) · 21.2 KB
/
add_python_operator.html
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
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
<!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Adding a new operator — SMAUG: Simulating Machine Learning Applications Using gem5-Aladdin</title>
<link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
<link rel="stylesheet" href="_static/pygments.css" type="text/css" />
<!--[if lt IE 9]>
<script src="_static/js/html5shiv.min.js"></script>
<![endif]-->
<script type="text/javascript" id="documentation_options" data-url_root="./" src="_static/documentation_options.js"></script>
<script src="_static/jquery.js"></script>
<script src="_static/underscore.js"></script>
<script src="_static/doctools.js"></script>
<script src="_static/language_data.js"></script>
<script type="text/javascript" src="_static/js/theme.js"></script>
<link rel="index" title="Index" href="genindex.html" />
<link rel="search" title="Search" href="search.html" />
<link rel="next" title="C++ API and Tutorials" href="cpp_docs.html" />
<link rel="prev" title="Run a model in gem5-Aladdin simulation" href="run_model_in_simulation.html" />
</head>
<body class="wy-body-for-nav">
<div class="wy-grid-for-nav">
<nav data-toggle="wy-nav-shift" class="wy-nav-side">
<div class="wy-side-scroll">
<div class="wy-side-nav-search" >
<a href="index.html" class="icon icon-home" alt="Documentation Home"> SMAUG
</a>
<div role="search">
<form id="rtd-search-form" class="wy-form" action="search.html" method="get">
<input type="text" name="q" placeholder="Search docs" />
<input type="hidden" name="check_keywords" value="yes" />
<input type="hidden" name="area" value="default" />
</form>
</div>
</div>
<div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
<p class="caption"><span class="caption-text">Python API and tutorials</span></p>
<ul class="current">
<li class="toctree-l1"><a class="reference internal" href="python_api.html">SMAUG Python APIs</a></li>
<li class="toctree-l1 current"><a class="reference internal" href="python_tutorials.html">Tutorials</a><ul class="current">
<li class="toctree-l2"><a class="reference internal" href="build_python_model.html">Build a SMAUG model with Python API</a></li>
<li class="toctree-l2"><a class="reference internal" href="run_model_in_simulation.html">Run a model in gem5-Aladdin simulation</a></li>
<li class="toctree-l2 current"><a class="current reference internal" href="#">Adding a new operator</a><ul>
<li class="toctree-l3"><a class="reference internal" href="#adding-a-simple-python-operator">Adding a simple Python operator</a></li>
<li class="toctree-l3"><a class="reference internal" href="#adding-an-operator-with-additional-parameters">Adding an operator with additional parameters</a></li>
<li class="toctree-l3"><a class="reference internal" href="#adding-a-complex-python-operator">Adding a complex Python operator</a></li>
</ul>
</li>
</ul>
</li>
</ul>
<p class="caption"><span class="caption-text">C++ docs</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="cpp_docs.html">C++ API and Tutorials</a></li>
</ul>
</div>
</div>
</nav>
<section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">
<nav class="wy-nav-top" aria-label="top navigation">
<i data-toggle="wy-nav-top" class="fa fa-bars"></i>
<a href="index.html">SMAUG</a>
</nav>
<div class="wy-nav-content">
<div class="rst-content">
<div role="navigation" aria-label="breadcrumbs navigation">
<ul class="wy-breadcrumbs">
<li><a href="index.html" class="icon icon-home"></a> »</li>
<li><a href="python_tutorials.html">Tutorials</a> »</li>
<li>Adding a new operator</li>
<li class="wy-breadcrumbs-aside">
<a href="_sources/add_python_operator.rst.txt" rel="nofollow"> View page source</a>
</li>
</ul>
<hr/>
</div>
<div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
<div itemprop="articleBody">
<div class="section" id="module-smaug.python.ops">
<span id="adding-a-new-operator"></span><h1>Adding a new operator<a class="headerlink" href="#module-smaug.python.ops" title="Permalink to this headline">¶</a></h1>
<p>To add a new operator to SMAUG, we first need to implement the actual operator
in C++. Please see the <a class="reference external" href="doxygen_html/index.html">C++ tutorial</a> for details.
Once that is done, all that is left is to expose a Python API which will add
the operator to the model.</p>
<p>Because at their core, Python operators simply add nodes to the graph, we can
build arbitrarily complex operators in Python that invoke other Python
operators. We will start simple, and then demonstrate more complex examples.</p>
<div class="section" id="adding-a-simple-python-operator">
<h2>Adding a simple Python operator<a class="headerlink" href="#adding-a-simple-python-operator" title="Permalink to this headline">¶</a></h2>
<p>A simple Python operator is one that adds only a single underlying SMAUG
operator. For the purposes of this tutorial, we’ll continue using the
<code class="code docutils literal notranslate"><span class="pre">MyCustomOperator</span></code> operator we built in the C++ tutorial, which
implements an elementwise add. This operator is as simple as it can get: it
takes just two input tensors and no additional parameters and produces a single
output tensor.</p>
<p>First, go to <code class="code docutils literal notranslate"><span class="pre">smaug/python/ops</span></code>. All operators are defined in Python
files here, organized by operator type. For example, <code class="code docutils literal notranslate"><span class="pre">math_ops</span></code> contains
arithmetic operators like add/multiply/etc, while <code class="code docutils literal notranslate"><span class="pre">array_ops</span></code> contains
operators for manipulating tensors by reshaping/concatentating/etc. In
practice, we would add a new operator to an existing file there. But for the
purposes of this tutorial, we’ll create a brand new file called
<cite>my_custom_operator.py</cite>.</p>
<p>A SMAUG Python operator takes some number of input tensors, additional operator
parameters (if appropriate), and returns some number of output tensors. Input
and output tensors are all represented as <a class="reference internal" href="smaug.html#smaug.Tensor" title="smaug.Tensor"><code class="xref py py-class docutils literal notranslate"><span class="pre">smaug.Tensor</span></code></a> objects. The
work that’s done is to add a properly formed <code class="code docutils literal notranslate"><span class="pre">NodeProto</span></code> to the
<code class="code docutils literal notranslate"><span class="pre">GraphProto</span></code> with the <a class="reference internal" href="internals.html#smaug.python.ops.common.add_node" title="smaug.python.ops.common.add_node"><code class="xref py py-func docutils literal notranslate"><span class="pre">smaug.python.ops.common.add_node()</span></code></a> API.
Here’s how the <cite>my_custom_operator.py</cite> file might look:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span> <span class="nn">smaug.core</span> <span class="kn">import</span> <span class="n">node_pb2</span><span class="p">,</span> <span class="n">types_pb2</span>
<span class="kn">from</span> <span class="nn">smaug.python.ops</span> <span class="kn">import</span> <span class="n">common</span>
<span class="k">def</span> <span class="nf">my_custom_operator</span><span class="p">(</span><span class="n">tensor_a</span><span class="p">,</span> <span class="n">tensor_b</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s2">"my_custom_operator"</span><span class="p">):</span>
<span class="k">if</span> <span class="n">tensor_a</span><span class="o">.</span><span class="n">shape</span><span class="o">.</span><span class="n">dims</span> <span class="o">!=</span> <span class="n">tensor_b</span><span class="o">.</span><span class="n">shape</span><span class="o">.</span><span class="n">dims</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
<span class="s2">"The input tensors to MyCustomOperator must be of the same shape"</span><span class="p">)</span>
<span class="k">return</span> <span class="n">common</span><span class="o">.</span><span class="n">add_node</span><span class="p">(</span>
<span class="n">name</span><span class="o">=</span><span class="n">name</span><span class="p">,</span>
<span class="n">op</span><span class="o">=</span><span class="n">types_pb2</span><span class="o">.</span><span class="n">MyCustomOperator</span><span class="p">,</span>
<span class="n">input_tensors</span><span class="o">=</span><span class="p">[</span><span class="n">tensor_a</span><span class="p">,</span> <span class="n">tensor_b</span><span class="p">],</span>
<span class="n">output_tensors_dims</span><span class="o">=</span><span class="p">[</span><span class="n">tensor_a</span><span class="o">.</span><span class="n">shape</span><span class="o">.</span><span class="n">dims</span><span class="p">],</span>
<span class="n">output_tensor_layout</span><span class="o">=</span><span class="n">tensor_a</span><span class="o">.</span><span class="n">shape</span><span class="o">.</span><span class="n">layout</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
</pre></div>
</div>
<p>The <code class="code docutils literal notranslate"><span class="pre">name</span></code> parameter is actually a prefix for the actual name of the node
that’s added to the graph. SMAUG will automatically generate a unique suffix to
ensure that no two nodes have the same name. Also, <a class="reference internal" href="internals.html#smaug.python.ops.common.add_node" title="smaug.python.ops.common.add_node"><code class="xref py py-func docutils literal notranslate"><span class="pre">common.add_node()</span></code></a>
returns a list of tensors, but in our case, we have only one, so to simplify
our API, we just return the first element. There are other optional parameters
to <a class="reference internal" href="internals.html#smaug.python.ops.common.add_node" title="smaug.python.ops.common.add_node"><code class="xref py py-func docutils literal notranslate"><span class="pre">common.add_node()</span></code></a>, but they aren’t needed in this basic scenario.</p>
<p>The final step is to expose this operator at the global <a class="reference internal" href="smaug.html#module-smaug" title="smaug"><code class="xref py py-mod docutils literal notranslate"><span class="pre">smaug</span></code></a> module
level. Open up <code class="file docutils literal notranslate"><span class="pre">smaug/__init__.py</span></code> and add the following line:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span> <span class="nn">smaug.python.ops</span> <span class="kn">import</span> <span class="n">my_custom_operator</span>
</pre></div>
</div>
<p>And that’s it! You’re now ready to use this new operator in a new model. Users
will refer to it as <code class="code docutils literal notranslate"><span class="pre">smaug.my_custom_operator.my_custom_operator</span></code>.
Obviously, the name for this small example is quite repetitive and
uninformative, but in practice, you would use a more descriptive module
and operator name.</p>
</div>
<div class="section" id="adding-an-operator-with-additional-parameters">
<h2>Adding an operator with additional parameters<a class="headerlink" href="#adding-an-operator-with-additional-parameters" title="Permalink to this headline">¶</a></h2>
<p>Some operators require additional parameters beyond just the input tensors. For
example, you may want to specify padding or stride lengths to a convolution. If
your operator needs additional parameters, you will need to add a custom
parameter protobuf message to store them. Open <a class="reference external" href="doxygen_html/node_8proto_source.html">smaug/core/node.proto</a>. This file contains the
<code class="code docutils literal notranslate"><span class="pre">NodeProto</span></code> definition along with various operator-specific parameters.
Then follow these steps.</p>
<ol class="arabic simple">
<li><p>Define a new message to store your operator’s parameters.</p></li>
<li><p>Add it as a <code class="code docutils literal notranslate"><span class="pre">oneof</span></code> field in the <code class="code docutils literal notranslate"><span class="pre">Params</span></code> message.</p></li>
<li><p>Build a <code class="code docutils literal notranslate"><span class="pre">Params</span></code> proto in your Python operator, populate it, and pass
it to <a class="reference internal" href="internals.html#smaug.python.ops.common.add_node" title="smaug.python.ops.common.add_node"><code class="xref py py-func docutils literal notranslate"><span class="pre">common.add_node()</span></code></a>.</p></li>
</ol>
<p>As an example, suppose our custom operator actually performed the operation A +
x*B, where x is a user-defined scalar. Then we would add a parameter message
like so:</p>
<div class="highlight-c notranslate"><div class="highlight"><pre><span></span><span class="hll"><span class="n">message</span> <span class="n">MyCustomOperatorParams</span> <span class="p">{</span>
</span><span class="hll"> <span class="kt">float</span> <span class="n">scale_factor</span> <span class="o">=</span> <span class="mi">1</span><span class="p">;</span>
</span><span class="hll"><span class="p">}</span>
</span>
<span class="n">message</span> <span class="n">Params</span> <span class="p">{</span>
<span class="n">oneof</span> <span class="n">value</span> <span class="p">{</span>
<span class="cp"># ... if we already have five other parameters already here...</span>
<span class="hll"> <span class="n">MyCustomOperatorParams</span> <span class="n">my_custom_operator_params</span> <span class="o">=</span> <span class="mi">6</span><span class="p">;</span>
</span> <span class="p">}</span>
<span class="cp"># ... anything else already here ...</span>
<span class="p">}</span>
</pre></div>
</div>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">my_custom_operator</span><span class="p">(</span><span class="n">tensor_a</span><span class="p">,</span> <span class="n">tensor_b</span><span class="p">,</span> <span class="n">scale_factor</span><span class="o">=</span><span class="mf">1.0</span> <span class="n">name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
<span class="k">if</span> <span class="n">tensor_a</span><span class="o">.</span><span class="n">shape</span><span class="o">.</span><span class="n">dims</span> <span class="o">!=</span> <span class="n">tensor_b</span><span class="o">.</span><span class="n">shape</span><span class="o">.</span><span class="n">dims</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
<span class="s2">"The input tensors to MyCustomOperator must be of the same shape"</span><span class="p">)</span>
<span class="n">params</span> <span class="o">=</span> <span class="n">node_pb2</span><span class="o">.</span><span class="n">Params</span><span class="p">()</span>
<span class="n">params</span><span class="o">.</span><span class="n">my_custom_operator_params</span><span class="o">.</span><span class="n">scale_factor</span> <span class="o">=</span> <span class="n">scale_factor</span>
<span class="k">return</span> <span class="n">common</span><span class="o">.</span><span class="n">add_node</span><span class="p">(</span>
<span class="n">name</span><span class="o">=</span><span class="n">name</span><span class="p">,</span>
<span class="n">op</span><span class="o">=</span><span class="n">types_pb2</span><span class="o">.</span><span class="n">MyCustomOperator</span><span class="p">,</span>
<span class="n">input_tensors</span><span class="o">=</span><span class="p">[</span><span class="n">tensor_a</span><span class="p">,</span> <span class="n">tensor_b</span><span class="p">],</span>
<span class="n">output_tensors_dims</span><span class="o">=</span><span class="p">[</span><span class="n">tensor_a</span><span class="o">.</span><span class="n">shape</span><span class="o">.</span><span class="n">dims</span><span class="p">],</span>
<span class="n">output_tensor_layout</span><span class="o">=</span><span class="n">tensor_a</span><span class="o">.</span><span class="n">shape</span><span class="o">.</span><span class="n">layout</span><span class="p">,</span>
<span class="n">params</span><span class="o">=</span><span class="n">params</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
</pre></div>
</div>
</div>
<div class="section" id="adding-a-complex-python-operator">
<h2>Adding a complex Python operator<a class="headerlink" href="#adding-a-complex-python-operator" title="Permalink to this headline">¶</a></h2>
<p>Since Python operators simply add nodes to the graph, we can call Python
operators from each other. As a very simple example, we can chain together
two instances of MyCustomOperator:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">my_custom_operator_chained</span><span class="p">(</span>
<span class="n">tensor_a</span><span class="p">,</span> <span class="n">tensor_b</span><span class="p">,</span> <span class="n">scale_factor</span><span class="o">=</span><span class="mf">1.0</span> <span class="n">name</span><span class="o">=</span><span class="s2">"my_custom_operator_chained"</span><span class="p">):</span>
<span class="k">if</span> <span class="n">tensor_a</span><span class="o">.</span><span class="n">shape</span><span class="o">.</span><span class="n">dims</span> <span class="o">!=</span> <span class="n">tensor_b</span><span class="o">.</span><span class="n">shape</span><span class="o">.</span><span class="n">dims</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
<span class="s2">"The input tensors to MyCustomOperator must be of the same shape"</span><span class="p">)</span>
<span class="n">params</span> <span class="o">=</span> <span class="n">node_pb2</span><span class="o">.</span><span class="n">Params</span><span class="p">()</span>
<span class="n">params</span><span class="o">.</span><span class="n">my_custom_operator_params</span><span class="o">.</span><span class="n">scale_factor</span> <span class="o">=</span> <span class="n">scale_factor</span>
<span class="n">output_tensor_1</span> <span class="o">=</span> <span class="n">my_custom_operator</span><span class="p">(</span>
<span class="n">tensor_a</span><span class="p">,</span> <span class="n">tensor_b</span><span class="p">,</span> <span class="n">scale_factor</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="n">name</span><span class="p">)</span>
<span class="k">return</span> <span class="n">my_custom_operator</span><span class="p">(</span>
<span class="n">output_tensor_1</span><span class="p">,</span> <span class="n">tensor_b</span><span class="p">,</span> <span class="n">scale_factor</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="n">name</span><span class="p">)</span>
</pre></div>
</div>
</div>
</div>
</div>
</div>
<footer>
<div class="rst-footer-buttons" role="navigation" aria-label="footer navigation">
<a href="cpp_docs.html" class="btn btn-neutral float-right" title="C++ API and Tutorials" accesskey="n" rel="next">Next <span class="fa fa-arrow-circle-right"></span></a>
<a href="run_model_in_simulation.html" class="btn btn-neutral float-left" title="Run a model in gem5-Aladdin simulation" accesskey="p" rel="prev"><span class="fa fa-arrow-circle-left"></span> Previous</a>
</div>
<hr/>
<div role="contentinfo">
<p>
© Copyright 2020, SMAUG Contributors
</p>
</div>
Built with <a href="http://sphinx-doc.org/">Sphinx</a> using a
<a href="https://github.com/rtfd/sphinx_rtd_theme">theme</a>
provided by <a href="https://readthedocs.org">Read the Docs</a>.
</footer>
</div>
</div>
</section>
</div>
<script type="text/javascript">
jQuery(function () {
SphinxRtdTheme.Navigation.enable(true);
});
</script>
</body>
</html>