はじめに

Tenstorrent の TT Architecture and Metalium Guide なるドキュメントが tt-metal github にあります。

TT Architecture and Metalium Guide

今回は、このドキュメントのTile (Tensix)に関する部分を眺めることにします。

Grayskull chip and Tensix core

下記の図は、上記の頁にあるもので、説明のために引用します。Tile が2次元に並んでいます。Tileの詳細が右側にあります。 Tensix core と読んでいるようです。

Tensix coreには、下記のものから構成されています。

5 small RISC-V processors (aka "Baby RISCVs") that run C/C++ kernels and dispatch instructions to the compute and data movement engines
1 MB SRAM memory, (aka L1) a scratch pad accessible by all RISCVs and engines within the core
Matrix engine (aka FPU) that performs Matrix multiplication, elementwise, and dot product operations on small matrices (or tiles) of shape 32x32 and similar
Vector engine (aka SFPU) for vectorized kernels such as Top-k, Sort and special functions such as GELU, Exp, and Sqrt
Data Movement engine connected to 2 Networks on Chip (NoCs)

5つの小さな RISC-V コアがメインです。これを Baby RISCV と呼んでいるようです。1MBのSRAMメモリは、基本的には scratch pad として使用しています。Matrix engine (FPU) と Vector engin (SFPU) は、3つのRISC-Vが使います。残りの2つのRISC-Vが Data Movement engine になり、Tensix の外部との接続部分(NoC-0/NoC-1) とのデータ移動を行います。

Grayskull には、Tensix が 12 x 10 個、DRAM(LPDDR4 x 8)、PCIE、ARCが載っています。

Tensix の中の処理

下図も上記の頁から説明のために引用します。

NoC-0からのデータは、ROUTER-0を経由して、SRAM内の buffer_0 に書き込まれます。ROUTER 0の制御は、RISC-V 1 上で動く Data movement kernel です。RISC-V 2/RISC-V 3/RISC-V 4 が Cimpute kernel 内で Matrix & Vectors engines を使って、SRAM 上の buffer_0 のデータを処理し、buffer_1 に書き込みます。RISC-V 5 上で動作する Data movement kernel が ROUTER 1 経由で NoC 1 に buffer_1 のデータを移動します。

下記の図は、"Reader -> Compute -> Write pipeline" を説明している図です。これも上記の頁から説明のために引用します。

Reader Kernel は、RISC-V 1 、Compute Kernel は、RISC-V 2/RISC-V 3/RISC-V 4、Writer Kernel が RISC-V 5 が実行します。CB A 16 tiles/CB B 16 tiles が上記の buffer_0 になり、CB C 2 tiles が buffer_1 になります。

Reader Kernel へのデータ(Buffer A/Buffer B)は、Global DRAM (or SRAM) から送られてきます。Writer Kernel からのデータ（Buffer C)は、Global DATA (or SRAM)に送られていきます。

Compute Kernel の例

下記が Compute Kernel の例です。Compute Kernel は、void MAIN {} の中に記述します。

namespace NAMESPACE {
void MAIN {
  mm_init();
  acquire_dst(tt::DstMode::Tile);

  cb_wait_front(tt::CB::c_in0, /* number of tiles */ 1);
  cb_wait_front(tt::CB::c_in1, /* number of tiles */ 1);

  matmul_tiles(tt::CB::c_in0, tt::CB::c_in1, 0, 0, 0, false);

  cb_pop_front(tt::CB::c_in1, /* number of tiles */ 1);
  cb_pop_front(tt::CB::c_in0, /* number of tiles */ 1);

  cb_reserve_back(tt::CB::c_out0, /* number of tiles */ 1);
  pack_tile(0, tt::CB::c_out0);
  cb_push_back(tt::CB::c_out0, /* number of tiles */ 1);

  release_dst(tt::DstMode::Tile);
}

下記の部分が、Compute Kernel のSRAM上の入力データを待っているところです。

  cb_wait_front(tt::CB::c_in0, /* number of tiles */ 1);
  cb_wait_front(tt::CB::c_in1, /* number of tiles */ 1);

下記で Compute Kernel を実行します。

  matmul_tiles(tt::CB::c_in0, tt::CB::c_in1, 0, 0, 0, false);

処理が終わったので、入力データを取り出します。

  cb_pop_front(tt::CB::c_in1, /* number of tiles */ 1);
  cb_pop_front(tt::CB::c_in0, /* number of tiles */ 1);

出力データをSRAMに書き出します。

  cb_reserve_back(tt::CB::c_out0, /* number of tiles */ 1);
  pack_tile(0, tt::CB::c_out0);
  cb_push_back(tt::CB::c_out0, /* number of tiles */ 1);

Tensix の RISC-V の名称

Tensix 内の5つのRISC-Vの名称ですが、

RISC-V 1 : BRISC
RISC-V 2 : TRISC0
RISC-V 3 : TRISC1
RISC-V 4 : TRISC2
RISC-V 5 : NCRISC

のようです。

おわりに

TT-Metalium の中身は、まだまだいろいろなものがありますので、少しづつ、調べていきたいと思います

Vengineerの妄想

人生を妄想しています。

Tenstorrent : TTArchitecture and Metalium Guide