How to use cudaSetDevice to select devices while using MIG on A100?


I am using A100 with Multi-Instance GPU (MIG). As described in the official document, it is straightforward to use environment variable CUDA_VISIBLE_DEVICES to select devices to run a single process. I am wondering if I could use CUDA runtime API cudaSetDevice or similar APIs to select devices for a thread in a multi-threading inference with CUDA? Thanks.

$ nvidia-smi -L
GPU 0: A100-SXM4-40GB (UUID: GPU-e86cb44c-6756-fd30-cd4a-1e6da3caf9b0)
  MIG 3g.20gb Device 0: (UUID: MIG-GPU-e86cb44c-6756-fd30-cd4a-1e6da3caf9b0/1/0)
  MIG 3g.20gb Device 1: (UUID: MIG-GPU-e86cb44c-6756-fd30-cd4a-1e6da3caf9b0/2/0)

$ CUDA_VISIBLE_DEVICES=MIG-GPU-e86cb44c-6756-fd30-cd4a-1e6da3caf9b0/1/0 ./BlackScholes &
$ CUDA_VISIBLE_DEVICES=MIG-GPU-e86cb44c-6756-fd30-cd4a-1e6da3caf9b0/2/0 ./BlackScholes &

Hello, this forum is dedicated to discussions related to using the sanitizer tools and API.
Questions related to CUDA can be raised at CUDA - NVIDIA Developer Forums