mindspore 2.0.0rc1__cp38-cp38-manylinux1_x86_64.whl → 2.2.0__cp38-cp38-manylinux1_x86_64.whl

mindspore/_akg/akg/tvm/autotvm/measure/measure_methods.py

@@ -33,8 +33,7 @@ import tempfile
 
 import numpy as np
 
-from ... import ir_pass, build, build_config, nd, TVMError, register_func, \
-    rpc as _rpc, target as _target
+from ... import ir_pass, build, build_config, nd, TVMError, register_func, target as _target
 from ...contrib import nvcc, ndk, tar
 
 from ..util import get_const_tuple
@@ -139,210 +138,6 @@ class LocalBuilder(Builder):
         return results
 
 
-class RPCRunner(Runner):
-    """Run generated code on remove devices.
-    This function will ask a RPC Tracker to get device for measurement.
-
-    Parameters
-    ----------
-    timeout: float
-        The timeout of a compilation
-    n_parallel: int
-        The number of tasks run in parallel. "None" will use all cpu cores
-    key: str
-        The key of the device registered in the tracker
-    host: str
-        The host address of RPC Tracker
-    port: int
-        The port of RPC Tracker
-    number: int
-        The number of times to run the generated code for taking average.
-        We call these runs as one `repeat` of measurement.
-    repeat : int, optional
-        The number of times to repeat the measurement.
-        In total, the generated code will be run (1 + number x repeat) times,
-        where the first "1" is warm up and will be discarded.
-        The returned result contains `repeat` costs,
-        each of which is an average of `number` costs.
-    min_repeat_ms: int, optional
-        The minimum duration of one `repeat` in milliseconds.
-        By default, one `repeat` contains `number` runs. If this parameter is set,
-        the parameters `number` will be dynamically adjusted to meet the
-        minimum duration requirement of one `repeat`.
-        i.e., When the run time of one `repeat` falls below this time, the `number` parameter
-        will be automatically increased.
-    cooldown_interval: float, optional
-        The cool down interval between two measurements.
-    check_correctness: bool, optional
-        Whether check correctness after measurement. This will use llvm cpu target to
-        call your template and get the reference output.
-        This can work for TOPI templates, but may not work for your custom template.
-    """
-    def __init__(self,
-                 key, host, port, priority=1,
-                 timeout=10, n_parallel=None,
-                 number=4, repeat=3, min_repeat_ms=0, cooldown_interval=0.1,
-                 check_correctness=False):
-        super(RPCRunner, self).__init__(timeout, n_parallel)
-
-        self.key = key
-        self.host = host
-        self.port = port
-        self.priority = priority
-        self.timeout = timeout
-
-        self.number = number
-        self.repeat = repeat
-        self.min_repeat_ms = min_repeat_ms
-
-        self.ref_input = None
-        self.ref_output = None
-        self.check_correctness = check_correctness
-        self.cooldown_interval = cooldown_interval
-
-        self.executor = LocalExecutor()
-
-    def set_task(self, task):
-        self.task = task
-
-        if check_remote(task.target, self.key, self.host, self.port):
-            logger.info("Get devices for measurement successfully!")
-        else:
-            raise RuntimeError("Cannot get remote devices from the tracker. "
-                               "Please check the status of tracker by "
-                               "'python -m tvm.exec.query_rpc_tracker --port [THE PORT YOU USE]' "
-                               "and make sure you have free devices on the queue status.")
-
-        if self.check_correctness:
-            # use llvm cpu to generate a reference input/output
-            # this option works for tuning topi, but might not work for you custom op
-            with _target.create("llvm"):
-                s, arg_bufs = task.instantiate(task.config_space.get(0))
-            self.ref_input = [np.random.uniform(size=get_const_tuple(x.shape)).astype(x.dtype)
-                              for x in arg_bufs]
-            func = build(s, arg_bufs, "llvm")
-            tvm_buf = [nd.array(x) for x in self.ref_input]
-            func(*tvm_buf)
-            self.ref_output = [x.asnumpy() for x in tvm_buf]
-
-    def get_build_kwargs(self):
-        kwargs = {}
-        if 'cuda' in self.task.target.keys or 'opencl' in self.task.target.keys or \
-           'rocm' in self.task.target.keys:
-            remote = request_remote(self.key, self.host, self.port)
-            ctx = remote.context(str(self.task.target), 0)
-            max_dims = ctx.max_thread_dimensions
-            kwargs['check_gpu'] = {
-                'max_shared_memory_per_block': ctx.max_shared_memory_per_block,
-                'max_threads_per_block': ctx.max_threads_per_block,
-                'max_thread_x': max_dims[0],
-                'max_thread_y': max_dims[1],
-                'max_thread_z': max_dims[2],
-            }
-
-            if 'cuda' in self.task.target.keys:
-                kwargs["cuda_arch"] = "sm_" + "".join(ctx.compute_version.split('.'))
-
-        return kwargs
-
-    def run(self, measure_inputs, build_results):
-        results = []
-        remote_args = (self.key, self.host, self.port, self.priority, self.timeout)
-
-        for i in range(0, len(measure_inputs), self.n_parallel):
-            futures = []
-            for measure_inp, build_res in zip(measure_inputs[i:i+self.n_parallel],
-                                              build_results[i:i+self.n_parallel]):
-                ret = self.executor.submit(run_through_rpc,
-                                           measure_inp,
-                                           build_res,
-                                           self.number,
-                                           self.repeat,
-                                           self.min_repeat_ms,
-                                           self.cooldown_interval,
-                                           remote_args,
-                                           self.ref_input,
-                                           self.ref_output)
-                futures.append(ret)
-
-            for future in futures:
-                res = future.get()
-                if isinstance(res, Exception):   # executor error or timeout
-                    results.append(MeasureResult((str(res),), MeasureErrorNo.RUN_TIMEOUT,
-                                                 self.timeout, time.time()))
-                else:
-                    results.append(res)
-
-        return results
-
-class LocalRunner(RPCRunner):
-    """Run generated code on local devices.
-
-    Parameters
-    ----------
-    timeout: float
-        The timeout of a compilation
-    number: int
-        The number of times to run the generated code for taking average.
-        We call these runs as one `repeat` of measurement.
-    repeat : int, optional
-        The number of times to repeat the measurement.
-        In total, the generated code will be run (1 + number x repeat) times,
-        where the first one is warm up and will be discarded.
-        The returned result contains `repeat` costs,
-        each of which is an average of `number` costs.
-    min_repeat_ms: int, optional
-        The minimum duration of one `repeat` in milliseconds.
-        By default, one `repeat` contains `number` runs. If this parameter is set,
-        the parameters `number` will be dynamically adjusted to meet the
-        minimum duration requirement of one `repeat`.
-        i.e., When the run time of one `repeat` falls below this time, the `number` parameter
-        will be automatically increased.
-    cooldown_interval: float, optional
-        The cool down interval between two measurements.
-    check_correctness: bool, optional
-        Whether check correctness after measurement. This will use llvm cpu target to
-        call your template and get the reference output.
-        This can work for TOPI templates, but may not work for your custom template.
-
-    Note
-    ----
-    This is a "fake" local mode. We start a silent rpc tracker and rpc server
-    for the user. In this way we reuse timeout/isolation mechanism in RPC infrastructure.
-    """
-    def __init__(self,
-                 timeout=10,
-                 number=4, repeat=3, min_repeat_ms=0, cooldown_interval=0.1,
-                 check_correctness=False):
-        super(LocalRunner, self).__init__('', None, None, 0,
-                                          timeout=timeout, n_parallel=1,
-                                          number=number, repeat=repeat,
-                                          min_repeat_ms=min_repeat_ms,
-                                          cooldown_interval=cooldown_interval,
-                                          check_correctness=check_correctness)
-        self.tracker = None
-        self.server = None
-
-    def set_task(self, task):
-        self.task = task
-
-        from ...rpc.tracker import Tracker
-        from ...rpc.server import Server
-
-        tracker = Tracker('0.0.0.0', port=9000, port_end=10000, silent=True)
-        device_key = '$local$device$%d' % tracker.port
-        server = Server('0.0.0.0', port=9000, port_end=10000,
-                        key=device_key,
-                        use_popen=True, silent=True,
-                        tracker_addr=(tracker.host, tracker.port))
-        self.key = device_key
-        self.host = tracker.host
-        self.port = tracker.port
-
-        super(LocalRunner, self).set_task(task)
-        return server, tracker
-
-
 def _build_func_common(measure_input, check_gpu=None, cuda_arch=None, build_option=None):
     """Common part for building a configuration"""
     target, task, config = measure_input
