migx ¶

@classmethod
def autoselect(cls, device_id: int = 0, **kwargs: Any) -> Backend:
    """
    Try to select the best backend for the current system.

    Args:
        device_id: The GPU device id.
        **kwargs: Additional arguments to pass to the backend.

    Returns:
        The selected backend.
    """

    gpu = get_gpu(device_id)
    vendor = (
        cast(str | None, gpu.vendor)
        if gpu
        else "apple"
        # macOS x86_64 is unsupported
        if platform.system().lower() == "darwin" and platform.machine() == "x86_64"
        else None
    )

    match vendor:
        # Windows & Linux
        case "nvidia":
            if hasattr(core, "trt"):
                backend = Backend.TRT
            elif hasattr(core, "trt_rtx"):
                backend = Backend.TRT_RTX
            elif platform.system().lower() == "windows" and hasattr(core, "ort"):
                backend = Backend.ORT_DML
            elif hasattr(core, "ort"):
                backend = Backend.ORT_CUDA
            elif hasattr(core, "ncnn"):
                backend = Backend.NCNN
            else:
                backend = Backend.OV_CPU
        # Windows & Linux
        case "amd":
            if platform.system().lower() == "windows" and hasattr(core, "ort"):
                backend = Backend.ORT_DML
            elif hasattr(core, "migx"):
                backend = Backend.MIGX
            elif hasattr(core, "ncnn"):
                backend = Backend.NCNN_VK
            else:
                backend = Backend.OV_CPU
        # Windows & Linux
        case "intel":
            # device-smi can't detect Intel NPUs in 0.5.6
            # https://github.com/ModelCloud/Device-SMI#roadmap
            if hasattr(core, "ov"):
                backend = Backend.OV_GPU
            elif platform.system().lower() == "windows" and hasattr(core, "ort"):
                backend = Backend.ORT_DML
            elif hasattr(core, "ncnn"):
                backend = Backend.NCNN_VK
            else:
                backend = Backend.OV_CPU
        # macOS ARM64 & x86_64
        case "apple":
            if hasattr(core, "ncnn"):
                backend = Backend.NCNN_VK
            elif hasattr(core, "ort"):
                backend = Backend.ORT_COREML
            else:
                backend = Backend.OV_CPU
        case _:
            backend = Backend.OV_CPU

    del gpu

    return backend(**kwargs)

def build_program(self, network_path: Path, channels: int, tilesize: Shape, input_name: str = "input") -> Path:
    if rocm_path := os.getenv("ROCM_PATH"):
        migraph_driver = f"{rocm_path}/bin/migraphx-driver"
    else:
        migraph_driver = shutil.which("migraphx-driver") or "migraphx-driver"

    dirname = get_artifacts_folder()
    dirname.mkdir(parents=True, exist_ok=True)
    identity = self.get_identity(network_path, tilesize)
    program_path = dirname / f"{identity}.mxr"

    command: list[Any] = [
        migraph_driver,
        "compile",
        "--onnx",
        network_path,
        "--gpu",
        "--optimize",
        "--binary",
        "--output",
        program_path,
    ]

    opt_shapes = self.opt_shapes or tilesize
    command.extend(["--input-dim", f"@{input_name}", "1", f"{channels}", f"{opt_shapes[1]}", f"{opt_shapes[0]}"])
    if self.fp16:
        command.append("--fp16")
    if self.bf16:
        command.append("--bf16")
    if not self.fast_math:
        command.append("--disable-fast-math")
    if self.exhaustive_tune:
        command.append("--exhaustive-tune")

    command.extend(self.custom_args)

    logger.debug("%s: Calling migraphx-driver with the command:", self)
    logger.debug(command)

    try:
        subprocess.run(command, env=os.environ | self.custom_env, check=True, stdout=sys.stderr)
    except subprocess.CalledProcessError as e:
        logger.debug("MIGraphx STDERR:\n%s", e.stderr)
        logger.debug("MIGraphx STDOUT:\n%s", e.stdout)
        raise CustomRuntimeError("The program compilation failed") from e

    return program_path

def get_args(self, clips: vs.VideoNode | Sequence[vs.VideoNode]) -> dict[str, Any]:
    return {"device_id": self.device_id, "num_streams": self.num_streams}

def get_identity(self, network_path: Path, tilesize: Shape) -> int:
    checksum = zlib.crc32(network_path.read_bytes())

    device_props = self.plugin.DeviceProperties(self.device_id)
    device = [device_props["name"].decode().replace(" ", "-"), device_props["driver_version"]]

    components = (
        str(self),
        str(self.version),
        network_path.name,
        f"{checksum:x}",
        str(tilesize),
        *map(str, device),
    )
    logger.debug("%s: Identity %r", self.get_identity, components)
    return zlib.crc32(bytes("|".join(components), "utf-8"))

@copy_signature(BackendAutoConvertFloat.inference)
def inference(
    self,
    clips: vs.VideoNode | Sequence[vs.VideoNode],
    network_path: str | os.PathLike[str],
    /,
    overlap: tuple[int, int],
    tilesize: tuple[int, int],
    *,
    flexible: bool = False,
    **kwargs: Any,
) -> vs.VideoNode | list[vs.VideoNode]:
    network_path = Path(network_path)
    channels = sum(clip.format.num_planes for clip in to_arr(clips))
    program_path = self.build_program(network_path, channels, tilesize)

    if self.fp16:
        # Clips must be in fp16 format is fp16 is enabled,
        # otherwise the MIGX plugin errors out.
        clips = [depth(c, 16, sample_type=vs.SampleType.FLOAT) for c in clips]
    else:
        clips = [depth(c, 32) for c in clips]

    return super().inference(clips, program_path, overlap, tilesize, flexible=flexible, **kwargs)