File: device_allocator.h

package info (click to toggle)
nvidia-cuda-toolkit 12.4.1-3
  • links: PTS, VCS
  • area: non-free
  • in suites: forky, sid
  • size: 18,505,836 kB
  • sloc: ansic: 203,477; cpp: 64,769; python: 34,699; javascript: 22,006; xml: 13,410; makefile: 3,085; sh: 2,343; perl: 352
file content (148 lines) | stat: -rw-r--r-- 4,164 bytes parent folder | download | duplicates (7)
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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
/*
 *  Copyright 2008-2018 NVIDIA Corporation
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */


/*! \file
 *  \brief An allocator which creates new elements in memory accessible by
 *  devices.
 */

#pragma once

#include <thrust/detail/config.h>

#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
#  pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
#  pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
#  pragma system_header
#endif // no system header
#include <thrust/device_ptr.h>
#include <thrust/mr/allocator.h>
#include <thrust/mr/device_memory_resource.h>

#include <limits>
#include <stdexcept>

THRUST_NAMESPACE_BEGIN

/** \addtogroup allocators Allocators
 *  \ingroup memory_management
 *  \{
 */

/*! Memory resource adaptor that turns any memory resource that returns a fancy
 *      with the same tag as \p device_ptr, and adapts it to a resource that returns
 *      a \p device_ptr.
 */
template<typename Upstream>
class device_ptr_memory_resource final
    : public thrust::mr::memory_resource<
        device_ptr<void>
    >
{
    typedef typename Upstream::pointer upstream_ptr;

public:
    /*! Initialize the adaptor with the global instance of the upstream resource. Obtains
     *      the global instance by calling \p get_global_resource.
     */
    __host__
    device_ptr_memory_resource() : m_upstream(mr::get_global_resource<Upstream>())
    {
    }

    /*! Initialize the adaptor with an upstream resource.
     *
     *  \param upstream the upstream memory resource to adapt.
     */
    __host__
    device_ptr_memory_resource(Upstream * upstream) : m_upstream(upstream)
    {
    }

    THRUST_NODISCARD __host__
    virtual pointer do_allocate(std::size_t bytes, std::size_t alignment = THRUST_MR_DEFAULT_ALIGNMENT) override
    {
        return pointer(m_upstream->do_allocate(bytes, alignment).get());
    }

    __host__
    virtual void do_deallocate(pointer p, std::size_t bytes, std::size_t alignment) override
    {
        m_upstream->do_deallocate(upstream_ptr(p.get()), bytes, alignment);
    }

private:
    Upstream * m_upstream;
};

/*! \brief An allocator which creates new elements in memory accessible by
 *         devices.
 *
 *  \see https://en.cppreference.com/w/cpp/named_req/Allocator
 */
template<typename T>
class device_allocator
    : public thrust::mr::stateless_resource_allocator<
        T,
        device_ptr_memory_resource<device_memory_resource>
    >
{
    typedef thrust::mr::stateless_resource_allocator<
        T,
        device_ptr_memory_resource<device_memory_resource>
    > base;

public:
    /*! The \p rebind metafunction provides the type of a \p device_allocator
     *  instantiated with another type.
     *
     *  \tparam U the other type to use for instantiation.
     */
    template<typename U>
    struct rebind
    {
        /*! The typedef \p other gives the type of the rebound \p device_allocator.
         */
        typedef device_allocator<U> other;
    };

    /*! Default constructor has no effect. */
    __host__ __device__
    device_allocator() {}

    /*! Copy constructor has no effect. */
    __host__ __device__
    device_allocator(const device_allocator& other) : base(other) {}

    /*! Constructor from other \p device_allocator has no effect. */
    template<typename U>
    __host__ __device__
    device_allocator(const device_allocator<U>& other) : base(other) {}

    device_allocator & operator=(const device_allocator &) = default;

    /*! Destructor has no effect. */
    __host__ __device__
    ~device_allocator() {}
};

/*! \} // allocators
 */

THRUST_NAMESPACE_END