Merge pull request #49325 from pallabidas/EMTF_setup_hls

[L1T] replacing tensorflow with hls4ml model for EMTF at setup

Author: cmsbuild

L1Trigger/L1TMuonEndCap/BuildFile.xml

@@ -8,3 +8,5 @@
 <use name="DataFormats/L1TMuon"/>
 <use name="L1Trigger/L1TMuon"/>
 <use name="PhysicsTools/TensorFlow"/>
+<use name="hls"/>
+<use name="hls4mlEmulatorExtras"/>

L1Trigger/L1TMuonEndCap/interface/EMTFSetup.h

@@ -16,6 +16,7 @@
 #include "L1Trigger/L1TMuonEndCap/interface/SectorProcessorLUT.h"
 #include "L1Trigger/L1TMuonEndCap/interface/PtAssignmentEngine.h"
 #include "L1Trigger/L1TMuonEndCap/interface/PtAssignmentEngineDxy.h"
+#include "hls4ml/emulator.h"
 
 class EMTFSetup {
 public:
@@ -72,6 +73,9 @@ class EMTFSetup {
   unsigned fw_ver_;
   unsigned pt_lut_ver_;
   unsigned pc_lut_ver_;
+
+  hls4mlEmulator::ModelLoader loader;
+  std::shared_ptr<hls4mlEmulator::Model> model;
 };
 
 #endif

L1Trigger/L1TMuonEndCap/interface/PtAssignment.h

@@ -23,8 +23,7 @@ class PtAssignment {
                  bool bugGMTPhi,
                  bool promoteMode7,
                  int modeQualVer,
-                 std::vector<int> promoteMode7Sectors,
-                 std::string pbFileName);
+                 std::vector<int> promoteMode7Sectors);
 
   void process(EMTFTrackCollection& best_tracks);
 

L1Trigger/L1TMuonEndCap/interface/PtAssignmentEngineDxy.h

@@ -10,33 +10,28 @@
 
 #include "L1Trigger/L1TMuonEndCap/interface/Common.h"
 #include "L1Trigger/L1TMuonEndCap/interface/PtAssignmentEngineAux2017.h"
-#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
 #include "FWCore/ParameterSet/interface/FileInPath.h"
+#include "ap_fixed.h"
+#include "hls4ml/emulator.h"
 
 class PtAssignmentEngineDxy {
 public:
-  explicit PtAssignmentEngineDxy();
+  explicit PtAssignmentEngineDxy(std::shared_ptr<hls4mlEmulator::Model> model);
   virtual ~PtAssignmentEngineDxy();
 
-  void configure(int verbose, const std::string pbFileNameDxy);
+  void configure(int verbose);
 
   const PtAssignmentEngineAux2017& aux() const;
 
   virtual void calculate_pt_dxy(const EMTFTrack& track, emtf::Feature& feature, emtf::Prediction& prediction) const;
 
   virtual void preprocessing_dxy(const EMTFTrack& track, emtf::Feature& feature) const;
 
-  virtual void call_tensorflow_dxy(const emtf::Feature& feature, emtf::Prediction& prediction) const;
+  virtual void call_hls_dxy(const emtf::Feature& feature, emtf::Prediction& prediction) const;
 
 protected:
   int verbose_;
-
-  tensorflow::GraphDef* graphDefDxy_;
-  tensorflow::Session* sessionDxy_;
-  std::string pbFileNameDxy_;
-  std::string pbFilePathDxy_;
-  std::string inputNameDxy_;
-  std::vector<std::string> outputNamesDxy_;
+  std::shared_ptr<hls4mlEmulator::Model> model_;
 };
 
 #endif

L1Trigger/L1TMuonEndCap/interface/VersionControl.h

