はじめに

Tenstorrent の TT-Metalium について、今年の2月に取り上げていて、

vengineer.hatenablog.com

次回は、TT-Metal のサンプルコードを見ていきたいと思います。

のままだったので、今日はサンプルコードを見てみます。

サンプルコード

TT-Metalium のサンプルコードは、

tt-metal/tt_metal/programming_examples at main · tenstorrent/tt-metal · GitHub

にあります。

今回はこの中の、

tt-metal/tt_metal/programming_examples/hello_world_compute_kernel at main · tenstorrent/tt-metal · GitHub

を見てみます。

ホスト側のプログラム : hello_world_compute_kernel.cpp
カーネルプログラム : kernels/compute/void_compute_kernel.cpp

があります。カーネルプログラムは、Tile (Tenisix) で動くプログラムです。

ホスト側のプログラム

ホスト側のプログラムは、OpenCL っぽい API を使っています。

#include "tt_metal/host_api.hpp"
#include "tt_metal/impl/device/device.hpp"

using namespace tt;
using namespace tt::tt_metal;

int main(int argc, char **argv) {

    // Initialize Program and Device

    constexpr CoreCoord core = {0, 0};
    int device_id = 0;
    Device *device = CreateDevice(device_id);
    CommandQueue& cq = device->command_queue();
    Program program = CreateProgram();

    // Configure and Create Void Kernel

    vector<uint32_t> compute_kernel_args = {};
    KernelHandle void_compute_kernel_id = CreateKernel(
        program,
        "tt_metal/programming_examples/hello_world_compute_kernel/kernels/compute/void_compute_kernel.cpp",
        core,
        ComputeConfig{
            .math_fidelity = MathFidelity::HiFi4,
            .fp32_dest_acc_en = false,
            .math_approx_mode = false,
            .compile_args = compute_kernel_args
        }
    );

    // Configure Program and Start Program Execution on Device

    SetRuntimeArgs(program, void_compute_kernel_id, core, {});
    EnqueueProgram(cq, program, false);
    printf("Hello, Core {0, 0} on Device 0, I am sending you a compute kernel. Standby awaiting communication.\n");

    // Wait Until Program Finishes, Print "Hello World!", and Close Device

    Finish(cq);
    printf("Thank you, Core {0, 0} on Device 0, for the completed task.\n");
    CloseDevice(device);

    return 0;

}

最初のおまじないが、

    constexpr CoreCoord core = {0, 0};
    int device_id = 0;
    Device *device = CreateDevice(device_id);
    CommandQueue& cq = device->command_queue();
    Program program = CreateProgram();

です。

CoreCoord
CreateDevice
CommandQueue
Program
CreateProgram

がおまじないとして必要なものです。

次は、カーネルプログラムをTileにロードする部分です。CreateKernel API を CreateProgram() によって生成した program にカーネルプログラムを紐付けます。カーネルプログラムのソースコードは、 "tt_metal/programming_examples/hello_world_compute_kernel/kernels/compute/void_compute_kernel.cpp"　で、Tileの番号 core( {0, 0} ) でカーネルプログラムを動かします。{0, 0} なので、1つのTile で動くことになります。

    vector<uint32_t> compute_kernel_args = {};
    KernelHandle void_compute_kernel_id = CreateKernel(
        program,
        "tt_metal/programming_examples/hello_world_compute_kernel/kernels/compute/void_compute_kernel.cpp",
        core,
        ComputeConfig{
            .math_fidelity = MathFidelity::HiFi4,
            .fp32_dest_acc_en = false,
            .math_approx_mode = false,
            .compile_args = compute_kernel_args
        }
    );

ComputeConfig が Tile (Tensix)の中の Compute コアの設定です。SetRuntimeArg API にて、progam にカーネルプログラム (void_compute_kernel_id) を core ({0, 0}) で動かすよと教えています。次の EnqueueProgram API でカーネルプログラムを実行します。3番目の引数が false になっているので、カーネルプログラムが終了するまで待ちます。

次は、カーネルプログラムの実行部分です。

    SetRuntimeArgs(program, void_compute_kernel_id, core, {});
    EnqueueProgram(cq, program, false);
    printf("Hello, Core {0, 0} on Device 0, I am sending you a compute kernel. Standby awaiting communication.\n");

EnqueueProgram API にてカーネルプログラムが終了したら、Finish と CloseDevice API を使って後処理をします。

    Finish(cq);
    printf("Thank you, Core {0, 0} on Device 0, for the completed task.\n");
    CloseDevice(device);

カーネルプログラム

カーネルプログラムは、void MAIN の中に書き込みます。このサンプルプログラムでは、DPRINT_MATHにて、”Hello, Master, I am running a void compute kernel." を表示しているだけです。このカーネルプログラムは、Compute コアが実行します。

#include "debug/dprint.h"  // required in all kernels using DPRINT
#include "compute_kernel_api.h"

namespace NAMESPACE {

    void MAIN {

        // Nothing to compute. Print respond message.
        // Make sure to export TT_METAL_DPRINT_CORES=0,0 before runtime.

        DPRINT_MATH(DPRINT << "Hello, Master, I am running a void compute kernel." << ENDL());

    }

}