paddlepaddle--paddle/python/paddle/device/cuda/graphs.py

# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import os
import warnings
from typing import NoReturn, overload

from paddle.base.core import (
    CUDAPlace,
    CustomPlace,
    XPUPlace,
    get_all_custom_device_type,
    is_compiled_with_cuda,
    is_compiled_with_custom_device,
    is_compiled_with_rocm,
    is_compiled_with_xpu,
)


def check_compiled_with_custom_device():
    custom_device_flag = False
    custom_devices_types = get_all_custom_device_type()
    for device_type in custom_devices_types:
        if is_compiled_with_custom_device(device_type):
            custom_device_flag = True
            break
    return custom_device_flag


if (
    is_compiled_with_cuda()
    or is_compiled_with_rocm()
    or check_compiled_with_custom_device()
    or is_compiled_with_xpu()
):
    from paddle.base.core import CUDAGraph as CoreCUDAGraph

    def is_cuda_graph_supported():
        return True
else:
    CoreCUDAGraph = None

    def is_cuda_graph_supported():
        return False


def current_expected_place():
    for device in get_all_custom_device_type():
        selected_devices = os.getenv(f"FLAGS_selected_{device}s", "0").split(
            ","
        )
        device_id = int(selected_devices[0])
        return CustomPlace(device, device_id)
    return None


ALL_MODES = ["global", "thread_local", "relaxed"]
cuda_graph_id = 0


