AlpakaCore/plugins/ModuleTypeResolverAlpaka.cc

0001 #include "FWCore/Framework/interface/ModuleTypeResolverMaker.h"
0002 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0003 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0004 #include "FWCore/Utilities/interface/EDMException.h"
0005 #include "FWCore/Utilities/interface/thread_safety_macros.h"
0006
0007 #include <cassert>
0008 #include <cstring>
0009 #include <string>
0010 #include <unordered_map>
0011
0012 class ModuleTypeResolverAlpaka : public edm::ModuleTypeResolverBase {
0013 public:
0014   ModuleTypeResolverAlpaka(std::string backendPrefix) : backendPrefix_(std::move(backendPrefix)) {}
0015
0016   std::pair<std::string, int> resolveType(std::string basename, int index) const final {
0017     assert(index == kInitialIndex);
0018     constexpr auto kAlpaka = "@alpaka";
0019     auto found = basename.find(kAlpaka);
0020     if (found != std::string::npos) {
0021       if (backendPrefix_.empty()) {
0022         throw edm::Exception(edm::errors::LogicError)
0023             << "AlpakaModuleTypeResolver encountered a module with type name " << basename
0024             << " but the backend prefix was empty. This should not happen. Please contact framework developers.";
0025       }
0026       basename.replace(found, std::strlen(kAlpaka), "");
0027       basename = backendPrefix_ + basename;
0028     }
0029     return {basename, kLastIndex};
0030   }
0031
0032 private:
0033   std::string const backendPrefix_;
0034 };
0035
0036 class ModuleTypeResolverMakerAlpaka : public edm::ModuleTypeResolverMaker {
0037 public:
0038   ModuleTypeResolverMakerAlpaka() {}
0039
0040   std::shared_ptr<edm::ModuleTypeResolverBase const> makeResolver(edm::ParameterSet const& modulePSet) const final {
0041     std::string prefix = "";
0042     if (modulePSet.existsAs<edm::ParameterSet>("alpaka", false)) {
0043       auto const& backend =
0044           modulePSet.getUntrackedParameter<edm::ParameterSet>("alpaka").getUntrackedParameter<std::string>("backend");
0045       prefix = fmt::format("alpaka_{}::", backend);
0046
0047       LogDebug("AlpakaModuleTypeResolver")
0048           .format("AlpakaModuleTypeResolver: module {} backend prefix {}",
0049                   modulePSet.getParameter<std::string>("@module_label"),
0050                   prefix);
0051     }
0052     auto found = cache_.find(prefix);
0053     if (found == cache_.end()) {
0054       bool inserted;
0055       std::tie(found, inserted) = cache_.emplace(prefix, std::make_shared<ModuleTypeResolverAlpaka>(prefix));
0056     }
0057     return found->second;
0058   }
0059
0060 private:
0061   // no protection needed because this object is used only in single-thread context
0062   CMS_SA_ALLOW mutable std::unordered_map<std::string, std::shared_ptr<ModuleTypeResolverAlpaka const>> cache_;
0063 };
0064
0065 #include "FWCore/Framework/interface/ModuleTypeResolverMakerFactory.h"
0066 DEFINE_EDM_PLUGIN(edm::ModuleTypeResolverMakerFactory, ModuleTypeResolverMakerAlpaka, "ModuleTypeResolverAlpaka");