Bring Your Own Datatypes

CMakeLists.txt

@@ -56,6 +56,7 @@ tvm_option(COMPILER_RT_PATH "Path to COMPILER-RT" "3rdparty/compiler-rt")
 tvm_option(PICOJSON_PATH "Path to PicoJSON" "3rdparty/picojson")
 
 # Contrib library options
+tvm_option(USE_BYOC_POSIT "Build with BYOC software emulated posit custom datatype" OFF)
 tvm_option(USE_BLAS "The blas library to be linked" none)
 tvm_option(USE_MKL "MKL root path when use MKL blas" OFF)
 tvm_option(USE_MKLDNN "Build with MKLDNN" OFF)
@@ -257,6 +258,7 @@ endif(USE_VM_PROFILER)
 
 file(GLOB DATATYPE_SRCS src/target/datatype/*.cc)
 list(APPEND COMPILER_SRCS ${DATATYPE_SRCS})
+list(APPEND COMPILER_SRCS "src/target/datatype/myfloat/myfloat.cc")
 
 file(GLOB RUNTIME_SRCS
   src/runtime/*.cc
@@ -279,8 +281,6 @@ if (INDEX_DEFAULT_I64)
   add_definitions(-DTVM_INDEX_DEFAULT_I64=1)
 endif()
 
-list(APPEND RUNTIME_SRCS 3rdparty/bfloat16/bfloat16.cc)
-
 if(USE_RPC)
   message(STATUS "Build with RPC support...")
   file(GLOB RUNTIME_RPC_SRCS src/runtime/rpc/*.cc)
@@ -334,6 +334,7 @@ include(cmake/modules/contrib/BLAS.cmake)
 include(cmake/modules/contrib/CODEGENC.cmake)
 include(cmake/modules/contrib/DNNL.cmake)
 include(cmake/modules/contrib/Random.cmake)
+include(cmake/modules/contrib/Posit.cmake)
 include(cmake/modules/contrib/MicroStandaloneRuntime.cmake)
 include(cmake/modules/contrib/Sort.cmake)
 include(cmake/modules/contrib/NNPack.cmake)

LICENSE

@@ -209,7 +209,6 @@ for text of these licenses.
 Apache Software Foundation License 2.0
 --------------------------------------
 
-3rdparty/bfloat16/bfloat16.cc
 3rdparty/dlpack
 3rdparty/dmlc-core
 

cmake/config.cmake

@@ -120,6 +120,9 @@ set(USE_LLVM OFF)
 #---------------------------------------------
 # Contrib libraries
 #---------------------------------------------
+# Whether to build with BYOC software emulated posit custom datatype
+set(USE_BYOC_POSIT OFF)
+
 # Whether use BLAS, choices: openblas, atlas, apple
 set(USE_BLAS none)
 

cmake/modules/LibInfo.cmake

@@ -53,6 +53,7 @@ function(add_lib_info src_file)
     TVM_INFO_HIDE_PRIVATE_SYMBOLS="${HIDE_PRIVATE_SYMBOLS}"
     TVM_INFO_USE_TF_TVMDSOOP="${USE_TF_TVMDSOOP}"
     TVM_INFO_USE_FALLBACK_STL_MAP="${USE_FALLBACK_STL_MAP}"
+    TVM_INFO_USE_BYOC_POSIT="${USE_BYOC_POSIT}"
     TVM_INFO_USE_BLAS="${USE_BLAS}"
     TVM_INFO_USE_MKL="${USE_MKL}"
     TVM_INFO_USE_MKLDNN="${USE_MKLDNN}"

cmake/modules/contrib/Posit.cmake

@@ -0,0 +1,26 @@
+# 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.
+
+if(USE_BYOC_POSIT)
+  message(STATUS "Build with contrib.posit")
+  if (NOT UNIVERSAL_PATH)
+    message(FATAL_ERROR "Fail to get Universal path")
+  endif(NOT UNIVERSAL_PATH)
+
+  include_directories(${UNIVERSAL_PATH}/include)
+  list(APPEND COMPILER_SRCS "src/target/datatype/posit/posit-wrapper.cc")
+endif(USE_BYOC_POSIT)

python/tvm/driver/build_module.py

@@ -264,6 +264,7 @@ def _build_for_device(input_mod, target, target_host):
             tvm.tir.transform.LowerWarpMemory(),
             tvm.tir.transform.Simplify(),
             tvm.tir.transform.LowerDeviceStorageAccessInfo(),
+            tvm.tir.transform.LowerCustomDatatypes(),
             tvm.tir.transform.LowerIntrin(),
         ]
     )
@@ -279,6 +280,7 @@ def _build_for_device(input_mod, target, target_host):
             tvm.tir.transform.Apply(lambda f: f.with_attr("target", target)),
             tvm.tir.transform.LowerTVMBuiltin(),
             tvm.tir.transform.LowerDeviceStorageAccessInfo(),
+            tvm.tir.transform.LowerCustomDatatypes(),
             tvm.tir.transform.LowerIntrin(),
             tvm.tir.transform.CombineContextCall(),
         ]

python/tvm/relay/backend/_backend.py

@@ -90,6 +90,9 @@ def _tensor_value_repr(tvalue):
 
 @tvm._ffi.register_func("relay._constant_repr")
 def _tensor_constant_repr(tvalue):
+    dtype = tvm.runtime.DataType(tvalue.data.dtype)
+    if tvm.target.datatype.get_type_registered(dtype.type_code):
+        return "custom tensor of type " + dtype.type_code
     return str(tvalue.data.asnumpy())
 
 

python/tvm/relay/frontend/__init__.py

@@ -34,3 +34,4 @@
 from .darknet import from_darknet
 from .pytorch import from_pytorch
 from .caffe import from_caffe
+from .change_datatype import ChangeDatatype

python/tvm/relay/frontend/change_datatype.py

@@ -0,0 +1,107 @@
+# 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.
+# pylint: disable=unused-argument
+"""Change Datatype Pass"""
+from ..function import Function
+from ..expr_functor import ExprMutator
+from ..transform.transform import function_pass
+from ..expr import var, bind
+
+
+@function_pass(opt_level=0)
+class ChangeDatatype(ExprMutator):
+    """Mutator for changing the datatype of Relay programs.
+
+    This pass should be useful for users of the Bring Your Own Datatypes
+    framework.
+    TODO(@gussmith23 @hypercubestart) Add link to documentation when it exists
+
+    Example:
+
+    .. code-block:: python
+
+        from tvm.relay.testing.inception_v3 import get_workload
+        mod, params = get_workload()
+
+        def change_dtype(mod, params, src, dst):
+            mod = ChangeDatatype(src, dst)(mod)
+            params = dict((p, tvm.nd.array(params[p].asnumpy().astype(dst))) for p in params)
+            return mod, params
+
+        mod, params = change_dtype(mod, params, "float32", "custom[posites2]32")
+
+    Parameters
+    ----------
+    src : String
+        The source datatype name, e.g. "float" or "posites2" (but not "float32"
+        or "custom[posites2]32").
+    dst : String
+        The destination datatype name, in the same format.
+
+    Returns
+    -------
+    mod : tvm.IRModule
+        Module where all nodes of dtype `src` have been changed to have dtype
+        `dst`.
+    """
+
+    def __init__(self, src, dst):
+        self.src = src
+        self.dst = dst
+        super().__init__()
+
+    def transform_function(self, func, mod, ctx):
+        return self.visit(func)
+
+    def visit_constant(self, const):
+        if const.data.dtype == self.src:
+            return const.astype(self.dst)
+        return const
+
+    def visit_function(self, fn):
+        new_params = []
+        binds = {}
+
+        for param in fn.params:
+            # Get the parameter's type annotation.
+            var_type = param.type_annotation
+
+            # See if we want to replace dtype.
+            if var_type.dtype == self.src:
+                dtype = self.dst
+            else:
+                dtype = var_type.dtype
+
+            # Generate new variable.
+            new_param = var(param.name_hint, shape=var_type.shape, dtype=dtype)
+
+            new_params.append(new_param)
+            binds[param] = new_param
+
+        new_body = self.visit(fn.body)
+        # Rewrite the body to use new parameters.
+        new_body = bind(new_body, binds)
+
+        # Construct the updated function and return.
+        return Function(
+            new_params,
+            new_body,
+            # You could change the return type, if you use None it will re-infer.
+            None,
+            type_params=fn.type_params,
+            attrs=fn.attrs,
+        )