[tmva][sofie] Drop the BLAS dependency of the generated code - #23600
guitargeek wants to merge 2 commits into
Conversation
dpiparo
left a comment
dpiparo
left a comment
There was a problem hiding this comment.
Thanks, this PR arrived so quickly!!
Test Results 24 files 24 suites 3d 18h 12m 23s ⏱️ For more details on these failures, see this check. Results for commit c938f8d. ♻️ This comment has been updated with latest results. |
The SOFIE tests are currently disabled on Windows CI (test_tmva_sofie=OFF in windows10.txt); enabling them surfaces two Windows-only failures: * RModelParser_ONNX::Parse split the model path at backslashes only, so a forward-slash path like dir/model.onnx (which the Windows file APIs happily accept) yielded an empty model directory, and the external weight data of such a model was looked up in the working directory instead of next to the model (TestSofieParser. ExternalDataLocationRelativeToModelDirectory). On Windows, split at the last of either separator; on POSIX a backslash stays a filename character. * Catching an exception thrown by cling-JIT code does not work on Windows: the runtime finds no handler and terminates the process with 0xE06D7363, so the safetensors tests whose whole point is checking that the generated session constructors throw can only run elsewhere (TestSafetensorsWeights dies in MissingFileThrows, and all the other throw-probing tests would follow). Skip them on MSVC; the success-path construction (ModerateNestingAccepted) stays enabled. 🤖 Done with the help of AI
The inference code SOFIE emits no longer calls any BLAS routines. Matrix multiplications are now done by Gemm_Ref, a self-contained reference implementation of the sgemm contract emitted into the generated header next to the other inference helpers, and bias accumulations by the corresponding Axpy_Ref (saxpy). Both keep the pointer-based BLAS calling convention, so the emission sites in the Gemm, Conv, ConvTranspose, Einsum, GRU, LSTM and RNN operators only change the function name. Gemm_Ref walks C column by column with unit stride so the compiler can auto-vectorize the innermost loop; this is close to BLAS performance for the per-event (gemv-shaped) inference SOFIE is used for in RDataFrame loops and RooFit, while removing - the link-time BLAS dependency of every generated model, - symbol collisions with whatever sgemm_ happens to be loaded in the process (e.g. a conda OpenBLAS in PyROOT sessions), - the implicit, environment-dependent multithreading of the BLAS implementation picked up at runtime, which interferes with RDataFrame implicit MT. Consequently deleted: the BLAS namespace emission in RModel::GenerateHeaderInfo, fNeededBlasRoutines / AddBlasRoutines / GetBlasRoutines, the extern "C" sgemm_/sgemv_/saxpy_ declarations (dead Gemv), the commented-out BLAS call sites in Conv/ConvTranspose, and the hand-rolled CMake BLAS handling of the SOFIE tests (BLAS_FOUND guards, BLAS::BLAS link libraries, the hard configure error in SearchInstalledSoftware). The hand-written Clad pullbacks/pushforwards for Gemm_Call are kept: they are pure C++ built on Gemm_Call itself, and are both more accurate and much cheaper than the tape Clad would generate for the inner loops. Also remove the test_tmva_sofie configuration option: with the BLAS requirement gone the tests need nothing beyond what any other ROOT test needs, so they now follow the global 'testing' option like all other subsystems. Verified: all SOFIE tests (TestSofieModels, TestCustomModelsFromONNX, TestCladAutodiff, TestGemmDerivative, TestSofieParser, ...), testRooONNXFunc and the SOFIE tutorials pass; a generated LSTM and ConvTranspose model header compiles and runs with plain g++ -std=c++17 with no ROOT and no BLAS. Performance (Zen5-class 12-core machine, emitted code compiled -O2, vs OpenBLAS 0.3.33 single-threaded -- ROOT's BLAS dependency here): Per-event inference (batch=1 GEMV shapes, the RDataFrame / RooFit SBI pattern) is at parity up to layer size ~512, both at kernel level (gemv 512x1x512: 3.5 vs 3.6 GFLOP/s) and end-to-end through Session::infer of a 4-layer MLP (D32: 2.1 vs 2.1 us, D128: 31 vs 27 us, D512: 464 vs 440 us per event). A regression sets in around layer size 1024 (weights ~4 MB): 3-5x kernel, 2.9x end-to-end (1899 vs 649 us). The transposed-A case (which is what per-event inference of row-major weights emits) needs the dot-product loop nest in Gemm_Ref: with the plain accumulation nest the strided A reads make it 27x slower than BLAS at D=1024 instead of 2.9x. Large *batched* evaluation regresses 20-100x end-to-end from D=128/batch=1000 up (D1024 B1000: 1.8 s vs 18 ms), matching the kernel gap for square GEMMs (~48x at 512^3, ~64-116x at 2048^3). Threaded OpenBLAS (12 threads) does not change the picture for typical SOFIE shapes: it is equal or slower than its single-threaded self below ~2048^3, and the shapes where it wins are exactly the ones SOFIE-emitted code does not target. 🤖 Done with the help of AI
b4a73d9 to
c938f8d
Compare
|
Sorry for jumping on this, but as you might have noticed (and discussed privately with @dpiparo), I am rather interested to SOFIE and I am keeping an eye on the development, since I still want to push upstream few improvements. I tested this on top of my current SOFIE branch with a Belle II GNN track-finding model (https://gitlab.desy.de/giacomo.depietro/catfinder-sofie) with dynamic input size (number of hits), single-threaded. Machine is an AMD EPYC 7452 (Zen 2, AVX2, no AVX-512), 512 KiB L2 per core, generated code compiled with gcc 11.5 -O3. Baseline is the same tree without this PR, linking OpenBLAS 0.3.29. Numerically everything agrees — outputs match ONNX Runtime to ~7 significant digits before and after, no NaN. But end-to-end inference gets substantially slower: With another model where GEMM is the dominating operator the slowdown is even larger. So SOFIE goes from being at parity with ONNX Runtime to being about 2x slower than it. The model emits 42 Those 24 calls are about 90% of the model's GEMM work. I understand the motivations for dropping the link-time BLAS dependency, the symbol collisions and the implicit threading. But batched inference with hundreds to thousands of rows is the normal case for us, not an edge case, and a 1.5–2x end-to-end regression there is hard to absorb. Would it be possible to keep an optional BLAS path, like emitting the |
|
Thanks for following this! And thanks for the systematic study, it's exactly the kind of input we need! Short answer first: yes, I'll add an opt-in BLAS path in this PR. The default generated code will be self-contained (no external symbols, as it is in this PR), and a generation option will make it emit the Some context on why the change happens at all, since it explains where SOFIE in ROOT is heading. It currently serves two mandates:
These mandates pull in opposite directions on performance, usability, and interface stability, and they require different interfaces to begin with. For the second usecase, it could even be a Python package with just an executable! The direction we plan to follow is:
The Belle II response to our Q3 questionnaire said SOFIE would "likely soon" be used for offline reconstruction given the GNN runtime performance on CPU, which is why the PR has stayed unmerged rather than landing the regression. But if and when the standalone tool materializes, the long-term future of mandate-2 usage patterns in ROOT itself is likely to evolve. The Q2/Q3 meeting notes ([1], [2]) have the details, and your input there (as the one experiment with an actual GNN about to go into production) is very important. What's your take on this? What part of ONNX tooling do you expect to be in ROOT, and what do you expect to be separate? [1] https://indico.cern.ch/event/1699702/ |
|
By the way, have you tried https://github.com/onnx/onnx-mlir ? That would make a great addition to your comparison table! And maybe the fact that it's also a code generation tool makes it easy to integrate in your pipeline? I have always been curious what are the differences in performance and user experience between SOFIE and ONNX-MLIR |
|
I need more time to formulate proper answers to your previous questions. Regarding |
This branch has not been deployed
3 participants
The inference code SOFIE emits no longer calls any BLAS routines.
Matrix multiplications are now done by Gemm_Ref, a self-contained reference implementation of the sgemm contract emitted into the generated header next to the other inference helpers, and bias accumulations by the corresponding Axpy_Ref (saxpy). Both keep the pointer-based BLAS calling convention, so the emission sites in the Gemm, Conv, ConvTranspose, Einsum, GRU, LSTM and RNN operators only change the function name. Gemm_Ref walks C column by column with unit stride so the compiler can auto-vectorize the innermost loop; this is close to BLAS performance for the per-event (gemv-shaped) inference SOFIE is used for in RDataFrame loops and RooFit, while removing
Consequently deleted: the BLAS namespace emission in RModel::GenerateHeaderInfo, fNeededBlasRoutines / AddBlasRoutines / GetBlasRoutines, the extern "C" sgemm_/sgemv_/saxpy_ declarations (dead Gemv), the commented-out BLAS call sites in Conv/ConvTranspose, and the hand-rolled CMake BLAS handling of the SOFIE tests (BLAS_FOUND guards, BLAS::BLAS link libraries, the hard configure error in SearchInstalledSoftware).
The hand-written Clad pullbacks/pushforwards for Gemm_Call are kept: they are pure C++ built on Gemm_Call itself, and are both more accurate and much cheaper than the tape Clad would generate for the inner loops.
Also remove the test_tmva_sofie configuration option: with the BLAS requirement gone the tests need nothing beyond what any other ROOT test needs, so they now follow the global 'testing' option like all other subsystems.
Verified: all SOFIE tests (TestSofieModels, TestCustomModelsFromONNX, TestCladAutodiff, TestGemmDerivative, TestSofieParser, ...), testRooONNXFunc and the SOFIE tutorials pass; a generated LSTM and ConvTranspose model header compiles and runs with plain g++ -std=c++17 with no ROOT and no BLAS.
Performance (Zen5-class 12-core machine, emitted code compiled -O2, vs OpenBLAS 0.3.33 single-threaded -- ROOT's BLAS dependency here):
Per-event inference (batch=1 GEMV shapes, the RDataFrame / RooFit SBI pattern) is at parity up to layer size ~512, both at kernel level (gemv 512x1x512: 3.5 vs 3.6 GFLOP/s) and end-to-end through Session::infer of a 4-layer MLP (D32: 2.1 vs 2.1 us, D128: 31 vs 27 us, D512: 464 vs 440 us per event). A regression sets in around layer size 1024 (weights ~4 MB): 3-5x kernel, 2.9x end-to-end (1899 vs 649 us).
The transposed-A case (which is what per-event inference of row-major weights emits) needs the dot-product loop nest in Gemm_Ref: with the plain accumulation nest the strided A reads make it 27x slower than BLAS at D=1024 instead of 2.9x.
Large batched evaluation regresses 20-100x end-to-end from D=128/batch=1000 up (D1024 B1000: 1.8 s vs 18 ms), matching the kernel gap for square GEMMs (~48x at 512^3, ~64-116x at 2048^3). Threaded OpenBLAS (12 threads) does not change the picture for typical SOFIE shapes: it is equal or slower than its single-threaded self below ~2048^3, and the shapes where it wins are exactly the ones SOFIE-emitted code does not target.
🤖 Done with the help of AI