@@ -413,172 +208,6 @@ def _wrap_build_func(build_func):
     return _wrapped
 
 
-def run_through_rpc(measure_input, build_result,
-                    number, repeat, min_repeat_ms, cooldown_interval,
-                    remote_args, ref_input=None, ref_output=None):
-    """Run a generated library through rpc
-
-    Parameters
-    ----------
-    measure_input: MeasureInput
-        The raw measure input
-    build_result: BuildResult
-        The result returned from Builder. This contains the path to the generated library.
-    number: int
-        The number of times to run the generated code for taking average.
-        We call these runs as one `repeat` of measurement.
-    repeat : int, optional
-        The number of times to repeat the measurement.
-        In total, the generated code will be run (1 + number x repeat) times,
-        where the first one is warm up and will be discarded.
-        The returned result contains `repeat` costs,
-        each of which is an average of `number` costs.
-    min_repeat_ms: int, optional
-        The minimum duration of one `repeat` in milliseconds.
-        By default, one `repeat` contains `number` runs. If this parameter is set,
-        the parameters `number` will be dynamically adjusted to meet the
-        minimum duration requirement of one `repeat`.
-        i.e., When the run time of one `repeat` falls below this time, the `number` parameter
-        will be automatically increased.
-    cooldown_interval: float
-        The cool down interval between two measurements
-    remote_args: Tuple
-        The argument for request_remote
-    ref_input: List of np.ndarray
-        The reference input used for checking correctness
-    ref_output: List of np.ndarray
-        The reference output used for checking correctness
-    """
-    if isinstance(build_result, MeasureResult):
-        return build_result
-
-    tic = time.time()
-    errno = MeasureErrorNo.NO_ERROR
-    try:
-        # upload built module
-        remote = request_remote(*remote_args)
-        # Program the FPGA every single time when targeting VTA
-        if hasattr(measure_input.target, 'device_name') and \
-            measure_input.target.device_name == 'vta':
-            from vta import program_fpga, reconfig_runtime
-            program_fpga(remote, None)
-            reconfig_runtime(remote)
-        remote.upload(build_result.filename)
-        func = remote.load_module(os.path.split(build_result.filename)[1])
-        ctx = remote.context(str(measure_input.target), 0)
-        time_f = func.time_evaluator(
-            func.entry_name, ctx, number=number, repeat=repeat, min_repeat_ms=min_repeat_ms)
-
-        # set input
-        if ref_input:
-            args = [nd.array(x, ctx=ctx) for x in ref_input]
-        else:
-            # create empty arrays on the remote device and copy them once.
-            # This can avoid some memory issues that make the measurement results unreliable.
-            args = [nd.empty(x[0], dtype=x[1], ctx=ctx) for x in build_result.arg_info]
-            args = [nd.array(x, ctx=ctx) for x in args]
-            ctx.sync()
-
-        costs = time_f(*args).results
-
-        # clean up remote files
-        remote.remove(build_result.filename)
-        remote.remove(os.path.splitext(build_result.filename)[0] + '.so')
-        remote.remove('')
-
-        if len(costs) > 2:  # remove largest and smallest value to reduce variance
-            costs = list(costs)
-            costs.sort()
-            costs = tuple(costs[1:-1])
-
-        # check correctness of output
-        if ref_output:
-            for expected, real in zip(ref_output, args):
-                if not np.allclose(expected, real.asnumpy(), rtol=1e-4):
-                    logger.warning("Wrong Answer!")
-                    errno = MeasureErrorNo.WRONG_ANSWER
-    except TVMError as exc:
-        msg = str(exc)
-        if "Stack trace returned" in msg:
-            msg = msg[:msg.index("Stack trace returned")]
-        if "CUDA Source" in msg:
-            msg = msg[:msg.index("CUDA Source")]
-        costs = (RuntimeError(msg[:1024]),)
-        errno = MeasureErrorNo.RUNTIME_DEVICE
-    tstamp = time.time()
-    time.sleep(cooldown_interval)
-    return MeasureResult(costs, errno, tstamp - tic + build_result.time_cost, tstamp)
-
-
-def request_remote(device_key, host=None, port=None, priority=1, timeout=60):
-    """Request a remote session
-
-    Parameters
-    ----------
-    device_key: string
-        The device key of registered device in tracker
-    host: host, optional
-        The host address of rpc tracker.
-        If is none, will use environment variable "TVM_TRACKER_HOST"
-    port: int, optional
-        The port of rpc tracker.
-        If is none, will use environment variable "TVM_TRACKER_PORT"
-    priority: int, optional
-        The priority of this request, larger is more prior
-    timeout: float, optional
-        The timeout of this session (units: second)
-
-    Returns
-    ------
-    session: RPCSession
-    """
-    # connect to the tracker
-    host = host or os.environ['TVM_TRACKER_HOST']
-    port = port or int(os.environ['TVM_TRACKER_PORT'])
-
-    tracker = _rpc.connect_tracker(host, port)
-    remote = tracker.request(device_key, priority=priority,
-                             session_timeout=timeout)
-    return remote
-
-
-def check_remote(target, device_key, host=None, port=None, priority=100, timeout=10):
-    """
-    Check the availability of a remote device
-
-    Parameters
-    ----------
-    target: Target
-        The wanted compilation target
-    device_key: string
-        device key of registered device in tracker
-    host: host, optional
-        The host address of rpc tracker.
-        If is none, will use environment variable "TVM_TRACKER_HOST"
-    port: int, optional
-        The port address of rpc tracker.
-        If is none, will use environment variable "TVM_TRACKER_PORT"
-    priority: int, optional
-        The priority of this request, larger is more prior
-    timeout: float, optional
-        The timeout of this check (units: seconds).
-
-    Returns
-    -------
-    available: bool
-        True if can find available device
-    """
-    def _check():
-        remote = request_remote(device_key, host, port, priority)
-        ctx = remote.context(str(target))
-        while not ctx.exist:  # wait until we get an available device
-            pass
-    t = threading.Thread(target=_check,)
-    t.start()
-    t.join(timeout)
-    return not t.is_alive()
-
-
 @register_func
 def tvm_callback_cuda_compile(code):
     """use nvcc to generate ptx code for better optimization"""

