[symbol_structure]: add mir op_class, add mir symbol_pass

Author: cortexlfj

python/mrt/mir/opclass.py

@@ -0,0 +1,188 @@
+import typing
+from dataclasses import dataclass
+from . import opns
+from . import symbol
+
+MRT_OP_MAP: typing.Dict[str, typing.Any] = {}
+
+#def _register_op_map_(op_name: str, clss:typing.Any=None):
+#    if len(op_name)>0 and clss!=None:
+#        if op_name not in MRT_OP_MAP:
+#            MRT_OP_MAP[op_name] = clss
+#    return MRT_OP_MAP
+
+def _register_op_map(op_name: str): #, clss:typing.Any=None):
+    def _wrapper(clss: typing.Any=None):
+        if len(op_name)>0 and clss!=None:
+            if op_name not in MRT_OP_MAP:
+                MRT_OP_MAP[op_name] = clss
+        return clss
+    return _wrapper
+
+@_register_op_map(opns.CONV2D)
+@dataclass(init=False)
+class Conv2D(symbol.Symbol):
+
+    op_name = opns.CONV2D
+
+    @property
+    def strides(self) -> typing.Tuple[int, int]:
+        default_val = (1,1)
+        return self._strides if self._strides else self.attrs['strides'] if 'strides' in self.attrs else default_val 
+
+    @property
+    def padding(self) -> typing.Tuple[int, int, int, int]:
+        default_val = (0,0,0,0)
+        return self._padding if self._padding else self.attrs['padding'] if 'padding' in self.attrs else default_val 
+
+    @property
+    def dilation(self) -> typing.Tuple[int, int]:
+        default_val = (1,1)
+        return self._ if self._ else self.attrs[''] if '' in self.attrs else default_val 
+
+    @property
+    def kernel_size(self) -> typing.Tuple[int, int]:
+        default_val = (3,3)
+        return self._kernel_size if self._kernel_size else self.attrs['kernel_size'] if 'kernel_size' in self.attrs else default_val 
+
+    def __init__(self, name:str, **kwargs):
+        self.name = name
+        self.args = kwargs.pop('args', [])
+        self.attrs = kwargs.pop('attrs', {})
+        self.extra_attrs = {}
+
+        # TODO: what if strides not in attrs?
+        self._strides = self.attrs['strides']
+        if 'padding' in self.attrs:
+            self._padding = self.attrs['padding']
+        if 'dilation' in self.attrs:
+            self._dilation = self.attrs['dilation']
+        if 'kernel_size' in self.attrs:
+            self._kernel_size = self.attrs['kernel_size']
+
+
+@_register_op_map(opns.DROP_OUT)
+@dataclass(init=False)
+class Dropout(symbol.Symbol):
+
+    op_name = opns.DROP_OUT
+
+    @property
+    def rate(self) -> float:
+        default_val = 0.0
+        return self._rate if self._rate else self.attrs['rate'] if 'rate' in self.attrs else default_val 
+    
+    def __init__(self, name:str, **kwargs):
+        self.name = name
+        self.args = kwargs.pop('args', [])
+        self.attrs = kwargs.pop('attrs', {})
+        self.extra_attrs = {}
+
+        self._rate = self.attrs['rate']
+
+@_register_op_map(opns.CLIP)
+@dataclass(init=False)
+class Clip(symbol.Symbol):
+
+    op_name = opns.CLIP
+    
+    @property
+    def min(self) -> float:
+        default_val = None
+        return self._min if self._min else self.attrs['min'] if 'min' in self.attrs else default_val
+
+    @property
+    def max(self) -> float:
+        default_val = None
+        return self._max if self._max else self.attrs['max'] if 'max' in self.attrs else default_val
+
+    def __init__(self, name:str, **kwargs):
+        self.name = name
+        self.args = kwargs.pop('args', [])
+        self.attrs = kwargs.pop('attrs', {})
+        self.extra_attrs = {}
+
+        self._min = self.attrs['min']
+        self._max = self.attrs['max']
+
+
+@_register_op_map(opns.BATCH_NORM)
+@dataclass(init=False)
+class BatchNorm(symbol.Symbol):
+
+    op_name = opns.BATCH_NORM
+
+    @property
+    def axis(self) -> float:
+        default_val = 1
+        return self._axis if self._axis else self.attrs['axis'] if 'axis' in self.attrs else default_val
+
+    @property
+    def epsilon(self) -> float:
+        default_val = 1e-5
+        return self._epsilon if self._epsilon else self.attrs['epsilon'] if 'epsilon' in self.attrs else default_val
+
+    @property
+    def center(self) -> float:
+        default_val = True
+        return self._center if self._center else self.attrs['center'] if 'center' in self.attrs else default_val
+
+    @property
+    def scale(self) -> float:
+        default_val = True
+        return self._scale if self._scale else self.attrs['scale'] if 'scale' in self.attrs else default_val
+
+    def __init__(self, name:str, **kwargs):
+        self.name = name
+        self.args = kwargs.pop('args', [])
+        self.attrs = kwargs.pop('attrs', {})
+        self.extra_attrs = {}
+
+        self._axis = self.attrs['axis']
+        self._epsilon = self.attrs['epsilon']
+        self._center = self.attrs['center']
+        self._scale = self.attrs['scale']
+
+@_register_op_map(opns.DENSE)
+@dataclass(init=False)
+class Dense(symbol.Symbol):
+
+    op_name = opns.DENSE
+    
+    def __init__(self, name:str, **kwargs):
+        self.name = name
+        self.args = kwargs.pop('args', [])
+        self.attrs = kwargs.pop('attrs', {})
+        self.extra_attrs = {}
+ 
+@_register_op_map(opns.TUPLE_GET_ITEM)
+@dataclass(init=False)
+class TupleGetItem(symbol.Symbol):
+
+    op_name = opns.TUPLE_GET_ITEM
+    
+    @property
+    def index(self) -> float:
+        default_val = 0
+        return self._index if self._index else self.attrs['index'] if 'index' in self.attrs else default_val
+
+    def __init__(self, name:str, **kwargs):
+        self.name = name
+        self.args = kwargs.pop('args', [])
+        self.attrs = kwargs.pop('attrs', {})
+        self.extra_attrs = {}
+
+        self._index  = self.attrs['index']
+
+@_register_op_map(opns.MUL)
+@dataclass(init=False)
+class Multiply(symbol.Symbol):
+
+    op_name = opns.MUL
+    
+    def __init__(self, name:str, **kwargs):
+        self.name = name
+        self.args = kwargs.pop('args', [])
+        self.attrs = kwargs.pop('attrs', {})
+        self.extra_attrs = {}
+ 

