swan改善案 - takgto/utokyo-chipathon2023 GitHub Wiki

swanを修正してみましょう。解析方法や改善方法の案を以下に示します。
こちらで用意した改善方法を以下に示しています。もっといい方法を思いついたら試してみてください。

このページではswanのソースコードの記述を変更する操作が出てきます。ファイルを直接開いてVSCode等で編集しても、Vitisで開いて修正してもいいです。
注意点があって、アプリケーションプロジェクトをビルドした時にビルドされるソースコードは元のGITでクローンしたswanのソースコードではなくて、プロジェクトディレクトリにコピーされたファイルです。ファイルを直接開いてVSCode等で編集したい時はプロジェクトディレクトリの中のソースコードを編集してください。カーネルのソースコードの場合は~/mkubos_swan/mkubos_swan/swan_kernels/srcディレクトリ、ホストのソースコードの場合は~/mkubos_swan/mkubos_swan/swan/srcディレクトリのファイルを編集してください。
Vitisのエディタで編集する場合は、Vitisの画面の左上のExplorerで編集したいソースコードをダブルクリックするとファイルが開けます。
個人的にはVSCodeの方が文字の色が変わって見やすかったり、コピーアンドペーストがやりやすかったりしますが、Vitisのエディタを使うとそこだけで完結できるので好みで使ってください。

0. 実行時間測定方法について

公開されているSwanのホストコードでは、実行時間の計測にclock()関数を使用していますが、clock()関数の場合、Kenelが実行している時間を正しく計測できていないです。
そこで、実行時間測定の精度をあげるため、std::chrono::high_resolution_clock 関数を使用して実行時間を測定するよう変更してください。

main.cppのファイルの１番上に以下の行を追加します。

#include <chrono>

330行目付近の以下の行を修正します。

 clock_t start_clk = clock();

コメントアウトか削除して、以下の行を追加します。

  auto start_time = std::chrono::high_resolution_clock::now();

ちなみに、コメントアウトは、文頭に//をつけるか、/*と*/で文字を挟めばできます。

387行目付近の以下の行を修正します。

  // 7. Print the time and speed.
  clock_t end_clk = clock();
  double decode_time = (double)(end_clk - start_clk) / CLOCKS_PER_SEC;
  std::cout << "Time : " << decode_time << "[s]" << std::endl
            << "Speed: " << args.max_seq / decode_time << "[tok/s]"
            << std::endl;

記述の場所も// 8. Flush OpenCL Device Memoryの前に書かれているので、400行目より下、#endif // USE_CPU_ONLYの次の行以降、return 0;の前に以下の記述を追加します。

  // 7. Print the time and speed.
  auto end_time = std::chrono::high_resolution_clock::now();
  double decode_time = std::chrono::duration<double, std::milli>(end_time - start_time).count();

  std::cout << "Total : " << decode_time / 1000.0 << "[s]" << std::endl
            << "Speed: " << args.max_seq / (decode_time/1000.0) << "[tok/s]"
            << std::endl;

1. Profiling

まず、各マクロを実行する時にどのくらいの時間が掛かっているのかを調べてみましょう。
MKUBOSのswanを実行する場所~に、xrt.iniというファイルを新規作成します。
xrt.iniの内容は以下のようにしてください。

[Debug]
profile=true
timeline_trace=true
data_transfer_trace=fine
opencl_trace=true
stall_trace=all

VSCodeでMKUBOSに接続できないので、自分のPCで作ってから以下のコマンドでMKUBOSに送ってください。

scp xrt.ini [email protected]:~

swanを実行すると、Profile機能をONにしてswanが実行されます。
Profile機能がONになっていると、実行時間が長くなる傾向があります。
xrt.run_summaryとsummary.csvとopencl_trace.csvというファイルがMKUBOSのswanの実行場所にできます。

できた３つのファイルを自PCに移動します。

ubuntu小

scp [email protected]:~/xrt.run_summary .
scp [email protected]:~/summary.csv .
scp [email protected]:~/opencl_trace.csv .

xrt.run_summarとsummary.csvを置いた場所で以下のコマンドを実行します。

ubuntu小

vitis_analyzer xrt.run_summary &

Summaryを確認しましょう。

Timeline Traceを確認しましょう。

上の図の赤枠のアイコンで見たい波形を見ることができます。
例えば、虫眼鏡に＋のアイコンを何回かクリックすると以下のように波形が拡大表示されます。拡大されるときの中心はマーカー（上の図の赤い矢印）です。マーカーはドラッグして動かせます。

2. 並列化

ProfileのSummaryを確認すると、matmul (行列・ベクトル乗算)に最も時間がかかっていることがわかります。
言語モデルを高位合成でFPGAに実装してみたにあるswanのブロック図を以下に示します。

これを見ると、以下の２か所が並列化できそうです。

QKV Attention機構のmatmul計算部分
SwiGLUの前段のmatmul計算部分

1の計算部分を並列化するやり方を以下に示します。
2の方をやってみたい場合は、1の並列化を参考に考えてみてください。

2.1 ハードウェアファンクションの複製

FPGAに実装されているmatmulファンクションを３つに増やします。
Vitisのswan_system_hw_linkタブを開き、kernel_matmulのCompute Unitsの数を3にします。

また、PS DDR4メモリのアクセスを分散させるために、connectivity.cfgファイルをコンパイルオプションに追加する必要があります。
connectivity.cfgというファイル名のファイルを新規作成し、内容は以下のようにしてください。

