gussmith23 commented on a change in pull request #5812:
URL: https://github.com/apache/incubator-tvm/pull/5812#discussion_r484567629
##########
File path: python/tvm/target/datatype.py
##########
@@ -14,73 +14,153 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Custom datatype functionality"""
-import tvm._ffi
+"""Bring Your Own Datatypes custom datatype framework
-import tvm.runtime._ffi_api
-from tvm.runtime import DataType
-import tvm.tir
-from tvm.tir.expr import Cast as _Cast, FloatImm as _FloatImm
+TODO(@gussmith23 @hypercubestart) link to BYODT docs when they exist"""
+import tvm
+from tvm.runtime import convert, DataType
+from tvm.tir.expr import (Call as _Call, Cast as _Cast,
+ FloatImm as _FloatImm, BinaryOpExpr as _BinaryOpExpr)
+from tvm.tir.op import call_pure_extern
+from tvm._ffi import register_func as _register_func
+from tvm.tir import call_intrin
def register(type_name, type_code):
"""Register a custom datatype with the given type name and type code
- Currently, the type code is manually allocated by the user, and the
- user must ensure that no two custom types share the same code.
- Generally, this should be straightforward, as the user will be
- manually registering all of their custom types.
+
+ Currently, the type code is manually allocated by the user, and the user
+ must ensure that no two custom types share the same code. Generally, this
+ should be straightforward, as the user will be manually registering all of
+ their custom types.
+
+ Example:
+
+ .. code-block:: python
+
+ # Register a dtype named 'posites2' under type code 130.
+ tvm.datatype.register('posites2', 130)
+
Parameters
----------
type_name : str
- The name of the custom datatype
+ The name of the custom datatype.
type_code : int
- The type's code, which should be >= kCustomBegin
+ The type's code, which should be >= kCustomBegin. See
+ include/tvm/runtime/data_type.h.
"""
tvm.runtime._ffi_api._datatype_register(type_name, type_code)
def get_type_name(type_code):
- """Get the type name from the type code
+ """Get the type name of a custom datatype from the type code.
+
+ Note that this only works for custom datatypes registered with
+ tvm.datatype.register(). It does not work for TVM-native types.
+
+ Example:
+
+ .. code-block:: python
+
+ tvm.datatype.register('posites2', 130)
+ assert tvm.datatype.get_type_name(130) == 'posites2'
Parameters
----------
type_code : int
- The type code
+ The type code of the custom datatype.
+
+ Returns
+ -------
+ type_name : String
+ The name of the custom datatype.
+
"""
return tvm.runtime._ffi_api._datatype_get_type_name(type_code)
def get_type_code(type_name):
- """Get the type code from the type name
+ """Get the type code of a custom datatype from its type name
+
+ Note that this only works for custom datatypes registered with
+ tvm.datatype.register(). It does not work for TVM-native types.
+
+ Example:
+
+ .. code-block:: python
+
+ tvm.datatype.register('posites2', 130)
+ assert tvm.datatype.get_type_code('posites2') == 130
Parameters
----------
type_name : str
The type name
+
+ Returns
+ -------
+ type_code : int
+ The type code of the custom datatype.
"""
return tvm.runtime._ffi_api._datatype_get_type_code(type_name)
def get_type_registered(type_code):
- """Get a boolean representing whether the type is registered
+ """Returns true if a custom datatype is registered under the given type
code
+
+ Example:
+
+ .. code-block:: python
+
+ tvm.datatype.register('posites2', 130)
+ assert tvm.datatype.get_type_registered(130)
Parameters
----------
type_code: int
The type code
+
+ Returns
+ -------
+ type_registered : bool
+ True if a custom datatype is registered under this type code, and false
+ otherwise.
"""
return tvm.runtime._ffi_api._datatype_get_type_registered(type_code)
-def register_op(lower_func, op_name, target, type_name, src_type_name=None):
- """Register an external function which computes the given op.
+def register_op(lower_func,
+ op_name,
+ target,
+ src_type_name,
+ dest_type_name=None,
+ intrinsic_name=None):
+ """Register a lowering function for a specific operator of a custom
datatype
+
+ At build time, Relay must lower operators over custom datatypes into
Review comment:
You're probably right...I'm going to leave this in so we can get some
comments from reviewers.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
