Skip to content

Commit

Permalink
Allow for custom ops in registry.
Browse files Browse the repository at this point in the history
  • Loading branch information
jvreca committed Sep 13, 2024
1 parent c5bd87f commit 171d038
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 2 deletions.
14 changes: 14 additions & 0 deletions src/qonnx/custom_op/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,20 @@

from qonnx.util.basic import get_preferred_onnx_opset

_QONNX_DOMAINS = ["finn", "qonnx.custom_op", "onnx.brevitas"]


def register_custom_domain(domain: str):
_QONNX_DOMAINS.append(domain)


def is_finn_op(op_type):
"Return whether given op_type string is a QONNX or FINN custom op"
is_finn = False
for domain in _QONNX_DOMAINS:
is_finn = is_finn or op_type.startswith(domain)
return is_finn


def getCustomOp(node, onnx_opset_version=get_preferred_onnx_opset(), brevitas_exception=True):
"Return a QONNX CustomOp instance for the given ONNX node, if it exists."
Expand Down
6 changes: 4 additions & 2 deletions src/qonnx/util/basic.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,9 @@
import string
import warnings

import qonnx
import qonnx.custom_op
import qonnx.custom_op.registry
from qonnx.core.datatype import DataType

# TODO solve by moving onnx-dependent fxns to onnx.py
Expand Down Expand Up @@ -63,8 +66,7 @@ def qonnx_make_model(graph_proto, **kwargs):


def is_finn_op(op_type):
"Return whether given op_type string is a QONNX or FINN custom op"
return op_type.startswith("finn") or op_type.startswith("qonnx.custom_op") or op_type.startswith("onnx.brevitas")
return qonnx.custom_op.registry.is_finn_op(op_type)


def get_num_default_workers():
Expand Down

0 comments on commit 171d038

Please sign in to comment.