rebase and update

Author: hypercubestart

python/tvm/target/datatype.py

@@ -17,9 +17,11 @@
 """Custom datatype functionality"""
 import tvm
 from tvm.runtime import convert, DataType
-from tvm.tir.expr import Call as _Call, Cast as _Cast, FloatImm as _FloatImm
+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
+from tvm.ir import Op
 
 
 def register(type_name, type_code):
@@ -154,17 +156,15 @@ def lower(op):
                 dtype += "x" + str(t.lanes)
         if isinstance(op, _Cast):
             src_bits = bit_length(op.value.dtype)
-            return _Call(dtype, extern_func_map[(src_bits, t.bits)], convert([op.value]),
-                         _Call.Extern)
+            return call_pure_extern(dtype, extern_func_map[(src_bits, t.bits)], op.value)
         elif isinstance(op, _FloatImm):
-            return _Call(dtype, extern_func_map[t.bits], convert([op.value]),
-                         _Call.Extern)
-        elif isinstance(op, _Call) and (op.call_type == _Call.Intrinsic or
-                                        op.call_type == _Call.PureIntrinsic):
-            return _Call(dtype, extern_func_map[t.bits], convert(op.args),
-                              _Call.Extern)
-        return _Call(dtype, extern_func_map[t.bits], convert([op.a, op.b]),
-                     _Call.Extern)
+            return call_pure_extern(dtype, extern_func_map[t.bits], op.value)
+        elif isinstance(op, _Call):
+            return call_pure_extern(dtype, extern_func_map[t.bits], *op.args)
+        elif isinstance(op, _BinaryOpExpr):
+            return call_pure_extern(dtype, extern_func_map[t.bits], op.a, op.b)
+
+        raise RuntimeError(f"lowering unsupported op: {op.astext()}")
 
     return lower
 
@@ -179,6 +179,6 @@ def lower_ite(ite_intrin):
     dtype = "uint" + str(t.bits)
     if t.lanes > 1:
         dtype += "x" + str(t.lanes)
-    return call_intrin(dtype, "tvm_if_then_else", convert(ite_intrin.args[0]),
+    return call_intrin(dtype, "tir.if_then_else", convert(ite_intrin.args[0]),
                        convert(ite_intrin.args[1]),
                        convert(ite_intrin.args[2]))

src/tir/transforms/lower_custom_datatypes.cc

@@ -23,6 +23,7 @@
 
 #include <tvm/runtime/registry.h>
 #include <tvm/target/target.h>
+#include <tvm/tir/op.h>
 #include <tvm/tir/stmt_functor.h>
 #include <tvm/tir/transform.h>
 
@@ -102,11 +103,12 @@ class CustomDatatypesLowerer : public StmtExprMutator {
     PrimExpr expr = StmtExprMutator::VisitExpr_(call);
     call = expr.as<CallNode>();
     if (toBeLowered) {
-      CHECK(call->call_type == CallNode::Intrinsic || call->call_type == CallNode::PureIntrinsic)
+      auto op = call->op.as<OpNode>();
+      CHECK(op != nullptr)
           << "Lowering non-intrinsic Calls not implemented";
-      auto lower = datatype::GetIntrinLowerFunc(target_, call->name, call->dtype.code());
+      auto lower = datatype::GetIntrinLowerFunc(target_, op->name, call->dtype.code());
       CHECK(lower) << "Intrinsic lowering function for target " << target_ << ", intrinsic name "
-                   << call->name << ", type " << static_cast<unsigned>(call->dtype.code())
+                   << op->name << ", type " << static_cast<unsigned>(call->dtype.code())
                    << " not found";
       return (*lower)(expr);
     }

tests/python/unittest/test_custom_datatypes.py

@@ -142,32 +142,32 @@ def setup():
         32: 'Posit32es2Sqrt',
         16: 'Posit16es2Sqrt',
         8: 'Posit8es2Sqrt'
-    }), "Call", "llvm", "posites2", intrinsic_name="sqrt")
+    }), "Call", "llvm", "posites2", intrinsic_name="tir.sqrt")
     register_op(lower_ite,
                 "Call",
                 "llvm",
                 "posites2",
-                intrinsic_name="tvm_if_then_else")
+                intrinsic_name="tir.if_then_else")
     register_op(create_lower_func({
         32: 'Posit32es2Exp',
         16: 'Posit16es2Exp',
         8: 'Posit8es2Exp'
-    }), "Call", "llvm", "posites2", intrinsic_name="exp")
+    }), "Call", "llvm", "posites2", intrinsic_name="tir.exp")
     register_op(create_lower_func({
         32: 'Posit32es2Log',
         16: 'Posit16es2Log',
         8: 'Posit8es2Log'
-    }), "Call", "llvm", "posites2", intrinsic_name="log")
+    }), "Call", "llvm", "posites2", intrinsic_name="tir.log")
     register_op(create_lower_func({
         32: 'Posit32es2Sigmoid',
         16: 'Posit16es2Sigmoid',
         8: 'Posit8es2Sigmoid'
-    }), "Call", "llvm", "posites2", intrinsic_name="sigmoid")
+    }), "Call", "llvm", "posites2", intrinsic_name="tir.sigmoid")
     register_op(create_lower_func({
         32: 'Posit32es2Tanh',
         16: 'Posit16es2Tanh',
         8: 'Posit8es2Tanh'
-    }), "Call", "llvm", "posites2", intrinsic_name="tanh")
+    }), "Call", "llvm", "posites2", intrinsic_name="tir.tanh")
     register_min_func(lambda num_bits: - (2 ** 2 ** 2) ** (num_bits - 2), "posites2")
 
 def run_ops(src_dtype, dst_dtype, rtol=1e-7, atol=1e-7):