GPUTPCNNClusterizerHost.h

Go to the documentation of this file.

// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
// All rights not expressly granted are reserved.
//
// This software is distributed under the terms of the GNU General Public
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.
 
 
#ifndef O2_GPUTPCNNCLUSTERIZERHOST_H
#define O2_GPUTPCNNCLUSTERIZERHOST_H
 
#include <string>
#include <unordered_map>
#include <vector>
#include "ML/OrtInterface.h"
 
using namespace o2::ml;
 
class OrtMemoryInfo;
class OrtAllocator;
struct MockedOrtAllocator;
namespace Ort
{
struct Env;
struct MemoryInfo;
} // namespace Ort
 
namespace o2::OrtDataType
{
struct Float16_t;
}
 
namespace o2::gpu
{
 
class GPUReconstruction;
class GPUTPCNNClusterizer;
struct GPUSettingsProcessingNNclusterizer;
 
class GPUTPCNNClusterizerHost
{
 public:
  GPUTPCNNClusterizerHost() = default;
  GPUTPCNNClusterizerHost(const GPUSettingsProcessingNNclusterizer& settings) { init(settings); }
 
  void init(const GPUSettingsProcessingNNclusterizer&);
  void initClusterizer(const GPUSettingsProcessingNNclusterizer&, GPUTPCNNClusterizer&);
 
  // ONNX
  void directOrtAllocator(Ort::Env*, Ort::MemoryInfo*, GPUReconstruction*, bool = false);
  MockedOrtAllocator* getMockedAllocator();
  const OrtMemoryInfo* getMockedMemoryInfo();
 
  std::unordered_map<std::string, std::string> mOrtOptions;
  o2::ml::OrtModel mModelClass, mModelReg1, mModelReg2;  // For splitting clusters
  std::vector<bool> mModelsUsed = {false, false, false}; // 0: class, 1: reg_1, 2: reg_2
  int32_t mDeviceId = -1;
  std::shared_ptr<MockedOrtAllocator> mMockedAlloc = nullptr;
}; // class GPUTPCNNClusterizerHost
 
} // namespace o2::gpu
 
#endif