Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update CUDA archs in ORT #291

Merged
merged 14 commits into from
Jan 15, 2025
3 changes: 2 additions & 1 deletion cmake/download_onnxruntime.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -28,4 +28,5 @@ if(DEFINED TRITON_ONNXRUNTIME_PACKAGE_URL)

endif(NOT DOWNLOAD_RESULT EQUAL 0)

endif(DEFINED TRITON_ONNXRUNTIME_PACKAGE_URL)
endif(DEFINED TRITON_ONNXRUNTIME_PACKAGE_URL)

14 changes: 5 additions & 9 deletions tools/gen_ort_dockerfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,7 @@ def dockerfile_for_linux(output_file):
if FLAGS.cudnn_home is not None:
ep_flags += ' --cudnn_home "{}"'.format(FLAGS.cudnn_home)
elif target_platform() == "igpu":
ep_flags += ' --cudnn_home "/usr/lib/aarch64-linux-gnu"'
ep_flags += ' --cudnn_home "/usr/include"'
if FLAGS.ort_tensorrt:
ep_flags += " --use_tensorrt"
if FLAGS.ort_version >= "1.12.1":
Expand All @@ -281,9 +281,9 @@ def dockerfile_for_linux(output_file):
ep_flags += (
" --skip_tests --cmake_extra_defines 'onnxruntime_BUILD_UNIT_TESTS=OFF'"
)
cuda_archs = "53;62;72;87"
cuda_archs = "87;101"
mc-nv marked this conversation as resolved.
Show resolved Hide resolved
else:
cuda_archs = "75;80;86;90"
cuda_archs = "75;80;86;89;90;100;120"

df += """
WORKDIR /workspace/onnxruntime
Expand Down Expand Up @@ -478,7 +478,7 @@ def dockerfile_for_windows(output_file):
WORKDIR /workspace/onnxruntime
ARG VS_DEVCMD_BAT="\BuildTools\VC\Auxiliary\Build\vcvars64.bat"
RUN powershell Set-Content 'build.bat' -value 'call %VS_DEVCMD_BAT%',(Get-Content 'build.bat')
RUN build.bat --cmake_generator "Visual Studio 17 2022" --config Release --cmake_extra_defines "CMAKE_CUDA_ARCHITECTURES=75;80;86;90" --skip_submodule_sync --parallel --build_shared_lib --compile_no_warning_as_error --skip_tests --update --build --build_dir /workspace/build {}
RUN build.bat --cmake_generator "Visual Studio 17 2022" --config Release --cmake_extra_defines "CMAKE_CUDA_ARCHITECTURES=75;80;86;90;100;120" --skip_submodule_sync --parallel --build_shared_lib --compile_no_warning_as_error --skip_tests --update --build --build_dir /workspace/build {}
""".format(
ep_flags
)
Expand Down Expand Up @@ -559,12 +559,8 @@ def preprocess_gpu_flags():
FLAGS.tensorrt_home = "/tensorrt"
else:
if "CUDNN_VERSION" in os.environ:
version = None
m = re.match(r"([0-9]\.[0-9])\.[0-9]\.[0-9]", os.environ["CUDNN_VERSION"])
if m:
version = m.group(1)
if FLAGS.cudnn_home is None:
FLAGS.cudnn_home = "/usr/local/cudnn-{}/cuda".format(version)
FLAGS.cudnn_home = "/usr"

if FLAGS.cuda_home is None:
FLAGS.cuda_home = "/usr/local/cuda"
Expand Down
Loading