Skip to content

Commit 3ba3ba4

Browse files
lukaszmichalskiichrisizeh
lukaszmichalskii
and
chrisizeh
committed
Beta support for efficient and portable ML inference with pytorch & alpaka
Co-authored-by: Christine Zeh <[email protected]>
1 parent dc931b8 commit 3ba3ba4

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

65 files changed

+4012
-334
lines changed

‎DataFormats/PyTorch/BuildFile.xml

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
<use name="rootcore"/>
2+
<use name="DataFormats/Common"/>
3+
<use name="DataFormats/Portable"/>
4+
<use name="DataFormats/SoATemplate"/>
5+
<use name="HeterogeneousCore/AlpakaInterface"/>
6+
<flags ALPAKA_BACKENDS="!serial"/>
7+
<export>
8+
<lib name="1"/>
9+
</export>

‎DataFormats/PyTorch/interface/Device.h

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
#ifndef DATA_FORMATS__PYTORCH__INTERFACE__DEVICE_H_
2+
#define DATA_FORMATS__PYTORCH__INTERFACE__DEVICE_H_
3+
4+
#include "DataFormats/Portable/interface/PortableDeviceCollection.h"
5+
#include "DataFormats/PyTorch/interface/Layout.h"
6+
7+
8+
namespace torchportable {
9+
10+
template <typename TDev>
11+
using ParticleCollectionDevice = PortableDeviceCollection<ParticleSoA, TDev>;
12+
13+
template <typename TDev>
14+
using ClassificationCollectionDevice = PortableDeviceCollection<ClassificationSoA, TDev>;
15+
16+
template <typename TDev>
17+
using RegressionCollectionDevice = PortableDeviceCollection<RegressionSoA, TDev>;
18+
19+
} // namespace torchportable
20+
21+
#endif // DATA_FORMATS__PYTORCH__INTERFACE__DEVICE_H_

‎DataFormats/PyTorch/interface/Host.h

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
#ifndef DATA_FORMATS__PYTORCH__INTERFACE__HOST_H_
2+
#define DATA_FORMATS__PYTORCH__INTERFACE__HOST_H_
3+
4+
#include "DataFormats/Portable/interface/PortableHostCollection.h"
5+
#include "DataFormats/PyTorch/interface/Layout.h"
6+
7+
namespace torchportable {
8+
9+
using ParticleCollectionHost = PortableHostCollection<ParticleSoA>;
10+
using ClassificationCollectionHost = PortableHostCollection<ClassificationSoA>;
11+
using RegressionCollectionHost = PortableHostCollection<RegressionSoA>;
12+
13+
} // namespace torchportable
14+
15+
#endif // DATA_FORMATS__PYTORCH__INTERFACE__HOST_H_

‎DataFormats/PyTorch/interface/Layout.h

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
#ifndef DATA_FORMATS__PYTORCH__INTERFACE__LAYOUT_H_
2+
#define DATA_FORMATS__PYTORCH__INTERFACE__LAYOUT_H_
3+
4+
#include "DataFormats/SoATemplate/interface/SoACommon.h"
5+
#include "DataFormats/SoATemplate/interface/SoALayout.h"
6+
#include "DataFormats/SoATemplate/interface/SoAView.h"
7+
8+
namespace torchportable {
9+
10+
GENERATE_SOA_LAYOUT(ParticleLayout,
11+
SOA_COLUMN(float, pt),
12+
SOA_COLUMN(float, eta),
13+
SOA_COLUMN(float, phi)
14+
)
15+
using ParticleSoA = ParticleLayout<>;
16+
17+
18+
GENERATE_SOA_LAYOUT(ClassificationLayout,
19+
SOA_COLUMN(float, c1),
20+
SOA_COLUMN(float, c2)
21+
)
22+
using ClassificationSoA = ClassificationLayout<>;
23+
24+
GENERATE_SOA_LAYOUT(RegressionLayout,
25+
SOA_COLUMN(float, reco_pt)
26+
)
27+
using RegressionSoA = RegressionLayout<>;
28+
29+
} // namespace torchportable
30+
31+
#endif // DATA_FORMATS__PYTORCH__INTERFACE__LAYOUT_H_

