NNgen と FPGA で作るニューラルネットワーク・アクセラレータ (2)

# --------------------
# (1) Represent a DNN model as a dataflow by NNgen operators
# --------------------

# data types
act_dtype = ng.int8
weight_dtype = ng.int8
bias_dtype = ng.int32
scale_dtype = ng.int8
batchsize = 1

# input
input_layer = ng.placeholder(dtype=act_dtype,
                             shape=(batchsize, 32, 32, 3),  # N, H, W, C
                             name='input_layer')

# layer 0: conv2d (with bias and scale (= batchnorm)), relu, max_pool
w0 = ng.variable(dtype=weight_dtype,
                 shape=(64, 3, 3, 3),  # Och, Ky, Kx, Ich
                 name='w0')
b0 = ng.variable(dtype=bias_dtype,
                 shape=(w0.shape[0],), name='b0')
s0 = ng.variable(dtype=scale_dtype,
                 shape=(w0.shape[0],), name='s0')

a0 = ng.conv2d(input_layer, w0,
               strides=(1, 1, 1, 1),
               bias=b0,
               scale=s0,
               act_func=ng.relu,
               dtype=act_dtype,
               sum_dtype=ng.int32)

a0p = ng.max_pool_serial(a0,
                         ksize=(1, 2, 2, 1),
                         strides=(1, 2, 2, 1))

# ...

# output
output_layer = ng.matmul(a2, w3,
                         bias=b3,
                         scale=s3,
                         transposed_b=True,
                         name='output_layer',
                         dtype=act_dtype,
                         sum_dtype=ng.int32)

targ = ng.to_ipxact([output_layer], 'hello_nngen', silent=silent,
                    config={'maxi_datawidth': axi_datawidth})

def to_ipxact(objs, name, ipname=None, config=None, silent=False):

    if ipname is None:
        ipname = name

    config = load_default_config(config)

    m = _to_veriloggen_module(objs, name, config,
                              silent=silent, where_from='to_ipxact', output=ipname)

    clk_name = config['clock_name']
    rst_name = config['reset_name']
    rst_polarity = ('ACTIVE_LOW'
                    if 'low' in config['reset_polarity'] else
                    'ACTIVE_HIGH')
    irq_name = config['interrupt_name']
    irq_sensitivity = 'LEVEL_HIGH'

    clk_ports = [(clk_name, (rst_name,))]
    rst_ports = [(rst_name, rst_polarity)]
    irq_ports = [(irq_name, irq_sensitivity)] if config['interrupt_enable'] else []

    ipxact.to_ipxact(m, ipname,
                     clk_ports=clk_ports,
                     rst_ports=rst_ports,
                     irq_ports=irq_ports)

    return m

def _to_veriloggen_module(objs, name, config=None,
                          silent=False, where_from=None, output=None):

    if not isinstance(objs, (list, tuple)):
        objs = [objs]

    (objs, num_storages,
     num_input_storages, num_output_storages) = analyze(config, objs)
    m, clk, rst, maxi, saxi = make_module(config, name, objs,
                                          num_storages, num_input_storages,
                                          num_output_storages)

    schedule_table = schedule(config, objs)

    header_info = make_header_addr_map(config, saxi)

    (ram_dict, substrm_dict, ram_set_cache,
     stream_cache, control_cache, main_fsm,
     global_map_info, global_mem_map) = allocate(config, m, clk, rst,
                                                 maxi, saxi, objs, schedule_table)

    reg_map = make_reg_map(config, global_map_info, header_info)

    if not silent:
        dump_config(config, where_from, output)
        dump_schedule_table(schedule_table)
        dump_rams(ram_dict)
        dump_substreams(substrm_dict)
        dump_streams(stream_cache)
        dump_controls(control_cache, main_fsm)
        dump_register_map(reg_map)
        dump_memory_map(global_mem_map)

    return m

def allocate(config, m, clk, rst, maxi, saxi, objs, schedule_table):
    set_storage_name(objs)
    merge_shared_attrs(objs)

    max_stream_rams = calc_max_stream_rams(config, schedule_table)
    ram_dict = make_rams(config, m, clk, rst, maxi, schedule_table, max_stream_rams)
    ram_set_cache = make_ram_sets(config, schedule_table, ram_dict, max_stream_rams)

    control_param_dict = make_control_params(config, schedule_table)

    substrm_dict = make_substreams(config, m, clk, rst, maxi, schedule_table)
    stream_cache = make_streams(config, schedule_table, ram_dict, substrm_dict)

    (global_addr_map, local_addr_map,
     global_map_info, global_mem_map) = make_addr_map(config, objs, saxi)

    if config['use_map_ram']:
        global_map_ram, local_map_ram = make_addr_map_rams(config, m, clk, rst, maxi,
                                                           global_addr_map, local_addr_map)
    else:
        global_map_ram = None
        local_map_ram = None

    control_cache, main_fsm = make_controls(
        config, m, clk, rst, maxi, saxi,
        schedule_table, control_param_dict,
        global_addr_map, local_addr_map,
        global_map_ram, local_map_ram)

    disable_unused_ram_ports(ram_dict)

    return (ram_dict, substrm_dict, ram_set_cache, stream_cache, control_cache,
            main_fsm, global_map_info, global_mem_map)