Gnn/src/DeviceOps.hpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #pragma once
0010
0011 #include "ActsPlugins/Gnn/Stages.hpp"
0012 #include "ActsPlugins/Gnn/Tensor.hpp"
0013
0014 #include <cstddef>
0015 #include <vector>
0016
0017 namespace ActsPlugins::detail {
0018
0019 /// Every operation in this plugin that needs a GPU, declared without naming a
0020 /// CUDA type: `cudaStream_t` here is the forward-declared `CUstream_st *` from
0021 /// Tensor.hpp, which is the type the CUDA runtime uses too.
0022 ///
0023 /// Exactly one implementation is linked -- the `.cu` translation units when the
0024 /// plugin is built with CUDA, DeviceOpsNoCuda.cpp otherwise -- so callers below
0025 /// dispatch on the runtime device and need no preprocessor branch. Both
0026 /// implementations are built by CI, which is what keeps them in step.
0027 ///
0028 /// Reaching any of these without CUDA means a CUDA tensor was constructed
0029 /// first, which cudaCreateTensorMemory refuses, so the no-CUDA definitions are
0030 /// unreachable rather than merely unsupported.
0031
0032 /// Allocate device memory for a tensor. @p ctx must carry a stream.
0033 TensorPtr cudaCreateTensorMemory(std::size_t nbytes,
0034                                  const ExecutionContext &ctx);
0035
0036 /// Copy @p nbytes between host and device in either direction. @p to must carry
0037 /// a stream.
0038 void cudaCopyTensorMemory(void *dst, const void *src, std::size_t nbytes,
0039                           Device from, const ExecutionContext &to);
0040
0041 /// Apply the logistic function to @p tensor in place.
0042 void cudaSigmoid(Tensor<float> &tensor, cudaStream_t stream);
0043
0044 /// Element-wise `scores > cut`.
0045 Tensor<bool> cudaScoreMask(const Tensor<float> &scores, float cut,
0046                            cudaStream_t stream);
0047
0048 /// Gather the rows of @p tensor selected by @p mask.
0049 template <typename T>
0050 Tensor<T> cudaSelectRows(const Tensor<T> &tensor, const Tensor<bool> &mask,
0051                          const ExecutionContext &execContext);
0052
0053 /// Gather the columns of @p tensor selected by @p mask.
0054 template <typename T>
0055 Tensor<T> cudaSelectCols(const Tensor<T> &tensor, const Tensor<bool> &mask,
0056                          const ExecutionContext &execContext);
0057
0058 /// Scale each column of @p src by the matching entry of @p scales.
0059 template <typename T>
0060 Tensor<T> cudaMulPerColumn(const Tensor<T> &src, const Tensor<T> &scales,
0061                            const ExecutionContext &execContext);
0062
0063 /// Drop nodes no edge refers to, renumbering the edge index to match.
0064 PipelineTensors cudaRemoveUnusedNodes(PipelineTensors &&tensors,
0065                                       std::vector<int> &spacePointIds,
0066                                       const ExecutionContext &execCtx);
0067
0068 }  // namespace ActsPlugins::detail