-
Notifications
You must be signed in to change notification settings - Fork 1
/
compute.clj
206 lines (162 loc) · 6.86 KB
/
compute.clj
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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
(ns tech.compute
(:require [tech.compute.driver :as drv]
[tech.compute.registry :as registry]
[tech.compute.context :as compute-ctx]
[tech.v2.datatype :as dtype]
[clojure.test :refer :all]
[tech.resource :as resource]))
(defn driver-names
"Get the names of the registered drivers."
[]
(registry/driver-names))
(defn driver
"Do a registry lookup to find a driver by its name."
[driver-name]
(registry/driver driver-name))
(defn ->driver
"Generically get a driver from a thing"
[item]
(drv/get-driver item))
(defn ->device
"Generically get a device from a thing"
[item]
(drv/get-device item))
(defn ->stream
"Generically get a stream from a thing"
[item]
(drv/get-stream item))
(defn driver-name
[driver]
(drv/driver-name driver))
(defn get-devices
[driver]
(drv/get-devices driver))
(defn default-device
[driver]
(first (get-devices driver)))
(defn allocate-host-buffer
"Allocate a host buffer. Usage type gives a hint as to the
intended usage of the buffer."
[driver elem-count elem-type & {:keys [usage-type]
:or {usage-type :one-time}
:as options}]
(drv/allocate-host-buffer driver elem-count
elem-type (assoc options
:usage-type usage-type)))
;; Device API
(defn supports-create-stream?
"Does this device support create-stream?"
[device]
(drv/supports-create-stream? device))
(defn default-stream
"All devices must have a default stream whether they support create or not."
[device]
(drv/default-stream device))
(defn create-stream
"Create a stream of execution. Streams are indepenent threads of execution. They can
be synchronized with each other and the main thread using events."
[device]
(drv/create-stream device))
(defn allocate-device-buffer
"Allocate a device buffer. This is the generic unit of data storage used for
computation. No options at this time."
[device elem-count elem-type & {:as options}]
(drv/allocate-device-buffer device elem-count elem-type options))
;;Stream API
(defn- check-legal-copy!
[src-buffer src-offset dst-buffer dst-offset elem-count]
(let [src-len (- (dtype/ecount src-buffer) (long src-offset))
dst-len (- (dtype/ecount dst-buffer) (long dst-offset))
elem-count (long elem-count)]
(when (> elem-count src-len)
(throw (ex-info "Copy out of range"
{:src-len src-len
:elem-count elem-count})))
(when (> elem-count dst-len)
(throw (ex-info "Copy out of range"
{:dst-len dst-len
:elem-count elem-count})))))
(defn- provided-or-default-stream
[stream device-buffer]
(or stream
(:stream compute-ctx/*context)
(default-stream (->device device-buffer))))
(defn device->device-copy-compatible?
[src-device dst-device]
(drv/device->device-copy-compatible? src-device dst-device))
(defn copy-host->device
"Copy from one device to another. If no stream is provided then the destination
buffer's device's default stream is used."
[host-buffer host-offset device-buffer device-offset elem-count & {:keys [stream]}]
(check-legal-copy! host-buffer host-offset device-buffer device-offset elem-count)
(let [stream (provided-or-default-stream stream device-buffer)]
(drv/copy-host->device stream host-buffer host-offset
device-buffer device-offset elem-count)))
(defn copy-device->host
"Copy from one device to another. If no stream is provided then the source
device buffer's device's default stream is used."
[device-buffer device-offset host-buffer host-offset elem-count & {:keys [stream]}]
(check-legal-copy! device-buffer device-offset host-buffer host-offset elem-count)
(let [stream (provided-or-default-stream stream device-buffer)]
(drv/copy-device->host stream device-buffer device-offset
host-buffer host-offset elem-count)))
(defn copy-device->device
"Copy from one device to another. If no stream is provided then the destination
buffer's device's default stream is used."
[dev-a dev-a-off dev-b dev-b-off elem-count & {:keys [stream]}]
(check-legal-copy! dev-a dev-a-off dev-b dev-b-off elem-count)
(let [stream (provided-or-default-stream stream dev-b)]
(drv/copy-device->device stream dev-a dev-a-off dev-b dev-b-off elem-count)))
(defn sync-with-host
"Block host until stream's queue is finished executing"
([stream]
(drv/sync-with-host stream))
([]
(sync-with-host (compute-ctx/default-stream))))
(defn sync-with-stream
"Create an event in src-stream's execution queue, then have dst stream wait on that
event. This allows dst-stream to ensure src-stream has reached a certain point of
execution before continuing. Both streams must be of the same driver."
[src-stream dst-stream & [options]]
(drv/sync-with-stream src-stream dst-stream))
(defn copy-host-data->device-buffer
"Robustly and synchronously make a device buffer with these elements in it.
1. Make host buffer of correct size.
2. Make device buffer of correct size.
3. Copy ary data into host buffer.
4. Copy host buffer into device buffer.
5. Wait until operation completes.
6. Release the host buffer.
7. Return device buffer."
[stream upload-ary & {:keys [datatype]
:or {datatype (dtype/get-datatype upload-ary)}
:as options}]
(let [device (->device stream)
driver (->driver device)
elem-count (dtype/ecount upload-ary)
device-buffer (allocate-device-buffer device elem-count datatype options)
upload-buffer (allocate-host-buffer driver elem-count datatype
:usage-type :one-time)]
(dtype/copy-raw->item! upload-ary upload-buffer 0 options)
(copy-host->device stream upload-buffer 0 device-buffer 0 elem-count)
;;Hold onto host buffer till this completes
(sync-with-host stream {:gc-root upload-buffer})
device-buffer))
(defn device-buffer->host-buffer
"Robustly and synchronously make a device buffer with these elements in it.
1. Make host buffer of correct size.
2. Copy device buffer into host buffer.
3. Wait until operation completes.
4. Return host buffer."
[stream device-buffer & {:keys [usage-type]
:or {usage-type :one-time}
:as options}]
(let [device (->device stream)
driver (->driver device)
elem-count (dtype/ecount device-buffer)
datatype (dtype/get-datatype device-buffer)
download-buffer (allocate-host-buffer driver elem-count datatype
(assoc options :usage-type usage-type))]
(copy-device->host stream device-buffer 0 download-buffer 0 elem-count)
(sync-with-host stream)
download-buffer))