Add TMVA SOFIE benchmarks

.github/workflows/ci.yaml

@@ -36,7 +36,7 @@ jobs:
     - name: Conda install dependencies
       shell: bash -l {0}
       run: |
-           conda create -n rootbench -y -c conda-forge root cmake pytest pytest-benchmark pytest-csv numpy numba
+           conda create -n rootbench -y -c conda-forge root cmake pytest pytest-benchmark pytest-csv numpy numba onnx
 
     - name: Configure and build
       shell: bash -l {0}

cmake/modules/FindONNXRuntime.cmake

@@ -0,0 +1,47 @@
+# Copyright (C) 1995-2019, Rene Brun and Fons Rademakers.
+# All rights reserved.
+#
+# For the licensing terms see $ROOTSYS/LICENSE.
+# For the list of contributors see $ROOTSYS/README/CREDITS.
+
+# Find the ONNXRuntime includes and library.
+#
+# This module defines
+# ONNXRuntime_INCLUDE_DIR, where to locate ONNXRuntime include file
+# ONNXRuntime_LIBRARIES, the libraries to link against to use ONNXRuntime
+# ONNXRuntime_FOUND.  If false, you cannot build anything that requires ONNXRuntime.
+# ONNXRuntime_LIBRARY, where to find the libONNXRuntime library.
+
+set(ONNXRuntime_FOUND 0)
+if(ONNXRuntime_LIBRARY AND ONNXRuntime_INCLUDE_DIR)
+  set(ONNXRuntime_FIND_QUIETLY TRUE)
+endif()
+
+find_path(ONNXRuntime_INCLUDE_DIR onnxruntime_cxx_api.h
+  $ENV{ONNXRuntime_DIR}/include
+  $ENV{ONNXRuntime} $ENV{ONNXRuntime}/include 
+  /usr/local/include
+  /usr/include
+  DOC "Specify the directory containing ONNXRuntime.h"
+)
+
+find_library(ONNXRuntime_LIBRARY NAMES onnxruntime  PATHS
+  $ENV{ONNXRuntime_DIR}/lib
+  $ENV{ONNXRuntime} $ENV{ONNXRuntime}/lib $ENV{ONNXRuntime}/.libs
+  /usr/local/lib
+  /usr/lib
+  /opt/ONNXRuntime/lib
+  DOC "Specify the ONNXRuntime library here."
+)
+
+if(ONNXRuntime_INCLUDE_DIR AND ONNXRuntime_LIBRARY)
+  set(ONNXRuntime_FOUND 1 )
+  if(NOT ONNXRuntime_FIND_QUIETLY)
+     message(STATUS "Found ONNXRuntime includes at ${ONNXRuntime_INCLUDE_DIR}")
+     message(STATUS "Found ONNXRuntime library at ${ONNXRuntime_LIBRARY}")
+  endif()
+endif()
+
+set(ONNXRuntime_LIBRARIES ${ONNXRuntime_LIBRARY})
+
+mark_as_advanced(ONNXRuntime_FOUND ONNXRuntime_LIBRARY ONNXRuntime_INCLUDE_DIR)

root/tmva/CMakeLists.txt

@@ -1 +1,2 @@
 add_subdirectory(tmva)
+add_subdirectory(sofie)

root/tmva/sofie/CMakeLists.txt