@@ -18,6 +18,7 @@ class VersionControl {
   int verbose() const { return verbose_; }
   bool useO2O() const { return useO2O_; }
   std::string era() const { return era_; }
+  std::string nnModelDxy() const { return nnModel_; }
 
   friend class SectorProcessor;  // allow access to private memebers
 
@@ -58,7 +59,7 @@ class VersionControl {
   bool bug9BitDPhi_, bugMode7CLCT_, bugNegPt_, bugGMTPhi_, promoteMode7_;
   int modeQualVer_;
   std::vector<int> promoteMode7Sectors_;  // Sectors to promote mode 7 tracks in Run 2 and Run 3, -1 for all sectors
-  std::string pbFileName_;
+  std::string nnModel_;
 };
 
 #endif

L1Trigger/L1TMuonEndCap/plugins/BuildFile.xml

@@ -1,4 +1,6 @@
 <library file="*.cc" name="L1TriggerL1TMuonEndCapPlugins">
   <use name="L1Trigger/L1TMuonEndCap"/>
+  <use name="hls"/>
+  <use name="hls4mlEmulatorExtras"/>
   <flags EDM_PLUGIN="1"/>
 </library>

L1Trigger/L1TMuonEndCap/python/simEmtfDigis_cfi.py

@@ -126,7 +126,7 @@
         PromoteMode7    = cms.bool(False), # Assign station 2-3-4 tracks with |eta| > 1.6 SingleMu quality
         ModeQualVer     = cms.int32(2),    # Version 2 contains modified mode-quality mapping for 2018
         PromoteMode7Sectors = cms.vint32(-1), # Sectors to promote mode 7 tracks in Run 2 and Run 3, -1 for all sectors
-        ProtobufFileName = cms.string('model_graph.displ.16.pb'), # Protobuf file name to be used by NN based pT assignment NNv16 is online since 26.06.2023
+        NNModel = cms.string('EMTFnn_v1'), # HLS model name to be used by NN based pT assignment NNv16 is online since 26.06.2023
     ),
 
 )

L1Trigger/L1TMuonEndCap/src/EMTFSetup.cc

