def run_compiler_on_subgraphs()

in python/graph_def_util.py [0:0]

• 61 lines of code
• 32 McCabe index (conditional complexity)

def run_compiler_on_subgraphs(graph_def, dumper):
    IOTensor = namedtuple('IOTensor', 'name, dtype, shape')
    for node in get_neuron_nodes(graph_def):
        is_compilable, reason = neuron_node_is_compilable(node)
        if not is_compilable:
            logging.warning('Not fusing subgraph {} because {}'.format(node.name, reason))
            continue

        # get graph_def and io tensors
        subgraph_def = get_subgraph_def(node)
        inputs = []
        outputs = []
        nal = lambda key: node.attr[key].list
        zip_inputs = zip(nal(knInputNames).s, nal(knInputDtypes).type, nal(knInputShapes).shape)
        zip_outputs = zip(nal(knOutputNames).s, nal(knOutputDtypes).type, nal(knOutputShapes).shape)
        for container, tensors in zip([inputs, outputs], [zip_inputs, zip_outputs]):
            for name, dtype_enum, shape in tensors:
                name = name.decode()
                dtype = dtypes.as_dtype(dtype_enum)
                tensor = IOTensor(name, dtype, shape)
                container.append(tensor)

        # remove attributes that are not recognized by neuron-cc
        for sg_node in subgraph_def.node:
            sg_node.attr.pop(kNeuronInferredShapes)

        # setup workdir and run neuron-cc
        executable, inputs, outputs = compile_savetemps(subgraph_def, inputs, outputs, node.name, dumper)
        if executable:
            node.attr[knExecutable].s = executable
            node.attr[knInputNames].list.s[:] = [ts.name.encode() for ts in inputs]
            node.attr[knOutputNames].list.s[:] = [ts.name.encode() for ts in outputs]
            try:
                input_shuffles = [inp.shuffle for inp in inputs]
            except AttributeError:
                input_shuffles = [None for inp in inputs]
            do_input_shuffles = any(shuffle is not None for shuffle in input_shuffles)
            if do_input_shuffles:
                for shuffle in input_shuffles:
                    idx_ts = node.attr['_input_shuffles'].list.tensor.add()
                    idx_ts.int64_val.extend(shuffle)
            try:
                input_batch_axis = [ts.batch_axis for ts in inputs]
                output_batch_axis = [ts.batch_axis for ts in outputs]
            except AttributeError:
                pass
            else:
                input_batch_axis = [vInvalidAxis if ax is None else ax for ax in input_batch_axis]
                output_batch_axis = [vInvalidAxis if ax is None else ax for ax in output_batch_axis]
                node.attr['input_batch_axis'].list.i[:] = input_batch_axis
                node.attr['output_batch_axis'].list.i[:] = output_batch_axis
                input_args = node.attr[knInputShapes].list.shape, inputs
                output_args = node.attr[knOutputShapes].list.shape, outputs
                for args in input_args, output_args:
                    for shape, ts in zip(*args):
                        if [dim.size for dim in shape.dim] != ts.shape:
                            for dim, size in zip(shape.dim, ts.shape):
                                dim.size = size
            try:
                input_can_use_shm = [ts.can_use_shm for ts in inputs]
                output_can_use_shm = [ts.can_use_shm for ts in outputs]
            except AttributeError:
                pass
            else:
                node.attr['_input_can_use_shm'].list.b[:] = input_can_use_shm
                node.attr['_output_can_use_shm'].list.b[:] = output_can_use_shm
    return graph_def