‎DataFormats/PyTorch/interface/alpaka/Collections.h

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
#ifndef DATA_FORMATS__PYTORCH__INTERFACE__ALPAKA__COLLECTIONS_H_
2+
#define DATA_FORMATS__PYTORCH__INTERFACE__ALPAKA__COLLECTIONS_H_
3+
4+
#include <alpaka/alpaka.hpp>
5+
6+
#include "DataFormats/Portable/interface/alpaka/PortableCollection.h"
7+
#include "DataFormats/PyTorch/interface/Device.h"
8+
#include "DataFormats/PyTorch/interface/Host.h"
9+
#include "DataFormats/PyTorch/interface/Layout.h"
10+
#include "HeterogeneousCore/AlpakaInterface/interface/config.h"
11+
#include "HeterogeneousCore/AlpakaInterface/interface/CopyToHost.h"
12+
13+
namespace ALPAKA_ACCELERATOR_NAMESPACE::torchportable {
14+
15+
/**
16+
* make the names from the top-level `torchportable` namespace visible for unqualified lookup
17+
* inside the `ALPAKA_ACCELERATOR_NAMESPACE::torchportable` namespace
18+
*/
19+
using namespace ::torchportable;
20+
using ::torchportable::ParticleCollectionHost;
21+
using ::torchportable::ClassificationCollectionHost;
22+
using ::torchportable::RegressionCollectionHost;
23+
using ::torchportable::ParticleCollectionDevice;
24+
using ::torchportable::ClassificationCollectionDevice;
25+
using ::torchportable::RegressionCollectionDevice;
26+
27+
using ParticleCollection =
28+
std::conditional_t<
29+
std::is_same_v<Device, alpaka::DevCpu>,
30+
ParticleCollectionHost,
31+
ParticleCollectionDevice<Device>>;
32+
33+
using ClassificationCollection =
34+
std::conditional_t<
35+
std::is_same_v<Device, alpaka::DevCpu>,
36+
ClassificationCollectionHost,
37+
ClassificationCollectionDevice<Device>>;
38+
39+
using RegressionCollection =
40+
std::conditional_t<
41+
std::is_same_v<Device, alpaka::DevCpu>,
42+
RegressionCollectionHost,
43+
RegressionCollectionDevice<Device>>;
44+
45+
} // namespace ALPAKA_ACCELERATOR_NAMESPACE::torchportable
46+
47+
ASSERT_DEVICE_MATCHES_HOST_COLLECTION(torchportable::ParticleCollection, torchportable::ParticleCollectionHost);
48+
ASSERT_DEVICE_MATCHES_HOST_COLLECTION(torchportable::ClassificationCollection, torchportable::ClassificationCollectionHost);
49+
ASSERT_DEVICE_MATCHES_HOST_COLLECTION(torchportable::RegressionCollection, torchportable::RegressionCollectionHost);
50+
51+
#endif // DATA_FORMATS__PYTORCH__INTERFACE__ALPAKA__COLLECTIONS_H_

‎DataFormats/PyTorch/src/alpaka/classes_cuda.h

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
#ifndef DATA_FORMATS__PYTORCH__SRC__ALPAKA__CLASSES_CUDA_H_
2+
#define DATA_FORMATS__PYTORCH__SRC__ALPAKA__CLASSES_CUDA_H_
3+
4+
#include "DataFormats/Common/interface/Wrapper.h"
5+
#include "DataFormats/Common/interface/DeviceProduct.h"
6+
#include "DataFormats/PyTorch/interface/Layout.h"
7+
#include "DataFormats/PyTorch/interface/alpaka/Collections.h"
8+
9+
#endif // DATA_FORMATS__PYTORCH__SRC__ALPAKA__CLASSES_CUDA_H_

‎DataFormats/PyTorch/src/alpaka/classes_cuda_def.xml

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
<lcgdict>
2+
<!-- Particle -->
3+
<class name="alpaka_cuda_async::PortableCollection<torchportable::ParticleSoA>" persistent="false"/>
4+
<class name="alpaka_cuda_async::torchportable::ParticleCollection" persistent="false"/>
5+
<class name="edm::DeviceProduct<alpaka_cuda_async::torchportable::ParticleCollection>" persistent="false"/>
6+
<class name="edm::Wrapper<edm::DeviceProduct<alpaka_cuda_async::torchportable::ParticleCollection>>" persistent="false"/>
7+
<!-- Classification -->
8+
<class name="alpaka_cuda_async::PortableCollection<torchportable::ClassificationSoA>" persistent="false"/>
9+
<class name="alpaka_cuda_async::torchportable::ClassificationCollection" persistent="false"/>
10+
<class name="edm::DeviceProduct<alpaka_cuda_async::torchportable::ClassificationCollection>" persistent="false"/>
11+
<class name="edm::Wrapper<edm::DeviceProduct<alpaka_cuda_async::torchportable::ClassificationCollection>>" persistent="false"/>
12+
<!-- Regression -->
13+
<class name="alpaka_cuda_async::PortableCollection<torchportable::RegressionSoA>" persistent="false"/>
14+
<class name="alpaka_cuda_async::torchportable::RegressionCollection" persistent="false"/>
15+
<class name="edm::DeviceProduct<alpaka_cuda_async::torchportable::RegressionCollection>" persistent="false"/>
16+
<class name="edm::Wrapper<edm::DeviceProduct<alpaka_cuda_async::torchportable::RegressionCollection>>" persistent="false"/>
17+
</lcgdict>