[connectivity]
sp=kernel_matmul_1.i_vec:HP0
sp=kernel_matmul_1.i_mat:HP0
sp=kernel_matmul_1.o_vec:HP0

sp=kernel_matmul_2.i_vec:HP1
sp=kernel_matmul_2.i_mat:HP1
sp=kernel_matmul_2.o_vec:HP1

sp=kernel_matmul_3.i_vec:HP2
sp=kernel_matmul_3.i_mat:HP2
sp=kernel_matmul_3.o_vec:HP2

connectivity.cfgファイルはmkubos_swan/swan_system_hw_linkディレクトリに置きます。swan_system_hw_linkを右クリックし、C/C++ Build Settingを開き、SettingでV++ Dernel LinkerのMiscellaneousをクリックします。右側にあるOther flagsの横の緑の＋のボタン（Addボタンです。下の画像を参照してください）を押すと、ウィンドウが出てくるので、空欄に以下のように記入してコンパイルオプションに追加してください。

--config ../connectivity.cfg

OKボタンを押してウィンドウを閉じます。Settingウィンドウの下の方にあるApply and Closeボタンを押してオプションを反映させます。

[swan_system_hw_link]ディレクトリを右クリックして、[Clean Project]を選択して、途中ファイルを削除してください。
MKUBOS上で言語モデルを動かすの「6.ビットストリームをビルドする」でやったように、左上の[Explorer]の[swan_system_hw_link/swan_system_hw_link.prj]をクリックしてから[Build] ボタン（金槌アイコン）をクリックしてビットストリームのビルドを行ってください。

2.2 ホストコード変更

ホストコードでは、Kernelを並列実行するために、QueueおよびBufferを複数用意します。

①tensor_fpga

ホストコードのtensor_fpga.cppを開きます。
matmulを並列で３つ動作させる関数を追加します。MatmulFPGAは並列化しない箇所で使うのでコメントアウトしないでください。
以下の関数をtensor_fpga.cppに追加します。MatmulFPGAの直後がいいかと思います。

// Matmul parallel 3
// Compute the matrix multiplication of two input tensors.
// Tensor1d [dim] . Tensor2dFFNB [dim, ffn_dim] = Tensor1dFFNB [ffn_dim]
// out[i] = w[i,j] . in[j]
// Compute the matrix multiplication of two input tensors.
// Tensor1d [dim] . Tensor2dAttn [dim, dim] = Tensor1d [dim]
// out[i] = w[i,j] . in[j]
void MatmulParaFPGA(Tensor1d& out1, const Tensor1d& in1, const Tensor2dAttn& w1,
				Tensor1d& out2, /*const Tensor1d& in2,*/ const Tensor2dAttn& w2,
				Tensor1d& out3, /*const Tensor1d& in3,*/ const Tensor2dAttn& w3,
                cl::CommandQueue& q1, cl::Kernel kernel_matmul_1, float* ptr_a1, float* ptr_b1, float* ptr_result1, cl::Buffer buffer_a1, cl::Buffer buffer_b1, cl::Buffer buffer_result1,
                cl::CommandQueue& q2, cl::Kernel kernel_matmul_2, float* ptr_a2, float* ptr_b2, float* ptr_result2, cl::Buffer buffer_a2, cl::Buffer buffer_b2, cl::Buffer buffer_result2,
                cl::CommandQueue& q3, cl::Kernel kernel_matmul_3, float* ptr_a3, float* ptr_b3, float* ptr_result3, cl::Buffer buffer_a3, cl::Buffer buffer_b3, cl::Buffer buffer_result3
                ) {
  for (int i = 0; i < kDim; i++) {
    ptr_a1[i] = in1[i];
    ptr_a2[i] = in1[i];
    ptr_a3[i] = in1[i];
  }
  for (int i = 0; i < kDim; i++) {
    for (int j = 0; j < kDim; j++) {
      ptr_b1[i * kDim + j] = w1[i][j];
      ptr_b2[i * kDim + j] = w2[i][j];
      ptr_b3[i * kDim + j] = w3[i][j];
    }
  }
  q1.enqueueMigrateMemObjects({buffer_a1, buffer_b1}, 0);
  q2.enqueueMigrateMemObjects({buffer_a2, buffer_b2}, 0);
  q3.enqueueMigrateMemObjects({buffer_a3, buffer_b3}, 0);
  kernel_matmul_1.setArg(3, kDim);
  kernel_matmul_1.setArg(4, kDim);
  kernel_matmul_2.setArg(3, kDim);
  kernel_matmul_2.setArg(4, kDim);
  kernel_matmul_3.setArg(3, kDim);
  kernel_matmul_3.setArg(4, kDim);

  cl::Event event1, event2, event3;
  q1.enqueueTask(kernel_matmul_1, nullptr, &event1);
  q2.enqueueTask(kernel_matmul_2, nullptr, &event2);
  q3.enqueueTask(kernel_matmul_3, nullptr, &event3);
  q1.enqueueMigrateMemObjects({buffer_result1}, CL_MIGRATE_MEM_OBJECT_HOST);
  q2.enqueueMigrateMemObjects({buffer_result2}, CL_MIGRATE_MEM_OBJECT_HOST);
  q3.enqueueMigrateMemObjects({buffer_result3}, CL_MIGRATE_MEM_OBJECT_HOST);
  event1.wait();
  event2.wait();
  event3.wait();
  for (int i = 0; i < kDim; i++) {
    out1[i] = ptr_result1[i];
    out2[i] = ptr_result2[i];
    out3[i] = ptr_result3[i];
  }
}

