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

compile_automap can use input/output types other than dictionaries. In this sample program, it receives torch.ones(3, 4) through inputs such as tuple, list, namedtuple, and dataclass, performs an addition (add) on MN-Core 2 processor, and directly returns the result as output. compile_automap supports these types (including nested structures) as long as they contain torch.Tensor elements.

Execution Method

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

Expected Output

• The codegen_dir containing the compiled results (/tmp/add_many_tensors)

• The computed results

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

Related Links

• mlsdk.Context.compile_automap

Sample Program

Listing 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()