‎DataFormats/PyTorch/src/alpaka/classes_rocm.h

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
#ifndef DATA_FORMATS__PYTORCH__SRC__ALPAKA__CLASSES_ROCM_H_
2+
#define DATA_FORMATS__PYTORCH__SRC__ALPAKA__CLASSES_ROCM_H_
3+
4+
#include "DataFormats/Common/interface/Wrapper.h"
5+
#include "DataFormats/Common/interface/DeviceProduct.h"
6+
#include "DataFormats/PyTorch/interface/Layout.h"
7+
#include "DataFormats/PyTorch/interface/alpaka/Collections.h"
8+
9+
#endif // DATA_FORMATS__PYTORCH__SRC__ALPAKA__CLASSES_ROCM_H_

‎DataFormats/PyTorch/src/alpaka/classes_rocm_def.xml

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
<lcgdict>
2+
<!-- Particle -->
3+
<class name="alpaka_rocm_async::PortableCollection<torchportable::ParticleSoA>" persistent="false"/>
4+
<class name="alpaka_rocm_async::torchportable::ParticleCollection" persistent="false"/>
5+
<class name="edm::DeviceProduct<alpaka_rocm_async::torchportable::ParticleCollection>" persistent="false"/>
6+
<class name="edm::Wrapper<edm::DeviceProduct<alpaka_rocm_async::torchportable::ParticleCollection>>" persistent="false"/>
7+
<!-- Classification -->
8+
<class name="alpaka_rocm_async::PortableCollection<torchportable::ClassificationSoA>" persistent="false"/>
9+
<class name="alpaka_rocm_async::torchportable::ClassificationCollection" persistent="false"/>
10+
<class name="edm::DeviceProduct<alpaka_rocm_async::torchportable::ClassificationCollection>" persistent="false"/>
11+
<class name="edm::Wrapper<edm::DeviceProduct<alpaka_rocm_async::torchportable::ClassificationCollection>>" persistent="false"/>
12+
<!-- Regression -->
13+
<class name="alpaka_rocm_async::PortableCollection<torchportable::RegressionSoA>" persistent="false"/>
14+
<class name="alpaka_rocm_async::torchportable::RegressionCollection" persistent="false"/>
15+
<class name="edm::DeviceProduct<alpaka_rocm_async::torchportable::RegressionCollection>" persistent="false"/>
16+
<class name="edm::Wrapper<edm::DeviceProduct<alpaka_rocm_async::torchportable::RegressionCollection>>" persistent="false"/>
17+
</lcgdict>

‎DataFormats/PyTorch/src/classes.cc

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
#include "DataFormats/Portable/interface/PortableHostCollectionReadRules.h"
2+
#include "DataFormats/PyTorch/interface/Host.h"
3+
4+
SET_PORTABLEHOSTCOLLECTION_READ_RULES(torchportable::ParticleCollectionHost);
5+
SET_PORTABLEHOSTCOLLECTION_READ_RULES(torchportable::ClassificationCollectionHost);
6+
SET_PORTABLEHOSTCOLLECTION_READ_RULES(torchportable::RegressionCollectionHost);

‎DataFormats/PyTorch/src/classes.h

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
#ifndef DATA_FORMATS__PYTORCH__SRC__CLASSES_H_
2+
#define DATA_FORMATS__PYTORCH__SRC__CLASSES_H_
3+
4+
#include "DataFormats/Common/interface/Wrapper.h"
5+
#include "DataFormats/PyTorch/interface/Host.h"
6+
#include "DataFormats/PyTorch/interface/Layout.h"
7+
8+
#endif // DATA_FORMATS__PYTORCH__SRC__CLASSES_H_

