SonicTriton/test/TritonGraphModules.cc

0001 #include "HeterogeneousCore/SonicTriton/interface/TritonEDProducer.h"
0002
0003 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0004 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0005 #include "FWCore/Framework/interface/MakerMacros.h"
0006
0007 #include <sstream>
0008 #include <string>
0009 #include <vector>
0010 #include <random>
0011
0012 class TritonGraphHelper {
0013 public:
0014   TritonGraphHelper(edm::ParameterSet const& cfg)
0015       : nodeMin_(cfg.getParameter<unsigned>("nodeMin")),
0016         nodeMax_(cfg.getParameter<unsigned>("nodeMax")),
0017         edgeMin_(cfg.getParameter<unsigned>("edgeMin")),
0018         edgeMax_(cfg.getParameter<unsigned>("edgeMax")),
0019         brief_(cfg.getParameter<bool>("brief")) {}
0020   void makeInput(edm::Event const& iEvent, TritonInputMap& iInput) const {
0021     //get event-based seed for RNG
0022     unsigned int runNum_uint = static_cast<unsigned int>(iEvent.id().run());
0023     unsigned int lumiNum_uint = static_cast<unsigned int>(iEvent.id().luminosityBlock());
0024     unsigned int evNum_uint = static_cast<unsigned int>(iEvent.id().event());
0025     std::uint32_t seed = (lumiNum_uint << 10) + (runNum_uint << 20) + evNum_uint;
0026     std::mt19937 rng(seed);
0027
0028     std::uniform_int_distribution<int> randint1(nodeMin_, nodeMax_);
0029     int nnodes = randint1(rng);
0030     std::uniform_int_distribution<int> randint2(edgeMin_, edgeMax_);
0031     int nedges = randint2(rng);
0032
0033     //set shapes
0034     auto& input1 = iInput.at("x__0");
0035     input1.setShape(0, nnodes);
0036     auto data1 = input1.allocate<float>();
0037     auto& vdata1 = (*data1)[0];
0038
0039     auto& input2 = iInput.at("edgeindex__1");
0040     input2.setShape(1, nedges);
0041     auto data2 = input2.allocate<int64_t>();
0042     auto& vdata2 = (*data2)[0];
0043
0044     //fill
0045     std::normal_distribution<float> randx(-10, 4);
0046     for (unsigned i = 0; i < input1.sizeShape(); ++i) {
0047       vdata1.push_back(randx(rng));
0048     }
0049
0050     std::uniform_int_distribution<int> randedge(0, nnodes - 1);
0051     for (unsigned i = 0; i < input2.sizeShape(); ++i) {
0052       vdata2.push_back(randedge(rng));
0053     }
0054
0055     // convert to server format
0056     input1.toServer(data1);
0057     input2.toServer(data2);
0058   }
0059   void makeOutput(const TritonOutputMap& iOutput, const std::string& debugName) const {
0060     //check the results
0061     const auto& output1 = iOutput.begin()->second;
0062     // convert from server format
0063     const auto& tmp = output1.fromServer<float>();
0064     if (brief_)
0065       edm::LogInfo(debugName) << "output shape: " << output1.shape()[0] << ", " << output1.shape()[1];
0066     else {
0067       edm::LogInfo msg(debugName);
0068       for (int i = 0; i < output1.shape()[0]; ++i) {
0069         msg << "output " << i << ": ";
0070         for (int j = 0; j < output1.shape()[1]; ++j) {
0071           msg << tmp[0][output1.shape()[1] * i + j] << " ";
0072         }
0073         msg << "\n";
0074       }
0075     }
0076   }
0077   static void fillPSetDescription(edm::ParameterSetDescription& desc) {
0078     desc.add<unsigned>("nodeMin", 100);
0079     desc.add<unsigned>("nodeMax", 4000);
0080     desc.add<unsigned>("edgeMin", 8000);
0081     desc.add<unsigned>("edgeMax", 15000);
0082     desc.add<bool>("brief", false);
0083   }
0084
0085 private:
0086   //members
0087   unsigned nodeMin_, nodeMax_;
0088   unsigned edgeMin_, edgeMax_;
0089   bool brief_;
0090 };
0091
0092 class TritonGraphProducer : public TritonEDProducer<> {
0093 public:
0094   explicit TritonGraphProducer(edm::ParameterSet const& cfg) : TritonEDProducer<>(cfg), helper_(cfg) {}
0095   void acquire(edm::Event const& iEvent, edm::EventSetup const& iSetup, Input& iInput) override {
0096     helper_.makeInput(iEvent, iInput);
0097   }
0098   void produce(edm::Event& iEvent, edm::EventSetup const& iSetup, Output const& iOutput) override {
0099     helper_.makeOutput(iOutput, debugName_);
0100   }
0101   ~TritonGraphProducer() override = default;
0102
0103   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0104     edm::ParameterSetDescription desc;
0105     TritonClient::fillPSetDescription(desc);
0106     TritonGraphHelper::fillPSetDescription(desc);
0107     //to ensure distinct cfi names
0108     descriptions.addWithDefaultLabel(desc);
0109   }
0110
0111 private:
0112   //member
0113   TritonGraphHelper helper_;
0114 };
0115
0116 DEFINE_FWK_MODULE(TritonGraphProducer);
0117
0118 #include "HeterogeneousCore/SonicTriton/interface/TritonEDFilter.h"
0119
0120 class TritonGraphFilter : public TritonEDFilter<> {
0121 public:
0122   explicit TritonGraphFilter(edm::ParameterSet const& cfg) : TritonEDFilter<>(cfg), helper_(cfg) {}
0123   void acquire(edm::Event const& iEvent, edm::EventSetup const& iSetup, Input& iInput) override {
0124     helper_.makeInput(iEvent, iInput);
0125   }
0126   bool filter(edm::Event& iEvent, edm::EventSetup const& iSetup, Output const& iOutput) override {
0127     helper_.makeOutput(iOutput, debugName_);
0128     return true;
0129   }
0130   ~TritonGraphFilter() override = default;
0131
0132   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0133     edm::ParameterSetDescription desc;
0134     TritonClient::fillPSetDescription(desc);
0135     TritonGraphHelper::fillPSetDescription(desc);
0136     //to ensure distinct cfi names
0137     descriptions.addWithDefaultLabel(desc);
0138   }
0139
0140 private:
0141   //member
0142   TritonGraphHelper helper_;
0143 };
0144
0145 DEFINE_FWK_MODULE(TritonGraphFilter);
0146
0147 #include "HeterogeneousCore/SonicTriton/interface/TritonOneEDAnalyzer.h"
0148
0149 class TritonGraphAnalyzer : public TritonOneEDAnalyzer<> {
0150 public:
0151   explicit TritonGraphAnalyzer(edm::ParameterSet const& cfg) : TritonOneEDAnalyzer<>(cfg), helper_(cfg) {}
0152   void acquire(edm::Event const& iEvent, edm::EventSetup const& iSetup, Input& iInput) override {
0153     helper_.makeInput(iEvent, iInput);
0154   }
0155   void analyze(edm::Event const& iEvent, edm::EventSetup const& iSetup, Output const& iOutput) override {
0156     helper_.makeOutput(iOutput, debugName_);
0157   }
0158   ~TritonGraphAnalyzer() override = default;
0159
0160   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0161     edm::ParameterSetDescription desc;
0162     TritonClient::fillPSetDescription(desc);
0163     TritonGraphHelper::fillPSetDescription(desc);
0164     //to ensure distinct cfi names
0165     descriptions.addWithDefaultLabel(desc);
0166   }
0167
0168 private:
0169   //member
0170   TritonGraphHelper helper_;
0171 };
0172
0173 DEFINE_FWK_MODULE(TritonGraphAnalyzer);