python/mrt/mir/opns.py

@@ -1,6 +1,14 @@
 """ MRT operator names """
+import typing
+
+MRT_OP_SET = set()
+def _register_op_list(*op_names: typing.List[str]):
+    for op_name in op_names:
+        if op_name not in MRT_OP_SET:
+            MRT_OP_SET.add(op_name)
 
 VAR = "var"
+_register_op_list(VAR)
 
 DROP_OUT = "nn.dropout"
 CONV2D = "nn.conv2d"
@@ -14,22 +22,29 @@
 ADAPTIVE_AVG_POOL2D = "nn.adaptive_avg_pool2d"
 AVG_POOL2D = "nn.avg_pool2d"
 MAX_POOL2D = "nn.max_pool2d"
+_register_op_list(DROP_OUT, CONV2D, DENSE, BATCH_NORM, RELU,
+                  HARDTANH, SILU, LEAKY_RELU, ADAPTIVE_AVG_POOL2D,
+                   AVG_POOL2D, MAX_POOL2D)
 
 SOFTMAX = "nn.softmax"
 LOG_SOFTMAX = "nn.log_softmax"
+_register_op_list(SOFTMAX, LOG_SOFTMAX)
 
 EXP = "exp"
 SIGMOID = "sigmoid"
+_register_op_list(EXP, SIGMOID)
 
 SUM = "sum"
 MEAN = "mean"
 MAX_AXIS = "max"
 MAXIMUM = "maximum"
 MINIMUM = "minimum"
+_register_op_list(SUM, MEAN, MAX_AXIS, MAXIMUM, MINIMUM)
 
 # =========== NON-CALC ops ===============
 TUPLE = "Tuple"
 TUPLE_GET_ITEM = "TupleGetItem"
+_register_op_list(TUPLE, TUPLE_GET_ITEM)
 
 REPEAT = "repeat"
 SQUEEZE = "squeeze"