tensor_fpga.cppを保存して閉じたら、tensor_fpga.hppを開きます。
ヘッダファイルtensor_fpga.hppにはtensor_fpga.cppにある関数の宣言が書いてあります。MatmulParaFPGAの関数の宣言もtensor_fpga.hppに追加します。
以下の宣言を追加してください。これも元の関数と同じ場所がいいので、MatmulFPGAの直後に追加するのがいいと思います。

void MatmulParaFPGA(Tensor1d& out1, const Tensor1d& in1, const Tensor2dAttn& w1,
				Tensor1d& out2, /*const Tensor1d& in2,*/ const Tensor2dAttn& w2,
				Tensor1d& out3, /*const Tensor1d& in3,*/ const Tensor2dAttn& w3,
                cl::CommandQueue& q1, cl::Kernel kernel_matmul_1, float* ptr_a1, float* ptr_b1, float* ptr_result1, cl::Buffer buffer_a1, cl::Buffer buffer_b1, cl::Buffer buffer_result1,
                cl::CommandQueue& q2, cl::Kernel kernel_matmul_2, float* ptr_a2, float* ptr_b2, float* ptr_result2, cl::Buffer buffer_a2, cl::Buffer buffer_b2, cl::Buffer buffer_result2,
                cl::CommandQueue& q3, cl::Kernel kernel_matmul_3, float* ptr_a3, float* ptr_b3, float* ptr_result3, cl::Buffer buffer_a3, cl::Buffer buffer_b3, cl::Buffer buffer_result3);

tensor_fpga.hppを保存して閉じます。

② decode

decode.cppを開きます。
引数を修正します。
cl::CommandQueue q,に追加して、３つのQueueを追加します。

cl::CommandQueue q1, cl::CommandQueue q2, cl::CommandQueue q3,

cl::Kernel kernel_matmul,は元のkernel_matmulは削除して、３つに増やします。

cl::Kernel kernel_matmul_1, cl::Kernel kernel_matmul_2, cl::Kernel kernel_matmul_3,

matmul用のバッファとその結果用のバッファを用意します。

float* ptr_a_1, float* ptr_b_1, float* ptr_result_1,
float* ptr_a_2, float* ptr_b_2, float* ptr_result_2,
float* ptr_a_3, float* ptr_b_3, float* ptr_result_3,

cl::Buffer buffer_a_1, cl::Buffer buffer_b_1, cl::Buffer buffer_result_1,
cl::Buffer buffer_a_2, cl::Buffer buffer_b_2, cl::Buffer buffer_result_2,
cl::Buffer buffer_a_3, cl::Buffer buffer_b_3, cl::Buffer buffer_result_3

引数の最後には,はつけてはいけません。なので、上のbuffer_result_3の後には,がないです。代わりに、buffer_result2の後に,をつけてください。
ヘッダファイルdecode.hppやmain.cppのdecode関数の記述箇所と引数の順序が違うと、FPGA上で動かした時にSegmentatiom faultという表示が出るので気を付けてください。

以上で引数の修正が完了しました。

matmulの関数呼び出しを変更します。58行目付近のMatmulFPGAを３回呼び出しているところを削除（コメントアウト）して、matmulを並列で行う関数MatmulParaFPGAを追加します。

	MatmulParaFPGA(ctx.attn_wqx[i_layer], ctx.attn_norm[i_layer], w.attn_wq[i_layer],
					ctx.attn_wkx[i_layer], /*ctx.attn_norm[i_layer],*/ w.attn_wk[i_layer],
					ctx.attn_wvx[i_layer], /*ctx.attn_norm[i_layer],*/ w.attn_wv[i_layer],
					q1, kernel_matmul_1, ptr_a_1, ptr_b_1, ptr_result_1,buffer_a_1, buffer_b_1, buffer_result_1,
					q2, kernel_matmul_2, ptr_a_2, ptr_b_2, ptr_result_2,buffer_a_2, buffer_b_2, buffer_result_2,
					q3, kernel_matmul_3, ptr_a_3, ptr_b_3, ptr_result_3,buffer_a_3, buffer_b_3, buffer_result_3);

上記の並列化したい箇所以外のmatmulFPGAを呼び出しているところで、kernel_matmulではなくkernel_matmul_1を使うように変更します。
例えば、以下のようにMatmulFPGAを呼び出しているところがあります。

    MatmulFPGA(ctx.attn_out[i_layer], ctx.attn_val[i_layer], w.attn_wo[i_layer],

４か所あります。以下のように引数を変更してください。

               q1, kernel_matmul_1, ptr_a_1, ptr_b_1, ptr_result_1, buffer_a_1, buffer_b_1,
               buffer_result_1);

decode.cppを保存して閉じます。

decode.hppを開きます。ここに関数Decodeの宣言があるので、decode.cppと合わせます。