@@ -17,7 +17,8 @@ EMTFSetup::EMTFSetup(const edm::ParameterSet& iConfig, edm::ConsumesCollector iC
       pt_assign_engine_dxy_(nullptr),
       fw_ver_(0),
       pt_lut_ver_(0),
-      pc_lut_ver_(0) {
+      pc_lut_ver_(0),
+      loader(version_control_.nnModelDxy()) {
   // Set pt assignment engine according to Era
   if (era() == "Run2_2016") {
     pt_assign_engine_ = std::make_unique<PtAssignmentEngine2016>();
@@ -29,7 +30,8 @@ EMTFSetup::EMTFSetup(const edm::ParameterSet& iConfig, edm::ConsumesCollector iC
   }
 
   // No era setup for displaced pT assignment engine
-  pt_assign_engine_dxy_ = std::make_unique<PtAssignmentEngineDxy>();
+  model = loader.load_model();
+  pt_assign_engine_dxy_ = std::make_unique<PtAssignmentEngineDxy>(model);
 
   emtf_assert(pt_assign_engine_ != nullptr);
   emtf_assert(pt_assign_engine_dxy_ != nullptr);

L1Trigger/L1TMuonEndCap/src/PtAssignment.cc

@@ -17,8 +17,7 @@ void PtAssignment::configure(PtAssignmentEngine* pt_assign_engine,
                              bool bugGMTPhi,
                              bool promoteMode7,
                              int modeQualVer,
-                             std::vector<int> promoteMode7Sectors,
-                             std::string pbFileName) {
+                             std::vector<int> promoteMode7Sectors) {
   emtf_assert(pt_assign_engine != nullptr);
   emtf_assert(pt_assign_engine_dxy != nullptr);
 
@@ -33,7 +32,7 @@ void PtAssignment::configure(PtAssignmentEngine* pt_assign_engine,
 
   pt_assign_engine_->configure(verbose_, readPtLUTFile, fixMode15HighPt, bug9BitDPhi, bugMode7CLCT, bugNegPt);
 
-  pt_assign_engine_dxy_->configure(verbose_, pbFileName);
+  pt_assign_engine_dxy_->configure(verbose_);
 
   bugGMTPhi_ = bugGMTPhi;
   promoteMode7_ = promoteMode7;

L1Trigger/L1TMuonEndCap/src/PtAssignmentEngineDxy.cc

@@ -6,35 +6,11 @@
 
 #include "helper.h"  // assert_no_abort
 
-PtAssignmentEngineDxy::PtAssignmentEngineDxy() : graphDefDxy_(nullptr), sessionDxy_(nullptr) {}
+PtAssignmentEngineDxy::PtAssignmentEngineDxy(std::shared_ptr<hls4mlEmulator::Model> model) { model_ = model; }
 
-PtAssignmentEngineDxy::~PtAssignmentEngineDxy() {
-  if (sessionDxy_ != nullptr) {
-    tensorflow::closeSession(sessionDxy_);
-  }
-  delete graphDefDxy_;
-}
-
-void PtAssignmentEngineDxy::configure(int verbose, const std::string pbFileNameDxy) {
-  verbose_ = verbose;
-
-  pbFileNameDxy_ = pbFileNameDxy;
-  std::string pbFilePathDxy_ = "L1Trigger/L1TMuon/data/emtf_luts/" + pbFileNameDxy_;
-
-  inputNameDxy_ = "input1";
-  outputNamesDxy_ = {"Identity"};
-
-  if (graphDefDxy_ == nullptr) {
-    graphDefDxy_ = tensorflow::loadGraphDef(edm::FileInPath(pbFilePathDxy_).fullPath());
-  }
-  emtf_assert(graphDefDxy_ != nullptr);
-
-  if (sessionDxy_ == nullptr) {
-    sessionDxy_ = tensorflow::createSession(graphDefDxy_);
-  }
+PtAssignmentEngineDxy::~PtAssignmentEngineDxy() {}
 
-  emtf_assert(sessionDxy_ != nullptr);
-}
+void PtAssignmentEngineDxy::configure(int verbose) { verbose_ = verbose; }
 
 const PtAssignmentEngineAux2017& PtAssignmentEngineDxy::aux() const {
   static const PtAssignmentEngineAux2017 instance;
@@ -46,7 +22,7 @@ void PtAssignmentEngineDxy::calculate_pt_dxy(const EMTFTrack& track,
                                              emtf::Prediction& prediction) const {
   // This is called for each track instead of for entire track collection as was done in Phase-2 implementation
   preprocessing_dxy(track, feature);
-  call_tensorflow_dxy(feature, prediction);
+  call_hls_dxy(feature, prediction);
   return;
 }
 
@@ -153,19 +129,24 @@ void PtAssignmentEngineDxy::preprocessing_dxy(const EMTFTrack& track, emtf::Feat
   return;
 }
 
-void PtAssignmentEngineDxy::call_tensorflow_dxy(const emtf::Feature& feature, emtf::Prediction& prediction) const {
-  tensorflow::Tensor input(tensorflow::DT_FLOAT, {1, emtf::NUM_FEATURES});
-  std::vector<tensorflow::Tensor> outputs;
+void PtAssignmentEngineDxy::call_hls_dxy(const emtf::Feature& feature, emtf::Prediction& prediction) const {
   emtf_assert(feature.size() == emtf::NUM_FEATURES);
 
-  float* d = input.flat<float>().data();
-  std::copy(feature.begin(), feature.end(), d);
-  tensorflow::run(sessionDxy_, {{inputNameDxy_, input}}, outputNamesDxy_, &outputs);
-  emtf_assert(outputs.size() == 1);
+  ap_uint<13> nn_input[29];
+  for (size_t i = 0; i < feature.size(); i++) {
+    nn_input[i] = feature[i];
+  }
+  ap_uint<8> nn_output[2];
+  model_->prepare_input(nn_input);
+  model_->predict();
+  model_->read_result(nn_output);
+  ap_uint<8> pT = nn_output[0];
+  ap_uint<7> dxy = (nn_output[1] > 127) ? ap_uint<7>(127) : ap_uint<7>(nn_output[1]);
+
   emtf_assert(prediction.size() == emtf::NUM_PREDICTIONS);
 
-  prediction.at(0) = outputs[0].matrix<float>()(0, 0);
-  prediction.at(1) = outputs[0].matrix<float>()(0, 1);
+  prediction.at(0) = pT.to_float();
+  prediction.at(1) = dxy.to_float();
 
   return;
 }