mindspore/_akg/akg/tvm/build_module.py

@@ -21,6 +21,7 @@ LoweredFunc and compiled Module.
 """
 # 2021.01.13 - Modify variable unroll_explicit to False.
 # 2021.03.17 - Add dump_mlir.
+# 2023.02.03 - Add activated_pass.
 
 from __future__ import absolute_import as _abs
 import warnings
@@ -182,6 +183,19 @@ class BuildConfig(NodeBase):
             add_lower_pass_args += [x[0], x[1]]
         _api_internal._BuildConfigSetAddLowerPass(self, *add_lower_pass_args)
 
+    @property
+    def activated_passes(self):
+        size = _api_internal._BuildConfigGetActivatedPassesInfo(self)
+        result = []
+        for i in range(size):
+            result.append(_api_internal._BuildConfigGetActivatedPassesInfo(self, i))
+        return result
+
+    @activated_passes.setter
+    def activated_passes(self, value):
+        activated_passes_args = [pass_name for pass_name in value]
+        _api_internal._BuildConfigSetActivatedPasses(self, *activated_passes_args)
+
     def __enter__(self):
         # pylint: disable=protected-access
         _api_internal._EnterBuildConfigScope(self)
@@ -266,7 +280,8 @@ def build_config(**kwargs):
 
     if "add_lower_pass" in kwargs:
         config.add_lower_pass = kwargs["add_lower_pass"]
-
+    if "activated_passes" in kwargs:
+        config.activated_passes = kwargs["activated_passes"]
     return config
 
 def get_binds(args, compact=False, binds=None):

mindspore/_akg/akg/tvm/contrib/graph_runtime.py

@@ -20,7 +20,6 @@ import numpy as np
 from .._ffi.base import string_types
 from .._ffi.function import get_global_func
 from .._ffi.runtime_ctypes import TVMContext
-from ..rpc import base as rpc_base
 
 def create(graph_json_str, libmod, ctx):
     """Create a runtime executor module given a graph and module.