void Decode(int tok, int pos, const Tensor1d& ctx_input,
            Tensor3dCache& ctx_k_cache, Tensor3dCache& ctx_v_cache,
            Tensor1d& ctx_final_norm, const Weights& w
#ifndef USE_CPU_ONLY
            ,
			cl::CommandQueue q,
			cl::CommandQueue q1, cl::CommandQueue q2, cl::CommandQueue q3,
			cl::Kernel kernel_matmul_1, cl::Kernel kernel_matmul_2, cl::Kernel kernel_matmul_3,
			cl::Kernel kernel_mul,
			cl::Kernel kernel_rmsnorm, cl::Kernel kernel_softmax,
			cl::Kernel kernel_add, cl::Kernel kernel_rope,
			float* ptr_a, float* ptr_b, float* ptr_c, float* ptr_d, float* ptr_result, float* ptr_result2,
			float* ptr_a_1, float* ptr_b_1, float* ptr_result_1,
			float* ptr_a_2, float* ptr_b_2, float* ptr_result_2,
			float* ptr_a_3, float* ptr_b_3, float* ptr_result_3,
			cl::Buffer buffer_a, cl::Buffer buffer_b,
			cl::Buffer buffer_c, cl::Buffer buffer_d, cl::Buffer buffer_result,
			cl::Buffer buffer_result2,
			cl::Buffer buffer_a_1, cl::Buffer buffer_b_1, cl::Buffer buffer_result_1,
			cl::Buffer buffer_a_2, cl::Buffer buffer_b_2, cl::Buffer buffer_result_2,
			cl::Buffer buffer_a_3, cl::Buffer buffer_b_3, cl::Buffer buffer_result_3
#endif // USE_CPU_ONLY
);

decode.hppを保存して閉じます。

③ main

main.cppを開きます。
161行目付近のcl::Kernel kernel_matmul;を削除（コメントアウト）して、代わりに３つに増やします。また、キューを追加します。

  cl::CommandQueue q1;
  cl::CommandQueue q2;
  cl::CommandQueue q3;
  cl::Kernel kernel_matmul_1;
  cl::Kernel kernel_matmul_2;
  cl::Kernel kernel_matmul_3;

216行目付近にOpenCLのコマンドキューを生成している記述があります。これをq1、q2、q3の分、追加します。
q1の例を示します。

    OCL_CHECK(err, q1 = cl::CommandQueue(context, device,
                                            CL_QUEUE_PROFILING_ENABLE, &err));

226行目付近にkernel_matmulを生成している記述があります。これを削除して、kernel_matmul_1、kernel_matmul_2、kernel_matmul_3を追加します。
kernel_matmul_1の例を示します。

      OCL_CHECK(err,
                kernel_matmul_1 = cl::Kernel(program, "kernel_matmul:{kernel_matmul_1}", &err));
      std::cout << "load : kernel_matmul_1" << std::endl;

251行目付近にバッファを生成している記述があります。Matmul用のバッファを追加で生成します。
kernel_matmul_1用のバッファの記述例を示します。kernel_matmul_2、kernel_matmul_3用の記述も作って追加してください。

  // Matmul Buffer 1
  OCL_CHECK(err, cl::Buffer buffer_a_1(context, CL_MEM_READ_ONLY, size_in_bytes,
                                       NULL, &err));
  OCL_CHECK(err, cl::Buffer buffer_b_1(context, CL_MEM_READ_ONLY,
                                  size_in_bytes * size_in_bytes, NULL, &err));
  OCL_CHECK(err, cl::Buffer buffer_result_1(context, CL_MEM_WRITE_ONLY,
                                            size_in_bytes, NULL, &err));

270行目付近にkernel_matmulの引数を設定している記述があるので、これを削除して代わりにkernel_matmul_1、kernel_matmul_2、kernel_matmul_3の引数を割り当てます。
kernel_matmul_1の記述例を示します。kernel_matmul_2、kernel_matmul_3用の記述も作って追加してください。

	OCL_CHECK(err, err = kernel_matmul_1.setArg(0, buffer_a_1));
	OCL_CHECK(err, err = kernel_matmul_1.setArg(1, buffer_b_1));
	OCL_CHECK(err, err = kernel_matmul_1.setArg(2, buffer_result_1));
	OCL_CHECK(err, err = kernel_matmul_1.setArg(3, swan::kDim));
	OCL_CHECK(err, err = kernel_matmul_1.setArg(4, swan::kDim));

296行目付近で、バッファのポインタを宣言しています。追加したバッファのポインタの宣言を追加します。
以下の記述を追加してください。

  float* ptr_a_1;
  float* ptr_b_1;
  float* ptr_result_1;
  float* ptr_a_2;
  float* ptr_b_2;
  float* ptr_result_2;
  float* ptr_a_3;
  float* ptr_b_3;
  float* ptr_result_3;

バッファをマップしてポインタを格納します。CPUがFPGAのメモリを読み書きできるようになります。
既に記述されている他のバッファを参考に、q1、q2、q3についてマップします。
例えば、q1についてはptr_a_1、ptr_b_1、ptr_result_1を設定します。
q1についての例を示します。

  OCL_CHECK(err, ptr_a_1 = (float*)q1.enqueueMapBuffer(
                     buffer_a_1, CL_TRUE, CL_MAP_WRITE, 0, size_in_bytes, NULL,
                     NULL, &err));
  OCL_CHECK(err, ptr_b_1 = (float*)q1.enqueueMapBuffer(
                     buffer_b_1, CL_TRUE, CL_MAP_WRITE, 0,
                     size_in_bytes * size_in_bytes, NULL, NULL, &err));
  OCL_CHECK(err, ptr_result_1 = (float*)q1.enqueueMapBuffer(
                     buffer_result_1, CL_TRUE, CL_MAP_READ, 0, size_in_bytes,
                     NULL, NULL, &err));

