PyTorchAlpaka: Batched inference with TensorCollections, eval and frozen model, FP16 convsersion support

DataFormats/Common/interface/FPX.h

@@ -0,0 +1,45 @@
+#ifndef FPX_h
+#define FPX_h
+
+#include <limits>
+
+#if defined ALPAKA_ACC_GPU_CUDA_ENABLED
+#include <cuda_fp16.h>
+#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<float>::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<FPX>::quiet_NaN(); }
+
+#endif
+
+#endif

DataFormats/PortableTestObjects/interface/ParticleHostCollection.h

@@ -7,6 +7,7 @@
 namespace portabletest {
 
   using ParticleHostCollection = PortableHostCollection<ParticleSoA>;
+  using ParticleHostCollectionFPX = PortableHostCollection<ParticleSoAFPX>;
 
 }  // namespace portabletest
 

DataFormats/PortableTestObjects/interface/ParticleSoA.h

@@ -5,13 +5,16 @@
 #include <Eigen/Dense>
 
 #include "DataFormats/Common/interface/StdArray.h"
+#include "DataFormats/Common/interface/FPX.h"
 #include "DataFormats/SoATemplate/interface/SoACommon.h"
 #include "DataFormats/SoATemplate/interface/SoALayout.h"
 
 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
 

DataFormats/PortableTestObjects/interface/SimpleNetHostCollection.h

@@ -7,6 +7,7 @@
 namespace portabletest {
 
   using SimpleNetHostCollection = PortableHostCollection<SimpleNetSoA>;
+  using SimpleNetHostCollectionFPX = PortableHostCollection<SimpleNetSoAFPX>;
 
 }  // namespace portabletest
 

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

DataFormats/PortableTestObjects/interface/alpaka/ParticleDeviceCollection.h

@@ -15,12 +15,15 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE {
     using namespace ::portabletest;
 
     using ParticleDeviceCollection = PortableCollection<ParticleSoA>;
+    using ParticleDeviceCollectionFPX = PortableCollection<ParticleSoAFPX>;
 
   }  // namespace portabletest
 
 }  // 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

DataFormats/PortableTestObjects/interface/alpaka/SimpleNetDeviceCollection.h

@@ -15,12 +15,15 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE {
     using namespace ::portabletest;
 
     using SimpleNetDeviceCollection = PortableCollection<SimpleNetSoA>;
+    using SimpleNetDeviceCollectionFPX = PortableCollection<SimpleNetSoAFPX>;
 
   }  // namespace portabletest
 
 }  // 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

DataFormats/PortableTestObjects/src/classes_def.xml

@@ -43,6 +43,14 @@
   <class name="portabletest::ParticleSoA::ConstView"/>
   <class name="edm::Wrapper<portabletest::ParticleHostCollection>" splitLevel="0"/>
 
+  <class name="portabletest::ParticleHostCollectionFPX" rntupleStreamerMode="true"/>
+  <!-- portabletest::ParticleHostCollection::Layout must be listed before the aliased-to type -->
+  <class name="portabletest::ParticleHostCollectionFPX::Layout"/>
+  <class name="portabletest::ParticleSoAFPX"/>
+  <class name="portabletest::ParticleSoAFPX::View"/>
+  <class name="portabletest::ParticleSoAFPX::ConstView"/>
+  <class name="edm::Wrapper<portabletest::ParticleHostCollectionFPX>" splitLevel="0"/>
+
   <class name="portabletest::SimpleNetHostCollection" rntupleStreamerMode="true"/>
   <!-- portabletest::SimpleNetHostCollection::Layout must be listed before the aliased-to type -->
   <class name="portabletest::SimpleNetHostCollection::Layout"/>
@@ -51,6 +59,14 @@
   <class name="portabletest::SimpleNetSoA::ConstView"/>
   <class name="edm::Wrapper<portabletest::SimpleNetHostCollection>" splitLevel="0"/>
 
+  <class name="portabletest::SimpleNetHostCollectionFPX" rntupleStreamerMode="true"/>
+  <!-- portabletest::SimpleNetHostCollection::Layout must be listed before the aliased-to type -->
+  <class name="portabletest::SimpleNetHostCollectionFPX::Layout"/>
+  <class name="portabletest::SimpleNetSoAFPX"/>
+  <class name="portabletest::SimpleNetSoAFPX::View"/>
+  <class name="portabletest::SimpleNetSoAFPX::ConstView"/>
+  <class name="edm::Wrapper<portabletest::SimpleNetHostCollectionFPX>" splitLevel="0"/>
+
   <class name="portabletest::MultiHeadNetHostCollection" rntupleStreamerMode="true"/>
   <!-- portabletest::MultiHeadNetHostCollection::Layout must be listed before the aliased-to type -->
   <class name="portabletest::MultiHeadNetHostCollection::Layout"/>

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

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 <cuda_fp16.h>
+#include <c10/core/ScalarType.h>
+
+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

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<SoA, Device> deviceCollection(batch_size, queue);
-PortableCollection<SoA_Result, Device> deviceResultCollection(batch_size, queue);
+PortableCollection<SoA, Device> deviceCollection(total_size, queue);
+PortableCollection<SoA_Result, Device> 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<SoA>("eigen_vector", records.a(), records.b());
 input.add<SoA>("eigen_matrix", records.c());
 input.add<SoA>("column", records.x(), records.y(), records.z());
 input.add<SoA>("scalar", records.type());
 input.change_order({"column", "scalar", "eigen_matrix", "eigen_vector"});
 
-TensorCollection output(batch_size);
+TensorCollection output(total_size);
 output.add<SoA>("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)

PhysicsTools/PyTorchAlpaka/interface/SoAConversion.h

@@ -24,10 +24,15 @@ namespace cms::torch::alpakatools::detail {
   }
 
   template <typename TQueue>
-  inline std::vector<::torch::IValue> convertInput(TensorCollection<TQueue>& inputs, ::torch::Device device) {
+  inline std::vector<::torch::IValue> convertInput(TensorCollection<TQueue>& 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;
   }