class CUDAGraph:
    """
    The native Paddle constructor takes ``place``, ``mode``, ``pool_id`` and
    ``enable_replace``; the PyTorch-compatible ``keep_graph`` keyword is
    accepted as well. ``capture_begin`` additionally accepts the PyTorch
    keywords ``pool`` and ``capture_error_mode`` so the same instance can be
    driven from either API style.
    """

    @overload
    def __init__(self, keep_graph: bool, /) -> None: ...

    @overload
    def __init__(
        self,
        place: CUDAPlace | XPUPlace | CustomPlace | None = None,
        mode: str = "thread_local",
        pool_id: int | None = None,
        enable_replace: bool = False,
        *,
        keep_graph: bool = False,
    ) -> None: ...

    def __init__(
        self,
        place=None,
        mode="thread_local",
        pool_id=None,
        enable_replace=False,
        *,
        keep_graph: bool = False,
    ):
        assert CoreCUDAGraph is not None, (
            "CUDA Graph is only supported on PaddlePaddle compiled with NVIDIA GPU."
        )

        if isinstance(place, bool):
            if keep_graph is not False:
                raise TypeError(
                    "keep_graph is specified both positionally and by keyword"
                )
            keep_graph = place
            place = None

        self._graph = None
        if place is None and check_compiled_with_custom_device():
            place = current_expected_place()
        elif place is None:
            if is_compiled_with_cuda():
                device_id = int(os.environ.get('FLAGS_selected_gpus', 0))
                place = CUDAPlace(device_id)
            elif is_compiled_with_xpu():
                device_id = int(os.environ.get('FLAGS_selected_xpus', 0))
                place = XPUPlace(device_id)
            else:
                raise RuntimeError("Not Supported devices")

        self._place = place
        assert mode in ALL_MODES
        self._mode = ALL_MODES.index(mode)
        self._pool_id = pool_id
        self._enable_replace = enable_replace
        self._keep_graph = keep_graph
        self._debug_mode = False

    def capture_begin(
        self, pool: int | None = None, capture_error_mode: str | None = None
    ) -> None:
        """Begin capturing CUDA work on the current stream.

        Args:
            pool (int, optional): A memory pool token from
                :func:`paddle.cuda.graph_pool_handle` or another graph's
                :meth:`pool`. When provided, this graph shares the indicated
                memory pool. Overrides ``pool_id`` from the constructor.
            capture_error_mode (str, optional): One of ``'global'``,
                ``'thread_local'``, ``'relaxed'`` (see :data:`ALL_MODES`).
                When ``None`` (default) the constructor's ``mode`` is used;
                otherwise it overrides the constructor for this capture and a
                :class:`UserWarning` is emitted to flag the precedence.
                Invalid values raise :class:`ValueError`.
        """
        if pool is not None:
            self._pool_id = pool
        elif self._pool_id is None:
            self._pool_id = CoreCUDAGraph.gen_new_memory_pool_id()

        if capture_error_mode is None:
            mode = self._mode
        else:
            if capture_error_mode not in ALL_MODES:
                raise ValueError(
                    f"capture_error_mode must be one of {ALL_MODES}, "
                    f"but got {capture_error_mode!r}."
                )
            mode = ALL_MODES.index(capture_error_mode)
            if mode != self._mode:
                warnings.warn(
                    f"capture_error_mode={capture_error_mode!r} differs from "
                    f"the constructor mode={ALL_MODES[self._mode]!r}; the "
                    f"explicit capture_error_mode takes precedence for this "
                    f"capture.",
                    stacklevel=2,
                )

        CoreCUDAGraph.begin_capture_with_pool_id(
            self._place, mode, self._pool_id, self._enable_replace
        )

    def capture_end(self):
        self._graph = CoreCUDAGraph.end_capture()

    def _require_captured(self) -> None:
        """Raise a clear error if no graph has been captured yet.

        ``self._graph`` is only populated by :meth:`capture_end`; methods that
        consume it (``replay`` / ``reset`` / ``debug_dump`` / ...) would
        otherwise raise ``AttributeError`` on ``NoneType`` when called too
        early. Centralizing the check produces a single, actionable message.
        """
        if self._graph is None:
            raise RuntimeError(
                "CUDAGraph has not been captured yet. "
                "Call capture_begin/capture_end first."
            )

    def instantiate(self) -> CoreCUDAGraph:
        """Return the instantiated core CUDA graph held by this wrapper.

        Paddle builds the executable graph eagerly inside :meth:`capture_end`,
        so by the time this is called the graph is already instantiated. It is
        kept for source compatibility with ``torch.cuda.CUDAGraph.instantiate``
        and returns the held core :class:`~paddle.base.core.CUDAGraph` produced
        by :meth:`capture_end`.
        """
        self._require_captured()
        return self._graph

    def replay(self):
        self._require_captured()
        self._graph.replay()

    def reset(self):
        self._require_captured()
        self._graph.reset()

    def pool(self) -> int:
        """Return an opaque integer token representing this graph's memory pool.

        The token can be passed as the ``pool`` argument to another graph's
        :meth:`capture_begin` (or to :class:`paddle.cuda.graph`) so the two
        graphs share the same memory pool.
        """
        if self._pool_id is None:
            self._pool_id = CoreCUDAGraph.gen_new_memory_pool_id()
        return self._pool_id

    def enable_debug_mode(self) -> None:
        """Enable debug mode so that :meth:`debug_dump` is permitted."""
        self._debug_mode = True

    def debug_dump(self, debug_path) -> None:
        """Dump the captured graph to ``debug_path`` for inspection.

        :meth:`enable_debug_mode` must be called first.
        """
        if not self._debug_mode:
            raise RuntimeError(
                "debug_dump requires debug mode to be enabled first. "
                "Call enable_debug_mode() before debug_dump()."
            )
        self._require_captured()
        self.print_to_dot_files(debug_path)

    def raw_cuda_graph(self) -> NoReturn:
        """Paddle does not expose the raw ``cudaGraph_t`` handle."""
        raise NotImplementedError(
            "raw_cuda_graph is not yet supported in Paddle CUDAGraph. "
            "The underlying cudaGraph_t handle is not exposed by the Python "
            "binding."
        )

    def raw_cuda_graph_exec(self) -> NoReturn:
        """Paddle does not expose the raw ``cudaGraphExec_t`` handle."""
        raise NotImplementedError(
            "raw_cuda_graph_exec is not yet supported in Paddle CUDAGraph. "
            "The underlying cudaGraphExec_t handle is not exposed by the "
            "Python binding."
        )

    def print_to_dot_files(self, dirname, flags=None):
        if not isinstance(dirname, (str, bytes)):
            dirname = dirname.name
        os.makedirs(name=dirname, exist_ok=True)
        assert os.path.isdir(dirname), (
            f"The dirname {dirname} should be a directory"
        )
        if flags is None:
            flags = 2047  # only all information. It can be any integer inside [1, 2048)
        self._graph.print_to_dot_files(dirname, flags)

    def replace_input_ptrs(self, old_ptrs, new_ptrs):
        self._graph.replace_input_ptrs(old_ptrs, new_ptrs)