Ryzen NPU の利用方法

#include <aie_api/aie.hpp>

extern "C"
void vadd_kernel(
    bfloat16* in0,
    bfloat16* in1,
    bfloat16* out
) {
    for (int i = 0; i < 1024; i++) {
        out[i] = in0[i] + in1[i];
    }
}

import numpy as np
from ml_dtypes import bfloat16

from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker
from aie.iron.placers import SequentialPlacer
from aie.iron.device import NPU2

# (1) 入出力データの型を定義
vec_type = np.ndarray[(1024,), np.dtype[bfloat16]]

# (2) カーネル関数定義
kernel = Kernel(
    "vadd_kernel",                  # カーネル関数名
    "vadd_kernel.o",                # コンパイル済みオブジェクト
    [vec_type, vec_type, vec_type], # 引数型
)

# (3) AIEコアで実行するメイン関数の定義
def core_fn(of_in0, of_in1, of_out, kernel):
    # メモリリソースを取得
    elemOut = of_out.acquire(1)
    elemIn0 = of_in0.acquire(1)
    elemIn1 = of_in1.acquire(1)
    # カーネルを実行
    kernel(elemIn0, elemIn1, elemOut)
    # メモリリソースを開放
    of_in0.release(1)
    of_in1.release(1)
    of_out.release(1)

# (4) データ移動のためのObject FIFOを作成
of_in0 = ObjectFifo(vec_type, name="in0")
of_in1 = ObjectFifo(vec_type, name="in1")
of_out = ObjectFifo(vec_type, name="out")

# (5) コアプログラムとObject FIFOを指定してWorkerを作成
worker = Worker(
    core_fn, [of_in0.cons(), of_in1.cons(), of_out.prod(), kernel],
)

# (6) 実行時のAIEアレイとDRAMの間のデータ移動を定義
rt = Runtime()
with rt.sequence(vec_type, vec_type, vec_type) as (in0, in1, out):
    rt.start(worker)
    rt.fill(of_in0.prod(), in0)
    rt.fill(of_in1.prod(), in1)
    rt.drain(of_out.cons(), out, wait=True)

# (7) リソースの配置を行いMLIRを生成
program = Program(NPU2(), rt)
module = program.resolve_program(SequentialPlacer())

# (8) MLIRを出力
print(module)

#include <fstream>
#include <iostream>
#include <stdfloat>
#include <string>
#include <vector>
#include <random>
#include <algorithm>

#include "xrt/xrt_bo.h"
#include "xrt/xrt_device.h"
#include "xrt/xrt_kernel.h"

int main(int argc, char** argv) {
    // (1) XRT環境の準備
    auto device = xrt::device(0);
    auto xclbin = xrt::xclbin(std::string("vadd.xclbin"));
    device.register_xclbin(xclbin);
    xrt::hw_context context(device, xclbin.get_uuid());
    auto kernel = xrt::kernel(context, "MLIR_AIE");

    // (2) 乱数生成の準備
    std::mt19937 rng(std::random_device{}());
    std::uniform_real_distribution<float> dist(-1.0f, 1.0f);

    // (3) 命令シーケンスの読み込み
    std::vector<uint32_t> inst;
    {
        std::ifstream ifs("vadd_inst.txt", std::ios::binary | std::ios::ate);
        inst.resize(ifs.tellg() / 4);
        ifs.seekg(0);
        ifs.read(reinterpret_cast<char*>(inst.data()), inst.size() * 4);
    }

    // (4) バッファオブジェクトの作成
    auto bo_inst = xrt::bo(device, inst.size() * sizeof(uint32_t), xrt::bo::flags::cacheable, kernel.group_id(1));
    auto bo_in0  = xrt::bo(device, 1024 * sizeof(std::bfloat16_t), xrt::bo::flags::host_only, kernel.group_id(3));
    auto bo_in1  = xrt::bo(device, 1024 * sizeof(std::bfloat16_t), xrt::bo::flags::host_only, kernel.group_id(4));
    auto bo_out  = xrt::bo(device, 1024 * sizeof(std::bfloat16_t), xrt::bo::flags::host_only, kernel.group_id(5));

    // (5) 命令シーケンスをコピー
    std::copy(inst.begin(), inst.end(), bo_inst.map<uint32_t*>());

    // (6) 入力ベクトルを乱数で初期化
    std::span in0 { bo_in0.map<std::bfloat16_t*>(), 1024 };
    for (auto& x : in0) x = static_cast<std::bfloat16_t>(dist(rng));

    std::span in1 { bo_in1.map<std::bfloat16_t*>(), 1024 };
    for (auto& x : in1) x = static_cast<std::bfloat16_t>(dist(rng));

    // (7) 出力ベクトルをゼロクリア
    std::span out { bo_out.map<std::bfloat16_t*>(), 1024 };
    std::fill(out.begin(), out.end(), 0);

    // (8) ホストからデバイスのメモリに同期
    bo_inst.sync(XCL_BO_SYNC_BO_TO_DEVICE);
    bo_in0 .sync(XCL_BO_SYNC_BO_TO_DEVICE);
    bo_in1 .sync(XCL_BO_SYNC_BO_TO_DEVICE);
    bo_out .sync(XCL_BO_SYNC_BO_TO_DEVICE);

    // (9) AIE実行
    unsigned int opcode = 3;
    auto run = kernel(opcode, bo_inst, inst.size(), bo_in0, bo_in1, bo_out);
    run.wait();

    // (10) デバイスからホストのメモリに同期
    bo_out.sync(XCL_BO_SYNC_BO_FROM_DEVICE);

    // (11) 誤差を許容して比較する関数
    auto close = [](std::bfloat16_t a, std::bfloat16_t b) {
        float diff = static_cast<float>(a) - static_cast<float>(b);
        float tol = 5e-4f + 8e-3f * std::fabs(b);
        return std::fabs(diff) < tol;
    };

    // (12) 出力結果をチェック
    int mismatch = 0;
    for (int i = 0; i < 1024; i++) {
        if (i < 4) {
            std::cout << "In0[" << i << "] = " << in0[i] << std::endl;
            std::cout << "In1[" << i << "] = " << in1[i] << std::endl;
            std::cout << "Out[" << i << "] = " << out[i] << std::endl;
        }
        if (!close(out[i], in0[i] + in1[i])) mismatch++;
    }

    if (mismatch == 0) {
        std::cout << "PASS" << std::endl;
    } else {
        std::cout << "FAIL : mismatch=" << mismatch << std::endl;
    }
}