diff --git a/DataFormats/Common/interface/FPX.h b/DataFormats/Common/interface/FPX.h new file mode 100644 index 0000000000000..6ff7ded9954ac --- /dev/null +++ b/DataFormats/Common/interface/FPX.h @@ -0,0 +1,45 @@ +#ifndef FPX_h +#define FPX_h + +#include + +#if defined ALPAKA_ACC_GPU_CUDA_ENABLED +#include +#endif + +/* + * Portable floating-point precision abstraction for CPU/GPU execution. + * + * This header defines a unified floating-point type `FPX` that maps to: + * - `__half` (FP16) when compiled with CUDA GPU support + * - `float` (FP32) otherwise + * + * It is designed to exploit half-precision acceleration on GPUs. + */ + +#if defined ALPAKA_ACC_GPU_CUDA_ENABLED +/* + * GPU (CUDA) implementation: + * Uses native IEEE-754 half precision (`__half`) type. + * Provides explicit conversion helpers between float and half. + */ +#define __F2H __float2half +#define __H2F __half2float +typedef __half FPX; + +__host__ __device__ inline FPX makeNaN() { return __float2half(std::numeric_limits::quiet_NaN()); } + +#else +/* + * CPU fallback implementation: + * Uses standard 32-bit floating point arithmetic. + */ +#define __F2H +#define __H2F +typedef float FPX; + +inline FPX makeNaN() { return std::numeric_limits::quiet_NaN(); } + +#endif + +#endif diff --git a/DataFormats/PortableTestObjects/interface/ParticleHostCollection.h b/DataFormats/PortableTestObjects/interface/ParticleHostCollection.h index f1253188f4c32..a6d919991066d 100644 --- a/DataFormats/PortableTestObjects/interface/ParticleHostCollection.h +++ b/DataFormats/PortableTestObjects/interface/ParticleHostCollection.h @@ -7,6 +7,7 @@ namespace portabletest { using ParticleHostCollection = PortableHostCollection; + using ParticleHostCollectionFPX = PortableHostCollection; } // namespace portabletest diff --git a/DataFormats/PortableTestObjects/interface/ParticleSoA.h b/DataFormats/PortableTestObjects/interface/ParticleSoA.h index f519d2484b03a..d7e02267f61b3 100644 --- a/DataFormats/PortableTestObjects/interface/ParticleSoA.h +++ b/DataFormats/PortableTestObjects/interface/ParticleSoA.h @@ -5,6 +5,7 @@ #include #include "DataFormats/Common/interface/StdArray.h" +#include "DataFormats/Common/interface/FPX.h" #include "DataFormats/SoATemplate/interface/SoACommon.h" #include "DataFormats/SoATemplate/interface/SoALayout.h" @@ -12,6 +13,8 @@ namespace portabletest { GENERATE_SOA_LAYOUT(ParticleLayout, SOA_COLUMN(float, pt), SOA_COLUMN(float, eta), SOA_COLUMN(float, phi)) using ParticleSoA = ParticleLayout<>; + GENERATE_SOA_LAYOUT(ParticleLayoutFPX, SOA_COLUMN(FPX, pt), SOA_COLUMN(FPX, eta), SOA_COLUMN(FPX, phi)) + using ParticleSoAFPX = ParticleLayoutFPX<>; } // namespace portabletest diff --git a/DataFormats/PortableTestObjects/interface/SimpleNetHostCollection.h b/DataFormats/PortableTestObjects/interface/SimpleNetHostCollection.h index 0694a2769edf0..00a39c4c13d62 100644 --- a/DataFormats/PortableTestObjects/interface/SimpleNetHostCollection.h +++ b/DataFormats/PortableTestObjects/interface/SimpleNetHostCollection.h @@ -7,6 +7,7 @@ namespace portabletest { using SimpleNetHostCollection = PortableHostCollection; + using SimpleNetHostCollectionFPX = PortableHostCollection; } // namespace portabletest diff --git a/DataFormats/PortableTestObjects/interface/SimpleNetSoA.h b/DataFormats/PortableTestObjects/interface/SimpleNetSoA.h index e12c369018984..75f9cd9e6556d 100644 --- a/DataFormats/PortableTestObjects/interface/SimpleNetSoA.h +++ b/DataFormats/PortableTestObjects/interface/SimpleNetSoA.h @@ -7,10 +7,13 @@ #include "DataFormats/Common/interface/StdArray.h" #include "DataFormats/SoATemplate/interface/SoACommon.h" #include "DataFormats/SoATemplate/interface/SoALayout.h" +#include "DataFormats/Common/interface/FPX.h" namespace portabletest { GENERATE_SOA_LAYOUT(SimpleNetLayout, SOA_COLUMN(float, reco_pt)) + GENERATE_SOA_LAYOUT(SimpleNetLayoutFPX, SOA_COLUMN(FPX, reco_pt)) + using SimpleNetSoAFPX = SimpleNetLayoutFPX<>; using SimpleNetSoA = SimpleNetLayout<>; } // namespace portabletest diff --git a/DataFormats/PortableTestObjects/interface/alpaka/ParticleDeviceCollection.h b/DataFormats/PortableTestObjects/interface/alpaka/ParticleDeviceCollection.h index d37fe13b3dc28..fcd9d9cbf68d7 100644 --- a/DataFormats/PortableTestObjects/interface/alpaka/ParticleDeviceCollection.h +++ b/DataFormats/PortableTestObjects/interface/alpaka/ParticleDeviceCollection.h @@ -15,6 +15,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE { using namespace ::portabletest; using ParticleDeviceCollection = PortableCollection; + using ParticleDeviceCollectionFPX = PortableCollection; } // namespace portabletest @@ -22,5 +23,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE { // heterogeneous ml data checks ASSERT_DEVICE_MATCHES_HOST_COLLECTION(portabletest::ParticleDeviceCollection, portabletest::ParticleHostCollection); +ASSERT_DEVICE_MATCHES_HOST_COLLECTION(portabletest::ParticleDeviceCollectionFPX, + portabletest::ParticleHostCollectionFPX); #endif // DataFormats_PortableTestObjects_interface_alpaka_ParticleDeviceCollection_h diff --git a/DataFormats/PortableTestObjects/interface/alpaka/SimpleNetDeviceCollection.h b/DataFormats/PortableTestObjects/interface/alpaka/SimpleNetDeviceCollection.h index 5c41648e68975..20e38b5b562fc 100644 --- a/DataFormats/PortableTestObjects/interface/alpaka/SimpleNetDeviceCollection.h +++ b/DataFormats/PortableTestObjects/interface/alpaka/SimpleNetDeviceCollection.h @@ -15,6 +15,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE { using namespace ::portabletest; using SimpleNetDeviceCollection = PortableCollection; + using SimpleNetDeviceCollectionFPX = PortableCollection; } // namespace portabletest @@ -22,5 +23,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE { // heterogeneous ml data checks ASSERT_DEVICE_MATCHES_HOST_COLLECTION(portabletest::SimpleNetDeviceCollection, portabletest::SimpleNetHostCollection); +ASSERT_DEVICE_MATCHES_HOST_COLLECTION(portabletest::SimpleNetDeviceCollectionFPX, + portabletest::SimpleNetHostCollectionFPX); #endif // DataFormats_PortableTestObjects_interface_alpaka_SimpleNetDeviceCollection_h diff --git a/DataFormats/PortableTestObjects/src/classes_def.xml b/DataFormats/PortableTestObjects/src/classes_def.xml index 50473ee0905b0..1a94317073322 100644 --- a/DataFormats/PortableTestObjects/src/classes_def.xml +++ b/DataFormats/PortableTestObjects/src/classes_def.xml @@ -43,6 +43,14 @@ + + + + + + + + @@ -51,6 +59,14 @@ + + + + + + + + diff --git a/PhysicsTools/PyTorch/interface/Model.h b/PhysicsTools/PyTorch/interface/Model.h index 57d19d2a5b3e6..3ebe1dc630d24 100644 --- a/PhysicsTools/PyTorch/interface/Model.h +++ b/PhysicsTools/PyTorch/interface/Model.h @@ -13,17 +13,42 @@ namespace cms::torch { // - https://docs.pytorch.org/cppdocs/api/classtorch_1_1nn_1_1_module.html#class-module class Model { public: - explicit Model(const std::string &model_path) : model_(cms::torch::load(model_path)), device_(::torch::kCPU) {} + explicit Model(const std::string &model_path, bool auto_freeze = true) + : model_(cms::torch::load(model_path)), device_(::torch::kCPU), auto_freeze_(auto_freeze) { + model_.eval(); + } + + explicit Model(const std::string &model_path, ::torch::Device dev, bool auto_freeze = true) + : model_(cms::torch::load(model_path, dev)), device_(dev), auto_freeze_(auto_freeze) { + model_.eval(); + } - explicit Model(const std::string &model_path, ::torch::Device dev) - : model_(cms::torch::load(model_path, dev)), device_(dev) {} + // Move model to specified device memory space and with the specified dtype. Async load by specifying `non_blocking` (in default stream if not overridden by the caller) + void to(::torch::Device dev, ::torch::Dtype dtype, const bool non_blocking = false) { + if (dev == device_) + return; + model_.to(dev, dtype, non_blocking); + device_ = dev; + } // Move model to specified device memory space. Async load by specifying `non_blocking` (in default stream if not overridden by the caller) void to(::torch::Device dev, const bool non_blocking = false) { if (dev == device_) return; + + assert(!is_frozen_ && "Model is frozen, cannot be moved to another device!"); model_.to(dev, non_blocking); device_ = dev; + if (auto_freeze_) { + freeze(); + } + } + + void freeze() { + if (!is_frozen_) { + model_ = ::torch::jit::freeze(model_); + is_frozen_ = true; + } } // Forward pass (inference) of model, returns torch::IValue (multi output support). Match native torchlib interface. @@ -39,6 +64,8 @@ namespace cms::torch { protected: ::torch::jit::script::Module model_; // underlying JIT model ::torch::Device device_; // device where the model is allocated (default CPU) + bool auto_freeze_; // flag to indicate if the model should be automatically frozen after loading or moving to device + bool is_frozen_ = false; // flag to indicate if the model is frozen }; } // namespace cms::torch diff --git a/PhysicsTools/PyTorch/interface/PyTorchFPXBridge.h b/PhysicsTools/PyTorch/interface/PyTorchFPXBridge.h new file mode 100644 index 0000000000000..817dc7671ce61 --- /dev/null +++ b/PhysicsTools/PyTorch/interface/PyTorchFPXBridge.h @@ -0,0 +1,23 @@ +#ifndef PhysicsTools_PyTorch_interface_PyTorchFPXBridge_h +#define PhysicsTools_PyTorch_interface_PyTorchFPXBridge_h + +#if defined(ALPAKA_ACC_GPU_CUDA_ENABLED) + +#include +#include + +namespace c10 { + + /* + * Map CUDA half precision type to PyTorch scalar type. + */ + template <> + struct CppTypeToScalarType<__half> { + static constexpr ScalarType value = ScalarType::Half; + }; + +} // namespace c10 + +#endif // ALPAKA_ACC_GPU_CUDA_ENABLED + +#endif // PhysicsTools_PyTorch_interface_PyTorchFPXBridge_h diff --git a/PhysicsTools/PyTorchAlpaka/README.md b/PhysicsTools/PyTorchAlpaka/README.md index a5f63b82f5a0b..87c9076412307 100644 --- a/PhysicsTools/PyTorchAlpaka/README.md +++ b/PhysicsTools/PyTorchAlpaka/README.md @@ -17,6 +17,15 @@ The interface provides a converter to dynamically wrap SoA data into one or more **Due to the lack of const correctness ensured by PyTorch, `const` data is currently being copied.** +## Model behavior + +The `Model` wrapper automatically sets the loaded TorchScript module to evaluation mode (`eval()`). + +Optionally, the model can be automatically frozen using `torch::jit::freeze` at construction time when a device is specified, or the first time it is moved. +**Important:** Once a model is frozen, it cannot be moved to another device. Attempting to do so will trigger a runtime assertion. + + + ### TensorCollection The structural information of the inputs/outputs SoA are stored in an `TensorCollection`. Which is a high level object to register column lists from which tensors are created @@ -41,10 +50,12 @@ GENERATE_SOA_LAYOUT(SoATemplate, GENERATE_SOA_LAYOUT(SoAOutputTemplate, SOA_COLUMN(int, cluster)); ``` + - **Get Metarecords from Portable Collections:** +If constructed with a single argument (`total_size`), the entire dataset is treated as a single batch. ```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 +64,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 +83,26 @@ 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 `(batch_size, total_size)` and internally manages batch offsets. + +**IMPORTANT:** the batchsize should be chosen carefully in order to respect the alignment. Otherwise, an assert will be trigged. + +**Constraints:** +- `total_size` must be divisible by `batch_size` +- All batches are assumed to be of equal size +- Partial (last) batches are currently not supported + +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`. + ## 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/SoAConversion.h b/PhysicsTools/PyTorchAlpaka/interface/SoAConversion.h index dd46a9a67dac4..54b7bd1c605ca 100644 --- a/PhysicsTools/PyTorchAlpaka/interface/SoAConversion.h +++ b/PhysicsTools/PyTorchAlpaka/interface/SoAConversion.h @@ -24,10 +24,15 @@ namespace cms::torch::alpakatools::detail { } template - inline std::vector<::torch::IValue> convertInput(TensorCollection& inputs, ::torch::Device device) { + inline std::vector<::torch::IValue> convertInput(TensorCollection& inputs, + ::torch::Device device, + bool to_half = false) { std::vector<::torch::IValue> tensors(inputs.size()); for (size_t i = 0; i < inputs.size(); i++) { - tensors[i] = cms::torch::alpakatools::detail::arrayToTensor(device, inputs[i]); + if (to_half) + tensors[i] = cms::torch::alpakatools::detail::arrayToTensor(device, inputs[i]).to(::torch::kHalf); + else + tensors[i] = cms::torch::alpakatools::detail::arrayToTensor(device, inputs[i]); } return tensors; } diff --git a/PhysicsTools/PyTorchAlpaka/interface/TensorCollection.h b/PhysicsTools/PyTorchAlpaka/interface/TensorCollection.h index 16494ee02e67e..705f4856e9390 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,29 @@ 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) {} + explicit TensorCollection(int batch_size, int total_size) : batch_size_(batch_size), total_size_(total_size) {} // 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_size(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 +130,16 @@ namespace cms::torch::alpakatools { else tensor_dims = {1 + sizeof...(Others), TSoAParamsImpl::ValueType::RowsAtCompileTime}; - emplace_tensor(name, SoALayout::alignment, ptr, batch_size, tensor_dims); + emplace_tensor(name, SoALayout::alignment, ptr, 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 +147,38 @@ 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_size(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; + emplace_tensor(name, SoALayout::alignment, ptr, 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 +209,26 @@ 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_batch_size(int batch_id) { + assert(batch_size_ > 0 && "Batch size must be positive!"); + assert(total_size_ > 0 && "Total size must be positive!"); + assert(batch_id >= 0 && "Batch id must be non-negative!"); + assert(total_size_ % batch_size_ == 0 && "Total size must be divisible by batch size!"); + assert((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..c972cdc0840a8 100644 --- a/PhysicsTools/PyTorchAlpaka/interface/TensorHandle.h +++ b/PhysicsTools/PyTorchAlpaka/interface/TensorHandle.h @@ -8,6 +8,7 @@ #include #include "PhysicsTools/PyTorch/interface/TorchInterface.h" +#include "PhysicsTools/PyTorch/interface/PyTorchFPXBridge.h" #include "PhysicsTools/PyTorchAlpaka/interface/Policy.h" // Forward declaration for friend @@ -21,7 +22,7 @@ namespace cms::torch::alpakatools::detail { template ::torch::ScalarType get_type() { - return ::torch::CppTypeToScalarType>(); + return c10::CppTypeToScalarType>::value; } inline int num_elements_per_column(const int n_elems, const size_t alignment, const size_t bytes) { @@ -86,13 +87,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 +133,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 = std::ceil(1.0 * total_size_ / per_bunch); // base stride initialization if (!dims_.is_scalar()) @@ -160,6 +163,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/PyTorchAlpaka/interface/alpaka/AlpakaModel.h b/PhysicsTools/PyTorchAlpaka/interface/alpaka/AlpakaModel.h index 0d43f212629ad..59f4e8daa8e65 100644 --- a/PhysicsTools/PyTorchAlpaka/interface/alpaka/AlpakaModel.h +++ b/PhysicsTools/PyTorchAlpaka/interface/alpaka/AlpakaModel.h @@ -36,7 +36,8 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torch { // Refer: PhysicsTools/PyTorch/interface/SoAConversion.h for details about wrapping memory layouts. void forward(Queue &queue, cms::torch::alpakatools::TensorCollection &inputs, - cms::torch::alpakatools::TensorCollection &outputs) { + cms::torch::alpakatools::TensorCollection &outputs, + bool to_half = false) { #ifdef ALPAKA_ACC_GPU_HIP_ENABLED inputs.copy(queue, cms::torch::alpakatools::detail::MemcpyKind::DeviceToHost); outputs.copy(queue, cms::torch::alpakatools::detail::MemcpyKind::DeviceToHost); @@ -48,7 +49,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torch { to(queue); } - auto input_tensor = cms::torch::alpakatools::detail::convertInput(inputs, device_); + auto input_tensor = cms::torch::alpakatools::detail::convertInput(inputs, device_, to_half); if (outputs.size() > 1) { auto output_tensors = model_.forward(input_tensor); cms::torch::alpakatools::detail::convertOutput(output_tensors, outputs, device_); @@ -61,8 +62,22 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torch { #endif // ALPAKA_ACC_GPU_HIP_ENABLED } - // Move model to specified device memory space. Async load (in default stream if not overridden by the caller) + // Move model to specified device memory space and with the specified dtype. Async load (in default stream if not overridden by the caller) // The caller should ensure the QueueGuard is instantiated and PyTorch stream context is properly set. + void to(const Device &dev, ::torch::Dtype dtype) { + if constexpr (std::is_same_v<::alpaka::Dev, ::alpaka::DevCpu>) { + this->Model::to(cms::torch::alpakatools::getDevice(dev), dtype); + return; + } +#ifdef ALPAKA_ACC_GPU_HIP_ENABLED + // ROCm/HIP not yet directly supported → fallback to CPU inference + this->Model::to(cms::torch::alpakatools::getDevice(dev), dtype); + return; +#endif // ALPAKA_ACC_GPU_HIP_ENABLED + // CUDA → keep async execution + this->Model::to(cms::torch::alpakatools::getDevice(dev), dtype, true); + } + void to(const Device &dev) { if constexpr (std::is_same_v<::alpaka::Dev, ::alpaka::DevCpu>) { this->Model::to(cms::torch::alpakatools::getDevice(dev)); @@ -78,6 +93,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torch { } // Overload for Queue to simplify the interface for the common case of async execution. + void to(const Queue &queue, ::torch::Dtype dtype) { this->AlpakaModel::to(::alpaka::getDev(queue), dtype); } void to(const Queue &queue) { this->AlpakaModel::to(::alpaka::getDev(queue)); } }; diff --git a/PhysicsTools/PyTorchAlpakaTest/plugins/InspectionSink.cc b/PhysicsTools/PyTorchAlpakaTest/plugins/InspectionSink.cc index 2238d780abff8..bdd3df8735d0e 100644 --- a/PhysicsTools/PyTorchAlpakaTest/plugins/InspectionSink.cc +++ b/PhysicsTools/PyTorchAlpakaTest/plugins/InspectionSink.cc @@ -32,26 +32,32 @@ namespace torchtest { : environment_{static_cast(params.getUntrackedParameter("environment"))}, particles_token_{consumes(params.getParameter("particles"))}, simple_net_token_{consumes(params.getParameter("simple_net"))}, + simple_net_batch_token_{consumes(params.getParameter("simple_net_batch"))}, 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_batch_token_{consumes(params.getParameter("resnet18_batch"))}, particles_backend_{consumes(getBackendTag(params.getParameter("particles")))}, simple_net_backend_{consumes(getBackendTag(params.getParameter("simple_net")))}, + simple_net_batch_backend_{consumes(getBackendTag(params.getParameter("simple_net_batch")))}, 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_batch_backend_{consumes(getBackendTag(params.getParameter("resnet18_batch")))} {} 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_batch"); desc.add("masked_net"); desc.add("multi_head_net"); desc.add("images"); desc.add("resnet18"); + desc.add("resnet18_batch"); descriptions.addWithDefaultLabel(desc); } @@ -68,10 +74,12 @@ namespace torchtest { // particles auto particles_handle = event.getHandle(particles_token_); auto simple_net_handle = event.getHandle(simple_net_token_); + auto simple_net_batch_handle = event.getHandle(simple_net_batch_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_batch_handle = event.getHandle(logits_batch_token_); // debug if (environment_ >= Environment::kDevelopment) { @@ -79,7 +87,8 @@ namespace torchtest { if (particles_handle.isValid()) { auto const& particles = *particles_handle; auto const particles_backend = static_cast(event.get(particles_backend_)); - if (simple_net_handle.isValid() || masked_net_handle.isValid() || multi_head_net_handle.isValid()) { + if (simple_net_handle.isValid() || simple_net_batch_handle.isValid() || masked_net_handle.isValid() || + multi_head_net_handle.isValid()) { print(particles.const_view(), cms::alpakatools::toString(particles_backend)); // assert ranges for (int32_t idx = 0; idx < particles.const_view().metadata().size(); idx++) { @@ -94,6 +103,26 @@ 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)); } + // simple_net_batch + if (simple_net_batch_handle.isValid()) { + auto const& simple_net_batch = *simple_net_batch_handle; + auto const simple_net_batch_backend = + static_cast(event.get(simple_net_batch_backend_)); + print(simple_net_batch.const_view(), cms::alpakatools::toString(simple_net_batch_backend)); + } + // if simple_net and simple_net_batch are both valid, assert they are producing the same results + if (simple_net_handle.isValid() && simple_net_batch_handle.isValid()) { + auto const& ref = *simple_net_handle; + auto const& batched = *simple_net_batch_handle; + + assert(ref.const_view().metadata().size() == batched.const_view().metadata().size()); + for (auto i = 0; i < ref.const_view().metadata().size(); i++) { + auto diff = std::abs(ref.const_view()[i].reco_pt() - batched.const_view()[i].reco_pt()) / + ref.const_view()[i].reco_pt(); + assert(diff < 1e-5 && "Results from simple_net and simple_net_batch do not match!"); + } + } + // masked_net if (masked_net_handle.isValid()) { auto const& masked_net = *masked_net_handle; @@ -145,6 +174,39 @@ namespace torchtest { assert(std::abs(sum - 1.0) < 1e-4); } } + if (images_handle.isValid() && logits_batch_handle.isValid()) { + auto const& images = *images_handle; + auto const images_backend = static_cast(event.get(images_backend_)); + print(images.const_view(), cms::alpakatools::toString(images_backend)); + + auto const& logits = *logits_batch_handle; + auto const logits_backend = static_cast(event.get(logits_batch_backend_)); + print(logits.const_view(), cms::alpakatools::toString(logits_backend)); + + const int dims = portabletest::LogitsType::RowsAtCompileTime; + for (int32_t idx = 0; idx < logits.const_view().metadata().size(); idx++) { + float sum = 0.0f; + const auto& logit = logits.const_view()[idx]; + for (int i = 0; i < dims; i++) { + sum += logit.logits()[i]; + } + assert(std::abs(sum - 1.0) < 1e-4); + } + } + if (logits_handle.isValid() && logits_batch_handle.isValid()) { + auto const& ref = *logits_handle; + auto const& batched = *logits_batch_handle; + + const int dims = portabletest::LogitsType::RowsAtCompileTime; + for (int32_t idx = 0; idx < ref.const_view().metadata().size(); idx++) { + const auto& ref_logit = ref.const_view()[idx]; + const auto& batched_logit = batched.const_view()[idx]; + for (int i = 0; i < dims; i++) { + auto diff = std::abs(ref_logit.logits()[i] - batched_logit.logits()[i]) / ref_logit.logits()[i]; + assert(diff < 1e-5 && "Results from logits and logits_batch do not match!"); + } + } + } } } @@ -169,17 +231,21 @@ namespace torchtest { const edm::EDGetTokenT particles_token_; const edm::EDGetTokenT simple_net_token_; + const edm::EDGetTokenT simple_net_batch_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_batch_token_; const edm::EDGetTokenT particles_backend_; const edm::EDGetTokenT simple_net_backend_; + const edm::EDGetTokenT simple_net_batch_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_batch_backend_; const int32_t kMaxView = 5; @@ -280,8 +346,9 @@ namespace torchtest { fmt::print("{}\n", fmt::to_string(buffer)); } - void print(const portabletest::SimpleNetHostCollection::ConstView& simple_net, - const std::string_view simple_net_backend, + template + void print(const ViewT& simple_net, + std::string_view simple_net_backend, const std::string& label = "SimpleNetCollection") { constexpr auto line = "+-------+---------+\n"; const auto size = simple_net.metadata().size(); @@ -289,19 +356,19 @@ namespace torchtest { // Header message fmt::format_to(std::back_inserter(buffer), "[DEBUG] {}[{}] ({}):\n", label, size, simple_net_backend); + fmt::format_to(std::back_inserter(buffer), "{}", line); fmt::format_to(std::back_inserter(buffer), "| {:>5} | {:>7} |\n", "index", "reco_pt"); fmt::format_to(std::back_inserter(buffer), "{}", line); - // Table rows (preview) int32_t range = (environment_ >= Environment::kTest) ? size : std::min(kMaxView, size); + for (int32_t i = 0; i < range; ++i) { fmt::format_to( std::back_inserter(buffer), "| {:5d} | {:7.2f} |\n", static_cast(i), simple_net[i].reco_pt()); } - // Ellipsis row if truncated - if (range < kMaxView) { + if (range < size) { // <-- small fix here (see below) fmt::format_to(std::back_inserter(buffer), "| {:>5} | {:>7} |\n", "...", "..."); } diff --git a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/CommonKernels.dev.cc b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/CommonKernels.dev.cc index 0b50040001dff..17ddf14e4be2e 100644 --- a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/CommonKernels.dev.cc +++ b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/CommonKernels.dev.cc @@ -6,7 +6,9 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest::kernels { - void randomFillParticleCollection(Queue& queue, portabletest::ParticleDeviceCollection& particles) { + void randomFillParticleCollection(Queue& queue, + portabletest::ParticleDeviceCollection& particles, + portabletest::ParticleDeviceCollectionFPX& particlesFPX) { const uint32_t threads_per_block = 64; const uint32_t blocks_per_grid = particles.view().metadata().size(); const auto grid = cms::alpakatools::make_workdiv(blocks_per_grid, threads_per_block); @@ -14,16 +16,23 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest::kernels { alpaka::exec( queue, grid, - [] ALPAKA_FN_ACC(Acc1D const& acc, portabletest::ParticleDeviceCollection::View particles_view) { + [] ALPAKA_FN_ACC(Acc1D const& acc, + portabletest::ParticleDeviceCollection::View particles_view, + portabletest::ParticleDeviceCollectionFPX::View particlesFPX_view) { for (int32_t thread_idx : cms::alpakatools::uniform_elements(acc, particles_view.metadata().size())) { auto rnd_gen = alpaka::rand::engine::createDefault(acc, 43, thread_idx); auto dist = alpaka::rand::distribution::createUniformReal(acc); particles_view[thread_idx].pt() = dist(rnd_gen); particles_view[thread_idx].eta() = dist(rnd_gen); particles_view[thread_idx].phi() = dist(rnd_gen); + + particlesFPX_view[thread_idx].pt() = particles_view[thread_idx].pt(); + particlesFPX_view[thread_idx].eta() = particles_view[thread_idx].eta(); + particlesFPX_view[thread_idx].phi() = particles_view[thread_idx].phi(); } }, - particles.view()); + particles.view(), + particlesFPX.view()); } struct RandomFillImageCollectionKernel { diff --git a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/CommonKernels.h b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/CommonKernels.h index d8eb6d835853f..a9772692a26fe 100644 --- a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/CommonKernels.h +++ b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/CommonKernels.h @@ -9,7 +9,9 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest::kernels { - void randomFillParticleCollection(Queue& queue, portabletest::ParticleDeviceCollection& particles); + void randomFillParticleCollection(Queue& queue, + portabletest::ParticleDeviceCollection& particles, + portabletest::ParticleDeviceCollectionFPX& particlesFPX); void randomFillImageCollection(Queue& queue, portabletest::ImageDeviceCollection& images); void fillMask(Queue& queue, portabletest::MaskDeviceCollection& mask); diff --git a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/DataSource.cc b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/DataSource.cc index 469d5cd93202d..d9f5fdb9648f8 100644 --- a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/DataSource.cc +++ b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/DataSource.cc @@ -19,6 +19,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest { DataSource(const edm::ParameterSet ¶ms) : FixedQueueEDProducer<>(params), particles_token_{produces()}, + particlesFPX_token_{produces()}, images_token_{produces()}, batch_size_(params.getParameter("batchSize")), environment_{static_cast<::torchtest::Environment>(params.getUntrackedParameter("environment"))} {} @@ -26,14 +27,16 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest { void produce(device::Event &event, const device::EventSetup &event_setup) override { // allocate data sources auto particles = portabletest::ParticleDeviceCollection(event.queue(), batch_size_); + auto particlesFPX = portabletest::ParticleDeviceCollectionFPX(event.queue(), batch_size_); auto images = portabletest::ImageDeviceCollection(event.queue(), batch_size_); // fill data - kernels::randomFillParticleCollection(event.queue(), particles); + kernels::randomFillParticleCollection(event.queue(), particles, particlesFPX); kernels::randomFillImageCollection(event.queue(), images); // put device-side data into event event.emplace(particles_token_, std::move(particles)); + event.emplace(particlesFPX_token_, std::move(particlesFPX)); event.emplace(images_token_, std::move(images)); } @@ -46,6 +49,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::torchtest { private: const device::EDPutToken particles_token_; + const device::EDPutToken particlesFPX_token_; const device::EDPutToken images_token_; const uint32_t batch_size_; const ::torchtest::Environment environment_; diff --git a/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNetMiniBatch.cc b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNetMiniBatch.cc new file mode 100644 index 0000000000000..8c7633a23c7aa --- /dev/null +++ b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/SimpleNetMiniBatch.cc @@ -0,0 +1,100 @@ +#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), + particlesFPX_token_(consumes(params.getParameter("particles"))), + simple_net_token_{produces()}, + model_(params.getParameter("model").fullPath()), + total_size_(params.getParameter("batchSize")), + batch_size_(params.getParameter("miniBatchSize")), + isFirstEvent_(false), + 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("miniBatchSize"); + 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 { + // if it's the first event and the device is 'cuda' move the model to device in FP16 precision + if (isFirstEvent_) { + isFirstEvent_ = false; +#if defined ALPAKA_ACC_GPU_CUDA_ENABLED + model_.to(event.queue(), ::torch::kHalf); +#endif + } + // in/out collections + const auto &particlesFPX = event.get(particlesFPX_token_); + auto regression_collectionFPX = portabletest::SimpleNetDeviceCollectionFPX(event.queue(), total_size_); + auto n_batches = total_size_ / batch_size_; + + // records + auto input_records = particlesFPX.const_view().records(); + auto output_records = regression_collectionFPX.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_collectionFPX)); + } + + private: + // event query tokens + const device::EDGetToken particlesFPX_token_; + const device::EDPutToken simple_net_token_; + // model + torch::AlpakaModel model_; + const int total_size_; + const int batch_size_; + bool isFirstEvent_; + // 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/TinyResNetMiniBatch.cc b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/TinyResNetMiniBatch.cc new file mode 100644 index 0000000000000..f6fa65aa6df80 --- /dev/null +++ b/PhysicsTools/PyTorchAlpakaTest/plugins/alpaka/TinyResNetMiniBatch.cc @@ -0,0 +1,90 @@ +#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()), + total_size_(params.getParameter("batchSize")), + batch_size_(params.getParameter("miniBatchSize")), + 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("miniBatchSize"); + 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_); + auto logits = portabletest::LogitsDeviceCollection(event.queue(), total_size_); + + auto n_batches = total_size_ / 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 total_size_; + 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..6606e06bdab52 100644 --- a/PhysicsTools/PyTorchAlpakaTest/python/options_cff.py +++ b/PhysicsTools/PyTorchAlpakaTest/python/options_cff.py @@ -38,10 +38,17 @@ def parse_args(): parser.add_argument( "-bs", "--batchSize", type=int, - default=8, + default=64, help="Batch size" ) + parser.add_argument( + "-mbs", "--miniBatchSize", + type=int, + default=32, + help="Mini-batch size" + ) + parser.add_argument( "-e", "--environment", type=int, @@ -57,6 +64,19 @@ def parse_args(): help="SimpleNet model (just-in-time compiled)" ) + parser.add_argument( + "--simpleNetBatch", + type=str, + default="PhysicsTools/PyTorchAlpakaTest/data/SimpleNet.pt", + help="SimpleNet model (just-in-time compiled)" + ) + + parser.add_argument( + "--compareBatch", + action="store_true", + help="Compare batched vs non-batched outputs" + ) + parser.add_argument( "--maskedNet", type=str, @@ -78,6 +98,13 @@ def parse_args(): help="TinyResNet model (just-in-time compiled)" ) + parser.add_argument( + "--tinyResNetBatch", + type=str, + default="PhysicsTools/PyTorchAlpakaTest/data/TinyResNet.pt", + help="TinyResNet model (just-in-time compiled)" + ) + parser.add_argument( "-o", "--only", nargs="+", diff --git a/PhysicsTools/PyTorchAlpakaTest/test/runPyTorchAlpakaTest.py b/PhysicsTools/PyTorchAlpakaTest/test/runPyTorchAlpakaTest.py index c8c5fe857d0f1..6eb997cdc99d4 100755 --- a/PhysicsTools/PyTorchAlpakaTest/test/runPyTorchAlpakaTest.py +++ b/PhysicsTools/PyTorchAlpakaTest/test/runPyTorchAlpakaTest.py @@ -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', @@ -50,7 +50,21 @@ ), environment = cms.untracked.int32(args.environment) ) + process.SimpleNetMiniBatch = torchtest_SimpleNetMiniBatch_alpaka( + model = cms.FileInPath(args.simpleNet), + batchSize = cms.int32(args.batchSize), + miniBatchSize = cms.int32(args.miniBatchSize), + particles = 'DataSource', + alpaka = cms.untracked.PSet( + backend = cms.untracked.string("serial_sync") + ), + environment = cms.untracked.int32(args.environment) + ) + process.path += process.SimpleNet + if args.compareBatch: + process.path += process.SimpleNetMiniBatch + # --only MultiHeadNet if "MultiHeadNet" in args.only: from PhysicsTools.PyTorchAlpakaTest.modules import torchtest_MultiHeadNet_alpaka @@ -77,7 +91,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', @@ -86,15 +100,30 @@ ), environment = cms.untracked.int32(args.environment) ) + process.TinyResNetMiniBatch = torchtest_TinyResNetMiniBatch_alpaka( + model = cms.FileInPath(args.tinyResNet), + batchSize = cms.int32(args.batchSize), + miniBatchSize = cms.int32(args.miniBatchSize), + images = 'DataSource', + alpaka = cms.untracked.PSet( + backend = cms.untracked.string(args.backend) + ), + environment = cms.untracked.int32(args.environment) + ) process.path += process.TinyResNet + if args.compareBatch: + process.path += process.TinyResNetMiniBatch + # debug (if --environment < 1 only assertions are checked) process.InspectionSink = torchtest_InspectionSink( particles = 'DataSource', simple_net = 'SimpleNet', + simple_net_batch = cms.InputTag('SimpleNetMiniBatch'), masked_net = 'MaskedNet', multi_head_net = 'MultiHeadNet', images = 'DataSource', resnet18 = 'TinyResNet', + resnet18_batch = cms.InputTag('TinyResNetMiniBatch'), environment = cms.untracked.int32(args.environment) ) process.path += process.InspectionSink