‎DataFormats/PyTorch/src/classes_def.xml

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
<lcgdict>
2+
<!-- Particle -->
3+
<class name="torchportable::ParticleSoA"/>
4+
<class name="torchportable::ParticleSoA::View"/>
5+
<class name="torchportable::ParticleSoA::ConstView"/>
6+
<class name="torchportable::ParticleCollectionHost"/>
7+
<class name="edm::Wrapper<torchportable::ParticleCollectionHost>" splitLevel="0"/>
8+
<!-- Classification -->
9+
<class name="torchportable::ClassificationSoA"/>
10+
<class name="torchportable::ClassificationSoA::View"/>
11+
<class name="torchportable::ClassificationSoA::ConstView"/>
12+
<class name="torchportable::ClassificationCollectionHost"/>
13+
<class name="edm::Wrapper<torchportable::ClassificationCollectionHost>" splitLevel="0"/>
14+
<!-- Regression -->
15+
<class name="torchportable::RegressionSoA"/>
16+
<class name="torchportable::RegressionSoA::View"/>
17+
<class name="torchportable::RegressionSoA::ConstView"/>
18+
<class name="torchportable::RegressionCollectionHost"/>
19+
<class name="edm::Wrapper<torchportable::RegressionCollectionHost>" splitLevel="0"/>
20+
</lcgdict>

‎PhysicsTools/PyTorch/BuildFile.xml

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
1+
<use name="eigen"/>
2+
<use name="pytorch"/>
3+
<use name="pytorch-cuda"/>
4+
<use name="DataFormats/PyTorch"/>
15
<use name="FWCore/Framework"/>
26
<use name="FWCore/MessageLogger"/>
37
<use name="FWCore/Utilities"/>
48
<use name="FWCore/ServiceRegistry"/>
5-
<use name="pytorch"/>
9+
<use name="HeterogeneousCore/AlpakaInterface"/>
610
<export>
711
<lib name="1"/>
812
</export>

‎PhysicsTools/PyTorch/README.md

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
# PyTorch Wrapper for C++ and Alpaka
2+
3+
The interface provides a converter to dynamically wrap SoA data into one or more `torch::tensors` without the need to copy data. This can be used directly with a PyTorch model. The result can also be dynamically placed into a SoA buffer.
4+
5+
## Metadata
6+
7+
The structual information of the input and output SoA are stored in an `SoAMetadata`. These two objects are then combined to a `ModelMetadata`, to be used by the `Converter`.
8+
9+
### Defining Metadata
10+
11+
The `SoAMetadata' can be defined by first initialising the object and then adding blocks to the metadata. Each block is transformed into a tensor whose size and type are derived from the columns provided.
12+
13+
#### Example SOA Template for Model Input:
14+
```cpp
15+
GENERATE_SOA_LAYOUT(SoATemplate,
16+
SOA_EIGEN_COLUMN(Eigen::Vector3d, a),
17+
SOA_EIGEN_COLUMN(Eigen::Vector3d, b),
18+
SOA_EIGEN_COLUMN(Eigen::Matrix2f, c),
19+
SOA_COLUMN(double, x),
20+
SOA_COLUMN(double, y),
21+
SOA_COLUMN(double, z),
22+
SOA_SCALAR(float, type),
23+
SOA_SCALAR(int, someNumber));
24+
```
25+
26+
#### Example SOA Template for Model Output:
27+
```cpp
28+
GENERATE_SOA_LAYOUT(SoAOutputTemplate,
29+
SOA_COLUMN(int, cluster));
30+
```
31+
32+
#### Metadata Definition (Automatic Approach):
33+
```cpp
34+
PortableCollection<SoA, Device> deviceCollection(batch_size, queue);
35+
PortableCollection<SoA_Result, Device> deviceResultCollection(batch_size, queue);
36+
fill(queue, deviceCollection);
37+
auto records = deviceCollection.view().records();
38+
auto result_records = deviceResultCollection.view().records();
39+
40+
SoAMetadata<SoA> input(batch_size);
41+
input.append_block("eigen_vector", records.a(), records.b());
42+
input.append_block("eigen_matrix", records.c());
43+
input.append_block("column", records.x(), records.y(), records.z());
44+
input.append_block("scalar", view.type());
45+
input.change_order({"column", "scalar", "eigen_matrix", "eigen_vector"});
46+
47+
SoAMetadata<SoA> output(batch_size);
48+
output.append_block("result", result_view.cluster());
49+
ModelMetadata metadata(input, output);
50+
```
51+
52+
### Ordering of Blocks
53+
54+
The function `change_order()` in the allows specifying the order in which the blocks should be processed. The order should match the expected input configuration of the PyTorch model.

0 commit comments

Comments
 (0)