From cc658c5f1e065a206555668b4aaa49b568839279 Mon Sep 17 00:00:00 2001 From: Brian Van Essen Date: Sun, 26 Sep 2021 06:50:09 -0700 Subject: [PATCH] lbann: add support for building with the ONNX C++ library (#26130) --- var/spack/repos/builtin/packages/lbann/package.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/var/spack/repos/builtin/packages/lbann/package.py b/var/spack/repos/builtin/packages/lbann/package.py index 2966d51bf8..650351938e 100644 --- a/var/spack/repos/builtin/packages/lbann/package.py +++ b/var/spack/repos/builtin/packages/lbann/package.py @@ -68,6 +68,7 @@ class Lbann(CMakePackage, CudaPackage, ROCmPackage): description='Builds with support for image processing data with OpenCV') variant('vtune', default=False, description='Builds with support for Intel VTune') variant('onednn', default=False, description='Support for OneDNN') + variant('onnx', default=False, description='Support for exporting models into ONNX format') variant('nvshmem', default=False, description='Support for NVSHMEM') variant('python', default=True, description='Support for Python extensions (e.g. Data Reader)') variant('pfe', default=True, description='Python Frontend for generating and launching models') @@ -237,6 +238,7 @@ class Lbann(CMakePackage, CudaPackage, ROCmPackage): depends_on('llvm-openmp', when='%apple-clang') depends_on('onednn cpu_runtime=omp gpu_runtime=none', when='+onednn') + depends_on('onnx', when='+onnx') depends_on('nvshmem', when='+nvshmem') depends_on('zstr') @@ -307,6 +309,7 @@ def cmake_args(self): '-DLBANN_WITH_NVSHMEM:BOOL=%s' % ('+nvshmem' in spec), '-DLBANN_WITH_FFT:BOOL=%s' % ('+fft' in spec), '-DLBANN_WITH_ONEDNN:BOOL=%s' % ('+onednn' in spec), + '-DLBANN_WITH_ONNX:BOOL=%s' % ('+onnx' in spec), '-DLBANN_WITH_EMBEDDED_PYTHON:BOOL=%s' % ('+python' in spec), '-DLBANN_WITH_PYTHON_FRONTEND:BOOL=%s' % ('+pfe' in spec), '-DLBANN_WITH_TBINF=OFF',