@@ -0,0 +1,160 @@
+# TMVA SOFIE inference benchmarks.
+# @author Federico Sossai (fsossai), Lorenzo Moneta
+
+# Check if SOFIE is available. ROOT versions up to 6.40 have a dedicated
+# tmva-sofie build option that is advertised as a ROOT feature, while in
+# later versions SOFIE is built unconditionally with TMVA and we probe for
+# its libraries directly: in ROOT-builtin builds the library targets exist,
+# and in standalone builds the libraries are found in the ROOT installation.
+set(RB_HAVE_SOFIE FALSE)
+if(ROOT_tmva-sofie_FOUND)
+  set(RB_HAVE_SOFIE TRUE)
+elseif(TARGET ROOTTMVASofie AND TARGET ROOTTMVASofieParser)
+  set(RB_HAVE_SOFIE TRUE)
+else()
+  find_library(RB_SOFIE_LIBRARY ROOTTMVASofie HINTS ${ROOT_LIBRARY_DIR})
+  find_library(RB_SOFIE_PARSER_LIBRARY ROOTTMVASofieParser HINTS ${ROOT_LIBRARY_DIR})
+  if(RB_SOFIE_LIBRARY AND RB_SOFIE_PARSER_LIBRARY)
+    set(RB_HAVE_SOFIE TRUE)
+  endif()
+endif()
+if(NOT (ROOT_tmva_FOUND AND RB_HAVE_SOFIE))
+  message(STATUS "TMVA SOFIE not found: disabling the SOFIE benchmarks")
+  return()
+endif()
+
+# The code generated by SOFIE uses BLAS for the matrix operations.
+find_package(BLAS)
+if(NOT BLAS_FOUND)
+  message(STATUS "BLAS not found: disabling the SOFIE benchmarks")
+  return()
+endif()
+
+# The ONNX input models are generated at build time by make_input_models.py,
+# which needs Python with the onnx and numpy packages.
+find_package(Python3 COMPONENTS Interpreter)
+if(Python3_FOUND)
+  execute_process(COMMAND ${Python3_EXECUTABLE} -c "import onnx, numpy"
+                  RESULT_VARIABLE onnx_missing OUTPUT_QUIET ERROR_QUIET)
+endif()
+if(NOT Python3_FOUND OR onnx_missing)
+  message(STATUS "Python with the onnx package not found: disabling the SOFIE benchmarks")
+  return()
+endif()
+
+# The models to benchmark. For each of them, an ONNX file is generated with
+# make_input_models.py, from which emitFromONNX then generates the inference
+# code compiled into the benchmarks. Everything happens at build time, so
+# the benchmarks always exercise the SOFIE version of the ROOT build they
+# run against.
+set(sofie_models
+  Conv3d_d32_L4_B1
+  ConvTModel_G4
+  ConvTrans2dModel_B1
+  Conv_d100_L14_B1
+  Conv_d100_L14_B32
+  Conv_d100_L1_B1
+  Generator_B1
+  Generator_B64
+  Linear_16
+  Linear_32
+  Linear_64
+  Linear_event
+  SimpleNN_Alice
+  higgs_model_dense)
+
+# Command line tool that generates the inference code for an ONNX model.
+add_executable(emitFromONNX EmitFromONNX.cxx)
+target_link_libraries(emitFromONNX Core ROOTTMVASofie ROOTTMVASofieParser)
+set_target_properties(emitFromONNX PROPERTIES POSITION_INDEPENDENT_CODE TRUE)
+
+set(model_dir ${CMAKE_CURRENT_BINARY_DIR}/input_models)
+set(model_generator ${CMAKE_CURRENT_SOURCE_DIR}/make_input_models.py)
+
+set(sofie_headers "")
+foreach(name ${sofie_models})
+  add_custom_command(
+    OUTPUT ${model_dir}/${name}.onnx
+    COMMAND ${Python3_EXECUTABLE} ${model_generator} --outdir ${model_dir} ${name}
+    DEPENDS ${model_generator}
+    COMMENT "Generating ONNX model ${name}")
+  add_custom_command(
+    OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/${name}.hxx ${CMAKE_CURRENT_BINARY_DIR}/${name}.dat
+    COMMAND emitFromONNX ${model_dir}/${name}.onnx
+    DEPENDS emitFromONNX ${model_dir}/${name}.onnx
+    WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}
+    COMMENT "Generating SOFIE inference code for ${name}")
+  list(APPEND sofie_headers ${CMAKE_CURRENT_BINARY_DIR}/${name}.hxx)
+endforeach()
+add_custom_target(SofieCompileModels DEPENDS ${sofie_headers})
+
+# Benchmark of the inference code emitted by SOFIE
+RB_ADD_GBENCHMARK(SOFIEInference
+  SOFIEInference.cxx
+  LABEL short
+  DEPENDS SofieCompileModels
+  LIBRARIES Core MathCore ROOTTMVASofie ${BLAS_LIBRARIES})
+
+# Benchmark of RSofieReader, which parses the model and JITs the generated
+# code at runtime
+RB_ADD_GBENCHMARK(SOFIEInference_Reader
+  SOFIEInference_Reader.cxx
+  LABEL short
+  DEPENDS SofieCompileModels
+  LIBRARIES Core Cling MathCore ROOTTMVASofie ${BLAS_LIBRARIES})
+
+# Benchmark of SOFIE inference inside an RDataFrame event loop
+RB_ADD_GBENCHMARK(RDF_SOFIE_Inference
+  RDF_SOFIE_Inference.cxx
+  LABEL short
+  DEPENDS SofieCompileModels
+  LIBRARIES Core Hist Imt RIO Tree TreePlayer ROOTDataFrame ROOTVecOps ROOTTMVASofie ${BLAS_LIBRARIES})
+
+# Compile the benchmarks with -O3 so that the generated inference code is
+# auto-vectorized like in an optimized user build. More aggressive options
+# (-march=native, -ffast-math) are deliberately not used: they would make the
+# results machine-dependent and change the numerical behavior of operators
+# that rely on infinities.
+target_compile_options(SOFIEInference PRIVATE -O3)
+target_compile_options(RDF_SOFIE_Inference PRIVATE -O3)
+
+# Optional comparison benchmark using ONNXRuntime on the same models. To help
+# CMake find ONNXRuntime, configure with
+#   -DONNXRuntime_INCLUDE_DIRS=<location>/include -DONNXRuntime_LIBRARIES=<location>/lib
+find_package(ONNXRuntime)
+if(ONNXRuntime_FOUND)
+  message(STATUS "Found ONNXRuntime (library is ${ONNXRuntime_LIBRARY}, libraries ${ONNXRuntime_LIBRARIES})")
+
+  # Generate one benchmark registration per model from
+  # ONNXRuntimeInference_Template.cxx.in
+  set(FUNC_NAME "BM_ONNXRuntime_Inference")
+  set(CAPTURE_STR "BENCHMARK_CAPTURE(${FUNC_NAME}, @1,\t@2)@3")
+  set(HEAD_COMMENT "Automatically configured by CMake")
+  set(ALL_CAPTURES "")
+  foreach(name ${sofie_models})
+    string(REPLACE "@1" ${name} cap ${CAPTURE_STR})
+    string(REPLACE "@2" "\"input_models/${name}.onnx\"" cap ${cap})
+    list(APPEND ALL_CAPTURES ${cap})
+  endforeach()
+  string(REPLACE ";" "\n" BENCHMARK_CAPTURES "${ALL_CAPTURES}")
+  string(REPLACE "@3" "->Unit(benchmark::kMillisecond);" BENCHMARK_CAPTURES "${BENCHMARK_CAPTURES}")
+  configure_file(ONNXRuntimeInference_Template.cxx.in ONNXRuntimeInference.cxx @ONLY)
+
+  RB_ADD_GBENCHMARK(ONNXRuntimeInference
+    ONNXRuntimeInference.cxx
+    LABEL short
+    DEPENDS SofieCompileModels
+    LIBRARIES Core ${ONNXRuntime_LIBRARIES})
+  target_link_directories(ONNXRuntimeInference PRIVATE ${ONNXRuntime_LIBRARIES})
+  target_include_directories(ONNXRuntimeInference PRIVATE ${ONNXRuntime_INCLUDE_DIR})
+
+  RB_ADD_GBENCHMARK(RDF_ONNXRuntime_Inference
+    RDF_ONNXRuntime_Inference.cxx
+    LABEL short
+    DEPENDS SofieCompileModels
+    LIBRARIES Core Hist Imt MathCore RIO Tree TreePlayer ROOTDataFrame ROOTVecOps ${ONNXRuntime_LIBRARIES})
+  target_link_directories(RDF_ONNXRuntime_Inference PRIVATE ${ONNXRuntime_LIBRARIES})
+  target_include_directories(RDF_ONNXRuntime_Inference PRIVATE ${ONNXRuntime_INCLUDE_DIR})
+else()
+  message(STATUS "ONNXRuntime not found: disabling the ONNXRuntime benchmarks")
+endif()

root/tmva/sofie/EmitFromONNX.cxx

@@ -0,0 +1,29 @@
+// Author: Federico Sossai
+// Last modified: 2021/07/30
+// Description:
+//    SOFIE command line compiler.
+//    This program is automatically run when the corresponding test target is built.
+//    Usage example: $./EmitFromONNX indir/mymodel.onnx outdir/myname.hxx
+
+#include <iostream>
+
+#include "TMVA/RModel.hxx"
+#include "TMVA/RModelParser_ONNX.hxx"
+
+using namespace TMVA::Experimental::SOFIE;
+
+int main(int argc, char *argv[]){
+   if (argc < 2) {
+      std::cerr << "ERROR: missing input file\n";
+      return -1;
+   }
+
+   std::string outname= (argc > 2) ? argv[2] : "";
+   RModelParser_ONNX parser;
+   std::cout << "Parsing file " << argv[1] << std::endl;
+   RModel model = parser.Parse(argv[1]);
+   model.Generate(Options::kDefault, 1);
+   model.PrintRequiredInputTensors();
+   model.OutputGenerated(outname);
+   return 0;
+}