create_min_lower_func

Author: hypercubestart

3rdparty/posit/posit-wrapper.cc

@@ -32,6 +32,7 @@
 #include "universal/posit/math/hyperbolic.hpp"
 #include "universal/posit/math/logarithm.hpp"
 #include "universal/posit/math/sqrt.hpp"
+#include "universal/posit/numeric_limits.hpp"
 
 TVM_DLL sw::unum::posit<8, 2> Uint8ToPosit8es2(uint8_t in) {
   sw::unum::bitblock<8> bb;
@@ -40,12 +41,15 @@ TVM_DLL sw::unum::posit<8, 2> Uint8ToPosit8es2(uint8_t in) {
 }
 
 extern "C" {
-TVM_DLL uint8_t RawPosit8es2(uint8_t in) { return in; }
-
 TVM_DLL uint8_t Posit8es2toUint8(sw::unum::posit<8, 2> in) {
   return static_cast<uint8_t>(in.get().to_ullong());
 }
 
+TVM_DLL uint8_t MinPosit8es2() {
+  auto min = std::numeric_limits<sw::unum::posit<8, 2>>::lowest();
+  return Posit8es2toUint8(min);
+}
+
 TVM_DLL float Posit8es2ToFloat(uint8_t in) { return Uint8ToPosit8es2(in).operator float(); }
 
 TVM_DLL uint8_t FloatToPosit8es2(float in) {
@@ -104,12 +108,15 @@ TVM_DLL sw::unum::posit<16, 2> Uint16ToPosit16es2(uint16_t in) {
 }
 
 extern "C" {
-TVM_DLL uint16_t RawPosit16es2(uint16_t in) { return in; }
-
 TVM_DLL uint16_t Posit16es2toUint16(sw::unum::posit<16, 2> in) {
   return static_cast<uint16_t>(in.get().to_ullong());
 }
 
+TVM_DLL uint8_t MinPosit16es2() {
+  auto min = std::numeric_limits<sw::unum::posit<16, 2>>::lowest();
+  return Posit16es2toUint16(min);
+}
+
 TVM_DLL float Posit16es2ToFloat(uint16_t in) { return Uint16ToPosit16es2(in).operator float(); }
 
 TVM_DLL uint16_t FloatToPosit16es2(float in) {
@@ -168,12 +175,15 @@ TVM_DLL sw::unum::posit<32, 2> Uint32ToPosit32es2(uint32_t in) {
 }
 
 extern "C" {
-TVM_DLL uint32_t RawPosit32es2(uint32_t in) { return in; }
-
 TVM_DLL uint32_t Posit32es2ToUint32(sw::unum::posit<32, 2> in) {
   return static_cast<uint32_t>(in.get().to_ullong());
 }
 
+TVM_DLL uint8_t MinPosit32es2() {
+  auto min = std::numeric_limits<sw::unum::posit<32, 2>>::lowest();
+  return Posit32es2ToUint32(min);
+}
+
 TVM_DLL float Posit32es2ToFloat(uint32_t in) { return Uint32ToPosit32es2(in).operator float(); }
 
 TVM_DLL uint32_t FloatToPosit32es2(float in) {

python/tvm/relay/frontend/change_datatype.py

@@ -21,8 +21,7 @@
 from ..transform.transform import function_pass
 from ..expr import var, bind
 
-
-@function_pass()
+@function_pass(opt_level=0)
 class ChangeDatatype(ExprMutator):
     """Mutator for changing the datatype of Relay programs.
 

python/tvm/target/datatype.py

@@ -214,6 +214,27 @@ def register_min_func(func, type_name):
     """
     _register_func("tvm.datatype.min." + type_name, func)
 
+def create_min_lower_func(extern_func_map, type_name):
+    """Returns a function which lowers the minimum value operation to a pure extern call.
+
+    Parameters
+    ----------
+    extern_func_map : map
+        A map from bit lengths to the external function name
+
+    type_name : string
+        The name of the custom datatype, e.g. posites2 (but not custom[posites2]32).
+    """
+    def lower(num_bits):
+        dtype = f'custom[{type_name}]{num_bits}'
+
+        if num_bits not in extern_func_map:
+            raise RuntimeError('missing minimum function for {dtype}')
+
+        return call_pure_extern(dtype, extern_func_map[num_bits])
+
+    return lower
+
 def create_lower_func(extern_func_map):
     """Returns a function which lowers an operation to a function call.
 

src/arith/rewrite_simplify.cc

@@ -30,6 +30,7 @@
 
 #include <algorithm>
 
+#include "../target/datatype/registry.h"
 #include "const_fold.h"
 #include "pattern_match.h"
 
@@ -460,6 +461,8 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const DivNode* op) {
 
   // x / 2.0 = x * 0.5
   if (const FloatImmNode* ptr = op->b.as<FloatImmNode>()) {
+    CHECK(op->dtype.is_float() ||
+          datatype::Registry::Global()->GetTypeRegistered(op->dtype.code()));
     return op->a * make_const(op->b.dtype(), 1.0 / ptr->value);
   }
 

src/tir/op/op.cc

@@ -182,7 +182,8 @@ PrimExpr min_value(const DataType& dtype) {
   if (datatype::Registry::Global()->GetTypeRegistered(dtype.code())) {
     auto f = datatype::GetMinFunc(dtype.code());
     CHECK(f) << "No minimum function registered for custom dtype " << (unsigned int)dtype.code();
-    // TODO(@hypercubestart) Document this change (and others associated with the overflowing floatimm min bug)
+    // TODO(@hypercubestart) Document this change (and others associated with the overflowing
+    // floatimm min bug)
     return (*f)(dtype.bits());
   } else if (dtype.is_int()) {
     if (dtype.bits() == 64) {

tests/python/unittest/test_custom_datatypes.py

@@ -27,7 +27,7 @@
 from tvm.relay.testing.resnet import get_workload as get_resnet
 from tvm.relay.testing.layers import batch_norm_infer
 from tvm.relay.testing.mobilenet import get_workload as get_mobilenet
-from tvm.target.datatype import register, register_min_func, register_op, create_lower_func, lower_ite, lower_call_pure_extern
+from tvm.target.datatype import register, register_min_func, register_op, create_lower_func, lower_ite, lower_call_pure_extern, create_min_lower_func
 from tvm.tir.op import call_pure_extern
 
 # we use a random seed to generate input_data
@@ -170,25 +170,11 @@ def setup():
         8: 'Posit8es2Tanh'
     }), "Call", "llvm", "posites2", intrinsic_name="tir.tanh")
 
-    def posit_min_func(num_bits):
-        # the minimum representable posit is all 1's in binary,
-        # here we encode the raw bit representation in an integer
-        # and use the extern function to simply interpret
-        # the integer as a posites2
-        # 
-        # another possible way is to create a FloatImm storing the value
-        # of the minimum as a float64 and then casting to `posites2`,
-        # but float imprecision makes this approach susceptible to hard-to-find bugs
-        value = np.dtype('int' + str(num_bits)).type(-1)
-        dtype = 'custom[posites2]' + str(num_bits)
-        func_map = {
-            32: 'RawPosit32es2',
-            16: 'RawPosit16es2',
-            8: 'RawPosit8es2'
-        }
-        return call_pure_extern(dtype, func_map[num_bits], value)
-    register_min_func(posit_min_func, "posites2")
-
+    register_min_func(create_min_lower_func({
+        32: 'MinPosit32es2',
+        16: 'MinPosit16es2',
+        8: 'MinPosit8es2'
+    }, "posites2"), "posites2")
 
 def run_ops(src_dtype, dst_dtype, rtol=1e-7, atol=1e-7):
     """Run the same op, but with two different datatypes"""