@@ -48,60 +47,9 @@ def create(graph_json_str, libmod, ctx):
         except AttributeError:
             raise ValueError("Type %s is not supported" % type(graph_json_str))
 
-    ctx, num_rpc_ctx, device_type_id = get_device_ctx(libmod, ctx)
-
-    if num_rpc_ctx == len(ctx):
-        hmod = rpc_base._ModuleHandle(libmod)
-        fcreate = ctx[0]._rpc_sess.get_function("tvm.graph_runtime.remote_create")
-        return GraphModule(fcreate(graph_json_str, hmod, *device_type_id))
-
     fcreate = get_global_func("tvm.graph_runtime.create")
     return GraphModule(fcreate(graph_json_str, libmod, *device_type_id))
 
-def get_device_ctx(libmod, ctx):
-    """Parse and validate all the device context(s).
-    Parameters
-    ----------
-    libmod : tvm.Module
-        The module of the corresponding function
-    ctx : TVMContext or list of TVMContext
-    Returns
-    -------
-    ctx : list of TVMContext
-    num_rpc_ctx : Number of rpc contexts
-    device_type_id : List of device type and device id
-    """
-
-    if isinstance(ctx, TVMContext):
-        ctx = [ctx]
-    elif not isinstance(ctx, (list, tuple)):
-        raise ValueError("ctx has to be the type of TVMContext or a list of "
-                         "TVMCTVMContext")
-    for cur_ctx in ctx:
-        if not isinstance(cur_ctx, TVMContext):
-            raise ValueError("ctx has to be the type of TVMContext or a list "
-                             "of TVMContext")
-
-    # device_type_id[0], device_type_id[1] are used as the primary/fallback
-    # context type and id. All other ones are used as device context for
-    # heterogeneous execution.
-    num_rpc_ctx = 0
-    device_type_id = []
-    for cur_ctx in ctx:
-        device_type = cur_ctx.device_type
-        if device_type >= rpc_base.RPC_SESS_MASK:
-            assert libmod.type_key == "rpc"
-            assert rpc_base._SessTableIndex(
-                libmod) == cur_ctx._rpc_sess._tbl_index
-            num_rpc_ctx += 1
-            device_type = cur_ctx.device_type % rpc_base.RPC_SESS_MASK
-        device_type_id.append(device_type)
-        device_type_id.append(cur_ctx.device_id)
-
-    if 0 < num_rpc_ctx < len(ctx):
-        raise ValueError("Either all or none of the contexts should be rpc.")
-    return ctx, num_rpc_ctx, device_type_id
-
 
 class GraphModule(object):
     """Wrapper runtime module.
@@ -149,7 +97,6 @@ class GraphModule(object):
             self._get_input(key).copyfrom(value)
 
         if params:
-            # upload big arrays first to avoid memory issue in rpc mode
             keys = list(params.keys())
             keys.sort(key=lambda x: -np.prod(params[x].shape))
             for k in keys:

mindspore/_akg/akg/tvm/hybrid/parser.py

@@ -16,6 +16,7 @@
 # under the License.
 """Hybrid Script Parser"""
 
+# 2023.08.16 - Support parser for python 3.9 and up.
 # 2022.02.15 - Support grid as new loop mode.
 # 2022.01.19 - Support negative extent for range.
 # 2021.12.15 - Support block_realize intrin in with scope.
@@ -304,7 +305,7 @@ class HybridParser(ast.NodeVisitor):
             _dtype = _buf.dtype
             _true = _api.convert(True)
             body = _make.Realize(_buf.op, _buf.value_index, _dtype, _domain, _true, body)
-            body = _make.AttrStmt(_buf.op, 'realize_scope',  _api.convert(_scope), body)
+            body = _make.AttrStmt(_buf.op, 'realize_scope',  _api.convert(_scope), body, None)
 
         for elem in to_pop:
             self.symbols.pop(elem)
@@ -315,7 +316,7 @@ class HybridParser(ast.NodeVisitor):
     def wrap_up_binds(self, body):
         for _, iter_var in self.binds.items():
             ext = iter_var.dom.extent
-            body = _make.AttrStmt(iter_var, 'thread_extent', ext, body)
+            body = _make.AttrStmt(iter_var, 'thread_extent', ext, body, None)
         self.binds = {}
         return body
 
@@ -650,7 +651,7 @@ class HybridParser(ast.NodeVisitor):
                     _apply_indices(self.closure_vars[node.value.id], args))
 
             buf = self.visit(node.value)
-            if isinstance(buf, Array):
+            if isinstance(buf, (Array, list)):
                 for i in args:
                     if isinstance(i, numbers.Integral):
                         buf = buf[i]
@@ -709,17 +710,17 @@ class HybridParser(ast.NodeVisitor):
         if context.func.id == "attr":
             args = [self.visit(i) for i in context.args]
             return _make.AttrStmt(_api._IterVar(None, block.loop_var.name, 0), args[0],
-                                  _api.convert(args[1]), block)
+                                  _api.convert(args[1]), block, None)
         elif context.func.id == "allocate":
             lhs = node.items[0].optional_vars
             rhs = self.visit(context)
             _ = self.generate_one_assign(lhs, rhs)
             return _make.AttrStmt(rhs, "type",
-                                  _api.convert("inline"), block)
+                                  _api.convert("inline"), block, None)
         elif context.func.id == "block_realize":
             args = [self.visit(i) for i in context.args]
             return _make.AttrStmt(args[0].op, "block_realize",
-                                  _api.convert(True), block)
+                                  _api.convert(True), block, None)
         else:
             raise ValueError("unsupported function in With scope")
 

mindspore/_akg/akg/tvm/ir_builder.py

@@ -174,7 +174,7 @@ class IRBuilder(object):
             node = _make.StringImm(node)
         if isinstance(value, string_types):
             value = _make.StringImm(value)
-        self.emit(lambda x: _make.AttrStmt(node, attr_key, value, x))
+        self.emit(lambda x: _make.AttrStmt(node, attr_key, value, x, None))
 
     def for_range(self, begin, end, name="i", dtype="int32", for_type="serial"):
         """Create a for iteration scope.

