Skip to content

Commit 10d616c

Browse files
GasoonjiaGasoonjia
committed
[executorch][PR] add cuda backend to backend test infra
Pull Request resolved: #17873 Integrate cuda backend into backend test infra; skipped the unsupported test for now ghstack-source-id: 349755553 @exported-using-ghexport Differential Revision: [D93019490](https://our.internmc.facebook.com/intern/diff/D93019490/)
1 parent f09bd55 commit 10d616c

8 files changed

Lines changed: 302 additions & 97 deletions

File tree

‎.ci/scripts/test_backend.sh‎

Lines changed: 33 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,32 @@ if [[ "$FLOW" == *vulkan* ]]; then
5656
EXTRA_BUILD_ARGS+=" -DEXECUTORCH_BUILD_VULKAN=ON"
5757
fi
5858

59+
if [[ "$FLOW" == *cuda* ]]; then
60+
# When running with the PyTorch test-infra Docker image (which has nvcc),
61+
# install executorch directly — it will auto-detect CUDA and install
62+
# CUDA-enabled PyTorch. Skip setup-linux.sh which expects the custom
63+
# Docker image with pre-built pinned-commit torch.
64+
echo "Installing ExecuTorch with CUDA support..."
65+
./install_executorch.sh --editable
66+
67+
# Verify PyTorch was installed with CUDA support
68+
python -c "import torch; assert torch.cuda.is_available(), 'PyTorch CUDA not available after reinstall'; print(f'PyTorch {torch.__version__} with CUDA {torch.version.cuda}')" || {
69+
echo "ERROR: PyTorch was not installed with CUDA support"
70+
exit 1
71+
}
72+
73+
# Fix libstdc++ GLIBCXX version for CUDA backend.
74+
# The embedded .so files in the CUDA blob require GLIBCXX_3.4.30
75+
# which the default conda libstdc++ doesn't have.
76+
echo "Installing newer libstdc++ for CUDA backend..."
77+
conda install -y -c conda-forge 'libstdcxx-ng>=12'
78+
export LD_LIBRARY_PATH="${CONDA_PREFIX}/lib:${LD_LIBRARY_PATH:-}"
79+
80+
source .ci/scripts/utils.sh
81+
CMAKE_ARGS="$EXTRA_BUILD_ARGS" build_executorch_runner cmake Release
82+
CUDA_SETUP_DONE=1
83+
fi
84+
5985
if [[ "$FLOW" == *arm* ]]; then
6086

6187
# Setup ARM deps.
@@ -78,12 +104,14 @@ if [[ "$FLOW" == *arm* ]]; then
78104
fi
79105
fi
80106

81-
if [[ $IS_MACOS -eq 1 ]]; then
82-
SETUP_SCRIPT=.ci/scripts/setup-macos.sh
83-
else
84-
SETUP_SCRIPT=.ci/scripts/setup-linux.sh
107+
if [[ "${CUDA_SETUP_DONE:-0}" != "1" ]]; then
108+
if [[ $IS_MACOS -eq 1 ]]; then
109+
SETUP_SCRIPT=.ci/scripts/setup-macos.sh
110+
else
111+
SETUP_SCRIPT=.ci/scripts/setup-linux.sh
112+
fi
113+
CMAKE_ARGS="$EXTRA_BUILD_ARGS" ${CONDA_RUN_CMD} $SETUP_SCRIPT --build-tool cmake --build-mode Release --editable true
85114
fi
86-
CMAKE_ARGS="$EXTRA_BUILD_ARGS" ${CONDA_RUN_CMD} $SETUP_SCRIPT --build-tool cmake --build-mode Release --editable true
87115

88116
GOLDEN_DIR="${ARTIFACT_DIR}/golden-artifacts"
89117
export GOLDEN_ARTIFACTS_DIR="${GOLDEN_DIR}"

