-
Notifications
You must be signed in to change notification settings - Fork 62
/
ctx.go
180 lines (154 loc) · 4.15 KB
/
ctx.go
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
// +build !debug
package cu
// #include <cuda.h>
import "C"
import (
"runtime"
"unsafe"
)
// Ctx is a standalone CUDA Context that is threadlocked.
type Ctx struct {
CUContext
work chan (func() error)
errChan chan error
err error
device Device
flags ContextFlags
locked bool
}
// NewContext creates a new context, and runs a listener locked to an OSThread. All work is piped through that goroutine
func NewContext(d Device, flags ContextFlags) *Ctx {
var cctx C.CUcontext
err := result(C.cuCtxCreate(&cctx, C.uint(flags), C.CUdevice(d)))
if err != nil {
panic(err)
}
ctx := newContext(CUContext{cctx})
ctx.device = d
ctx.flags = flags
errChan := make(chan error)
go ctx.Run(errChan)
if err := <-errChan; err != nil {
panic(err)
}
return ctx
}
// NewManuallyManagedContext creates a new context, but the Run() method which locks a goroutine to an OS thread, has to be manually run
func NewManuallyManagedContext(d Device, flags ContextFlags) *Ctx {
var cctx C.CUcontext
err := result(C.cuCtxCreate(&cctx, C.uint(flags), C.CUdevice(d)))
if err != nil {
panic(err)
}
ctx := newContext(CUContext{cctx})
ctx.device = d
ctx.flags = flags
return ctx
}
// CtxFromCUContext is another way of buildinga *Ctx.
//
// Typical example:
// cuctx, err := dev.MakeContext(SchedAuto)
// if err != nil {
// ..error handling..
// }
// ctx := CtxFroMCUContext(d, cuctx)
func CtxFromCUContext(d Device, cuctx CUContext, flags ContextFlags) *Ctx {
ctx := newContext(cuctx)
ctx.device = d
ctx.flags = flags
return ctx
}
func newContext(c CUContext) *Ctx {
ctx := &Ctx{
CUContext: c,
work: make(chan func() error),
errChan: make(chan error),
}
runtime.SetFinalizer(ctx, finalizeCtx)
return ctx
}
// Close destroys the CUDA context and associated resources that has been created. Additionally, all channels of communications will be closed.
func (ctx *Ctx) Close() error {
var empty C.CUcontext
if ctx.CUContext.ctx == empty {
return nil
}
if ctx.errChan != nil {
close(ctx.errChan)
ctx.errChan = nil
}
if ctx.work != nil {
close(ctx.work)
ctx.work = nil
}
err := result(C.cuCtxDestroy(C.CUcontext(unsafe.Pointer(ctx.CUContext.ctx))))
ctx.CUContext.ctx = empty
return err
}
// Do does one function at a time.
func (ctx *Ctx) Do(fn func() error) error {
ctx.work <- fn
return <-ctx.errChan
}
// CUDAContext returns the CUDA Context
func (ctx *Ctx) CUDAContext() CUContext { return ctx.CUContext }
// Error returns the errors that may have occured during the calls.
func (ctx *Ctx) Error() error { return ctx.err }
// Work returns the channel where work will be passed in. In most cases you don't need this. Use Run instead.
func (ctx *Ctx) Work() <-chan func() error { return ctx.work }
// ErrChan returns the internal error channel used
func (ctx *Ctx) ErrChan() chan error { return ctx.errChan }
// Run locks the goroutine to the OS thread and ties the CUDA context to the OS thread. For most cases, this would suffice
//
// Note: errChan that is passed in should NOT be the same errChan as the one used internally for signalling.
// The main reasoning for passing in an error channel is to support two different kinds of run modes:
//
// The typical use example is as such:
//
/*
func A() {
ctx := NewContext(d, SchedAuto)
errChan := make(chan error)
go ctx.Run(errChan)
if err := <- errChan; err != nil {
// handleError
}
doSomethingWithCtx(ctx)
}
*/
// And yet another run mode supported is running of the context in the main thread:
//
/*
func main() {
ctx := NewContext(d, SchedAuto)
go doSomethingWithCtx(ctx)
if err := ctx.Run(nil); err != nil{
// handle error
}
}
*/
func (ctx *Ctx) Run(errChan chan error) error {
runtime.LockOSThread()
// set current, which locks the context to the OS thread
if err := SetCurrentContext(ctx.CUContext); err != nil {
if errChan != nil {
errChan <- err
} else {
return err
}
return nil
}
close(errChan)
// wait for Do()s
for w := range ctx.work {
ctx.errChan <- w()
}
runtime.UnlockOSThread()
return nil
}
func finalizeCtx(ctx *Ctx) { ctx.Close() }
/* Manually Written Methods */
func (ctx *Ctx) ResetL2Cache() {
ctx.Do(ctx.CUContext.ResetL2Cache)
}