339行目付近にDecode関数があるので、引数を書き換えます。

    swan::Decode(token, pos, ctx_input, ctx_k_cache, ctx_v_cache,
                 ctx_final_norm, weights
#ifndef USE_CPU_ONLY
                ,
                q,
		q1, q2, q3,
		kernel_matmul_1, kernel_matmul_2, kernel_matmul_3,
		kernel_mul,
		kernel_rmsnorm, kernel_softmax,
                kernel_add, kernel_rope,
		ptr_a, ptr_b, ptr_c, ptr_d, ptr_result, 
                ptr_result2,
		ptr_a_1, ptr_b_1, ptr_result_1,
		ptr_a_2, ptr_b_2, ptr_result_2,
		ptr_a_3, ptr_b_3, ptr_result_3,
		buffer_a, buffer_b,
                buffer_c, buffer_d, buffer_result,
		buffer_result2,
		buffer_a_1, buffer_b_1, buffer_result_1,
		buffer_a_2, buffer_b_2, buffer_result_2,
		buffer_a_3, buffer_b_3, buffer_result_3
#endif // USE_CPU_ONLY
    );

396行目付近のメモリをアンマップしているところも更新します。以下のような記述があります。

  // 8. Flush OpenCL Device Memory
  OCL_CHECK(err, err = q.enqueueUnmapMemObject(buffer_a, ptr_a));
  OCL_CHECK(err, err = q.enqueueUnmapMemObject(buffer_b, ptr_b));
  OCL_CHECK(err, err = q.enqueueUnmapMemObject(buffer_result, ptr_result));
  OCL_CHECK(err, err = q.finish());

ここに、以下の記述を追加します。

  OCL_CHECK(err, err = q1.enqueueUnmapMemObject(buffer_a_1, ptr_a_1));
  OCL_CHECK(err, err = q1.enqueueUnmapMemObject(buffer_b_1, ptr_b_1));
  OCL_CHECK(err, err = q1.enqueueUnmapMemObject(buffer_result_1, ptr_result_1));

  OCL_CHECK(err, err = q2.enqueueUnmapMemObject(buffer_a_2, ptr_a_2));
  OCL_CHECK(err, err = q2.enqueueUnmapMemObject(buffer_b_2, ptr_b_2));
  OCL_CHECK(err, err = q2.enqueueUnmapMemObject(buffer_result_2, ptr_result_2));

  OCL_CHECK(err, err = q3.enqueueUnmapMemObject(buffer_a_3, ptr_a_3));
  OCL_CHECK(err, err = q3.enqueueUnmapMemObject(buffer_b_3, ptr_b_3));
  OCL_CHECK(err, err = q3.enqueueUnmapMemObject(buffer_result_3, ptr_result_3));

  OCL_CHECK(err, err = q1.finish());
  OCL_CHECK(err, err = q2.finish());
  OCL_CHECK(err, err = q3.finish());

main.cppを保存すれば、ソースコードの変更は完了です。

MKUBOS上で言語モデルを動かすの「7. ホストプログラムをビルドする」と同様に、左上のExplorerのswan/swan.prjをダブルクリックした後、[Build] ボタン（金槌アイコン）をクリックしてホストプログラムをビルドしてください。

2.3 実行

ビルドが完了したら、更新したビットストリームと実行ファイルをMKUBOSに送ります。

mkubos_swan/swan/Hardware/swan
mkubos_swan/swan_system_hw_link/Hardware/binary_container_1.xclbin をbinary_container_1.bin にリネームしたもの

swan.dtboとmodelディレクトリは変更の必要はありません。
ハードウェアを生成後、swan on FPGAを実行し、速度を比較してみてください。

3. matmul分割

さらなる効率化のために、以下のようなPartitioningという考え方を使って、行列ｘベクトル乗算内部の並列化を行い、高速化を目指します。
https://github.com/Xilinx/xup_high_level_synthesis_design_flow/blob/main/source/matmult/notebook/matmul_part2.ipynb

以下にPartitioningの簡単な概念図を示します。Matrixを分割し、ベクトルを複製することで、内部の累積加算動作の並列化を行います。

本件では、FPGA内部のリソースの限界から、以下の２つの行列ｘベクトル乗算のPartitioningを行いました。

(288,288) matrix × (288,) vector
(768,288) matrix × (288,) vector
※Swan内部には、(288,768) matrix × (768,) vectorの乗算もありますが、FPGA内部リソースの限界のため、この演算はCPUに実行させます。

3.1 コンパイルオプションの削除

並列化で追加したハードウェアファンクションの複製を削除します。
swan_system_hw_linkを右クリックし、C/C++ Build Settingを開き、SettingでV++ Dernel LinkerのMiscellaneousをクリックします。右側にある--config ../connectivity.cfgを選択して先程のAdd...ボタンの隣に赤いバツのDeleteボタンがあります。これを押して設定を削除します。Apply and Closeを押して設定を反映させます。

3.2 kernel_matmul_pt_288x

新しいmatmulのソースコードを作ります。
kernel_matmul_pt_288x.cppというファイルを新規作成します。
kernel_matmul_pt_288x.cppというファイルをswan_kernels/srcに作ると、自動的にProjectにインポートされます。または、Projectとは無関係のディレクトリでソースコードを作成してインポートします。どちらでもいいですが、直接swan_kernels/srcに作るとVitis上でファイルを削除すると元のファイルも削除されるので気を付けてください。
kernel_matmul_pt_288x.cppは内容を以下のようにしてください。

