lbann: add support for building with the ONNX C++ library (#26130)

This commit is contained in:
Brian Van Essen 2021-09-26 06:50:09 -07:00 committed by GitHub
parent 1f972f90c8
commit cc658c5f1e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -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',