pygfx.renderers.wgpu.select_adapter

pygfx.renderers.wgpu.select_adapter(adapter)

Select a specific adapter / GPU.

Select an adapter as obtained via wgpu.gpu.enumerate_adapters(), which can be useful in multi-gpu environments.

For example:

adapters = wgpu.gpu.enumerate_adapters()
adapters_tesla = [a for a in adapters if "Tesla" in a.summary]
adapters_discrete = [a for a in adapters if "DiscreteGPU" in a.summary]
pygfx.renderers.wgpu.select_adapter(adapters_discrete[0])

Note that using this function reduces the portability of your code, because it’s highly specific for your current machine/environment.

The order of the adapters returned by wgpu.gpu.enumerate_adapters() is such that Vulkan adapters go first, then Metal, then D3D12, then OpenGL. Within each category, the order as provided by the particular backend is maintained. Note that the same device may be present via multiple backends (e.g. vulkan/opengl).

We cannot make guarantees about whether the order of the adapters matches the order as reported by e.g. nvidia-smi. We have found that on a Linux multi-gpu cluster, the order does match, but we cannot promise that this is always the case. If you want to make sure, do some testing by allocating big buffers and checking memory usage using nvidia-smi.

Example to allocate and check GPU mem usage:

import subprocess

import wgpu
import torch

def allocate_gpu_mem_with_wgpu(idx):
    a = wgpu.gpu.enumerate_adapters()[idx]
    d = a.request_device()
    b = d.create_buffer(size=10*2**20, usage=wgpu.BufferUsage.COPY_DST)
    return b

def allocate_gpu_mem_with_torch(idx):
    d = torch.device(f"cuda:{idx}")
    return torch.ones([2000, 10], dtype=torch.float32, device=d)

def show_mem_usage():
    print(subprocess.run(["nvidia-smi"]))

See https://github.com/pygfx/wgpu-py/issues/482 for more details.