TensorFlow から bfloat16 のデータをつくって FPGA で使う

TensorFlow からごくごく簡単な MNIST のモデルを作り bfloat16 に変換して、最終的には FPGA(Zybo) で動くようにしてみた。

MNIST のデータ

散々書かれているので省略。
https://github.com/ryos36/polyphony-with-tf-mnist
にソース等を置いた。とにかく浮動小数点のデータを作ればよい

bfloat16 に変換

Session で eval とすれば NumPy 形式に落としてくれる。tf.cast でキャストすればよい。

bfloat16.py

import numpy as np
import tensorflow as tf

w_np = np.loadtxt('w_value.txt', delimiter=',')
b_np = np.loadtxt('b_value.txt', delimiter=',')

print(type(w_np), type(w_np[0][0]))

w_b = tf.cast(w_np, tf.bfloat16)
b_b = tf.cast(b_np, tf.bfloat16)

print(type(w_b), type(w_b[0][0]))

with tf.Session() as sess:
    w_b_np = w_b.eval()
    b_b_np = b_b.eval()

print(type(w_b_np[0][0]))

NumPy でセーブしたものをちょちょいと加工

ヘッダ付きのバイナリ形式でセーブされる。pickle?

save.py

np.save('w_b_value.npy', w_b_np)
np.save('b_b_value.npy', b_b_np)

dd で 128バイトスキップする。

Python で予測用プログラムを作る

mnist7

def do_mnist7_mem(a:List[bit16], _mem:List[bit16], lst_len = LEN):
    rom_w = W_PARAM
    rom_b = B_PARAM
    mem = [0] * 10

    xi = 0
    for i in range(lst_len):
        x = a[i]
        for j in range(10):
            mem[j] = bfloat.mul_add(x, rom_w[xi + j], mem[j])

        xi += 10

    for j in range(10):
        _mem[j] = bfloat.add(mem[j], rom_b[j])

コンパイルし Zybo にもっていく

持っていくときに少し細工。これで Vivado の BRAM の I/F とリセット極性が設定される。あと BRAM 用にクロックも追加。

インタフェース

    input wire clk,
(* X_INTERFACE_INFO = "xilinx.com:signal:reset:1.0 rst RST" *)
(* X_INTERFACE_PARAMETER = "POLARITY ACTIVE_HIGH" *)
    input wire rst,
    input wire do_mnist7_mem_ready,
    input wire do_mnist7_mem_accept,
    output reg do_mnist7_mem_valid,
(* X_INTERFACE_INFO = "xilinx.com:interface:bram:1.0 bram_in CLK" *)
    output wire do_mnist7_in_a_clk,
(* X_INTERFACE_INFO = "xilinx.com:interface:bram:1.0 bram_in DOUT" *)
    input wire signed [15:0] do_mnist7_mem_in_a_q,
    input wire [10:0] do_mnist7_mem_in_a_len,
(* X_INTERFACE_INFO = "xilinx.com:interface:bram:1.0 bram_in ADDR" *)
    output wire signed [10:0] do_mnist7_mem_in_a_addr,
(* X_INTERFACE_INFO = "xilinx.com:interface:bram:1.0 bram_in DIN" *)
    output wire signed [15:0] do_mnist7_mem_in_a_d,
(* X_INTERFACE_INFO = "xilinx.com:interface:bram:1.0 bram_in WE" *)
    output wire do_mnist7_mem_in_a_we,
(* X_INTERFACE_INFO = "xilinx.com:interface:bram:1.0 bram_in EN" *)
    output wire do_mnist7_mem_in_a_req,

<中略>

  assign do_mnist7_in_a_clk = clk;

おそるべし Polyphony の最適化

かなり最適化が効いてほとんど assign になる。ただ DSP をつかってくれない。
そこで、ちょっとこれも Verilog を細工して DSP を使うようにしてみる。

  assign new_n_inl1_inl13 = (t622_inl1_inl1 * t623_inl1_inl1);

の部分を reg にしてステートを一つ入れてみた。たしかに DSP を使うようにはなった。ただ、結局 fmax のボトルネックはメモリアクセスなのでこの改造は意味がなかった。

unroll してみる

これは失敗に終わる。よーく考えたらメモリアクセスを同時にできないと意味がない。ここは将来的な検討項目。

システム全体はこんな感じ

VIO で確認

10進だからわかりづらい。

今後の展開

Retro FORTH と連携する。AXI Stream と連携する。unroll 対応をする。といったところか。
現時点でも使えているので、パラメタのリアルタイムの可視化とかできそうな気がしている。