Migrate PyTorch alpaka modules to stream::FixedQueueEDProducer

This ensures that PyTorch sees only a limited number of device streams,
reducing the overall device memory utilisation.

Author: fwyzard

PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/DataSource.cc

@@ -7,17 +7,17 @@
 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/Event.h"
 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/EventSetup.h"
 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/MakerMacros.h"
-#include "HeterogeneousCore/AlpakaCore/interface/alpaka/stream/EDProducer.h"
+#include "HeterogeneousCore/AlpakaCore/interface/alpaka/stream/FixedQueueEDProducer.h"
 #include "HeterogeneousCore/AlpakaInterface/interface/config.h"
 #include "PhysicsTools/PyTorchAlpakaTest/interface/Environment.h"
 #include "PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/CommonKernels.h"
 
 namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest {
 
-  class DataSource : public stream::EDProducer<> {
+  class DataSource : public stream::FixedQueueEDProducer<> {
   public:
     DataSource(const edm::ParameterSet &params)
-        : EDProducer<>(params),
+        : FixedQueueEDProducer<>(params),
           particles_token_{produces()},
           images_token_{produces()},
           batch_size_(params.getParameter<uint32_t>("batchSize")),

PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/MaskedNet.cc

@@ -9,7 +9,7 @@
 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/Event.h"
 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/EventSetup.h"
 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/MakerMacros.h"
-#include "HeterogeneousCore/AlpakaCore/interface/alpaka/stream/EDProducer.h"
+#include "HeterogeneousCore/AlpakaCore/interface/alpaka/stream/FixedQueueEDProducer.h"
 #include "HeterogeneousCore/AlpakaInterface/interface/config.h"
 #include "PhysicsTools/PyTorchAlpaka/interface/TensorCollection.h"
 #include "PhysicsTools/PyTorchAlpaka/interface/alpaka/AlpakaModel.h"
@@ -18,10 +18,10 @@
 
 namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest {
 
-  class MaskedNet : public stream::EDProducer<> {
+  class MaskedNet : public stream::FixedQueueEDProducer<> {
   public:
     MaskedNet(const edm::ParameterSet &params)
-        : EDProducer<>(params),
+        : FixedQueueEDProducer<>(params),
           particles_token_(consumes(params.getParameter<edm::InputTag>("particles"))),
           masked_net_token_{produces()},
           model_(params.getParameter<edm::FileInPath>("model").fullPath()),

PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/MultiHeadNet.cc

@@ -8,18 +8,18 @@
 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/Event.h"
 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/EventSetup.h"
 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/MakerMacros.h"
-#include "HeterogeneousCore/AlpakaCore/interface/alpaka/stream/EDProducer.h"
+#include "HeterogeneousCore/AlpakaCore/interface/alpaka/stream/FixedQueueEDProducer.h"
 #include "HeterogeneousCore/AlpakaInterface/interface/config.h"
 #include "PhysicsTools/PyTorchAlpaka/interface/TensorCollection.h"
 #include "PhysicsTools/PyTorchAlpaka/interface/alpaka/AlpakaModel.h"
 #include "PhysicsTools/PyTorchAlpakaTest/interface/Environment.h"
 
 namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest {
 
-  class MultiHeadNet : public stream::EDProducer<> {
+  class MultiHeadNet : public stream::FixedQueueEDProducer<> {
   public:
     MultiHeadNet(const edm::ParameterSet &params)
-        : EDProducer<>(params),
+        : FixedQueueEDProducer<>(params),
           particles_token_(consumes(params.getParameter<edm::InputTag>("particles"))),
           multi_head_net_token_{produces()},
           model_(params.getParameter<edm::FileInPath>("model").fullPath()),

PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNet.cc

@@ -8,18 +8,18 @@
 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/Event.h"
 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/EventSetup.h"
 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/MakerMacros.h"
-#include "HeterogeneousCore/AlpakaCore/interface/alpaka/stream/EDProducer.h"
+#include "HeterogeneousCore/AlpakaCore/interface/alpaka/stream/FixedQueueEDProducer.h"
 #include "HeterogeneousCore/AlpakaInterface/interface/config.h"
 #include "PhysicsTools/PyTorchAlpaka/interface/TensorCollection.h"
 #include "PhysicsTools/PyTorchAlpaka/interface/alpaka/AlpakaModel.h"
 #include "PhysicsTools/PyTorchAlpakaTest/interface/Environment.h"
 
 namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest {
 
-  class SimpleNet : public stream::EDProducer<> {
+  class SimpleNet : public stream::FixedQueueEDProducer<> {
   public:
     SimpleNet(const edm::ParameterSet &params)
-        : EDProducer<>(params),
+        : FixedQueueEDProducer<>(params),
           particles_token_(consumes(params.getParameter<edm::InputTag>("particles"))),
           simple_net_token_{produces()},
           model_(params.getParameter<edm::FileInPath>("model").fullPath()),

PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/TinyResNet.cc

@@ -8,18 +8,18 @@
 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/Event.h"
 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/EventSetup.h"
 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/MakerMacros.h"
-#include "HeterogeneousCore/AlpakaCore/interface/alpaka/stream/EDProducer.h"
+#include "HeterogeneousCore/AlpakaCore/interface/alpaka/stream/FixedQueueEDProducer.h"
 #include "HeterogeneousCore/AlpakaInterface/interface/config.h"
 #include "PhysicsTools/PyTorchAlpaka/interface/TensorCollection.h"
 #include "PhysicsTools/PyTorchAlpaka/interface/alpaka/AlpakaModel.h"
 #include "PhysicsTools/PyTorchAlpakaTest/interface/Environment.h"
 
 namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest {
 
-  class TinyResNet : public stream::EDProducer<> {
+  class TinyResNet : public stream::FixedQueueEDProducer<> {
   public:
     TinyResNet(const edm::ParameterSet &params)
-        : EDProducer<>(params),
+        : FixedQueueEDProducer<>(params),
           images_token_(consumes(params.getParameter<edm::InputTag>("images"))),
           logits_token_{produces()},
           model_(params.getParameter<edm::FileInPath>("model").fullPath()),