‎.github/workflows/test-backend-cuda.yml‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
name: Test CUDA Backend
2+
3+
on:
4+
schedule:
5+
- cron: 0 2 * * *
6+
push:
7+
branches:
8+
- release/*
9+
tags:
10+
- ciflow/nightly/*
11+
pull_request:
12+
paths:
13+
- .github/workflows/test-backend-cuda.yml
14+
- .ci/scripts/test_backend.sh
15+
workflow_dispatch:
16+
17+
concurrency:
18+
group: ${{ github.workflow }}--${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
19+
cancel-in-progress: true
20+
21+
jobs:
22+
test-cuda:
23+
strategy:
24+
fail-fast: false
25+
matrix:
26+
suite: [models, operators]
27+
28+
uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
29+
with:
30+
ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
31+
runner: linux.g5.4xlarge.nvidia.gpu
32+
gpu-arch-type: cuda
33+
gpu-arch-version: '12.6'
34+
use-custom-docker-registry: false
35+
submodules: recursive
36+
timeout: 120
37+
upload-artifact: test-report-cuda-${{ matrix.suite }}
38+
script: |
39+
set -eux
40+
41+
source .ci/scripts/test_backend.sh "${{ matrix.suite }}" "cuda" "${RUNNER_ARTIFACT_DIR}"

‎backends/cuda/test/tester.py‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
from typing import Any, List, Optional, Tuple
8+
9+
import executorch
10+
import executorch.backends.test.harness.stages as BaseStages
11+
import torch
12+
from executorch.backends.cuda.cuda_backend import CudaBackend
13+
from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
14+
from executorch.backends.test.harness import Tester as TesterBase
15+
from executorch.backends.test.harness.stages import StageType
16+
from executorch.exir import EdgeCompileConfig
17+
from executorch.exir.backend.partitioner import Partitioner
18+
19+
20+
def _create_default_partitioner() -> CudaPartitioner:
21+
"""Create a CudaPartitioner with default compile specs."""
22+
compile_specs = [CudaBackend.generate_method_name_compile_spec("forward")]
23+
return CudaPartitioner(compile_specs)
24+
25+
26+
class ToEdgeTransformAndLower(BaseStages.ToEdgeTransformAndLower):
27+
"""CUDA-specific ToEdgeTransformAndLower stage."""
28+
29+
def __init__(
30+
self,
31+
partitioners: Optional[List[Partitioner]] = None,
32+
edge_compile_config: Optional[EdgeCompileConfig] = None,
33+
):
34+
if partitioners is None:
35+
partitioners = [_create_default_partitioner()]
36+
37+
super().__init__(
38+
default_partitioner_cls=_create_default_partitioner,
39+
partitioners=partitioners,
40+
edge_compile_config=edge_compile_config
41+
or EdgeCompileConfig(_check_ir_validity=False),
42+
)
43+
44+
45+
class CudaTester(TesterBase):
46+
"""
47+
Tester subclass for CUDA backend.
48+
49+
This tester defines the recipe for lowering models to the CUDA backend
50+
using AOTInductor compilation.
51+
"""
52+
53+
def __init__(
54+
self,
55+
module: torch.nn.Module,
56+
example_inputs: Tuple[torch.Tensor],
57+
dynamic_shapes: Optional[Tuple[Any]] = None,
58+
):
59+
stage_classes = (
60+
executorch.backends.test.harness.Tester.default_stage_classes()
61+
| {
62+
StageType.TO_EDGE_TRANSFORM_AND_LOWER: ToEdgeTransformAndLower,
63+
}
64+
)
65+
66+
super().__init__(
67+
module=module,
68+
stage_classes=stage_classes,
69+
example_inputs=example_inputs,
70+
dynamic_shapes=dynamic_shapes,
71+
)

‎backends/test/harness/stages/serialize.py‎

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,9 @@
11
import copy
22
import logging
3-
4-
from typing import Optional
3+
from typing import Dict, Optional
54

65
from executorch.backends.test.harness.stages.stage import Stage, StageType
76
from executorch.exir import ExecutorchProgramManager
8-
97
from torch.utils._pytree import tree_flatten
108

119
logger = logging.getLogger(__name__)
@@ -23,12 +21,15 @@
2321
class Serialize(Stage):
2422
def __init__(self):
2523
self.buffer = None
24+
self.data_files: Dict[str, bytes] = {}
2625

2726
def stage_type(self) -> StageType:
2827
return StageType.SERIALIZE
2928

3029
def run(self, artifact: ExecutorchProgramManager, inputs=None) -> None:
3130
self.buffer = artifact.buffer
31+
# Capture external data files (e.g., .ptd files for CUDA backend)
32+
self.data_files = artifact.data_files
3233

3334
@property
3435
def artifact(self) -> bytes:
@@ -40,8 +41,29 @@ def graph_module(self) -> None:
4041

4142
def run_artifact(self, inputs):
4243
inputs_flattened, _ = tree_flatten(inputs)
44+
45+
# Combine all external data files into a single buffer for data_map_buffer
46+
# Most backends have at most one external data file, but we concatenate
47+
# in case there are multiple (though this may not be fully supported)
48+
data_map_buffer = None
49+
if self.data_files:
50+
# If there's exactly one data file, use it directly
51+
# Otherwise, log a warning - multiple external files may need special handling
52+
if len(self.data_files) == 1:
53+
data_map_buffer = list(self.data_files.values())[0]
54+
else:
55+
# For multiple files, we use the first one and warn
56+
# This is a limitation - proper handling would need runtime support
57+
logger.warning(
58+
f"Multiple external data files found ({list(self.data_files.keys())}). "
59+
f"Using the first one. This may not work correctly for all backends."
60+
)
61+
data_map_buffer = list(self.data_files.values())[0]
62+
4363
executorch_module = _load_for_executorch_from_buffer(
44-
self.buffer, program_verification=Verification.Minimal
64+
self.buffer,
65+
data_map_buffer=data_map_buffer,
66+
program_verification=Verification.Minimal,
4567
)
4668
executorch_output = copy.deepcopy(
4769
executorch_module.run_method("forward", tuple(inputs_flattened))

‎backends/test/suite/conftest.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33

44
import pytest
55
import torch
6-
76
from executorch.backends.test.suite.flow import all_flows, TestFlow
87
from executorch.backends.test.suite.reporting import _sum_op_counts
98
from executorch.backends.test.suite.runner import run_test
@@ -103,7 +102,14 @@ def lower_and_run_model(
103102
ids=str,
104103
)
105104
def test_runner(request):
106-
return TestRunner(request.param, request.node.name, request.node.originalname)
105+
flow = request.param
106+
test_name = request.node.name
107+
108+
# Check if this test should be skipped based on the flow's skip_patterns
109+
if flow.should_skip_test(test_name):
110+
pytest.skip(f"Test '{test_name}' matches skip pattern for flow '{flow.name}'")
111+
112+
return TestRunner(flow, test_name, request.node.originalname)
107113

108114

109115
@pytest.hookimpl(optionalhook=True)

0 commit comments

Comments
 (0)