mindspore/_akg/akg/tvm/module.py

@@ -206,7 +206,7 @@ class Module(ModuleBase):
             The ProfileResult reports `repeat` time costs in seconds.
         """
         try:
-            feval = _RPCTimeEvaluator(
+            feval = _TimeEvaluator(
                 self, func_name, ctx.device_type, ctx.device_id, number, repeat, min_repeat_ms,
                 cooldown_interval_ms, repeats_to_cooldown, preproc_name)
 
@@ -268,7 +268,6 @@ def load(path, fmt=""):
     cc.create_shared if the path is in format .o or .tar
     """
     # High level handling for .o and .tar file.
-    # We support this to be consistent with RPC module load.
     if path.endswith(".o"):
         _cc.create_shared(path + ".so", path)
         path += ".so"

mindspore/_akg/akg/tvm/stmt.py

@@ -235,9 +235,9 @@ class AttrStmt(Stmt):
     body : Stmt
         The body statement.
     """
-    def __init__(self, node, attr_key, value, body):
+    def __init__(self, node, attr_key, value, body, bounds=None):
         self.__init_handle_by_constructor__(
-            _make.AttrStmt, node, attr_key, value, body)
+            _make.AttrStmt, node, attr_key, value, body, bounds)
 
 
 @register_node

mindspore/_akg/akg/utils/composite_op_helper.py

@@ -95,7 +95,7 @@ def compare_tensor(acu_output, exp_output, rtol=1.e-5, atol=1.e-8, equal_nan=Fal
 
 def get_rtol_atol(op_name, dtype, rtol=5e-03, atol=5e-03):
     run_mode = os.environ.get('RUNTIME_MODE')
-    if run_mode in ("rpc_cloud", "air_cloud"):
+    if run_mode in ("air_cloud",):
         if dtype == "float16":
             rtol = atol = 1e-03
         else:
@@ -272,11 +272,7 @@ def _gen_uniq_file_name(op_name):
 
 def _collect_inplace_assign_infos(op, infos, sum_out):
     """Collect inplace assign infos."""
-    if op["name"] != "InplaceAssign":
-        return
-    fake_output = get_attr(op["attr"], "fake_output")
-    if fake_output:
-        infos["fake_output_tensors"].append(op["output_desc"][0]["tensor_name"])
+    infos["fake_output_tensors"].append(op["output_desc"][0]["tensor_name"])
     input0, input1 = op["input_desc"][0][0], op["input_desc"][1][0]
     if input1["tensor_name"] in sum_out:
         infos["clean_input"].append(input0["tensor_name"])
@@ -302,9 +298,8 @@ def _collect_infos(desc, infos):
                                                                         indices_shape=input1["shape"],
                                                                         indices_dtype=input1["data_type"],
                                                                         attrs=get_attr(op["attr"], "num_segments"))
-        elif op["name"] in ["InplaceAssign", "Assign"]:
-            if op["name"] == "InplaceAssign":
-                _collect_inplace_assign_infos(op, infos, sum_out)
+        elif op["name"] == "Assign":
+            _collect_inplace_assign_infos(op, infos, sum_out)
             infos["inplace_assign_write"].append(op["input_desc"][0][0]["tensor_name"])
         elif op["name"] in ["TensorScatterAdd", "Gather", "GatherNd"]:
             input0, input1 = op["input_desc"][0][0], op["input_desc"][1][0]
@@ -543,7 +538,11 @@ def _update_workspace_data(kernel_name, input_for_mod, output_indexes):
 def gen_json_data(op_desc, with_compute=True, input_for_mod=None):
     """Generating test data for composite json"""
     desc = json.loads(op_desc)
-
+    from akg.ms.info_version_adapt import InfoVersionAdapt
+    info_adapter = InfoVersionAdapt(desc)
+    ret = info_adapter.run()
+    if not ret:
+        raise RuntimeError(info_adapter.msg)
     output_indexes = []
     expect = []
     infos = {"gen_data": False,