7.1.8. Example: Load codegen_dir

A sample program that directly reuses the compilation result from Example: Adding Two Vectors (/tmp/add_two_tensors)

Execution Method

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

Expected Output

• Computation result (should be equal to the result of add.py)

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

Related Links

• Directly loading codegen_dir

Sample Program

Listing 7.8 /opt/pfn/pfcomp/codegen/MLSDK/examples/load_codegen_dir.py

import torch
from mlsdk import Context, MNDevice, storage

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

# Load the previously compiled function
loaded_add = context.load_codegen_dir(storage.path("/tmp/add_two_tensors"))

# Now you can use it directly
result = loaded_add({"x": torch.ones(3, 4), "y": torch.ones(3, 4)})
result_on_cpu = result["out"].cpu()
print(f"Result from loaded function: {result_on_cpu=}")
assert torch.allclose(result_on_cpu, torch.ones(3, 4) * 2)