8.1.5. Example: Using non-dict inputs and outputs with compile_automap

compile_automap は入出力として辞書以外を用いることができます。このサンプルプログラムでは、入力として tuple, list, namedtuple, dataclass を通して torch.ones(3, 4) を受け取り、MN-Core 2 上で加算 (add) してそれを直接出力します。 compile_automap は torch.Tensor を要素として持つこれらの型 (のネスト) に対応しています。

実行方法

$ cd /opt/pfn/pfcomp/codegen/MLSDK/examples/
$ ./exec_with_env.sh python3 add_automap.py

想定出力

• コンパイル結果を格納した codegen_dir (/tmp/add_many_tensors)

• 計算結果

result_on_cpu=tensor([[4., 4., 4., 4.],
        [4., 4., 4., 4.],
        [4., 4., 4., 4.]])

関連リンク

• mlsdk.Context.compile_automap

サンプルプログラム

リスト 8.5 /opt/pfn/pfcomp/codegen/MLSDK/examples/add_automap.py

from dataclasses import dataclass
from typing import List, NamedTuple, Tuple

import torch
from mlsdk import (
    Context,
    MNDevice,
    register_pytree_dataclass,
    register_pytree_namedtuple,
    storage,
)


@register_pytree_namedtuple
class NT(NamedTuple):
    x: torch.Tensor


@register_pytree_dataclass
@dataclass
class DC:
    x: torch.Tensor


def run_add():
    device = MNDevice("mncore2:auto")
    context = Context(device)
    Context.switch_context(context)

    def add(
        arg1: Tuple[torch.Tensor],
        arg2: NT,
        *,
        kwarg_ls: List[torch.Tensor],
        kwarg_dc: DC,
    ) -> torch.Tensor:
        return arg1[0] + arg2.x + kwarg_ls[0] + kwarg_dc.x

    arg1 = (torch.randn(3, 4),)
    arg2 = NT(x=torch.randn(3, 4))
    kwarg_ls = [torch.randn(3, 4)]
    kwarg_dc = DC(x=torch.randn(3, 4))

    compiled_add = context.compile_automap(
        add,
        (arg1, arg2),
        {"kwarg_ls": kwarg_ls, "kwarg_dc": kwarg_dc},
        storage.path("/tmp/add_many_tensors"),
        options={"float_dtype": "float"},
    )
    result = compiled_add(
        (torch.ones(3, 4),),
        NT(x=torch.ones(3, 4)),
        kwarg_ls=[torch.ones(3, 4)],
        kwarg_dc=DC(x=torch.ones(3, 4)),
    )
    result_on_cpu = result.cpu()
    print(f"{result_on_cpu=}")
    assert torch.allclose(result_on_cpu, torch.ones(3, 4) * 4)


if __name__ == "__main__":
    run_add()