@@ -40,16 +55,20 @@
 SPLIT = "split"
 TRANSPOSE = "transpose"
 BROADCAST_TO = "broadcast_to"
+_register_op_list(REPEAT, SQUEEZE, FLATTEN, BATCH_FLATTEN, RESHAPE,
+                  CONCAT, SPLIT, TRANSPOSE, BROADCAST_TO, )
 
 EXPAND_DIMS = "expand_dims"
 TILE = "tile"
+_register_op_list(EXPAND_DIMS, TILE)
 
 WHERE = "where"
 GREATER = "greater"
 STRIDED_SLICE = "strided_slice"
 SLICE_LIKE = "slice_like"
 GET_VALID_COUNT = "vision.get_valid_counts"
 NON_MAX_SUPRESSION = "vision.non_max_suppression"
+_register_op_list(WHERE, GREATER, STRIDED_SLICE, SLICE_LIKE, GET_VALID_COUNT, NON_MAX_SUPRESSION)
 
 # relax clip attrs from a_min/a_max to min/max
 CLIP = "clip"
@@ -58,11 +77,14 @@
 # relax support astype instead of cast
 AS_TYPE = "astype"
 #  CAST = "cast"
+_register_op_list(CLIP, CEIL, RIGHT_SHIFT, AS_TYPE)
 
 ADV_INDEX = "adv_index"
+_register_op_list(ADV_INDEX)
 
 CALL_TIR = "call_tir"
 CALL_DPS_PACKED = "call_dps_packed"
+_register_op_list(CALL_TIR, CALL_DPS_PACKED)
 
 # ======= binary ops =============
 
@@ -71,6 +93,7 @@
 MUL = "multiply"
 MATMUL = "matmul"
 DIV = "divide"
+_register_op_list(ADD, SUB, MUL, MATMUL, DIV)
 
 # ======= unary ops ==============
 
@@ -81,14 +104,17 @@
 POW = "pow"
 
 PASS = "pass"
+_register_op_list(NEGATIVE, ABS, LOG, SQRT, POW, PASS)
 # ======= auto generate op =========
 ARANGE = "arange"
 ZEROS_LIKE = "zeros_like"
 ONES_LIKE = "ones_like"
+_register_op_list(ARANGE, ZEROS_LIKE, ONES_LIKE)
 
 # ======= control flow op ===========
 IF = "if"
 ARGWHERE = "argwhere"
+_register_op_list(IF, ARGWHERE)
 
 # ======= mrt requant op ==========
 REQUANT = "mrt.requant"
@@ -98,4 +124,9 @@
 """ right shift precision clip """
 LUT = "mrt.lut"
 """ look up table, equals adv_index in tvm """
+_register_op_list(REQUANT, PCLIP, RS_PCLIP, LUT)
+
 
+def Opname2Funcname(op_name: str):
+    return op_name.replace('.', '_')
+#print('MRT_OP_SET:', MRT_OP_SET)

python/mrt/mir/symbol.py

@@ -11,8 +11,7 @@
 
 # from . import config
 # from .utils import *
-#  from .types import *
-from .opns import *
+from . import opns
 
 __ALL__ = [
         "Symbol",
@@ -277,34 +276,6 @@ def __hash__(self) -> int:
     def hash(self) -> int:
         return hash(str(self))
 
-# class Convolution2D(Symbol):
-#     strides: typing.Tuple[int, int]
-
-# class Dropout(Symbol):
-#     eps: float = 1e-5
-
-# class Pass:
-#     symbol: Symbol
-
-#     def visit(self, op: Symbol):
-#         env: typing.Dict[Symbol, Symbol] = {}
-#         for sym in sym2list(self.symbol):
-#             out = getattr(self, f"visit_{op.op_name}")(op) or op
-#             assert isinstance(sym, Symbol)
-#             env[sym] = out
-#         return env[op]
-
-# def _default_visit_op(op):
-#     return op
-
-# for op in op_list:
-#     setattr(Pass, f"visit_{op.op_name}", _default_visit_op)
-
-# class FuseDropoutPass(Pass):
-#     def visit_dropout(self, op: Dropout):
-#         op.eps
-#         return op.args[0]
-
 def _topo_sort(symbol: Symbol, sym_list: typing.List[Symbol]):
     assert isinstance(symbol, Symbol), \
             f"({type(symbol).__name__}){str(symbol)}"