| Name | S Akash |
| Organisation | CERN HSF (Root Project) |
| Mentor | Sanjiban Sengupta, Dr. Lorenzo Moneta |
| Project | TMVA SOFIE - GPU Support for Machine Learning Inference |
ROOT's TMVA( Toolkit for Multi-Variate Analysis ) has SOFIE (System for Optimized Fast Inference code Emit) which offers a parser capable of converting ML models trained in Keras, PyTorch, or ONNX format into its own Intermediate Representation, and generates C++ functions that can be easily invoked for fast inference of trained neural networks. It is currently implemented for CPU inference along with a SYCL implementation. This project aims to explore different GPU stacks (such as CUDA, ROCm, ALPAKA) and implement GPU-based inference functionalities in SOFIE. Although there exist CPU implementations, working with HEP applications, need for GPU inferences become important.
Any model built with TensorFlow, PyTorch, or Keras, once converted to ONNX, can be parsed by SOFIE. Once parsed, header files in c++ are generated for fast inference in two files, a .hxx header file containing the inference code and a .dat file holding the weights.
SOFIE reads an ONNX model, translates it into its own IR (RModel), and then the generator produces a .hxx header plus a .dat weights file.
Once we compile that header file and call for inference, it runs as a standalone C++ with minimal dependencies, without a separate framework runtime unlike ONNX which uses onnx runtime, numpy, etc.
Internally similar to ONNX it still builds a graph of operators, but the key difference is that this graph is compiled into simple, fast C++ specialized for your model.
RModel is the general code generator, it writes the header guards, includes, opens the namespace, allocates the intermediate memory pool, lays out tensors, builds the infer() function signature/return, and makes the generated code usable.
ROperator is per-node, for each operators that is encountered in a topological manner, the codes for those specific kernels are generated by it.
Below is a example of how a add.hxx inference file is generated and from where it is actually generated.
| Generated code snippet | Generator & location |
|---|---|
//Code generated automatically …#ifndef SOFIE_ADD#define SOFIE_ADD#include <vector>#include "SOFIE/SOFIE_common.hxx"namespace SOFIE_Add{ |
Produced by RModel_Base::GenerateHeaderInfo. It writes the banner comment, header guards, needed <vector> include, SOFIE_common.hxx, and opens the model’s namespace |
//--- Allocating session memory pool…char* fIntermediateMemoryPool = new char[8]; |
Emitted in RModel::GenerateIntermediateMemoryPool, which sizes and allocates the shared block for intermediate tensors |
// --- Positioning intermediate tensor memory -- |
Start of the allocation pass assembled in RModel::GenerateSessionCode before placement of individual tensor pointers |
// Allocating memory for intermediate tensor 2 with size 8 bytesfloat* tensor_2 = reinterpret_cast<float*>(fIntermediateMemoryPool + 0); |
Generated by RModel::AllocateIntermediateMemory, which emits per‑tensor comments and pointer assignments into the pool |
std::vector<float> infer(float* tensor_onnxAdd_0, float* tensor_onnxAdd_1){ |
Function signature built by RModel::GenerateOutput (return type, name, and parameters come from model metadata and GenerateInferSignature) |
for (size_t id = 0; id < 2 ; id++){ tensor_2[id] = tensor_onnxAdd_0[id] + tensor_onnxAdd_1[id] ; } |
The Add loop is emitted by ROperator_BasicBinary::Generate for the Add operator instance |
return {std::vector<float>(tensor_2, tensor_2 + 2)}; |
Return statement pieced together by RModel::GenerateOutput, which calls createOutputTensor to wrap the tensor pointer into an std::vector |
} //SOFIE_Add#endif // SOFIE_ADD |
Closing lines appended at the end of RModel::Generate after the session code is generated |
-
Investigate CUDA, ROCm, and ALPAKA frameworks for GPU acceleration. Evaluate pros and cons of each framework in terms of performance and portability.
-
Develop GPU kernels for operations in SOFIE using ALPAKA as an abstraction layer for use of CPU, CUDA or ROCm.
-
Integrate the operators within SOFIE.
-
LayerNormFlat1DKernel_kernel.hpp— LayerNorm 1D . -
relu_kernel.hpp/sigmoid_kernel.hpp/tanh_kernel.hpp/selu_kernel.hpp— activation suite. -
binary_kernel.hpp— arithmetic kernels likeAddKernel,SubKernel,MulKernelandDivKernel. -
transpose_kernel.hpp— 2D transpose. -
unary_kernel.hpp— unary operations likeNegKernel,AbsKernel,ReciprocalKernel,SquareKernel,SqrtKernel,ExpKernelandLogKernel. -
template_kernel.hpp— boilerplate for new kernels.
Example of how a kernel looks in alpaka
struct TemplateKernel
{
// Operator that Alpaka calls on device
template <typename TAcc, typename T>
ALPAKA_FN_ACC void operator()(TAcc const& acc, T* data, int numElements) const
{
// Loop over elements this thread is responsible for
for (auto i : alpaka::uniformElements(acc, numElements))
{
// Example operation replace this with your kernel logic
data[i] = data[i];
}
}
};These kernels needs unit tests to be done by making a test file.
To test a kernel, the test files are structured in this way for proper execution and makes :
- Platforms & devices
auto platAcc = alpaka::Platform<Acc>{};
auto platHost = alpaka::PlatformCpu{};
auto devAcc = alpaka::getDevByIdx(platAcc, 0);
auto devHost = alpaka::getDevByIdx(platHost, 0);
alpaka::Queue<Acc, alpaka::Blocking> queue{devAcc};- Allocate buffers
auto dData = alpaka::allocBuf<T, Idx>(devAcc, extent);
auto hData = alpaka::allocBuf<T, Idx>(devHost, extent);- Copy INPUT into host buffer
{
T* p = alpaka::getPtrNative(hData);
for (Idx i = 0; i < N; ++i) p[i] = INPUT[i];
}- Host to Device
alpaka::memcpy(queue, dData, hData, extent);
alpaka::wait(queue);- Work division
auto workDiv = alpaka::WorkDivMembers<Dim, Idx>(
alpaka::Vec<Dim, Idx>::all(gridSize),
alpaka::Vec<Dim, Idx>::all(blockSize),
alpaka::Vec<Dim, Idx>::all(1)
);- Run the kernel
TemplateKernel kernel; // replace with your kernel name
alpaka::exec<Acc>(queue, workDiv, kernel, alpaka::getPtrNative(dData), N);
alpaka::wait(queue);- Device to Host
alpaka::memcpy(queue, hData, dData, extent);
alpaka::wait(queue);kernels/LayerNormFlat1DKernel_kernel.hpp
#pragma once
#include <alpaka/alpaka.hpp>
namespace alpaka_kernels {
struct LayerNormFlat1DKernel
{
// Operator that Alpaka calls on device
template <typename TAcc, typename T>
ALPAKA_FN_ACC void operator()(TAcc const& acc,
float const epsilon,
int const axis,
std::size_t stashType,
T const* nameX,
T const* nameScale,
T const* nameB,
T* nameY,
int N) const
{
// computing mean
T mean = T{0};
for(auto i : alpaka:: uniformElements(acc,N)) mean = mean + nameX[i];
mean = mean / N;
T variance = T{0};
for(auto i : alpaka:: uniformElements(acc,N)){
T xi = nameX[i] - mean ;
variance = variance + xi*xi;
}
// use ? /N or (N-1) might break at N=1
variance = variance/(N-1);
// Addition of epsilon value
variance = variance + epsilon;
// computing standard deviation using 1/sqrt(varianceiance)
T std_dev = T{1}/alpaka::math::sqrt(acc, variance);
// final equation
for( auto i : alpaka:: uniformElements(acc,N)){
nameY[i] = (nameX[i] - mean) * std_dev * nameScale[i] + nameB[i];
}
}
};
} // namespace alpaka_kernelstests/norm1D_test.cpp
#include <alpaka/alpaka.hpp>
#include <iostream>
#include <vector>
#include <cstddef>
#include <LayerNormFlat1DKernel_kernel.hpp>
using Dim = alpaka::DimInt<1>;
using Idx = std::size_t;
#if defined(USE_CPU_OMP)
using Acc = alpaka::AccCpuOmp2Blocks<Dim, Idx>;
#elif defined(USE_CPU_THREADS)
using Acc = alpaka::AccCpuThreads<Dim, Idx>;
#else
using Acc = alpaka::AccGpuCudaRt<Dim, Idx>; // default to CUDA
#endif
int main() {
using namespace alpaka_kernels;
using T = float;
// input here
std::vector<T> INPUT = { 1.f, 2.f, 3.f};
const int N = INPUT.size();
auto extent = alpaka::Vec<Dim, Idx>::all(N);
// platforms & devices
auto platAcc = alpaka::Platform<Acc>{};
auto platHost = alpaka::PlatformCpu{};
auto devAcc = alpaka::getDevByIdx(platAcc, 0);
auto devHost = alpaka::getDevByIdx(platHost, 0);
alpaka::Queue<Acc, alpaka::Blocking> queue{devAcc};
// buffers
auto dx = alpaka::allocBuf<T, Idx>(devAcc, extent);
auto dy = alpaka::allocBuf<T, Idx>(devAcc, extent);
auto dg = alpaka::allocBuf<T, Idx>(devAcc, extent);
auto db = alpaka::allocBuf<T, Idx>(devAcc, extent);
// host
auto hx = alpaka::allocBuf<T, Idx>(devHost, extent);
auto hy = alpaka::allocBuf<T, Idx>(devHost, extent);
auto hg = alpaka::allocBuf<T, Idx>(devHost, extent);
auto hb = alpaka::allocBuf<T, Idx>(devHost, extent);
// copy INPUT -> host buffer
{
T* p = alpaka::getPtrNative(hx);
for (Idx i = 0; i < N; ++i) p[i] = INPUT[i];
T* pg = alpaka::getPtrNative(hg);
T* pb = alpaka::getPtrNative(hb);
for (Idx i = 0; i < N; ++i) {
pg[i] = T{1}; // gamma_i = 1
pb[i] = T{1}; // beta_i = 0
}
}
std::cout << "Input: ";
for (auto v : INPUT) std::cout << v << ' ';
std::cout << '\n';
// H2D
alpaka::memcpy(queue, dx, hx);
alpaka::memcpy(queue, dg, hg);
alpaka::memcpy(queue, db, hb);
alpaka::wait(queue);
// work division
constexpr Idx blockSize = 256;
Idx gridSize = (N + blockSize - 1) / blockSize;
if (gridSize == 0) gridSize = 1;
auto workDiv = alpaka::WorkDivMembers<Dim, Idx>(
alpaka::Vec<Dim, Idx>::all(1),
alpaka::Vec<Dim, Idx>::all(1),
alpaka::Vec<Dim, Idx>::all(1)
);
// launch norm
LayerNormFlat1DKernel kernel;
constexpr T eps = 1e-5f;
T axis = T{0};
size_t stashType = 0;
alpaka::exec<Acc>(queue, workDiv, kernel,
eps, axis, stashType,
alpaka::getPtrNative(dx),
alpaka::getPtrNative(dg),
alpaka::getPtrNative(db),
alpaka::getPtrNative(dy),N);
alpaka::wait(queue);
// D2H
alpaka::memcpy(queue, hy, dy);
alpaka::wait(queue);
// print output
std::vector<T> OUT(N);
{
T* p = alpaka::getPtrNative(hy);
for (Idx i = 0; i < N; ++i) OUT[i] = p[i];
}
std::cout << "Output: ";
for (auto v : OUT) std::cout << v << ' ';
std::cout << '\n';
return 0;
}As most of neural networks and ML models deal with linear algebra at their core, we are completely wanting to abstract the BLAS libraries like cuBLAS and rocBLAS on our backend kernel instead of writing them on our own.
We are currently still working on developing sofieBLAS.
| PR (title/feature) | PR number | code path(s) | status |
|---|---|---|---|
| ALPAKA support in SOFIE | - | alpaka_kernels | In progress |
| sofieBLAS development | - | example | Planned |
- Faced difficulty while setting up ROOT project with SOFIE enabled.
- Navigated through SOFIE's complex code base and working.
- Got hands on experince with ALPAKA, SOFIE and overall structure of its workings.
- Learned how to write ALPAKA kernels with concise and modular unit tests.
I will be contributing to SOFIE beyond the GSoC period and will be primarily be working on integrating existing kernels to be generated within SOFIE and to develop sofieBLAS.
I am thankful for my mentors, Sanjiban Sengupta and Dr. Lorenzo Moneta, who were really patient and helped me through every step in the project. I am really thankful for the guidance and our bi-weekly discussions where I got to learn more each time. I am really happy and delighted to have such supportive mentors and community to interact with all summer. I look forward to continuing to work with the community beyond the GSoC timeline and am delighted to do so.
A heartfelt thanks as well to the ALPAKA community and Andrea Bocci for your timely feedback and pointers while we iterated on the SofieBLAS prototype and for including me in their meetings. I look forward to more such discussions in the future.
I also thank the Google Summer of Code for granting me such a fantastic opportunity this summer!
Heartfult Thanks,
S Akash