#ifndef USE_CPU_ONLY

#include <hls_stream.h>
#include <stdint.h>
#include <hls_vector.h>

#define BLK_ROW 96
#define COLSIZE 288
#define MAX_ROW_SIZE 768
#define MAX_COL_SIZE 288
#define MAX_BLK_NUM  MAX_ROW_SIZE/BLK_ROW	// 8
#define MAX_BLK_NUM_HALF MAX_BLK_NUM/2	// 4
#define MAX_BLK_NUM_URAM 1	// 1
#define MAX_BLK_NUM_BRAM (MAX_BLK_NUM-MAX_BLK_NUM_URAM)	// 8-1=6

extern "C" {
void kernel_matmul_pt_288x(float* i_vec, float* i_mat, float* o_vec, int row) {
#pragma HLS INTERFACE m_axi port = i_vec bundle = gmem0 // max_read_burst_length=256 max_write_burst_length=256 depth=1024
#pragma HLS INTERFACE m_axi port = i_mat bundle = gmem1 // max_read_burst_length=256 max_write_burst_length=256 depth=1024*1024
#pragma HLS INTERFACE m_axi port = o_vec bundle = gmem1 // max_read_burst_length=256 max_write_burst_length=256 depth=1024
#pragma HLS INTERFACE s_axilite port=row  bundle=CTRL
#pragma HLS INTERFACE s_axilite port=return bundle=CTRL

	int blk_num = row / BLK_ROW;	// 3 or 8
	// Buffers
	float vec_buf[MAX_BLK_NUM][MAX_COL_SIZE];
	float split_buf[MAX_BLK_NUM][BLK_ROW][MAX_COL_SIZE];
	float out_buf[BLK_ROW][MAX_BLK_NUM];
//
#pragma HLS BIND_STORAGE variable=vec_buf type=RAM_T2P impl=uram
#pragma HLS BIND_STORAGE variable=out_buf type=RAM_T2P impl=uram

	SWEEP_VEC:
	// Load Vector (replicate)
	for (int i=0; i<blk_num; i++){
//#pragma HLS DATAFLOW
		// ------ Read Vector ------
		LOAD_VEC:
		//for (int v_i=0; v_i<col; v_i++){
		for (int v_i=0; v_i<COLSIZE; v_i++){
#pragma HLS PIPELINE
			vec_buf[i][v_i] = i_vec[v_i];
		}
		// Load Matrix (Split)
		// ------ Read Matrix ------
		READ_MAT:
		for (int mr_i=0; mr_i<BLK_ROW; mr_i++){
			for (int m_i=0; m_i<COLSIZE; m_i++){
#pragma HLS PIPELINE
				split_buf[i][mr_i][m_i] = i_mat[(i*BLK_ROW + mr_i)*COLSIZE + m_i];
			}
		}
		// Compute multiply
		// ------ Compute ------
		COMPUTE_MULT:
		for (int c_i=0; c_i<BLK_ROW; c_i++){
			float c_tmp = 0;
#pragma HLS PIPELINE
			COMPUTE_ELE:
			for (int c_j=0; c_j<COLSIZE; c_j++){
#pragma HLS UNROLL
				c_tmp += split_buf[i][c_i][c_j] * vec_buf[i][c_j];
			}
			out_buf[c_i][i] = c_tmp;
		}
		// Output
		// ------ Write result ------
		WRITE_RESULT:
		for (int w_i=0; w_i<BLK_ROW; w_i++){
#pragma HLS PIPELINE
			int new_col = i*BLK_ROW + w_i;
			o_vec[new_col] = out_buf[w_i][i];
		}
	}
}
}

#endif // USE_CPU_ONLY

kernel_matmul_pt_288x.cppを保存します。

FPGAはkernel_matmul_pt_288xで容量がいっぱいになってしまうので、カーネルをkernel_matmul_pt_288xだけにします。他の機能はCPUで行います。
[swan_kernels]タブで、現在追加されているカーネルをクリックして選択します。Shiftキーで複数選択できます。バツのアイコン[Remove]を押してカーネルを削除します。

カーネルが消えたら、丸いアイコンの[Add Hardware Functions]を押して、kernel_matmul_pt_288xを追加します。[Hardware Functions]のところに出てきたら、[swan_system_hw_link]タブに移動して、kernel_matmul_pt_288xの[Compute Unit]が1になっていることを確認します。万が一違っていたら1にしてください。
[swan_system_hw_link]ディレクトリを右クリックして、[Clean Project]を選択して、途中ファイルを削除します。その後、swan_system_hw_link.prjをビルドしてください。

3.3 ホストコード変更

① tensor_fpga

tensor_fpga.cppを開き、2. で追加した関数MatmulParaFPGAの次、232行目付近に、以下のような関数MatmulPt288x288とMatmulPt288x768を追加します。

