Skip to content

Commit 4994ccc

Browse files
guschmueguschmue
authored
python package for webgpu (#26477)
1 parent 423a03f commit 4994ccc

File tree

7 files changed

+374
-2
lines changed

7 files changed

+374
-2
lines changed

‎setup.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,9 @@ def parse_arg_remove_string(argv, arg_name_equal):
9191
is_qnn = True
9292
package_name = "onnxruntime-qnn"
9393
qnn_version = parse_arg_remove_string(sys.argv, "--qnn_version=")
94+
elif parse_arg_remove_boolean(sys.argv, "--use_webgpu"):
95+
package_name = "onnxruntime-webgpu"
96+
9497

9598
# PEP 513 defined manylinux1_x86_64 and manylinux1_i686
9699
# PEP 571 defined manylinux2010_x86_64 and manylinux2010_i686

‎tools/ci_build/build.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1939,6 +1939,7 @@ def build_python_wheel(
19391939
use_acl,
19401940
use_armnn,
19411941
use_dml,
1942+
use_webgpu,
19421943
use_cann,
19431944
use_azure,
19441945
use_qnn,
@@ -1996,6 +1997,8 @@ def build_python_wheel(
19961997
args.append("--use_armnn")
19971998
elif use_dml:
19981999
args.append("--wheel_name_suffix=directml")
2000+
elif use_webgpu:
2001+
args.append("--wheel_name_suffix=webgpu")
19992002
elif use_cann:
20002003
args.append("--use_cann")
20012004
elif use_qnn:
@@ -2632,6 +2635,7 @@ def main():
26322635
args.use_acl,
26332636
args.use_armnn,
26342637
args.use_dml,
2638+
args.use_webgpu,
26352639
args.use_cann,
26362640
args.use_azure,
26372641
args.use_qnn,

‎tools/ci_build/github/azure-pipelines/py-webgpu-packaging-pipeline.yml‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
trigger: none
2+
# The `resources` specify the location and version of the 1ES PT.
3+
resources:
4+
repositories:
5+
- repository: 1esPipelines
6+
type: git
7+
name: 1ESPipelineTemplates/1ESPipelineTemplates
8+
ref: refs/tags/release
9+
parameters:
10+
- name: cmake_build_type
11+
type: string
12+
default: 'Release'
13+
values:
14+
- Debug
15+
- Release
16+
- RelWithDebInfo
17+
- MinSizeRel
18+
19+
extends:
20+
# The pipeline extends the 1ES PT which will inject different SDL and compliance tasks.
21+
# For non-production pipelines, use "Unofficial" as defined below.
22+
# For productions pipelines, use "Official".
23+
template: v1/1ES.Official.PipelineTemplate.yml@1esPipelines
24+
parameters:
25+
sdl:
26+
componentgovernance:
27+
ignoreDirectories: '$(Build.Repository.LocalPath)/cmake/external/emsdk/upstream/emscripten/tests,$(Build.Repository.LocalPath)/cmake/external/onnx/third_party/benchmark,$(Build.Repository.LocalPath)/cmake/external/onnx/third_party/pybind11,$(Build.Repository.LocalPath)/cmake/external/onnx/third_party/pybind11/tests,$(Build.Repository.LocalPath)/cmake/external/onnxruntime-extensions,$(Build.Repository.LocalPath)/js/react_native/e2e/node_modules,$(Build.Repository.LocalPath)/js/node_modules,$(Build.Repository.LocalPath)/onnxruntime-inference-examples,$(Build.SourcesDirectory)/cmake/external/emsdk/upstream/emscripten/tests,$(Build.SourcesDirectory)/cmake/external/onnx/third_party/benchmark,$(Build.SourcesDirectory)/cmake/external/onnx/third_party/pybind11,$(Build.SourcesDirectory)/cmake/external/onnx/third_party/pybind11/tests,$(Build.SourcesDirectory)/cmake/external/onnxruntime-extensions,$(Build.SourcesDirectory)/js/react_native/e2e/node_modules,$(Build.SourcesDirectory)/js/node_modules,$(Build.SourcesDirectory)/onnxruntime-inference-examples,$(Build.BinariesDirectory)'
28+
alertWarningLevel: High
29+
failOnAlert: false
30+
verbosity: Normal
31+
timeout: 3600
32+
tsa:
33+
enabled: true
34+
codeSignValidation:
35+
enabled: true
36+
break: true
37+
policheck:
38+
enabled: true
39+
exclusionsFile: '$(Build.SourcesDirectory)\tools\ci_build\policheck_exclusions.xml'
40+
codeql:
41+
compiled:
42+
enabled: false
43+
justificationForDisabling: 'CodeQL is taking nearly 6 hours resulting in timeouts in our production pipelines'
44+
pool:
45+
name: 'onnxruntime-Win-CPU-VS2022-Latest' # Name of your hosted pool
46+
os: windows # OS of the image. This value cannot be a variable. Allowed values: windows, linux, macOS
47+
48+
stages:
49+
- template: stages/py-webgpu-packaging-stage.yml
50+
parameters:
51+
cmake_build_type: ${{ parameters.cmake_build_type }}

‎tools/ci_build/github/azure-pipelines/stages/py-linux-webgpu-stage.yml‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
parameters:
2+
- name: arch
3+
type: string
4+
5+
- name: machine_pool
6+
type: string
7+
8+
- name: extra_build_arg
9+
type: string
10+
default: ''
11+
12+
- name: cmake_build_type
13+
type: string
14+
default: 'Release'
15+
values:
16+
- Debug
17+
- Release
18+
- RelWithDebInfo
19+
- MinSizeRel
20+
21+
- name: docker_base_image
22+
type: string
23+
24+
stages:
25+
- stage: Linux_py_webgpu_Wheels_${{ parameters.arch }}
26+
dependsOn: []
27+
jobs:
28+
- job: Linux_py_webgpu_Wheels_${{ parameters.arch }}
29+
timeoutInMinutes: 240
30+
workspace:
31+
clean: all
32+
pool:
33+
name: ${{ parameters.machine_pool }}
34+
os: linux
35+
templateContext:
36+
outputs:
37+
- output: pipelineArtifact
38+
targetPath: $(Build.ArtifactStagingDirectory)/dist
39+
artifactName: onnxruntime_webgpu
40+
- output: pipelineArtifact
41+
targetPath: $(Build.ArtifactStagingDirectory)/${{ parameters.cmake_build_type }}
42+
artifactName: linux_webgpu_wheel_${{ parameters.arch }}
43+
variables:
44+
- name: extra_build_args
45+
${{ if ne(parameters.extra_build_arg, '') }}:
46+
value: -x ${{ parameters.extra_build_arg }}
47+
${{ if eq(parameters.extra_build_arg, '') }}:
48+
value: ''
49+
- template: ../templates/common-variables.yml
50+
steps:
51+
- checkout: self
52+
clean: true
53+
submodules: recursive
54+
55+
- template: ../templates/set-nightly-build-option-variable-step.yml
56+
57+
- template: ../templates/get-docker-image-steps.yml
58+
parameters:
59+
Dockerfile: tools/ci_build/github/linux/docker/inference/${{ parameters.arch }}/python/cuda/Dockerfile
60+
Context: tools/ci_build/github/linux/docker/inference/${{ parameters.arch }}/python/cuda
61+
DockerBuildArgs: "--build-arg BASEIMAGE=${{ parameters.docker_base_image }} --build-arg BUILD_UID=$( id -u )"
62+
Repository: onnxruntimewebgpu
63+
64+
- task: Bash@3
65+
displayName: 'Build Python Wheel'
66+
inputs:
67+
targetType: filePath
68+
filePath: tools/ci_build/github/linux/run_python_dockerbuild.sh
69+
arguments: -i onnxruntimewebgpu -d "WEBGPU" -c ${{ parameters.cmake_build_type }} $(extra_build_args)
70+
71+
- script: |
72+
set -e -x
73+
mv $(Build.BinariesDirectory)/${{ parameters.cmake_build_type }} ./${{ parameters.cmake_build_type }}
74+
mv $(Build.BinariesDirectory)/dist ./dist
75+
pushd dist
76+
find . -name \*.whl -exec unzip -qq -o {} \;
77+
rm -r onnxruntime
78+
popd
79+
pushd ${{ parameters.cmake_build_type }}
80+
find . -name \*.whl -exec unzip -qq -o {} \;
81+
popd
82+
workingDirectory: '$(Build.ArtifactStagingDirectory)'
83+
displayName: 'Move files'

‎tools/ci_build/github/azure-pipelines/stages/py-webgpu-packaging-stage.yml‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
parameters:
2+
- name: build_py_parameters
3+
displayName: >
4+
Extra parameters to pass to build.py. Don't put newlines in here.
5+
type: string
6+
default: ''
7+
8+
# TODO: Now the Windows jobs use a different cmake build type. Consider to merge it.
9+
- name: cmake_build_type
10+
type: string
11+
displayName: 'Linux packages cmake build type. Linux Only.'
12+
default: 'Release'
13+
values:
14+
- Debug
15+
- Release
16+
- RelWithDebInfo
17+
- MinSizeRel
18+
19+
- name: PythonVersions
20+
type: object
21+
displayName: 'Python versions to build'
22+
default:
23+
- '3.11'
24+
- '3.12'
25+
- '3.13'
26+
27+
stages:
28+
- ${{ each python_version in parameters.PythonVersions }}:
29+
- template: py-win-webgpu-stage.yml
30+
parameters:
31+
PYTHON_VERSION: ${{ python_version }}
32+
EP_BUILD_FLAGS: --enable_lto --use_webgpu
33+
34+
- ${{ each python_version in parameters.PythonVersions }}:
35+
- stage: MacOS_py_webgpu_Wheels_${{ replace(python_version,'.','_') }}_Build
36+
dependsOn: []
37+
jobs:
38+
- template: ../templates/py-macos.yml
39+
parameters:
40+
arch: 'arm64'
41+
python_version: ${{ python_version }}
42+
extra_build_arg: --use_webgpu ${{ parameters.build_py_parameters }}
43+
cmake_build_type: ${{ parameters.cmake_build_type }}
44+
45+
- template: py-linux-webgpu-stage.yml
46+
parameters:
47+
arch: 'x86_64'
48+
machine_pool: 'onnxruntime-Ubuntu2404-AMD-CPU'
49+
extra_build_arg: ${{ parameters.build_py_parameters }}
50+
cmake_build_type: ${{ parameters.cmake_build_type }}
51+
docker_base_image: onnxruntimebuildcache.azurecr.io/internal/azureml/onnxruntime/build/cuda12_x64_almalinux8_gcc14:20251017.1

0 commit comments

Comments
 (0)