Project
Loading...
Searching...
No Matches
GPUTPCNNClusterizerHost.h
Go to the documentation of this file.
1// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
2// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
3// All rights not expressly granted are reserved.
4//
5// This software is distributed under the terms of the GNU General Public
6// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
7//
8// In applying this license CERN does not waive the privileges and immunities
9// granted to it by virtue of its status as an Intergovernmental Organization
10// or submit itself to any jurisdiction.
11
14
15#ifndef O2_GPUTPCNNCLUSTERIZERHOST_H
16#define O2_GPUTPCNNCLUSTERIZERHOST_H
17
18#include <string>
19#include <unordered_map>
20#include <vector>
21#include "ML/OrtInterface.h"
22
23using namespace o2::ml;
24
25class OrtMemoryInfo;
26class OrtAllocator;
28namespace Ort
29{
30struct Env;
31struct MemoryInfo;
32} // namespace Ort
33
34namespace o2::OrtDataType
35{
36struct Float16_t;
37}
38
39namespace o2::gpu
40{
41
42class GPUReconstruction;
43class GPUTPCNNClusterizer;
44struct GPUSettingsProcessingNNclusterizer;
45
47{
48 public:
50 GPUTPCNNClusterizerHost(const GPUSettingsProcessingNNclusterizer& settings) { init(settings); }
51
52 void init(const GPUSettingsProcessingNNclusterizer&);
53 void initClusterizer(const GPUSettingsProcessingNNclusterizer&, GPUTPCNNClusterizer&);
54
55 // ONNX
56 void volatileOrtAllocator(Ort::Env*, Ort::MemoryInfo*, GPUReconstruction*, bool = false);
58 const OrtMemoryInfo* getMockedMemoryInfo();
59
60 std::unordered_map<std::string, std::string> mOrtOptions;
62 std::vector<bool> mModelsUsed = {false, false, false}; // 0: class, 1: reg_1, 2: reg_2
63 int32_t mDeviceId = -1;
64 std::shared_ptr<MockedOrtAllocator> mMockedAlloc = nullptr;
65}; // class GPUTPCNNClusterizerHost
66
67} // namespace o2::gpu
68
69#endif
A header library for loading ONNX models and inferencing them on CPU and GPU.
void volatileOrtAllocator(Ort::Env *, Ort::MemoryInfo *, GPUReconstruction *, bool=false)
std::unordered_map< std::string, std::string > mOrtOptions
void initClusterizer(const GPUSettingsProcessingNNclusterizer &, GPUTPCNNClusterizer &)
std::shared_ptr< MockedOrtAllocator > mMockedAlloc
GPUTPCNNClusterizerHost(const GPUSettingsProcessingNNclusterizer &settings)
void init(const GPUSettingsProcessingNNclusterizer &)