-
Notifications
You must be signed in to change notification settings - Fork 8
/
TensorflowCompute.cuh
50 lines (43 loc) · 2.43 KB
/
TensorflowCompute.cuh
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
// Copyright (c) 2020 HOOMD-TF Developers
#ifndef m_TENSORFLOW_COMPUTE_CUH_
#define m_TENSORFLOW_COMPUTE_CUH_
#include "hoomd/HOOMDMath.h"
#include "hoomd/ParticleData.cuh"
#include "hoomd/Index1D.h"
#include "hoomd/HOOMDMath.h"
/*! \file TensorflowCompute.cuh
* \brief Declarations of GPU version of some TensorflowCompute methods
*/
//! Add Scalar4 type arrays in parallel on GPU by invoking the kernel.
extern "C" cudaError_t htf_gpu_add_scalar4(Scalar4 *dest,
Scalar4 *src,
unsigned int N,
cudaStream_t stream);
//! Add up virial arrays in parallel on GPU by invoking the kernel.
extern "C" cudaError_t htf_gpu_add_virial(Scalar *dest,
Scalar *src,
unsigned int N,
unsigned int pitch,
cudaStream_t stream);
/*! Reshape neighbor list to fit into GPU memory in an agreeable way for TensorFlow.
* Attempts to use texture memory first, where possible, or will
* use device main memory if not. Invokes kernel function.
*/
extern "C" cudaError_t htf_gpu_reshape_nlist(Scalar4 *dest,
const Scalar4 *d_pos,
const unsigned int N,
const unsigned int NN,
const unsigned int offset,
const unsigned int batch_size,
const unsigned int n_ghost,
const BoxDim &box,
const unsigned int *d_n_neigh,
const unsigned int *d_nlist,
const unsigned int *d_head_list,
const unsigned int size_nlist,
const unsigned int block_size,
const unsigned int compute_capability,
const unsigned int max_tex1d_width,
double rmax,
cudaStream_t stream);
#endif // m_TENSORFLOW_COMPUTE_CUH_