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
23class OrtMemoryInfo;
24class OrtAllocator;
26namespace Ort
27{
28struct Env;
29struct MemoryInfo;
30} // namespace Ort
31
32namespace o2::OrtDataType
33{
34struct Float16_t;
35}
36
37namespace o2::gpu
38{
39
40class GPUReconstruction;
41class GPUTPCNNClusterizer;
42struct GPUSettingsProcessingNNclusterizer;
43
45{
46 public:
48 GPUTPCNNClusterizerHost(const GPUSettingsProcessingNNclusterizer& settings) { init(settings); }
49
50 void init(const GPUSettingsProcessingNNclusterizer&);
51 void initClusterizer(const GPUSettingsProcessingNNclusterizer&, GPUTPCNNClusterizer&);
52
53 // ONNX
54 void directOrtAllocator(Ort::Env*, Ort::MemoryInfo*, GPUReconstruction*, bool = false);
56 const OrtMemoryInfo* getMockedMemoryInfo();
57
58 std::unordered_map<std::string, std::string> mOrtOptions;
60 std::vector<bool> mModelsUsed = {false, false, false}; // 0: class, 1: reg_1, 2: reg_2
61 int32_t mDeviceId = -1;
62 std::shared_ptr<MockedOrtAllocator> mMockedAlloc = nullptr;
63}; // class GPUTPCNNClusterizerHost
64
65} // namespace o2::gpu
66
67#endif
A header library for loading ONNX models and inferencing them on CPU and GPU.
std::unordered_map< std::string, std::string > mOrtOptions
void initClusterizer(const GPUSettingsProcessingNNclusterizer &, GPUTPCNNClusterizer &)
std::shared_ptr< MockedOrtAllocator > mMockedAlloc
void directOrtAllocator(Ort::Env *, Ort::MemoryInfo *, GPUReconstruction *, bool=false)
GPUTPCNNClusterizerHost(const GPUSettingsProcessingNNclusterizer &settings)
void init(const GPUSettingsProcessingNNclusterizer &)