[RELAY][BYOC] Register pattern tables from external codegens

This adds utility functions to support registering
and retrieving pattern tables used by MergeComposite for
external codegens.

Change-Id: I5be165a321440e48b15ff6aff4970e0c67496aaa

Author: mbaret

python/tvm/relay/op/contrib/__init__.py

@@ -16,4 +16,6 @@
 # under the License.
 # pylint: disable=wildcard-import
 """Contrib modules."""
+from .register import get_pattern_table, register_pattern_table
+
 from .dnnl import *

python/tvm/relay/op/contrib/register.py

@@ -0,0 +1,49 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Register utilities for external codegen."""
+pattern_tables = {}
+
+
+def register_pattern_table(compiler, table=None):
+    """Register a pattern table for an external compiler.
+
+    Pattern tables are used to create composite functions.
+    See the MergeComposite pass.
+
+    Parameters
+    ----------
+    compiler : str
+        The name of compiler
+
+    table : function, optional
+        A function that returns the pattern table
+
+    Returns
+    -------
+    fregister : function
+        Register function if value is not specified.
+    """
+    def _register(t):
+        """internal register function"""
+        pattern_tables[compiler] = t()
+        return t
+    return _register(table) if table is not None else _register
+
+
+def get_pattern_table(compiler):
+    """Get the pattern table associated with a compiler (if it's registered)."""
+    return pattern_tables[compiler] if compiler in pattern_tables else None

tests/python/relay/test_pattern_table.py

@@ -0,0 +1,45 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit test for pattern table registry (BYOC)."""
+from tvm.relay.op.contrib import get_pattern_table, register_pattern_table
+from tvm import relay
+
+
+@register_pattern_table("test_pattern_table")
+def pattern_table():
+    def _make_add_relu_pattern():
+        x = relay.var('x')
+        y = relay.var('y')
+        add_node = relay.add(x, y)
+        r = relay.nn.relu(add_node)
+        return r
+
+    def _check_add_relu_pattern():
+        return True
+
+    return [
+        ("test_pattern_table.add_relu", _make_add_relu_pattern(), _check_add_relu_pattern)
+    ]
+
+
+def test_retrieve_pattern_table():
+    table = get_pattern_table("test_pattern_table")
+    assert table[0][0] == "test_pattern_table.add_relu"
+
+
+if __name__ == "__main__":
+    test_retrieve_pattern_table()