Skip to content
Projects
Groups
Snippets
Help
This project
Loading...
Sign in / Register
Toggle navigation
T
tic
Overview
Overview
Details
Activity
Cycle Analytics
Repository
Repository
Files
Commits
Branches
Tags
Contributors
Graph
Compare
Charts
Issues
0
Issues
0
List
Board
Labels
Milestones
Merge Requests
0
Merge Requests
0
CI / CD
CI / CD
Pipelines
Jobs
Schedules
Charts
Wiki
Wiki
Snippets
Snippets
Members
Members
Collapse sidebar
Close sidebar
Activity
Graph
Charts
Create a new issue
Jobs
Commits
Issue Boards
Open sidebar
wenyuanbo
tic
Commits
2475556a
Commit
2475556a
authored
Sep 29, 2018
by
Steven S. Lyubomirsky
Committed by
Tianqi Chen
Sep 29, 2018
Browse files
Options
Browse Files
Download
Email Patches
Plain Diff
[Relay][DOC] Add tutorial for adding an operator to Relay (#1778)
parent
9319b6f9
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
148 additions
and
0 deletions
+148
-0
docs/dev/relay_add_op.rst
+148
-0
No files found.
docs/dev/relay_add_op.rst
0 → 100644
View file @
2475556a
Adding an Operator to Relay
===========================
In order to use TVM operators from within the Relay IR, the
operators need to be registered in Relay in order to ensure
that they will be integrated into Relay's type system.
Registering an operator requires three steps:
- Using the ``RELAY_REGISTER_OP`` macro in C++ to
register the operator's arity and type information
- Defining a C++ function to produce a call node for the
operator and registering a Python API hook for the function
- Wrapping the above Python API hook in a neater interface
The file ``src/relay/op/tensor/elemwise.cc`` provides
examples of the first two steps, while
``python/tvm/relay/op/tensor.py`` gives examples of the
last.
Registering an Operator
-----------------------
TVM already has an operator registry, but Relay cannot properly
incorporate TVM operators without additional type information.
To allow for flexibility in registering operators and greater
expressivity and granularity in expressing types in Relay, operators
are typed using relations between input and output types. These relations
are represented as functions that take in a list of input types and
output types (any of these types may be incomplete) and return a list
of input and output types that satisfies the relation. Essentially, a
relation for an operator can enforce all the necessary typing rules
(namely by inspecting the input types) in addition to computing the
output type.
For example, see ``src/relay/op/type_relations.h`` and their
implementations. E.g., ``BroadcastRel`` takes two input types and an
output type, checks that they are all tensor types with the same underlyin
data type, and finally ensures that the shape of the output type is the
broadcast of the input types' shapes.
It may be necessary to add another type relation to ``type_relations.h``
if the existing ones do not capture the behavior of the desired operator.
The ``RELAY_REGISTER_OP`` macro in C++ allows a developer
to specify the following information about an operator in Relay:
- Arity (number of arguments)
- Names and descriptions for positional arguments
- Support level (1 indicating an internal intrinsic, higher numbers
indicating operators that are not as integral to the framework or are
supported externally)
- A type relation for the operator
The below example is from ``elemwise.cc`` and uses a broadcasting
add for tensors:
.. code:: c
RELAY_REGISTER_OP("add")
.set_num_inputs(2)
.add_argument("lhs", "Tensor", "The left hand side tensor.")
.add_argument("rhs", "Tensor", "The right hand side tensor.")
.set_support_level(1)
.add_type_rel("Broadcast", BroadcastRel);
Creating a Call Node
--------------------
This step requires simply writing a function that takes
the arguments to the operator (as Relay expressions) and
returning a call node to the operator (i.e., the node that
should be placed into the Relay AST where the call to the
operator is intended).
At present call attributes and type arguments (the last two fields)
are not supported, so it suffices to use ``Op::Get`` to fetch
the operator's information from the operator registry and pass in
the arguments to the call node, as below.
.. code:: c
TVM_REGISTER_API("relay.op._make.add")
.set_body_typed<Expr(Expr, Expr)>([](Expr lhs, Expr rhs) {
static const Op& op = Op::Get("add");
return CallNode::make(op, {lhs, rhs}, Attrs(), {});
});
Including a Python API Hook
---------------------------
It is generally the convention in Relay, that functions exported
through ``TVM_REGISTER_API`` should be wrapped in a separate
Python function rather than called directly in Python. In the case
of the functions that produce calls to operators, it may be convenient
to bundle them, as in ``python/tvm/relay/op/tensor.py``, where
elementwise operators on tensors are all provided. For example,
the following is how the add function from the previous section is
exposed in Python:
.. code:: python
def add(lhs, rhs):
"""Elementwise addition.
Parameters
----------
lhs : relay.Expr
The left hand side input data
rhs : relay.Expr
The right hand side input data
Returns
-------
result : relay.Expr
The computed result.
"""
return _make.add(lhs, rhs)
Note that these Python wrappers might also be good opportunities to
provide an easier interface to the operator. For example, the
``concat`` operator is registered as taking only one operator,
namely a tuple with the tensors to be concatenated, but the Python
wrapper takes the tensors as arguments and combines them into a tuple
before producing the call node:
.. code:: python
def concat(*args):
"""Concatenate the input tensors along the zero axis.
Parameters
----------
args: list of Tensor
Returns
-------
tensor: The concatenated tensor.
"""
tup = Tuple(list(args))
return _make.concat(tup)
Summary
-------
- A TVM operator can be registered in Relay using a relation to express
the appropriate type information.
- Using an operator in Relay requires a function to produce a
call node for the operator.
- It is best to have a simple Python wrapper for producing the call node.
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment