Skip to content

Commit 06241d2

Browse files
author
Naren Dasan
committed
fix: Fix the CUDAGraphs C++ runtime implementation
Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
1 parent 655ed6b commit 06241d2

21 files changed

+838
-542
lines changed

‎.github/workflows/build-test-linux.yml

Lines changed: 47 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,9 @@ on:
88
- nightly
99
- release/*
1010
tags:
11-
# NOTE: Binary build pipelines should only get triggered on release candidate builds
12-
# Release candidate tags look like: v1.11.0-rc1
13-
- v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
11+
# NOTE: Binary build pipelines should only get triggered on release candidate builds
12+
# Release candidate tags look like: v1.11.0-rc1
13+
- v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
1414
workflow_dispatch:
1515

1616
jobs:
@@ -84,9 +84,9 @@ jobs:
8484
popd
8585
pushd .
8686
cd tests/py/ts
87-
python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/ts_api_test_results.xml api/
88-
python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/ts_models_test_results.xml models/
89-
python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/ts_integrations_test_results.xml integrations/
87+
python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/ts_api_test_results.xml api/
88+
python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/ts_models_test_results.xml models/
89+
python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/ts_integrations_test_results.xml integrations/
9090
popd
9191
9292
tests-py-dynamo-converters:
@@ -114,7 +114,7 @@ jobs:
114114
export USE_HOST_DEPS=1
115115
pushd .
116116
cd tests/py/dynamo
117-
python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_converters_test_results.xml -n 10 conversion/
117+
python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_converters_test_results.xml -n 8 conversion/
118118
popd
119119
120120
tests-py-dynamo-fe:
@@ -142,8 +142,8 @@ jobs:
142142
export USE_HOST_DEPS=1
143143
pushd .
144144
cd tests/py/dynamo
145-
python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_fe_test_results.xml --ir dynamo models/test_models_export.py
146-
python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/dyn_models_export.xml --ir dynamo models/test_dyn_models.py
145+
python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_fe_test_results.xml --ir dynamo models/test_models_export.py
146+
python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dyn_models_export.xml --ir dynamo models/test_dyn_models.py
147147
popd
148148
149149
tests-py-dynamo-serde:
@@ -171,7 +171,7 @@ jobs:
171171
export USE_HOST_DEPS=1
172172
pushd .
173173
cd tests/py/dynamo
174-
python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/export_serde_test_results.xml --ir dynamo models/test_export_serde.py
174+
python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/export_serde_test_results.xml --ir dynamo models/test_export_serde.py
175175
popd
176176
177177
tests-py-torch-compile-be:
@@ -199,9 +199,9 @@ jobs:
199199
export USE_HOST_DEPS=1
200200
pushd .
201201
cd tests/py/dynamo
202-
python -m pytest -n 10 --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_compile_be_test_results.xml backend/
203-
python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_complete_be_e2e_test_results.xml --ir torch_compile models/test_models.py
204-
python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_compile_dyn_models_export.xml --ir torch_compile models/test_dyn_models.py
202+
python -m pytest -ra -n 10 --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_compile_be_test_results.xml backend/
203+
python -m pytest -ra -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_complete_be_e2e_test_results.xml --ir torch_compile models/test_models.py
204+
python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_compile_dyn_models_export.xml --ir torch_compile models/test_dyn_models.py
205205
popd
206206
207207
tests-py-dynamo-core:
@@ -229,9 +229,39 @@ jobs:
229229
export USE_HOST_DEPS=1
230230
pushd .
231231
cd tests/py/dynamo
232-
python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_runtime_test_results.xml runtime/
233-
python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_partitioning_test_results.xml partitioning/
234-
python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_lowering_test_results.xml lowering/
232+
python -m pytest -ra -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_runtime_test_results.xml --ignore runtime/test_002_cudagraphs_py.py --ignore runtime/test_002_cudagraphs_cpp.py runtime/
233+
python -m pytest -ra -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_partitioning_test_results.xml partitioning/
234+
python -m pytest -ra -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_lowering_test_results.xml lowering/
235+
popd
236+
237+
tests-py-dynamo-cudagraphs:
238+
name: Test dynamo cudagraphs [Python]
239+
needs: [generate-matrix, build]
240+
strategy:
241+
fail-fast: false
242+
matrix:
243+
include:
244+
- repository: pytorch/tensorrt
245+
package-name: torch_tensorrt
246+
pre-script: packaging/pre_build_script.sh
247+
post-script: packaging/post_build_script.sh
248+
smoke-test-script: packaging/smoke_test_script.sh
249+
uses: ./.github/workflows/linux-test.yml
250+
with:
251+
job-name: tests-py-dynamo-cudagraphs
252+
repository: "pytorch/tensorrt"
253+
ref: ""
254+
test-infra-repository: pytorch/test-infra
255+
test-infra-ref: main
256+
build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
257+
pre-script: ${{ matrix.pre-script }}
258+
script: |
259+
export USE_HOST_DEPS=1
260+
pushd .
261+
cd tests/py/dynamo
262+
nvidia-smi
263+
python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_runtime_cudagraphs_cpp_test_results.xml runtime/test_002_cudagraphs_cpp.py || true
264+
python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_runtime_cudagraphs_py_test_results.xml runtime/test_002_cudagraphs_py.py || true
235265
popd
236266
237267
tests-py-core:
@@ -259,7 +289,7 @@ jobs:
259289
export USE_HOST_DEPS=1
260290
pushd .
261291
cd tests/py/core
262-
python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_core_test_results.xml .
292+
python -m pytest -ra -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_core_test_results.xml .
263293
popd
264294
265295
concurrency:

‎.github/workflows/build-test-windows.yml

Lines changed: 38 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,9 @@ on:
88
- nightly
99
- release/*
1010
tags:
11-
# NOTE: Binary build pipelines should only get triggered on release candidate builds
12-
# Release candidate tags look like: v1.11.0-rc1
13-
- v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
11+
# NOTE: Binary build pipelines should only get triggered on release candidate builds
12+
# Release candidate tags look like: v1.11.0-rc1
13+
- v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
1414
workflow_dispatch:
1515

1616
jobs:
@@ -192,8 +192,8 @@ jobs:
192192
export USE_HOST_DEPS=1
193193
pushd .
194194
cd tests/py/dynamo
195-
python -m pytest -n 10 --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_compile_be_test_results.xml backend/
196-
python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_complete_be_e2e_test_results.xml --ir torch_compile models/test_models.py
195+
python -m pytest -ra -n 10 --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_compile_be_test_results.xml backend/
196+
python -m pytest -ra -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_complete_be_e2e_test_results.xml --ir torch_compile models/test_models.py
197197
python -m pytest --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_compile_dyn_models_export.xml --ir torch_compile models/test_dyn_models.py
198198
popd
199199
@@ -219,9 +219,38 @@ jobs:
219219
export USE_HOST_DEPS=1
220220
pushd .
221221
cd tests/py/dynamo
222-
python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_runtime_test_results.xml runtime/
223-
python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_partitioning_test_results.xml partitioning/
224-
python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_lowering_test_results.xml lowering/
222+
python -m pytest -ra -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_runtime_test_results.xml --ignore runtime/test_002_cudagraphs_py.py --ignore runtime/test_002_cudagraphs_cpp.py runtime/
223+
python -m pytest -ra -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_partitioning_test_results.xml partitioning/
224+
python -m pytest -ra -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_lowering_test_results.xml lowering/
225+
popd
226+
227+
tests-py-dynamo-cudagraphs:
228+
name: Test dynamo cudagraphs [Python]
229+
needs: [generate-matrix, build]
230+
strategy:
231+
fail-fast: false
232+
matrix:
233+
include:
234+
- repository: pytorch/tensorrt
235+
package-name: torch_tensorrt
236+
pre-script: packaging/pre_build_script.sh
237+
post-script: packaging/post_build_script.sh
238+
smoke-test-script: packaging/smoke_test_script.sh
239+
uses: ./.github/workflows/linux-test.yml
240+
with:
241+
job-name: tests-py-dynamo-cudagraphs
242+
repository: "pytorch/tensorrt"
243+
ref: ""
244+
test-infra-repository: pytorch/test-infra
245+
test-infra-ref: main
246+
build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
247+
pre-script: ${{ matrix.pre-script }}
248+
script: |
249+
export USE_HOST_DEPS=1
250+
pushd .
251+
cd tests/py/dynamo
252+
python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_runtime_cudagraphs_cpp_test_results.xml runtime/test_002_cudagraphs_cpp.py
253+
#python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_runtime_cudagraphs_py_test_results.xml runtime/test_002_cudagraphs_py.py
225254
popd
226255
227256
tests-py-core:
@@ -246,7 +275,7 @@ jobs:
246275
export USE_HOST_DEPS=1
247276
pushd .
248277
cd tests/py/core
249-
python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_core_test_results.xml .
278+
python -m pytest -ra -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_core_test_results.xml .
250279
popd
251280
252281
concurrency:

‎core/conversion/var/Var.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@ bool Var::isITensorList() {
153153
// Unpack the Var as a List and check if each entry is a custom class since
154154
// ITensors are stored in CustomClassHolder
155155
auto ival_list = ptr_.ivalue->toList();
156-
for (int i = 0; i < ival_list.size(); i++) {
156+
for (size_t i = 0; i < ival_list.size(); i++) {
157157
if (!ival_list.get(i).isCustomClass()) {
158158
return false;
159159
}
@@ -167,7 +167,7 @@ std::vector<nvinfer1::ITensor*> Var::unwrapToITensorList() {
167167
TORCHTRT_CHECK(isITensorList(), "Expected IValue to be an ITensorList");
168168
auto ivalue_list = ptr_.ivalue->toList();
169169
std::vector<nvinfer1::ITensor*> outputs;
170-
for (int i = 0; i < ivalue_list.size(); i++) {
170+
for (size_t i = 0; i < ivalue_list.size(); i++) {
171171
auto element = ivalue_list.get(i).toCustomClass<TensorContainer>()->tensor();
172172
outputs.push_back(std::move(element));
173173
}

‎core/runtime/TRTEngine.cpp

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -71,15 +71,6 @@ TRTEngine::TRTEngine(
7171
multi_gpu_device_check();
7272
set_rt_device(device_info);
7373

74-
// Set active stream to non-default stream
75-
auto current_stream = c10::cuda::getCurrentCUDAStream(device_info.id);
76-
if (current_stream == c10::cuda::getDefaultCUDAStream(device_info.id)) {
77-
active_stream = c10::cuda::getStreamFromPool(false, device_info.id);
78-
c10::cuda::setCurrentCUDAStream(active_stream);
79-
} else {
80-
active_stream = current_stream;
81-
}
82-
8374
rt = make_trt(nvinfer1::createInferRuntime(util::logging::get_logger()));
8475

8576
name = slugify(mod_name);
@@ -205,6 +196,7 @@ TRTEngine::TRTEngine(
205196
}
206197

207198
TRTEngine::~TRTEngine() {
199+
cudagraph.reset();
208200
trt_engine_profiler.reset();
209201
exec_ctx.reset();
210202
cuda_engine.reset();
@@ -253,6 +245,7 @@ void TRTEngine::set_profiling_paths() {
253245
enqueue_profile_path = std::filesystem::path{profile_path_prefix + "/" + name + "_enqueue_profile.trace"}.string();
254246
trt_engine_profile_path =
255247
std::filesystem::path{profile_path_prefix + "/" + name + "_engine_exectuion_profile.trace"}.string();
248+
cuda_graph_debug_path = std::filesystem::path{profile_path_prefix + "/" + name + "_cudagraph.dot"}.string();
256249
}
257250

258251
std::string TRTEngine::to_str() const {

‎core/runtime/TRTEngine.h

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,8 @@ struct TRTEngine : torch::CustomClassHolder {
7070

7171
// CUDAGraph-Related Functionality
7272
at::cuda::CUDAGraph cudagraph = {};
73-
at::cuda::CUDAStream active_stream = c10::cuda::getDefaultCUDAStream();
73+
at::cuda::CUDAStream engine_stream = c10::cuda::getDefaultCUDAStream();
74+
at::cuda::CUDAStream caller_stream = c10::cuda::getDefaultCUDAStream();
7475
std::vector<at::Tensor> input_buffers = {};
7576
std::vector<at::Tensor> output_buffers = {};
7677
std::string shape_key;
@@ -89,6 +90,7 @@ struct TRTEngine : torch::CustomClassHolder {
8990
std::string output_profile_path;
9091
std::string enqueue_profile_path;
9192
std::string trt_engine_profile_path;
93+
std::string cuda_graph_debug_path;
9294
std::mutex mu;
9395
std::unique_ptr<TRTEngineProfiler> trt_engine_profiler;
9496
};

0 commit comments

Comments
 (0)