CUDA backend: a program crashes with an illegal memory access when CUDA device 1 is current #24887

import os, subprocess, sys, tempfile

import torch
from executorch.backends.cuda.cuda_backend import CudaBackend
from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
from executorch.exir import to_edge_transform_and_lower
from executorch.runtime import Runtime


def make_model():
    torch.manual_seed(0)
    return torch.nn.Sequential(torch.nn.Linear(8, 16), torch.nn.ELU(), torch.nn.Linear(16, 4)).eval()


if len(sys.argv) == 1:
    folder = tempfile.mkdtemp()
    program = to_edge_transform_and_lower(
        torch.export.export(make_model(), (torch.zeros(1, 8),)),
        partitioner=[CudaPartitioner([CudaBackend.generate_method_name_compile_spec("forward")])],
    ).to_executorch()
    with open(os.path.join(folder, "model.pte"), "wb") as f:
        program.write_to_file(f)
    program.write_tensor_data_to_file(folder)
    # Run each case in a fresh process, because a failed run leaves the CUDA context broken.
    for device in ("0", "1"):
        subprocess.run([sys.executable, __file__, folder, device])
else:
    folder, device = sys.argv[1], int(sys.argv[2])
    torch.cuda.set_device(device)
    method = Runtime.get().load_program(
        os.path.join(folder, "model.pte"), data_path=os.path.join(folder, "aoti_cuda_blob.ptd")
    ).load_method("forward")
    x = torch.randn(1, 8)
    try:
        out = method.execute([x])[0]
        print(f"current device {device}: OK, max error {(out - make_model()(x)).abs().max().item():.1e}")
    except RuntimeError as e:
        print(f"current device {device}: FAILED: {e}")