wehub-resource-sync
133 lines
3.8 KiB
Python
133 lines
3.8 KiB
Python
# Copyright (c) 2025 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.
|
|
|
|
import unittest
|
|
from contextlib import contextmanager
|
|
|
|
import numpy as np
|
|
from dygraph_to_static_utils import Dy2StTestBase
|
|
|
|
import paddle
|
|
from paddle.jit.dy2static.utils import CUDAGraphState
|
|
|
|
SEED = 2025
|
|
np.random.seed(2025)
|
|
GLOBAL_GRAPH_WITH_BUFFER = None
|
|
|
|
|
|
class GraphWithBuffer:
|
|
def __init__(self, inputs, outputs):
|
|
self.inputs_buffer = inputs
|
|
self.outputs_buffer = outputs
|
|
|
|
def set_inputs_buffer(self, inputs):
|
|
assert len(self.inputs_buffer) == len(inputs)
|
|
for i, _ in enumerate(inputs):
|
|
self.inputs_buffer[i][:] = inputs[i]
|
|
|
|
def get_inputs(self):
|
|
return self.inputs_buffer
|
|
|
|
def get_real_outputs(self):
|
|
return self.outputs_buffer
|
|
|
|
def get_outputs(self):
|
|
return [out.clone() for out in self.outputs_buffer]
|
|
|
|
|
|
def capture_run_impl(original_run_impl, inputs, parameters, attrs):
|
|
prog_attrs, cuda_graph_attrs = attrs
|
|
cuda_graph_attrs |= {
|
|
"cuda_graph_state": CUDAGraphState.CAPTURE,
|
|
"cuda_graph_dispatch_key": inputs[0].shape[0],
|
|
}
|
|
outputs = original_run_impl(
|
|
inputs, parameters, (prog_attrs, cuda_graph_attrs)
|
|
)
|
|
|
|
global GLOBAL_GRAPH_WITH_BUFFER
|
|
if GLOBAL_GRAPH_WITH_BUFFER is None:
|
|
GLOBAL_GRAPH_WITH_BUFFER = GraphWithBuffer(inputs, outputs)
|
|
|
|
return outputs
|
|
|
|
|
|
def replay_run_impl(original_run_impl, inputs, parameters, attrs):
|
|
prog_attrs, cuda_graph_attrs = attrs
|
|
cuda_graph_attrs |= {
|
|
"cuda_graph_state": CUDAGraphState.REPLAY,
|
|
"cuda_graph_dispatch_key": inputs[0].shape[0],
|
|
}
|
|
global GLOBAL_GRAPH_WITH_BUFFER
|
|
assert GLOBAL_GRAPH_WITH_BUFFER is not None
|
|
GLOBAL_GRAPH_WITH_BUFFER.set_inputs_buffer(inputs)
|
|
|
|
_ = original_run_impl(
|
|
GLOBAL_GRAPH_WITH_BUFFER.get_inputs(),
|
|
parameters,
|
|
(prog_attrs, cuda_graph_attrs),
|
|
)
|
|
|
|
return GLOBAL_GRAPH_WITH_BUFFER.get_outputs()
|
|
|
|
|
|
@contextmanager
|
|
def capture_run_impl_guard():
|
|
with paddle.jit.dy2static.pir_partial_program.replace_run_impl_guard(
|
|
capture_run_impl,
|
|
):
|
|
yield
|
|
|
|
|
|
@contextmanager
|
|
def replay_run_impl_guard():
|
|
with paddle.jit.dy2static.pir_partial_program.replace_run_impl_guard(
|
|
replay_run_impl,
|
|
):
|
|
yield
|
|
|
|
|
|
@unittest.skipIf(
|
|
(not paddle.is_compiled_with_cuda()) or paddle.is_compiled_with_rocm(),
|
|
"Skipped on non-GPU devices and ROCm devices(DCU) as this test requires NVIDIA CUDA Graph.",
|
|
)
|
|
class TestCUDAGraph(Dy2StTestBase):
|
|
def initialize(self):
|
|
global GLOBAL_GRAPH_WITH_BUFFER
|
|
GLOBAL_GRAPH_WITH_BUFFER = None
|
|
|
|
def func(x, y):
|
|
return x + y
|
|
|
|
self.fn = func
|
|
self.static_fn = paddle.jit.to_static(func)
|
|
|
|
def test_capture_replay(self):
|
|
self.initialize()
|
|
x = paddle.randn([2, 2, 3, 3], dtype='float32')
|
|
y = paddle.randn([2, 2, 3, 3], dtype='float32')
|
|
with capture_run_impl_guard():
|
|
_ = self.static_fn(x, y)
|
|
|
|
a = paddle.randn([2, 2, 3, 3], dtype='float32')
|
|
b = paddle.randn([2, 2, 3, 3], dtype='float32')
|
|
with replay_run_impl_guard():
|
|
c = self.static_fn(a, b)
|
|
|
|
np.testing.assert_allclose(self.fn(a, b), c)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|