// Compute the matrix multiplication of two input tensors.
// Tensor1d [dim] . Tensor2dAttn [dim, dim] = Tensor1d [dim]
// out[i] = w[i,j] . in[j]
void MatmulPt288x288(Tensor1d& out, const Tensor1d& in, const Tensor2dAttn& w,
                cl::CommandQueue q, cl::Kernel kernel_matmul_pt_288x, float* ptr_a,
                float* ptr_b, float* ptr_result, cl::Buffer buffer_a,
                cl::Buffer buffer_b, cl::Buffer buffer_result) {
  for (int i = 0; i < kDim; i++) {
    ptr_a[i] = in[i];
  }
  for (int i = 0; i < kDim; i++) {
    for (int j = 0; j < kDim; j++) {
      ptr_b[i * kDim + j] = w[i][j];
    }
  }
  q.enqueueMigrateMemObjects({buffer_a, buffer_b}, 0);
  kernel_matmul_pt_288x.setArg(3, kDim);
  // kernel_matmul.setArg(4, kDim);
  q.enqueueTask(kernel_matmul_pt_288x);
  q.enqueueMigrateMemObjects({buffer_result}, CL_MIGRATE_MEM_OBJECT_HOST);
  q.finish();
  for (int i = 0; i < kDim; i++) {
    out[i] = ptr_result[i];
  }
}
// Compute the matrix multiplication of two input tensors.
// Tensor1dFFNB [ffn_dim] . Tensor2dFFNA [ffn_dim, dim] = Tensor1dFFNB [ffn_dim]
// out[i] = w[i,j] . in[j]
void MatmulPt288x768(Tensor1dFFNB& out, const Tensor1d& in, const Tensor2dFFNA& w,
                cl::CommandQueue q, cl::Kernel kernel_matmul_pt_288x, float* ptr_a,
                float* ptr_b, float* ptr_result, cl::Buffer buffer_a,
                cl::Buffer buffer_b, cl::Buffer buffer_result) {
  for (int i = 0; i < kDim; i++) {
    ptr_a[i] = in[i];
  }
  for (int i = 0; i < kFFNDim; i++) {
    for (int j = 0; j < kDim; j++) {
      ptr_b[i * kDim + j] = w[i][j];
    }
  }
  q.enqueueMigrateMemObjects({buffer_a, buffer_b}, 0);
  kernel_matmul_pt_288x.setArg(3, kFFNDim);
  // kernel_matmul.setArg(4, kFFNDim);
  q.enqueueTask(kernel_matmul_pt_288x);
  q.enqueueMigrateMemObjects({buffer_result}, CL_MIGRATE_MEM_OBJECT_HOST);
  q.finish();
  for (int i = 0; i < kFFNDim; i++) {
    out[i] = ptr_result[i];
  }
}

tensor_fpga.hppにMatmulPt288x288とMatmulPt288x768を追加します。

// Added
void MatmulPt288x288(Tensor1d& out, const Tensor1d& in, const Tensor2dAttn& w,
                cl::CommandQueue q, cl::Kernel kernel_matmul_pt_288x, float* ptr_a,
                float* ptr_b, float* ptr_result, cl::Buffer buffer_a,
                cl::Buffer buffer_b, cl::Buffer buffer_result);
void MatmulPt288x768(Tensor1dFFNB& out, const Tensor1d& in, const Tensor2dFFNA& w,
                cl::CommandQueue q, cl::Kernel kernel_matmul_pt_288x, float* ptr_a,
                float* ptr_b, float* ptr_result, cl::Buffer buffer_a,
                cl::Buffer buffer_b, cl::Buffer buffer_result);

② decode

decode.cppを開きます。
まずは、関数Decodeの引数を変更します。
#ifndef USE_CPU_ONLYから#endif // USE_CPU_ONLYまでの間の引数は、以下のものだけ残してコメントアウトか削除してください。

cl::CommandQueue q
float* ptr_a
float* ptr_b
float* ptr_result
cl::Buffer buffer_a
cl::Buffer buffer_b
cl::Buffer buffer_result

そして、kernel_matmul_pt_288xを追加します。

cl::Kernel kernel_matmul_pt_288x,

追加場所は#ifndef USE_CPU_ONLYから#endif // USE_CPU_ONLYまでの間ならどこでもいいですが、kernelを宣言していた辺り（qの下、kernel_~をコメントアウトしている場合はその直後）に追加するのがいいです。

79行目付近にあるMatmulParaFPGAをコメントアウトします。
代わりに、MatmulPt288x288を３つ追加してください。
MatmulPt288x288の最初の３つの引数Tensor1d& out、const Tensor1d& in、const Tensor2dAttn& wは、対応するMatmulFPGAやCPUのMatmulの引数を参考に記述してみてください。
例として、1つめのMatmulPt288x288の引数は以下のようになります。

	MatmulPt288x288(ctx.attn_wqx[i_layer], ctx.attn_norm[i_layer],
               w.attn_wq[i_layer], q, kernel_matmul_pt_288x, ptr_a, ptr_b, ptr_result,
               buffer_a, buffer_b, buffer_result);

144行目のMatmulFPGAもコメントアウトしてMatmulPt288x288を代わりに追加します。

    MatmulPt288x288(ctx.attn_out[i_layer], ctx.attn_val[i_layer], w.attn_wo[i_layer],
               q, kernel_matmul_pt_288x, ptr_a, ptr_b, ptr_result, buffer_a, buffer_b,
               buffer_result);

173行目と182行目のMatmulFPGAはコメントアウトしたらMatmulPt288x768を代わりに追加します。
173行目の方のMatmulPt288x768は以下のようにしてください。

    MatmulPt288x768(ctx.ffn_w1x[i_layer], ctx.ffn_norm[i_layer], w.ffn_w1[i_layer],
               q, kernel_matmul_pt_288x, ptr_a, ptr_b, ptr_result, buffer_a, buffer_b,
               buffer_result);

