mirror of
https://github.com/PaddlePaddle/FastDeploy.git
synced 2025-09-26 20:41:53 +08:00
[CudaGraph] support cudagraph use shared pool (#4199)
Some checks failed
CE Compile Job / ce_job_pre_check (push) Has been cancelled
CE Compile Job / print_ce_job_pre_check_outputs (push) Has been cancelled
CE Compile Job / FD-Clone-Linux (push) Has been cancelled
CE Compile Job / Show Code Archive Output (push) Has been cancelled
CE Compile Job / BUILD_SM8090 (push) Has been cancelled
CE Compile Job / BUILD_SM8689 (push) Has been cancelled
CE Compile Job / CE_UPLOAD (push) Has been cancelled
Some checks failed
CE Compile Job / ce_job_pre_check (push) Has been cancelled
CE Compile Job / print_ce_job_pre_check_outputs (push) Has been cancelled
CE Compile Job / FD-Clone-Linux (push) Has been cancelled
CE Compile Job / Show Code Archive Output (push) Has been cancelled
CE Compile Job / BUILD_SM8090 (push) Has been cancelled
CE Compile Job / BUILD_SM8689 (push) Has been cancelled
CE Compile Job / CE_UPLOAD (push) Has been cancelled
* support cudagraph use shared pool * add envs * change CUDAGRAPH_POOL_ID to int * change CUDAGRAPH_POOL_ID to use_memory_pool * unify use_unique_memory_pool * fix use_unique_memory_pool
This commit is contained in:
@@ -588,6 +588,9 @@ class GraphOptimizationConfig:
|
|||||||
Thus this flag cannot be used together with splitting_ops."""
|
Thus this flag cannot be used together with splitting_ops."""
|
||||||
self.full_cuda_graph: bool = True
|
self.full_cuda_graph: bool = True
|
||||||
|
|
||||||
|
""" Whether to use shared memory pool for multi capture_size """
|
||||||
|
self.use_unique_memory_pool: bool = False
|
||||||
|
|
||||||
self.max_capture_size: int = None
|
self.max_capture_size: int = None
|
||||||
self.real_shape_to_captured_size: dict[int, int] = None
|
self.real_shape_to_captured_size: dict[int, int] = None
|
||||||
# CINN Config ...
|
# CINN Config ...
|
||||||
|
@@ -20,6 +20,7 @@ from typing import Callable, Dict, List, Optional
|
|||||||
|
|
||||||
import paddle.jit.dy2static.utils as jit_utils
|
import paddle.jit.dy2static.utils as jit_utils
|
||||||
import paddle.nn.layer
|
import paddle.nn.layer
|
||||||
|
from paddle.base.core import CUDAGraph
|
||||||
from paddle.device.cuda import graphs
|
from paddle.device.cuda import graphs
|
||||||
|
|
||||||
from fastdeploy import envs
|
from fastdeploy import envs
|
||||||
@@ -85,17 +86,14 @@ class Dy2StCudaGraphManager:
|
|||||||
class CudaGraphPiecewiseBackend:
|
class CudaGraphPiecewiseBackend:
|
||||||
"""Manage the capture and replay of CUDA graphs at the subgraph level."""
|
"""Manage the capture and replay of CUDA graphs at the subgraph level."""
|
||||||
|
|
||||||
def __init__(
|
def __init__(self, fd_config: FDConfig, runnable: Callable):
|
||||||
self,
|
|
||||||
fd_config: FDConfig,
|
|
||||||
runnable: Callable,
|
|
||||||
):
|
|
||||||
self.fd_config = fd_config
|
self.fd_config = fd_config
|
||||||
self.runnable = runnable
|
self.runnable = runnable
|
||||||
self.cudagraph_capture_sizes = fd_config.graph_opt_config.cudagraph_capture_sizes
|
self.cudagraph_capture_sizes = fd_config.graph_opt_config.cudagraph_capture_sizes
|
||||||
self.warm_up_size = fd_config.graph_opt_config.cudagraph_num_of_warmups
|
self.warm_up_size = fd_config.graph_opt_config.cudagraph_num_of_warmups
|
||||||
self.real_shape_to_captured_size = fd_config.graph_opt_config.real_shape_to_captured_size
|
self.real_shape_to_captured_size = fd_config.graph_opt_config.real_shape_to_captured_size
|
||||||
|
if self.fd_config.graph_opt_config.use_unique_memory_pool:
|
||||||
|
self.unique_memory_pool_id = CUDAGraph.gen_new_memory_pool_id()
|
||||||
self._create_entry_dict()
|
self._create_entry_dict()
|
||||||
|
|
||||||
self.cuda_graph_manager = None
|
self.cuda_graph_manager = None
|
||||||
@@ -168,7 +166,11 @@ class CudaGraphPiecewiseBackend:
|
|||||||
input_addresses = [x.data_ptr() for (_, x) in kwargs.items() if isinstance(x, paddle.Tensor)]
|
input_addresses = [x.data_ptr() for (_, x) in kwargs.items() if isinstance(x, paddle.Tensor)]
|
||||||
entry.input_addresses = input_addresses
|
entry.input_addresses = input_addresses
|
||||||
|
|
||||||
new_grpah = graphs.CUDAGraph()
|
new_grpah = (
|
||||||
|
graphs.CUDAGraph(pool_id=self.unique_memory_pool_id)
|
||||||
|
if self.fd_config.graph_opt_config.use_unique_memory_pool
|
||||||
|
else graphs.CUDAGraph()
|
||||||
|
)
|
||||||
paddle.device.synchronize()
|
paddle.device.synchronize()
|
||||||
|
|
||||||
# Capture
|
# Capture
|
||||||
|
Reference in New Issue
Block a user