src/alpaka/NeuralNetwork.h

0001 #ifndef RecoTracker_LSTCore_src_alpaka_NeuralNetwork_h
0002 #define RecoTracker_LSTCore_src_alpaka_NeuralNetwork_h
0003
0004 #include "FWCore/Utilities/interface/CMSUnrollLoop.h"
0005
0006 #include "RecoTracker/LSTCore/interface/alpaka/Common.h"
0007 #include "RecoTracker/LSTCore/interface/MiniDoubletsSoA.h"
0008
0009 #include "NeuralNetworkWeights.h"
0010
0011 namespace ALPAKA_ACCELERATOR_NAMESPACE::lst::t5dnn {
0012
0013   template <int FEATURES>
0014   ALPAKA_FN_ACC ALPAKA_FN_INLINE void relu_activation(float (&input)[FEATURES]) {
0015     CMS_UNROLL_LOOP
0016     for (unsigned int col = 0; col < FEATURES; ++col) {
0017       input[col] = (input[col] > 0.f) ? input[col] : 0.f;
0018     }
0019   }
0020
0021   template <typename TAcc>
0022   ALPAKA_FN_ACC ALPAKA_FN_INLINE float sigmoid_activation(TAcc const& acc, const float x) {
0023     return alpaka::math::exp(acc, x) / (alpaka::math::exp(acc, x) + 1.f);
0024   }
0025
0026   template <int IN_FEATURES, int OUT_FEATURES>
0027   ALPAKA_FN_ACC ALPAKA_FN_INLINE void linear_layer(const float (&input)[IN_FEATURES],
0028                                                    float (&output)[OUT_FEATURES],
0029                                                    const float (&weights)[IN_FEATURES][OUT_FEATURES],
0030                                                    const float (&biases)[OUT_FEATURES]) {
0031     CMS_UNROLL_LOOP
0032     for (unsigned int i = 0; i < OUT_FEATURES; ++i) {
0033       output[i] = biases[i];
0034       CMS_UNROLL_LOOP
0035       for (int j = 0; j < IN_FEATURES; ++j) {
0036         output[i] += input[j] * weights[j][i];
0037       }
0038     }
0039   }
0040
0041   ALPAKA_FN_ACC ALPAKA_FN_INLINE float delta_phi(const float phi1, const float phi2) {
0042     float delta = phi1 - phi2;
0043     // Adjust delta to be within the range [-M_PI, M_PI]
0044     if (delta > kPi) {
0045       delta -= 2 * kPi;
0046     } else if (delta < -kPi) {
0047       delta += 2 * kPi;
0048     }
0049
0050     return delta;
0051   }
0052
0053   template <typename TAcc>
0054   ALPAKA_FN_ACC ALPAKA_FN_INLINE bool runInference(TAcc const& acc,
0055                                                    MiniDoubletsConst mds,
0056                                                    const unsigned int mdIndex1,
0057                                                    const unsigned int mdIndex2,
0058                                                    const unsigned int mdIndex3,
0059                                                    const unsigned int mdIndex4,
0060                                                    const unsigned int mdIndex5,
0061                                                    const float innerRadius,
0062                                                    const float outerRadius,
0063                                                    const float bridgeRadius) {
0064     // Constants
0065     constexpr unsigned int kinputFeatures = 23;
0066     constexpr unsigned int khiddenFeatures = 32;
0067
0068     float eta1 = alpaka::math::abs(acc, mds.anchorEta()[mdIndex1]);  // inner T3 anchor hit 1 eta (t3_0_eta)
0069     float eta2 = alpaka::math::abs(acc, mds.anchorEta()[mdIndex2]);  // inner T3 anchor hit 2 eta (t3_2_eta)
0070     float eta3 = alpaka::math::abs(acc, mds.anchorEta()[mdIndex3]);  // inner T3 anchor hit 3 eta (t3_4_eta)
0071     float eta4 = alpaka::math::abs(acc, mds.anchorEta()[mdIndex4]);  // outer T3 anchor hit 4 eta (t3_2_eta)
0072     float eta5 = alpaka::math::abs(acc, mds.anchorEta()[mdIndex5]);  // outer T3 anchor hit 5 eta (t3_4_eta)
0073
0074     float phi1 = mds.anchorPhi()[mdIndex1];  // inner T3 anchor hit 1 phi (t3_0_phi)
0075     float phi2 = mds.anchorPhi()[mdIndex2];  // inner T3 anchor hit 2 phi (t3_2_phi)
0076     float phi3 = mds.anchorPhi()[mdIndex3];  // inner T3 anchor hit 3 phi (t3_4_phi)
0077     float phi4 = mds.anchorPhi()[mdIndex4];  // outer T3 anchor hit 4 phi (t3_2_phi)
0078     float phi5 = mds.anchorPhi()[mdIndex5];  // outer T3 anchor hit 5 phi (t3_4_phi)
0079
0080     float z1 = alpaka::math::abs(acc, mds.anchorZ()[mdIndex1]);  // inner T3 anchor hit 1 z (t3_0_z)
0081     float z2 = alpaka::math::abs(acc, mds.anchorZ()[mdIndex2]);  // inner T3 anchor hit 2 z (t3_2_z)
0082     float z3 = alpaka::math::abs(acc, mds.anchorZ()[mdIndex3]);  // inner T3 anchor hit 3 z (t3_4_z)
0083     float z4 = alpaka::math::abs(acc, mds.anchorZ()[mdIndex4]);  // outer T3 anchor hit 4 z (t3_2_z)
0084     float z5 = alpaka::math::abs(acc, mds.anchorZ()[mdIndex5]);  // outer T3 anchor hit 5 z (t3_4_z)
0085
0086     float r1 = mds.anchorRt()[mdIndex1];  // inner T3 anchor hit 1 r (t3_0_r)
0087     float r2 = mds.anchorRt()[mdIndex2];  // inner T3 anchor hit 2 r (t3_2_r)
0088     float r3 = mds.anchorRt()[mdIndex3];  // inner T3 anchor hit 3 r (t3_4_r)
0089     float r4 = mds.anchorRt()[mdIndex4];  // outer T3 anchor hit 4 r (t3_2_r)
0090     float r5 = mds.anchorRt()[mdIndex5];  // outer T3 anchor hit 5 r (t3_4_r)
0091
0092     // Build the input feature vector using pairwise differences after the first hit
0093     float x[kinputFeatures] = {
0094         eta1 / kEta_norm,                          // inner T3: First hit eta normalized
0095         alpaka::math::abs(acc, phi1) / kPhi_norm,  // inner T3: First hit phi normalized
0096         z1 / kZ_max,                               // inner T3: First hit z normalized
0097         r1 / kR_max,                               // inner T3: First hit r normalized
0098
0099         eta2 - eta1,                        // inner T3: Difference in eta between hit 2 and 1
0100         delta_phi(phi2, phi1) / kPhi_norm,  // inner T3: Difference in phi between hit 2 and 1
0101         (z2 - z1) / kZ_max,                 // inner T3: Difference in z between hit 2 and 1 normalized
0102         (r2 - r1) / kR_max,                 // inner T3: Difference in r between hit 2 and 1 normalized
0103
0104         eta3 - eta2,                        // inner T3: Difference in eta between hit 3 and 2
0105         delta_phi(phi3, phi2) / kPhi_norm,  // inner T3: Difference in phi between hit 3 and 2
0106         (z3 - z2) / kZ_max,                 // inner T3: Difference in z between hit 3 and 2 normalized
0107         (r3 - r2) / kR_max,                 // inner T3: Difference in r between hit 3 and 2 normalized
0108
0109         eta4 - eta3,                        // outer T3: Difference in eta between hit 4 and 3
0110         delta_phi(phi4, phi3) / kPhi_norm,  // inner T3: Difference in phi between hit 4 and 3
0111         (z4 - z3) / kZ_max,                 // outer T3: Difference in z between hit 4 and 3 normalized
0112         (r4 - r3) / kR_max,                 // outer T3: Difference in r between hit 4 and 3 normalized
0113
0114         eta5 - eta4,                        // outer T3: Difference in eta between hit 5 and 4
0115         delta_phi(phi5, phi4) / kPhi_norm,  // inner T3: Difference in phi between hit 5 and 4
0116         (z5 - z4) / kZ_max,                 // outer T3: Difference in z between hit 5 and 4 normalized
0117         (r5 - r4) / kR_max,                 // outer T3: Difference in r between hit 5 and 4 normalized
0118
0119         alpaka::math::log10(acc, innerRadius),   // T5 inner radius (t5_innerRadius)
0120         alpaka::math::log10(acc, bridgeRadius),  // T5 bridge radius (t5_bridgeRadius)
0121         alpaka::math::log10(acc, outerRadius)    // T5 outer radius (t5_outerRadius)
0122     };
0123
0124     float x_1[khiddenFeatures];  // Layer 1 output
0125     float x_2[khiddenFeatures];  // Layer 2 output
0126     float x_3[1];                // Layer 3 linear output
0127
0128     // Layer 1: Linear + Relu
0129     linear_layer<kinputFeatures, khiddenFeatures>(x, x_1, wgtT_layer1, bias_layer1);
0130     relu_activation<khiddenFeatures>(x_1);
0131
0132     // Layer 2: Linear + Relu
0133     linear_layer<khiddenFeatures, khiddenFeatures>(x_1, x_2, wgtT_layer2, bias_layer2);
0134     relu_activation<khiddenFeatures>(x_2);
0135
0136     // Layer 3: Linear + Sigmoid
0137     linear_layer<khiddenFeatures, 1>(x_2, x_3, wgtT_output_layer, bias_output_layer);
0138     float x_5 = sigmoid_activation(acc, x_3[0]);
0139
0140     // Get the bin index based on abs(eta) of first hit and t5_pt
0141     float t5_pt = innerRadius * lst::k2Rinv1GeVf * 2;
0142
0143     uint8_t pt_index = (t5_pt > 5);
0144     uint8_t bin_index = (eta1 > 2.5f) ? (kEtaBins - 1) : static_cast<unsigned int>(eta1 / 0.25f);
0145
0146     // Compare x_5 to the cut value for the relevant bin
0147     return x_5 > kWp[pt_index][bin_index];
0148   }
0149 }  // namespace ALPAKA_ACCELERATOR_NAMESPACE::lst::t5dnn
0150
0151 #endif