182行目の方は自分で考えてみてください。

203行目のMatmulFPGAはFPGA内部リソース不足のため、CPUでやります。
コメントアウトして、Matmulを代わりに追加してください。
#elseの下にMatmulという関数を呼び出しています。これは一番最初に試したswanをCPUのみで動かすときに実行される関数です。
今はFPGAで動作させるので#ifndefの方が呼ばれています。ここに#elseと同じMatmulを書けば、MatmulはCPUで実行されます。

これでMatmulFPGAの変更は終わりましたが、内部リソースが足りないため、他の計算はCPUでやります。
以下の関数を203行目のMatmulと同様に変更してください。

61行目付近のRMSNormFPGA
94行目付近のRoPE
122行目付近のMul
130行目付近のSoftmax
153行目付近のAdd
164行目付近のRMSNorm
194行目付近のMul
212行目付近のAdd
222行目付近のRMSNorm

decode.hppのDecodeの引数も、decode.cppと一致するように変更してください。

③ main

main.cppを変更します。並列化後のmain.cppを開きます。

162行目付近の、並列化で追加したCommandQueueのq1、q2、q3の宣言をコメントアウトします。
その直後の追加したカーネルの宣言kernel_matmul_1、kernel_matmul_2、kernel_matmul_3もコメントアウトします。
CPUで行うkernel_mul、kernel_rmsnorm、kernel_softmax、kernel_add、kernel_ropeもコメントアウトします。
代わりに、kernel_matmul_pt_288xの宣言を追加します。

  cl::Kernel kernel_matmul_pt_288x;

225行目付近のCommandQueueの初期化もコメントアウトします。
コメントアウトの例です。

//    OCL_CHECK(err, q1 = cl::CommandQueue(context, device,
//                                            // CL_QUEUE_PROFILING_ENABLE, &err));
//                                            CL_QUEUE_PROFILING_ENABLE  /*| CL_QUEUE_OUT_OF_ORDER_EXEC_MODE_ENABLE*/, &err));

q1、q2、q3をコメントアウトします。

241行目付近のカーネルの初期化もコメントアウトします。
kernel_matmul_1、kernel_matmul_2、kernel_matmul_3とkernel_mul、kernel_rmsnorm、kernel_softmax、kernel_add、kernel_ropeをコメントアウトします。
代わりに、kernel_matmul_pt_288xの初期化を追加します。

      OCL_CHECK(err,
                kernel_matmul_pt_288x = cl::Kernel(program, "kernel_matmul_pt_288x", &err));
      std::cout << "load : kernel_matmul_pt_288x" << std::endl;

278行目付近のバッファの生成の不要なものはコメントアウトします。
buffer_a、buffer_b、buffer_result以外はコメントアウトします。

310行目付近のカーネルに引数を指定しているところは、今記述してあるものは全てコメントアウトして、kernel_matmul_pt_288xのものを追加します。

	OCL_CHECK(err, err = kernel_matmul_pt_288x.setArg(0, buffer_a));
	OCL_CHECK(err, err = kernel_matmul_pt_288x.setArg(1, buffer_b));
	OCL_CHECK(err, err = kernel_matmul_pt_288x.setArg(2, buffer_result));
	OCL_CHECK(err, err = kernel_matmul_pt_288x.setArg(3, swan::kDim));

359行目付近からのポインタの宣言と、375行目付近からの初期化も、ptr_a、ptr_b、ptr_result以外はコメントアウトします。

443行目付近にDecodeの引数があるので、不要なものを削除します。decode.cppで書いたDecodeに合わせて、#ifndef USE_CPU_ONLYから#endif // USE_CPU_ONLYまでの間の引数は、以下のものだけ残してコメントアウトか削除してください。

q
ptr_a
ptr_b
ptr_result
buffer_a
buffer_b
buffer_result

Decodeの引数と同じ位置にkernel_matmul_pt_288xの引数を追加してください。

kernel_matmul_pt_288x,

508行目付近のバッファ開放を、q以外コメントアウトします。
525行目付近、finishもq以外コメントアウトします。

main.cppを保存すれば、ソースコードの変更は完了です。

MKUBOS上で言語モデルを動かすの「7. ホストプログラムをビルドする」や並列化と同様に、左上のExplorerのswan/swan.prjをダブルクリックした後、[Build] ボタン（金槌アイコン）をクリックしてホストプログラムをビルドしてください。

3.4 実行

ビルドが完了したら、更新したビットストリームと実行ファイルをMKUBOSに送ります。並列化と同じです。

mkubos_swan/swan/Hardware/swan
mkubos_swan/swan_system_hw_link/Hardware/binary_container_1.xclbin をbinary_container_1.bin にリネームしたもの

swan.dtboとmodelディレクトリは変更の必要はありません。
ハードウェアを生成後、swan on FPGAを実行し、速度を比較してみてください。

swan改善案 - takgto/utokyo-chipathon2023 GitHub Wiki

0. 実行時間測定方法について

1. Profiling

2. 並列化

2.1 ハードウェアファンクションの複製

2.2 ホストコード変更

①tensor_fpga

② decode

③ main

2.3 実行

3. matmul分割

3.1 コンパイルオプションの削除

3.2 kernel_matmul_pt_288x

3.3 ホストコード変更

① tensor_fpga

② decode

③ main

3.4 実行

⚠️ **GitHub.com Fallback** ⚠️

⚠️ GitHub.com Fallback ⚠️