From 43bafa26768fee910359eb8dd23f7850c896769d Mon Sep 17 00:00:00 2001 From: Emanuele Coradin Date: Thu, 23 Apr 2026 10:19:03 +0200 Subject: [PATCH] Add PyTorchAlpaka mini-batching support --- PhysicsTools/PyTorchAlpaka/README.md | 27 +++++- .../interface/TensorCollection.h | 83 ++++++++++------ .../PyTorchAlpaka/interface/TensorHandle.h | 7 +- .../plugins/InspectionSink.cc | 48 +++++++++- .../plugins/alpaka/DataSource.cc | 10 +- .../plugins/alpaka/MaskedNet.cc | 12 +-- .../plugins/alpaka/MultiHeadNet.cc | 8 +- .../plugins/alpaka/SimpleNet.cc | 8 +- .../plugins/alpaka/SimpleNetMiniBatch.cc | 95 +++++++++++++++++++ .../plugins/alpaka/TinyResNet.cc | 8 +- .../plugins/alpaka/TinyResNetMiniBatch.cc | 93 ++++++++++++++++++ .../PyTorchAlpakaTest/python/options_cff.py | 28 +++++- .../test/runPyTorchAlpakaTest.py | 31 +++++- 13 files changed, 395 insertions(+), 63 deletions(-) create mode 100644 PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNetMiniBatch.cc create mode 100644 PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/TinyResNetMiniBatch.cc diff --git a/PhysicsTools/PyTorchAlpaka/README.md b/PhysicsTools/PyTorchAlpaka/README.md index a5f63b82f5a0b..4f2cd391ec2a8 100644 --- a/PhysicsTools/PyTorchAlpaka/README.md +++ b/PhysicsTools/PyTorchAlpaka/README.md @@ -8,8 +8,10 @@ All Pytorch based modules should add `PyTorchService` to disable internal torchl Examples demonstrating the interoperability of PyTorch with Alpaka in the CMSSW environment can be found in the [PyTorchAlpakaTest](../PyTorchAlpakaTest) directory. The basic test pipeline includes: - *SimpleNet* composed with few Dense layers, that operate on SoA style portable data structures +- *SimpleNetMiniBatch*, providing and example of inference perfomed in mini-batches - *MaskedNet* shows how to use multiple input data with `Eigen::Vector` and `SOA_SCALAR` - *TinyResNet* emulate more complex scenario with `Eigen::Matrix` and how one can implement image-like Tensor implementation +- *TinyResNetMiniBatch* to test the inference in mini-batches in a more complex scenario - *MulitHeadNet* handle networks that return more than one output tensor ## Direct Inference on SoA @@ -43,8 +45,8 @@ GENERATE_SOA_LAYOUT(SoAOutputTemplate, ``` - **Get Metarecords from Portable Collections:** ```cpp -PortableCollection deviceCollection(batch_size, queue); -PortableCollection deviceResultCollection(batch_size, queue); +PortableCollection deviceCollection(total_size, queue); +PortableCollection deviceResultCollection(total_size, queue); fill(queue, deviceCollection); auto records = deviceCollection.view().records(); auto result_records = deviceResultCollection.view().records(); @@ -53,14 +55,14 @@ auto result_records = deviceResultCollection.view().records(); **IMPORTANT:** continuity of memory is a strict requirement! ``` -TensorCollection input(batch_size); +TensorCollection input(total_size); input.add("eigen_vector", records.a(), records.b()); input.add("eigen_matrix", records.c()); input.add("column", records.x(), records.y(), records.z()); input.add("scalar", records.type()); input.change_order({"column", "scalar", "eigen_matrix", "eigen_vector"}); -TensorCollection output(batch_size); +TensorCollection output(total_size); output.add("result", result_view.cluster()); ``` @@ -72,6 +74,23 @@ After adding all the blocks to the `TensorCollection`, the order of the blocks f More examples about usage can be found in [PyTorchAlpakaTest](../PyTorchAlpakaTest). +### Batching semantics + +When using batched inference, `TensorCollection` is constructed with `(total_size, total_size)` and internally manages batch offsets. + +**IMPORTANT:** the batchsize should be chosen carefully in order to respect the alignment (typically a multiple of 32). Otherwise, an assert will be trigged. + +The `batch_id` passed to `add()` selects which batch slice is exposed to the model. + +Runtime checks are performed to ensure: +- valid batch indices +- consistency between batch size and total size +- memory contiguity between columns + +These checks rely on `assert`. + +Look at [SimpleNetMiniBatch](PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNetMiniBatch.cc) to have an example. + ## Limitations - Current implementation supports `SerialSync` and `CudaAsync` backends only. `ROCmAsync` backend is supported via SerialSync fallback mechanism due to missing `pytorch-hip` library in CMSSW (see: https://github.com/pytorch/pytorch/blob/main/aten/CMakeLists.txt#L75), with explicit `alpaka::wait()` call to copy data to host and back to device. - Const correctness and thread-safety relies on `torch::from_blob()` mechanism which currently does not ensure that data will not be modified internally. There is ongoing work to support COW tensors but until this support will be integrated in mainstream PyTorch the provided solution materialises (copies) the tensors if passed registry points to `const` memory. For more information please check [Const correctness and thread-safety of torch::from_blob with external memory](https://discuss.pytorch.org/t/const-correctness-and-thread-safety-of-torch-from-blob-with-external-memory/223521) and [pytorch:#97856](https://github.com/pytorch/pytorch/issues/97856) diff --git a/PhysicsTools/PyTorchAlpaka/interface/TensorCollection.h b/PhysicsTools/PyTorchAlpaka/interface/TensorCollection.h index 16494ee02e67e..71c206f64d3c0 100644 --- a/PhysicsTools/PyTorchAlpaka/interface/TensorCollection.h +++ b/PhysicsTools/PyTorchAlpaka/interface/TensorCollection.h @@ -5,6 +5,7 @@ #include #include #include +#include #include #include @@ -69,11 +70,18 @@ namespace cms::torch::alpakatools { // SOA_COLUMN(float, phi)) // // can register the following: - // TensorCollection registry(batch_size); + // + // TensorCollection registry(batch_size, total_size); + // registry.add("features", batch_id, records.pt(), records.eta(), records.phi()); + // + // In the above example, the add function automatically computes the offset for the batch and ensures the provided columns are contiguous in memory. + // If the user wants to perform inference on the entire dataset without batching, he can simply register by passing just the total size: + // + // TensorCollection registry(total_size); // registry.add("features", records.pt(), records.eta(), records.phi()); // - // but if want to use only pt() and phi() then below will not work as pt() and phi() are not contiguous: - // TensorCollection registry(batch_size); + // If the user wants to use only pt() and phi() then below will not work as pt() and phi() are not contiguous: + // TensorCollection registry(batch_size, total_size); // registry.add("features", records.pt(), records.phi()); // // potential solution would be to arrange layout dependent on model requirements @@ -91,24 +99,31 @@ namespace cms::torch::alpakatools { friend class alpaka_rocm_async::torch::AlpakaModel; friend class alpaka_serial_sync::torch::AlpakaModel; - explicit TensorCollection(int batch_size) : batch_size_(batch_size) {} + explicit TensorCollection(int total_size) : batch_size_(total_size), total_size_(total_size) { assert_sizes(); } + explicit TensorCollection(int batch_size, int total_size) : batch_size_(batch_size), total_size_(total_size) { + assert_sizes(); + } // SOA_EIGEN_COLUMN template requires(SameValueType && TSoAParamsImpl::columnType == cms::soa::SoAColumnType::eigen) void add(const std::string& name, - int batch_size, + int batch_id, std::tuple column, std::tuple... others) { using DataType = typename TSoAParamsImpl::ScalarType; + assert_batch_id(batch_id); + int offset = batch_id * batch_size_; auto ptr = std::get<0>(column).data(); int n_elems = - cms::torch::alpakatools::detail::num_elements_per_column(batch_size, SoALayout::alignment, sizeof(DataType)); + cms::torch::alpakatools::detail::num_elements_per_column(total_size_, SoALayout::alignment, sizeof(DataType)); assert_location( n_elems * TSoAParamsImpl::ValueType::RowsAtCompileTime * TSoAParamsImpl::ValueType::ColsAtCompileTime, ptr, std::get<0>(others).data()...); + ptr += offset; + std::vector tensor_dims; if constexpr (TSoAParamsImpl::ValueType::ColsAtCompileTime > 1) tensor_dims = {1 + sizeof...(Others), @@ -117,16 +132,18 @@ namespace cms::torch::alpakatools { else tensor_dims = {1 + sizeof...(Others), TSoAParamsImpl::ValueType::RowsAtCompileTime}; - emplace_tensor(name, SoALayout::alignment, ptr, batch_size, tensor_dims); + // Handle the case in which the last batch contains less elements + auto effective_batch_size = std::min(batch_size_, total_size_ - offset); + emplace_tensor(name, SoALayout::alignment, ptr, effective_batch_size, total_size_, tensor_dims); } - // SOA_EIGEN_COLUMN with default batch size + // SOA_EIGEN_COLUMN with default batch size = default size template requires(SameValueType && TSoAParamsImpl::columnType == cms::soa::SoAColumnType::eigen) void add(const std::string& name, std::tuple column, std::tuple... others) { - add(name, batch_size_, column, others...); + add(name, 0, column, others...); } // SOA_COLUMN @@ -134,43 +151,39 @@ namespace cms::torch::alpakatools { requires(SameScalarType && TSoAParamsImpl::columnType == cms::soa::SoAColumnType::column) void add(const std::string& name, - int batch_size, + int batch_id, std::tuple column, std::tuple... others) { using DataType = typename TSoAParamsImpl::ScalarType; - int n_elems = - cms::torch::alpakatools::detail::num_elements_per_column(batch_size, SoALayout::alignment, sizeof(DataType)); - assert_location(n_elems, std::get<0>(column).data(), std::get<0>(others).data()...); + assert_batch_id(batch_id); + int offset = batch_id * batch_size_; auto ptr = std::get<0>(column).data(); - emplace_tensor(name, SoALayout::alignment, ptr, batch_size, {1 + sizeof...(Others)}); + int n_elems = + cms::torch::alpakatools::detail::num_elements_per_column(total_size_, SoALayout::alignment, sizeof(DataType)); + assert_location(n_elems, ptr, std::get<0>(others).data()...); + + ptr += offset; + auto effective_batch_size = std::min(batch_size_, total_size_ - offset); + emplace_tensor(name, SoALayout::alignment, ptr, effective_batch_size, total_size_, {1 + sizeof...(Others)}); } - // SOA_COLUMN with default batch size + // SOA_COLUMN with default batch size = total size template requires(SameScalarType && TSoAParamsImpl::columnType == cms::soa::SoAColumnType::column) void add(const std::string& name, std::tuple column, std::tuple... others) { - add(name, batch_size_, column, others...); + add(name, 0, column, others...); } // SOA_SCALAR template requires(std::is_arithmetic_v && column_t == cms::soa::SoAColumnType::scalar) void add(const std::string& name, - int batch_size, std::tuple, cms::soa::size_type> column) { auto ptr = std::get<0>(column).data(); - emplace_tensor(name, SoALayout::alignment, ptr, batch_size, {1}, true); - } - - // SOA_SCALAR with default batch size - template - requires(std::is_arithmetic_v && column_t == cms::soa::SoAColumnType::scalar) - void add(const std::string& name, - std::tuple, cms::soa::size_type> column) { - add(name, batch_size_, column); + emplace_tensor(name, SoALayout::alignment, ptr, batch_size_, total_size_, {1}, true); } // The order is defined by the order `add()` is called. @@ -201,16 +214,32 @@ namespace cms::torch::alpakatools { size_t alignment, Tptr ptr, int batch_size, + int total_size, std::vector dims = {1}, const bool is_scalar = false) { using T = std::remove_pointer_t; registry_.try_emplace(name, std::make_unique>( - alignment, sizeof(T), ptr, batch_size, std::move(dims), is_scalar)); + alignment, sizeof(T), ptr, batch_size, total_size, std::move(dims), is_scalar)); order_.push_back(name); } + void assert_sizes() { + assert(total_size_ >= 0 && "Total size must be positive!"); + if (batch_size_ == 0) { + assert(total_size_ == 0 && "Batch size 0 only allowed when total size is 0"); + return; + } + assert(batch_size_ > 0 && "Batch size must be positive!"); + } + + void assert_batch_id(int batch_id) { + assert(batch_id >= 0 && "Batch id must be non-negative!"); + assert((total_size_ == 0 || (batch_id * batch_size_ < total_size_)) && "Batch id is out of bounds!"); + } + int batch_size_; + int total_size_; std::vector order_; std::unordered_map>> registry_; }; diff --git a/PhysicsTools/PyTorchAlpaka/interface/TensorHandle.h b/PhysicsTools/PyTorchAlpaka/interface/TensorHandle.h index 894f617a367ef..eaef77ddac10c 100644 --- a/PhysicsTools/PyTorchAlpaka/interface/TensorHandle.h +++ b/PhysicsTools/PyTorchAlpaka/interface/TensorHandle.h @@ -86,13 +86,15 @@ namespace cms::torch::alpakatools::detail { const size_t bytes, T* data, const int batch_size, + const int total_size, const std::vector dims, const bool is_scalar = false) : alignment_(alignment), bytes_(bytes), data_(data), + total_size_(total_size), dims_(batch_size, dims, is_scalar), - policy_(data, dims_.volume() * num_elements_per_column(batch_size, alignment, bytes)) { + policy_(data, dims_.volume() * num_elements_per_column(total_size, alignment, bytes)) { init_sizes(); init_strides(); } @@ -130,7 +132,7 @@ namespace cms::torch::alpakatools::detail { strides_ = std::vector(N); int per_bunch = alignment_ / bytes_; - int bunches = std::ceil(1.0 * dims_.batch_size() / per_bunch); + int bunches = (total_size_ + per_bunch - 1) / per_bunch; // base stride initialization if (!dims_.is_scalar()) @@ -160,6 +162,7 @@ namespace cms::torch::alpakatools::detail { const size_t alignment_; const size_t bytes_; T* data_; + const int total_size_; const Dims dims_; std::vector strides_; diff --git a/PhysicsTools/PyTorchAlpakaTest/plugins/InspectionSink.cc b/PhysicsTools/PyTorchAlpakaTest/plugins/InspectionSink.cc index 2238d780abff8..783f9768798a4 100644 --- a/PhysicsTools/PyTorchAlpakaTest/plugins/InspectionSink.cc +++ b/PhysicsTools/PyTorchAlpakaTest/plugins/InspectionSink.cc @@ -32,26 +32,34 @@ namespace torchtest { : environment_{static_cast(params.getUntrackedParameter("environment"))}, particles_token_{consumes(params.getParameter("particles"))}, simple_net_token_{consumes(params.getParameter("simple_net"))}, + simple_net_minibatch_token_{consumes(params.getParameter("simple_net_minibatch"))}, masked_net_token_{consumes(params.getParameter("masked_net"))}, multi_head_net_token_{consumes(params.getParameter("multi_head_net"))}, images_token_{consumes(params.getParameter("images"))}, logits_token_{consumes(params.getParameter("resnet18"))}, + logits_minibatch_token_{consumes(params.getParameter("resnet18_minibatch"))}, particles_backend_{consumes(getBackendTag(params.getParameter("particles")))}, simple_net_backend_{consumes(getBackendTag(params.getParameter("simple_net")))}, + simple_net_minibatch_backend_{ + consumes(getBackendTag(params.getParameter("simple_net_minibatch")))}, masked_net_backend_{consumes(getBackendTag(params.getParameter("masked_net")))}, multi_head_net_backend_{consumes(getBackendTag(params.getParameter("multi_head_net")))}, images_backend_{consumes(getBackendTag(params.getParameter("images")))}, - logits_backend_{consumes(getBackendTag(params.getParameter("resnet18")))} {} + logits_backend_{consumes(getBackendTag(params.getParameter("resnet18")))}, + logits_minibatch_backend_{consumes(getBackendTag(params.getParameter("resnet18_minibatch")))} { + } static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) { edm::ParameterSetDescription desc; desc.addUntracked("environment", static_cast(Environment::kProduction)); desc.add("particles"); desc.add("simple_net"); + desc.add("simple_net_minibatch"); desc.add("masked_net"); desc.add("multi_head_net"); desc.add("images"); desc.add("resnet18"); + desc.add("resnet18_minibatch"); descriptions.addWithDefaultLabel(desc); } @@ -68,10 +76,12 @@ namespace torchtest { // particles auto particles_handle = event.getHandle(particles_token_); auto simple_net_handle = event.getHandle(simple_net_token_); + auto simple_net_minibatch_handle = event.getHandle(simple_net_minibatch_token_); auto masked_net_handle = event.getHandle(masked_net_token_); auto multi_head_net_handle = event.getHandle(multi_head_net_token_); auto images_handle = event.getHandle(images_token_); auto logits_handle = event.getHandle(logits_token_); + auto logits_minibatch_handle = event.getHandle(logits_minibatch_token_); // debug if (environment_ >= Environment::kDevelopment) { @@ -94,6 +104,12 @@ namespace torchtest { auto const simple_net_backend = static_cast(event.get(simple_net_backend_)); print(simple_net.const_view(), cms::alpakatools::toString(simple_net_backend)); } + if (simple_net_minibatch_handle.isValid()) { + auto const& simple_net_minibatch = *simple_net_minibatch_handle; + auto const simple_net_minibatch_backend = + static_cast(event.get(simple_net_minibatch_backend_)); + print(simple_net_minibatch.const_view(), cms::alpakatools::toString(simple_net_minibatch_backend)); + } // masked_net if (masked_net_handle.isValid()) { auto const& masked_net = *masked_net_handle; @@ -145,6 +161,12 @@ namespace torchtest { assert(std::abs(sum - 1.0) < 1e-4); } } + if (logits_minibatch_handle) { + auto const& logits_minibatch = *logits_minibatch_handle; + auto const logits_minibatch_backend = + static_cast(event.get(logits_minibatch_backend_)); + print(logits_minibatch.const_view(), cms::alpakatools::toString(logits_minibatch_backend)); + } } } @@ -169,21 +191,29 @@ namespace torchtest { const edm::EDGetTokenT particles_token_; const edm::EDGetTokenT simple_net_token_; + const edm::EDGetTokenT simple_net_minibatch_token_; const edm::EDGetTokenT masked_net_token_; const edm::EDGetTokenT multi_head_net_token_; const edm::EDGetTokenT images_token_; const edm::EDGetTokenT logits_token_; + const edm::EDGetTokenT logits_minibatch_token_; const edm::EDGetTokenT particles_backend_; const edm::EDGetTokenT simple_net_backend_; + const edm::EDGetTokenT simple_net_minibatch_backend_; const edm::EDGetTokenT masked_net_backend_; const edm::EDGetTokenT multi_head_net_backend_; const edm::EDGetTokenT images_backend_; const edm::EDGetTokenT logits_backend_; + const edm::EDGetTokenT logits_minibatch_backend_; const int32_t kMaxView = 5; void print(const portabletest::LogitsHostCollection::ConstView& logits, const std::string_view logits_backend) { + if (logits.metadata().size() == 0) { + fmt::print("[DEBUG] LogitsCollection[0]: empty\n"); + return; + } const int rows = portabletest::LogitsType::RowsAtCompileTime; constexpr auto line = "+------+------+------+------+------+------+------+------+------+------+\n"; fmt::memory_buffer buffer; @@ -200,6 +230,10 @@ namespace torchtest { void print(const portabletest::ImageHostCollection::ConstView& images, const std::string_view images_backend) { const auto size = images.metadata().size(); + if (size == 0) { + fmt::print("[DEBUG] ImageCollection[0]: empty\n"); + return; + } const int rows = portabletest::ColorChannel::RowsAtCompileTime; const int cols = portabletest::ColorChannel::ColsAtCompileTime; constexpr auto line = "+-------+-------+-------+-------+-------+-------+-------+-------+-------+\n"; @@ -240,6 +274,10 @@ namespace torchtest { const std::string_view multi_head_net_backend) { constexpr auto line = "+-------+-----------------+-------+-------+-------+\n"; const auto size = multi_head_net.metadata().size(); + if (size == 0) { + fmt::print("[DEBUG] MultiHeadNetCollection[0]: empty\n"); + return; + } fmt::memory_buffer buffer; // Header message @@ -285,6 +323,10 @@ namespace torchtest { const std::string& label = "SimpleNetCollection") { constexpr auto line = "+-------+---------+\n"; const auto size = simple_net.metadata().size(); + if (size == 0) { + fmt::print("[DEBUG] SimpleNetCollection[0]: empty\n"); + return; + } fmt::memory_buffer buffer; // Header message @@ -313,6 +355,10 @@ namespace torchtest { const std::string_view particles_backend) { constexpr auto line = "+-------+---------+---------+---------+\n"; const auto size = particles.metadata().size(); + if (size == 0) { + fmt::print("[DEBUG] ParticleCollection[0]: empty\n"); + return; + } fmt::memory_buffer buffer; // Header message diff --git a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/DataSource.cc b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/DataSource.cc index 469d5cd93202d..feeb99603686a 100644 --- a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/DataSource.cc +++ b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/DataSource.cc @@ -20,13 +20,13 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest { : FixedQueueEDProducer<>(params), particles_token_{produces()}, images_token_{produces()}, - batch_size_(params.getParameter("batchSize")), + total_size_(params.getParameter("totalSize")), environment_{static_cast<::torchtest::Environment>(params.getUntrackedParameter("environment"))} {} void produce(device::Event &event, const device::EventSetup &event_setup) override { // allocate data sources - auto particles = portabletest::ParticleDeviceCollection(event.queue(), batch_size_); - auto images = portabletest::ImageDeviceCollection(event.queue(), batch_size_); + auto particles = portabletest::ParticleDeviceCollection(event.queue(), total_size_); + auto images = portabletest::ImageDeviceCollection(event.queue(), total_size_); // fill data kernels::randomFillParticleCollection(event.queue(), particles); @@ -39,7 +39,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest { static void fillDescriptions(edm::ConfigurationDescriptions &descriptions) { edm::ParameterSetDescription desc; - desc.add("batchSize"); + desc.add("totalSize"); desc.addUntracked("environment", static_cast(::torchtest::Environment::kProduction)); descriptions.addWithDefaultLabel(desc); } @@ -47,7 +47,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest { private: const device::EDPutToken particles_token_; const device::EDPutToken images_token_; - const uint32_t batch_size_; + const uint32_t total_size_; const ::torchtest::Environment environment_; }; diff --git a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/MaskedNet.cc b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/MaskedNet.cc index 74d8ae73bb62a..0ce91bb4b87fc 100644 --- a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/MaskedNet.cc +++ b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/MaskedNet.cc @@ -38,14 +38,14 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest { void produce(device::Event &event, const device::EventSetup &event_setup) override { // in/out collections const auto &particles = event.get(particles_token_); - const auto batch_size = particles.const_view().metadata().size(); - auto masked_net_output = portabletest::SimpleNetDeviceCollection(event.queue(), batch_size); + const auto total_size = particles.const_view().metadata().size(); + auto masked_net_output = portabletest::SimpleNetDeviceCollection(event.queue(), total_size); // mask - auto mask = portabletest::MaskDeviceCollection(event.queue(), batch_size); + auto mask = portabletest::MaskDeviceCollection(event.queue(), total_size); kernels::fillMask(event.queue(), mask); // note that scalar mask can be used to mask out entire batch at once (scalars are broadcasted) - // auto scalar_mask = ScalarMaskDeviceCollection(batch_size, event.queue()); + // auto scalar_mask = ScalarMaskDeviceCollection(total_size, event.queue()); // scalar_mask.zeroInitialise(event.queue()); // records @@ -54,14 +54,14 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest { // auto scalar_mask_records = scalar_mask.view().records(); auto output_records = masked_net_output.view().records(); // input tensor definition - cms::torch::alpakatools::TensorCollection inputs(batch_size); + cms::torch::alpakatools::TensorCollection inputs(total_size); inputs.add( "particles", particle_records.pt(), particle_records.eta(), particle_records.phi()); // note override of default `ParticleSoA` layout with `MaskSoA` inputs.add("mask", mask_records.mask()); // inputs.add("scalar_mask", scalar_mask_records.scalar_mask()); // output tensor definition - cms::torch::alpakatools::TensorCollection outputs(batch_size); + cms::torch::alpakatools::TensorCollection outputs(total_size); outputs.add("regression_head", output_records.reco_pt()); // metadata for automatic tensor conversion // ModelMetadata metadata(inputs, outputs); diff --git a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/MultiHeadNet.cc b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/MultiHeadNet.cc index a449fda460a07..322362e3289cd 100644 --- a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/MultiHeadNet.cc +++ b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/MultiHeadNet.cc @@ -36,17 +36,17 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest { void produce(device::Event &event, const device::EventSetup &event_setup) override { // in/out collections const auto &particles = event.get(particles_token_); - const auto batch_size = particles.const_view().metadata().size(); - auto multi_head_output = portabletest::MultiHeadNetDeviceCollection(event.queue(), batch_size); + const auto total_size = particles.const_view().metadata().size(); + auto multi_head_output = portabletest::MultiHeadNetDeviceCollection(event.queue(), total_size); // records auto input_records = particles.const_view().records(); auto output_records = multi_head_output.view().records(); // input tensor definition - cms::torch::alpakatools::TensorCollection inputs(batch_size); + cms::torch::alpakatools::TensorCollection inputs(total_size); inputs.add("particles", input_records.pt(), input_records.eta(), input_records.phi()); // output tensor definition - cms::torch::alpakatools::TensorCollection outputs(batch_size); + cms::torch::alpakatools::TensorCollection outputs(total_size); outputs.add("regression_head", output_records.regression_head()); outputs.add("classification_head", output_records.classification_head()); diff --git a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNet.cc b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNet.cc index ec99a096256bf..bdb23913edf69 100644 --- a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNet.cc +++ b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNet.cc @@ -36,17 +36,17 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest { void produce(device::Event &event, const device::EventSetup &event_setup) override { // in/out collections const auto &particles = event.get(particles_token_); - const auto batch_size = particles.const_view().metadata().size(); - auto regression_collection = portabletest::SimpleNetDeviceCollection(event.queue(), batch_size); + const auto total_size = particles.const_view().metadata().size(); + auto regression_collection = portabletest::SimpleNetDeviceCollection(event.queue(), total_size); // records auto input_records = particles.const_view().records(); auto output_records = regression_collection.view().records(); // input tensor definition - cms::torch::alpakatools::TensorCollection inputs(batch_size); + cms::torch::alpakatools::TensorCollection inputs(total_size); inputs.add("particles", input_records.pt(), input_records.eta(), input_records.phi()); // output tensor definition - cms::torch::alpakatools::TensorCollection outputs(batch_size); + cms::torch::alpakatools::TensorCollection outputs(total_size); outputs.add("regression_head", output_records.reco_pt()); model_.forward(event.queue(), inputs, outputs); diff --git a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNetMiniBatch.cc b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNetMiniBatch.cc new file mode 100644 index 0000000000000..4c39b51d0cece --- /dev/null +++ b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNetMiniBatch.cc @@ -0,0 +1,95 @@ +#include +#include + +#include "HeterogeneousCore/AlpakaCore/interface/alpaka/EDPutToken.h" +#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/AlpakaInterface/interface/config.h" +#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h" +#include "FWCore/ParameterSet/interface/ParameterSet.h" +#include "FWCore/ParameterSet/interface/ParameterSetDescription.h" +#include "DataFormats/PortableTestObjects/interface/TestSoA.h" +#include "DataFormats/PortableTestObjects/interface/alpaka/ParticleDeviceCollection.h" +#include "DataFormats/PortableTestObjects/interface/alpaka/SimpleNetDeviceCollection.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 { + + struct BatchIO { + cms::torch::alpakatools::TensorCollection inputs; + cms::torch::alpakatools::TensorCollection outputs; + }; + + class SimpleNetMiniBatch : public stream::EDProducer<> { + public: + SimpleNetMiniBatch(const edm::ParameterSet ¶ms) + : EDProducer<>(params), + particles_token_(consumes(params.getParameter("particles"))), + simple_net_token_{produces()}, + model_(params.getParameter("model").fullPath()), + batch_size_(params.getParameter("batchSize")), + environment_{static_cast<::torchtest::Environment>(params.getUntrackedParameter("environment"))} {} + + static void fillDescriptions(edm::ConfigurationDescriptions &descriptions) { + edm::ParameterSetDescription desc; + desc.add("model"); + desc.add("batchSize"); + desc.add("particles"); + desc.addUntracked("environment", static_cast(::torchtest::Environment::kProduction)); + descriptions.addWithDefaultLabel(desc); + } + + void produce(device::Event &event, const device::EventSetup &event_setup) override { + // in/out collections + const auto &particles = event.get(particles_token_); + const auto total_size = particles.const_view().metadata().size(); + auto regression_collection = portabletest::SimpleNetDeviceCollection(event.queue(), total_size); + + int n_batches; + if (batch_size_ == 0) { + assert(total_size == 0 && "Batch size can be 0 only if the total size is 0"); + n_batches = 1; + } else + n_batches = (total_size + batch_size_ - 1) / batch_size_; + + // records + auto input_records = particles.const_view().records(); + auto output_records = regression_collection.view().records(); + + // input and output tensor definitions + std::deque batches; + for (int i_batch = 0; i_batch < n_batches; ++i_batch) { + BatchIO batch{cms::torch::alpakatools::TensorCollection(batch_size_, total_size), + cms::torch::alpakatools::TensorCollection(batch_size_, total_size)}; + + batch.inputs.add( + "particles", i_batch, input_records.pt(), input_records.eta(), input_records.phi()); + + batch.outputs.add("regression_head", i_batch, output_records.reco_pt()); + batches.push_back(std::move(batch)); + } + // forward pass on mini-batches + for (auto &batch : batches) { + model_.forward(event.queue(), batch.inputs, batch.outputs); + } + // put device-side product into event + event.emplace(simple_net_token_, std::move(regression_collection)); + } + + private: + // event query tokens + const device::EDGetToken particles_token_; + const device::EDPutToken simple_net_token_; + // model + torch::AlpakaModel model_; + const int batch_size_; + // debug mode flag + const ::torchtest::Environment environment_; + }; +} // namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest + +DEFINE_FWK_ALPAKA_MODULE(torchtest::SimpleNetMiniBatch); diff --git a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/TinyResNet.cc b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/TinyResNet.cc index 30bb682379d93..88e02cc41cf12 100644 --- a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/TinyResNet.cc +++ b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/TinyResNet.cc @@ -36,17 +36,17 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest { void produce(device::Event &event, const device::EventSetup &event_setup) override { // in/out collections const auto &images = event.get(images_token_); - const auto batch_size = images.const_view().metadata().size(); - auto logits = portabletest::LogitsDeviceCollection(event.queue(), batch_size); + const auto total_size = images.const_view().metadata().size(); + auto logits = portabletest::LogitsDeviceCollection(event.queue(), total_size); // records auto input_records = images.const_view().records(); auto output_records = logits.view().records(); // input tensor definition - cms::torch::alpakatools::TensorCollection inputs(batch_size); + cms::torch::alpakatools::TensorCollection inputs(total_size); inputs.add("images", input_records.r(), input_records.g(), input_records.b()); // output tensor definition - cms::torch::alpakatools::TensorCollection outputs(batch_size); + cms::torch::alpakatools::TensorCollection outputs(total_size); outputs.add("logits", output_records.logits()); model_.forward(event.queue(), inputs, outputs); diff --git a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/TinyResNetMiniBatch.cc b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/TinyResNetMiniBatch.cc new file mode 100644 index 0000000000000..29dda17541ea4 --- /dev/null +++ b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/TinyResNetMiniBatch.cc @@ -0,0 +1,93 @@ +#include "DataFormats/PortableTestObjects/interface/TestSoA.h" +#include "DataFormats/PortableTestObjects/interface/alpaka/ImageDeviceCollection.h" +#include "DataFormats/PortableTestObjects/interface/alpaka/LogitsDeviceCollection.h" +#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h" +#include "FWCore/ParameterSet/interface/ParameterSet.h" +#include "FWCore/ParameterSet/interface/ParameterSetDescription.h" +#include "HeterogeneousCore/AlpakaCore/interface/alpaka/EDPutToken.h" +#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/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 { + struct BatchIO { + cms::torch::alpakatools::TensorCollection inputs; + cms::torch::alpakatools::TensorCollection outputs; + }; + + class TinyResNetMiniBatch : public stream::EDProducer<> { + public: + TinyResNetMiniBatch(const edm::ParameterSet ¶ms) + : EDProducer<>(params), + images_token_(consumes(params.getParameter("images"))), + logits_token_{produces()}, + model_(params.getParameter("model").fullPath()), + batch_size_(params.getParameter("batchSize")), + environment_{static_cast<::torchtest::Environment>(params.getUntrackedParameter("environment"))} {} + + static void fillDescriptions(edm::ConfigurationDescriptions &descriptions) { + edm::ParameterSetDescription desc; + desc.add("model"); + desc.add("batchSize"); + desc.add("images"); + desc.addUntracked("environment", static_cast(::torchtest::Environment::kProduction)); + descriptions.addWithDefaultLabel(desc); + } + + void produce(device::Event &event, const device::EventSetup &event_setup) override { + // in/out collections + const auto &images = event.get(images_token_); + const auto total_size = images.const_view().metadata().size(); + auto logits = portabletest::LogitsDeviceCollection(event.queue(), total_size); + + int n_batches; + if (batch_size_ == 0) { + assert(total_size == 0 && "Batch size can be 0 only if the total size is 0"); + n_batches = 1; + } else + n_batches = (total_size + batch_size_ - 1) / batch_size_; + // records + auto input_records = images.const_view().records(); + auto output_records = logits.view().records(); + + // input and output tensor definitions + std::deque batches; + for (int i_batch = 0; i_batch < n_batches; ++i_batch) { + BatchIO batch{cms::torch::alpakatools::TensorCollection(batch_size_, total_size), + cms::torch::alpakatools::TensorCollection(batch_size_, total_size)}; + + batch.inputs.add( + "images", i_batch, input_records.r(), input_records.g(), input_records.b()); + + batch.outputs.add("logits", i_batch, output_records.logits()); + batches.push_back(std::move(batch)); + } + + // forward pass on mini-batches + for (auto &batch : batches) { + model_.forward(event.queue(), batch.inputs, batch.outputs); + } + + // put device-side product into event + event.emplace(logits_token_, std::move(logits)); + } + + private: + // event query tokens + const device::EDGetToken images_token_; + const device::EDPutToken logits_token_; + // model + torch::AlpakaModel model_; + const int batch_size_; + // debug mode flag + const ::torchtest::Environment environment_; + }; + +} // namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest + +DEFINE_FWK_ALPAKA_MODULE(torchtest::TinyResNetMiniBatch); diff --git a/PhysicsTools/PyTorchAlpakaTest/python/options_cff.py b/PhysicsTools/PyTorchAlpakaTest/python/options_cff.py index c428d90506e5e..590847f2be9b8 100644 --- a/PhysicsTools/PyTorchAlpakaTest/python/options_cff.py +++ b/PhysicsTools/PyTorchAlpakaTest/python/options_cff.py @@ -35,10 +35,17 @@ def parse_args(): help="Accelerator backend" ) + parser.add_argument( + "-ts", "--totalSize", + type=int, + default=35, + help="Total size" + ) + parser.add_argument( "-bs", "--batchSize", type=int, - default=8, + default=32, help="Batch size" ) @@ -57,6 +64,14 @@ def parse_args(): help="SimpleNet model (just-in-time compiled)" ) + parser.add_argument( + "--simpleNetMiniBatch", + type=str, + default="PhysicsTools/PyTorchAlpakaTest/data/SimpleNet.pt", + help="SimpleNetMiniBatch model (just-in-time compiled)" + ) + + parser.add_argument( "--maskedNet", type=str, @@ -78,11 +93,18 @@ def parse_args(): help="TinyResNet model (just-in-time compiled)" ) + parser.add_argument( + "--tinyResNetMiniBatch", + type=str, + default="PhysicsTools/PyTorchAlpakaTest/data/TinyResNet.pt", + help="TinyResNetMiniBatch model (just-in-time compiled)" + ) + parser.add_argument( "-o", "--only", nargs="+", - default=["SimpleNet", "MultiHeadNet", "MaskedNet", "TinyResNet"], - choices=["SimpleNet", "MaskedNet", "MultiHeadNet", "TinyResNet"], + default=["SimpleNet", "SimpleNetMiniBatch", "MultiHeadNet", "MaskedNet", "TinyResNet", "TinyResNetMiniBatch"], + choices=["SimpleNet", "SimpleNetMiniBatch", "MultiHeadNet", "MaskedNet", "TinyResNet", "TinyResNetMiniBatch"], help="Run selected test(s). Default: all modules run in parallel." ) diff --git a/PhysicsTools/PyTorchAlpakaTest/test/runPyTorchAlpakaTest.py b/PhysicsTools/PyTorchAlpakaTest/test/runPyTorchAlpakaTest.py index c8c5fe857d0f1..d2d2e861548cf 100755 --- a/PhysicsTools/PyTorchAlpakaTest/test/runPyTorchAlpakaTest.py +++ b/PhysicsTools/PyTorchAlpakaTest/test/runPyTorchAlpakaTest.py @@ -32,7 +32,7 @@ process.path = cms.Path() # data provider process.DataSource = torchtest_DataSource_alpaka( - batchSize = cms.uint32(args.batchSize if args.batchSize > 1 else 1), + totalSize = cms.uint32(args.totalSize if args.totalSize >= 0 else 0), alpaka = cms.untracked.PSet( backend = cms.untracked.string(args.backend) ), @@ -41,7 +41,7 @@ process.path += process.DataSource # --only SimpleNet if "SimpleNet" in args.only: - from PhysicsTools.PyTorchAlpakaTest.modules import torchtest_SimpleNet_alpaka + from PhysicsTools.PyTorchAlpakaTest.modules import torchtest_SimpleNet_alpaka, torchtest_SimpleNetMiniBatch_alpaka process.SimpleNet = torchtest_SimpleNet_alpaka( model = cms.FileInPath(args.simpleNet), particles = 'DataSource', @@ -51,6 +51,17 @@ environment = cms.untracked.int32(args.environment) ) process.path += process.SimpleNet +if "SimpleNetMiniBatch" in args.only: + process.SimpleNetMiniBatch = torchtest_SimpleNetMiniBatch_alpaka( + model = cms.FileInPath(args.simpleNet), + batchSize = cms.int32(args.batchSize), + particles = 'DataSource', + alpaka = cms.untracked.PSet( + backend = cms.untracked.string("serial_sync") + ), + environment = cms.untracked.int32(args.environment) + ) + process.path += process.SimpleNetMiniBatch # --only MultiHeadNet if "MultiHeadNet" in args.only: from PhysicsTools.PyTorchAlpakaTest.modules import torchtest_MultiHeadNet_alpaka @@ -77,7 +88,7 @@ process.path += process.MaskedNet # --only TinyResNet if "TinyResNet" in args.only: - from PhysicsTools.PyTorchAlpakaTest.modules import torchtest_TinyResNet_alpaka + from PhysicsTools.PyTorchAlpakaTest.modules import torchtest_TinyResNet_alpaka, torchtest_TinyResNetMiniBatch_alpaka process.TinyResNet = torchtest_TinyResNet_alpaka( model = cms.FileInPath(args.tinyResNet), images = 'DataSource', @@ -87,14 +98,28 @@ environment = cms.untracked.int32(args.environment) ) process.path += process.TinyResNet +if "TinyResNetMiniBatch" in args.only: + process.TinyResNetMiniBatch = torchtest_TinyResNetMiniBatch_alpaka( + model = cms.FileInPath(args.tinyResNet), + batchSize = cms.int32(args.batchSize), + images = 'DataSource', + alpaka = cms.untracked.PSet( + backend = cms.untracked.string(args.backend) + ), + environment = cms.untracked.int32(args.environment) + ) + process.path += process.TinyResNetMiniBatch + # debug (if --environment < 1 only assertions are checked) process.InspectionSink = torchtest_InspectionSink( particles = 'DataSource', simple_net = 'SimpleNet', + simple_net_minibatch = 'SimpleNetMiniBatch', masked_net = 'MaskedNet', multi_head_net = 'MultiHeadNet', images = 'DataSource', resnet18 = 'TinyResNet', + resnet18_minibatch = 'TinyResNetMiniBatch', environment = cms.untracked.int32(args.environment) ) process.path += process.InspectionSink