Package: compilerinterface
ML-Compiler-BridgeML-Compiler-Bridgeis a compiler agnostic library to aid in ML-Enabled Compiler Optimizations. ML-Compiler-Bridge supports both training and inference scenarios. Library exposes Python and C/C++ APIs to interface
with the Python-based ML models and a C/C++ compiler. This design allows ML model development within a traditional Python framework while making end-to-end integration with an optimizing compiler possible and efficient.This repo contains the source code and relevant information described in our paper,"The Next 700 ML-Enabled Compiler Optimizations"(arxiv).
Please seeherefor documentation and other details.The Next 700 ML-Enabled Compiler Optimizations, S. VenkataKeerthy, Siddharth Jain, Umesh Kalvakuntla, Pranav Sai Gorantla, Rajiv Shailesh Chitale, Eugene Brevdo, Albert Cohen, Mircea Trofin and Ramakrishna Upadrasta. CC 2024.FeaturesUnified Framework:Comes with a suite of two inter-process and two in-process model runners and three serialization-deserialization mechanisms to support interleaved and non-interleaved communication between models and compiler.Multi-language Support:Exposes C++ and C APIs to interface model runners and serializers with the compilers and Python APIs to interface inter-process model runners with ML models.Compiler and ML-Framework Independence:Provides compiler and ML-Framework independent APIs, and supports easier integration with compilers like LLVM, MLIR, and Pluto and ML Frameworks like TensorFlow, PyTorch, JAX, etc.Deeper Integration:Enables deeper integration of ML models within the compiler in a framework-independent manner to support easier inference in case of ML driven compiler optimizations.Requirementscmake (>= 3.10)GNU Make (4.2.1)LLVM (10.X) -src,releasePython (3.10), C++17gRPC v1.34 and protobuf v3.13 - for gRPC Model RunnerBuilding GRPC from Source: Please followBuild GRPC with cmakev1.34 (protobuf v3.13) to build GRPC from source.In the above tutorial settingDCMAKE_INSTALL_PREFIXmay not be necessary and the default install prefix can be used.The following dependencies will be required for Python:pip install grpcio-tools.ONNXRuntimev1.13.1TensorFlow - for TF Model Runner (AOT flow)Tested with TensorFlow 2.13.0Other python requirements are available inmlbridge.ymlConda/Anaconda based virtual environment is assumed(Experiments are done on an Ubuntu 20.04 machine)SetupML-Compiler-Bridgecan be built as a stand-alone library to generate.afiles that can in turn be linked with any compiler.mkdir build && cd buildcmake [-DCMAKE_BUILD_TYPE=Release|Debug] [-DCMAKE_INSTALL_PREFIX=<Install_path>] [-DMLBRIDGE_ENABLE_TEST=ON|OFF] -DONNXRUNTIME_ROOTDIR=<Path to ONNX install dir> -DPROTOS_DIRECTORY=<Path to protobuf files> -DTENSORFLOW_AOT_PATH=<Path to TensorFlow pip install dir> ../make -j [&& make install]This process would generatelibMLCompilerBridge.aandlibMLCompilerBridgeC.alibraries underbuild/libdirectory, required headers underbuild/includedirectory.libMLCompilerBridgeC.aexposes C APIs for using with C-based compilers like Pluto, where aslibMLCompilerBridge.aexposes C++ APIs that can be used with any compiler written in C++.Python end points are available underCompilerInterface.To ensure the correctness, runmake verify-all. This would need enabling tests in cmake (-DMLBRIDGE_ENABLE_TEST=ON) andPROTOS_DIRECTORYshould point totest/protos.Using ML-Compiler-Bridge with LLVM and MLIRML-Compiler-Bridgecan be integrated and built along with the LLVM project. This can be done by adding this repository as a new project and settingLLVM_MLBRIDGEoption toON.You can check theCMakeLists.txtof theml-llvm-projectrepository which demonstrates such an integration.The passes that need to make use of this library can then just link withLLVMMLBridge.ExampleCMakeLists.txtof an LLVM pass that would use the library is shown below.add_llvm_component_library(LLVMMLPass
pass.cpp
ml.cpp

ADDITIONAL_HEADER_DIRS
  ${CMAKE_CURRENT_SOURCE_DIR}/includes

DEPENDS
  LLVMMLBridge
  intrinsics_gen
)
target_link_libraries(LLVMMLPass PRIVATE LLVMMLBridge)To use TensorFlow AOT Model Runner, you need to make use oftf_find_and_compilemethod exposed incmake/modules/TensorFlowCompile.cmakein the CMakeLists.txt of your pass with appropriate arguments. An example of integrating TF AOT Model with inlining pass is shownhere.ArtifactsLibraries are autogenerated for every relevant check-in with GitHub actions. Such generated artifacts are tagged along with the successful runs ofPublishaction.Citation@inproceedings{venkatakeerthy-2024-MLCompilerBridge,
author = {VenkataKeerthy, S. and Jain, Siddharth and Kalvakuntla, Umesh and Gorantla, Pranav Sai and Chitale, Rajiv Shailesh and Brevdo, Eugene and Cohen, Albert and Trofin, Mircea and Upadrasta, Ramakrishna},
title = {The Next 700 ML-Enabled Compiler Optimizations},
year = {2024},
isbn = {9798400705076},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3640537.3641580},
doi = {10.1145/3640537.3641580},
booktitle = {Proceedings of the 33rd ACM SIGPLAN International Conference on Compiler Construction},
pages = {238–249},
numpages = {12},
keywords = {Machine Learning for Compiler Optimizations, ONNX, Pipes, TensorFlow AOT, gRPC},
location = {<conf-loc>, <city>Edinburgh</city>, <country>United Kingdom</country>, </conf-loc>},
series = {CC 2024}
}ContributionsPlease feel free to raise issues to file a bug, pose a question, or initiate any related discussions. Pull requests are welcome :)LicenseML-Compiler-Bridge is released under Apache 2.0 license with LLVM Exceptions. See LICENSE file for more details.