Use Stream wrapping instead of casting

Author: egonelbre

addressing.go

@@ -81,7 +81,7 @@ func (d DevicePtr) MemAdvise(count int64, advice MemAdvice, dev Device) error {
 func (d DevicePtr) MemPrefetchAsync(count int64, dst Device, hStream Stream) error {
 	devPtr := C.CUdeviceptr(d)
 	cc := C.size_t(count)
-	str := C.CUstream(unsafe.Pointer(uintptr(hStream)))
+	str := hStream.s
 	dv := C.CUdevice(dst)
 	return result(C.cuMemPrefetchAsync(devPtr, cc, dv, str))
 }

batch.go

@@ -413,7 +413,7 @@ func (ctx *BatchedContext) LaunchKernel(function Function, gridDimX, gridDimY, g
 		blockDimY:      C.uint(blockDimY),
 		blockDimZ:      C.uint(blockDimZ),
 		sharedMemBytes: C.uint(sharedMemBytes),
-		stream:         C.CUstream(unsafe.Pointer(uintptr(stream))),
+		stream:         stream.c(),
 		kernelParams:   (*unsafe.Pointer)(argp),
 		extra:          (*unsafe.Pointer)(unsafe.Pointer(uintptr(0))),
 	}

batch_test.go

@@ -64,7 +64,7 @@ func TestBatchContext(t *testing.T) {
 		}
 
 		bctx.MemcpyHtoD(memB, unsafe.Pointer(&b[0]), size)
-		bctx.LaunchKernel(fn, 1, 1, 1, len(a), 1, 1, 0, Stream(0), args)
+		bctx.LaunchKernel(fn, 1, 1, 1, len(a), 1, 1, 0, Stream{}, args)
 		bctx.Synchronize()
 		bctx.MemcpyDtoH(unsafe.Pointer(&a[0]), memA, size)
 		bctx.MemcpyDtoH(unsafe.Pointer(&b[0]), memB, size)
@@ -161,7 +161,7 @@ func TestLargeBatch(t *testing.T) {
 			}
 
 			bctx.MemcpyHtoD(memB, unsafe.Pointer(&b[0]), size)
-			bctx.LaunchKernel(fn, 1, 1, 1, len(a), 1, 1, 0, Stream(0), args)
+			bctx.LaunchKernel(fn, 1, 1, 1, len(a), 1, 1, 0, Stream{}, args)
 			bctx.Synchronize()
 
 			if i%13 == 0 {
@@ -273,7 +273,7 @@ func BenchmarkNoBatching(bench *testing.B) {
 				bench.Fatalf("Failed to copy memory from b: %v", err)
 			}
 
-			if err = fn.LaunchAndSync(100, 10, 1, 1000, 1, 1, 1, Stream(0), args); err != nil {
+			if err = fn.LaunchAndSync(100, 10, 1, 1000, 1, 1, 1, Stream{}, args); err != nil {
 				bench.Errorf("Launch and Sync Failed: %v", err)
 			}
 
@@ -353,7 +353,7 @@ func BenchmarkBatching(bench *testing.B) {
 			default:
 				bctx.MemcpyHtoD(memA, unsafe.Pointer(&a[0]), size)
 				bctx.MemcpyHtoD(memB, unsafe.Pointer(&b[0]), size)
-				bctx.LaunchKernel(fn, 100, 10, 1, 1000, 1, 1, 0, Stream(0), args)
+				bctx.LaunchKernel(fn, 100, 10, 1, 1000, 1, 1, 0, Stream{}, args)
 				bctx.Synchronize()
 				bctx.MemcpyDtoH(unsafe.Pointer(&a[0]), memA, size)
 				bctx.MemcpyDtoH(unsafe.Pointer(&b[0]), memB, size)

batchedPatterns.go

@@ -53,7 +53,7 @@ func (fn Function) LaunchAndSync(gridDimX, gridDimY, gridDimZ, blockDimX, blockD
 		C.uint(blockDimY),
 		C.uint(blockDimZ),
 		C.uint(sharedMemBytes),
-		C.CUstream(unsafe.Pointer(uintptr(stream))),
+		stream.c(),
 		(*unsafe.Pointer)(argp),
 		(*unsafe.Pointer)(unsafe.Pointer(uintptr(0)))))
 	return err

batchedPatterns_test.go

@@ -101,7 +101,7 @@ func TestLaunchAndSync(t *testing.T) {
 		unsafe.Pointer(&size),
 	}
 
-	if err = fn.LaunchAndSync(1, 1, 1, len(a), 1, 1, 1, Stream(0), args); err != nil {
+	if err = fn.LaunchAndSync(1, 1, 1, len(a), 1, 1, 1, Stream{}, args); err != nil {
 		t.Errorf("Launch and Sync Failed: %v", err)
 	}
 

cucontext_test.go

@@ -173,7 +173,7 @@ func TestMultipleContextSingleHostThread(t *testing.T) {
 		unsafe.Pointer(&size),
 	}
 
-	if err = fn0.LaunchAndSync(1, 1, 1, len(data), 1, 1, 0, Stream(0), args); err != nil {
+	if err = fn0.LaunchAndSync(1, 1, 1, len(data), 1, 1, 0, Stream{}, args); err != nil {
 		t.Errorf("Failed to launcj add32: %v", err)
 	}
 
@@ -204,7 +204,7 @@ func TestMultipleContextSingleHostThread(t *testing.T) {
 		unsafe.Pointer(&size),
 	}
 
-	if err = fn1.LaunchAndSync(1, 1, 1, len(data), 1, 1, 0, Stream(0), args); err != nil {
+	if err = fn1.LaunchAndSync(1, 1, 1, len(data), 1, 1, 0, Stream{}, args); err != nil {
 		t.Errorf("Failed to launcj add32: %v", err)
 	}
 
@@ -221,7 +221,7 @@ func TestMultipleContextSingleHostThread(t *testing.T) {
 		unsafe.Pointer(&size),
 	}
 
-	if err = fn0.LaunchAndSync(1, 1, 1, len(data), 1, 1, 0, Stream(0), args); err == nil {
+	if err = fn0.LaunchAndSync(1, 1, 1, len(data), 1, 1, 0, Stream{}, args); err == nil {
 		t.Errorf("Expected error when launching a kernel defined in a different context")
 	}
 	t.Log(err)
@@ -230,7 +230,7 @@ func TestMultipleContextSingleHostThread(t *testing.T) {
 	if err = SetCurrentContext(ctx0); err != nil {
 		t.Errorf("Failed to swtch to ctx0 %v", err)
 	}
-	if err = fn0.LaunchAndSync(1, 1, 1, len(data), 1, 1, 0, Stream(0), args); err != nil {
+	if err = fn0.LaunchAndSync(1, 1, 1, len(data), 1, 1, 0, Stream{}, args); err != nil {
 		t.Errorf("fn0 errored while using memory declared in ctx1: %v", err)
 	}
 	t.Log(err)

execution.go

@@ -35,7 +35,7 @@ func (fn Function) LaunchKernel(gridDimX, gridDimY, gridDimZ int, blockDimX, blo
 		C.uint(blockDimY),
 		C.uint(blockDimZ),
 		C.uint(sharedMemBytes),
-		C.CUstream(unsafe.Pointer(uintptr(stream))),
+		stream.c(),
 		(*unsafe.Pointer)(argp),
 		(*unsafe.Pointer)(unsafe.Pointer(uintptr(0)))))
 	return err
@@ -67,7 +67,7 @@ func (ctx *Ctx) LaunchKernel(fn Function, gridDimX, gridDimY, gridDimZ int, bloc
 			C.uint(blockDimY),
 			C.uint(blockDimZ),
 			C.uint(sharedMemBytes),
-			C.CUstream(unsafe.Pointer(uintptr(stream))),
+			stream.c(),
 			(*unsafe.Pointer)(argp),
 			(*unsafe.Pointer)(unsafe.Pointer(uintptr(0)))))
 	}

module_test.go

@@ -53,7 +53,7 @@ func TestModule(t *testing.T) {
 	grid := DivUp(N, block)
 	shmem := 0
 	args := []unsafe.Pointer{unsafe.Pointer(&A), unsafe.Pointer(&value), unsafe.Pointer(&n)}
-	if err = f.LaunchKernel(grid, 1, 1, block, 1, 1, shmem, 0, args); err != nil {
+	if err = f.LaunchKernel(grid, 1, 1, block, 1, 1, shmem, Stream{}, args); err != nil {
 		t.Fatal(err)
 	}
 

stream.go

@@ -3,27 +3,25 @@ package cu
 // #include <cuda.h>
 import "C"
 import (
-	"unsafe"
-
 	"github.com/pkg/errors"
 )
 
 // Stream represents a CUDA stream.
-type Stream uintptr
+type Stream struct {
+	s C.CUstream
+}
 
-func makeStream(s C.CUstream) Stream { return Stream(uintptr(unsafe.Pointer(s))) }
-func (s Stream) c() C.CUstream       { return C.CUstream(unsafe.Pointer(uintptr(s))) }
+func makeStream(s C.CUstream) Stream { return Stream{s} }
+func (s Stream) c() C.CUstream       { return s.s }
 
 // C is the exported version of the c method
 func (s Stream) C() C.CUstream { return s.c() }
 
 // MakeStream creates a stream. The flags determines the behaviors of the stream.
 func MakeStream(flags StreamFlags) (Stream, error) {
-	var s C.CUstream
-	if err := result(C.cuStreamCreate(&s, C.uint(flags))); err != nil {
-		return 0, err
-	}
-	return makeStream(s), nil
+	var s Stream
+	err := result(C.cuStreamCreate(&s.s, C.uint(flags)))
+	return s, err
 }
 
 // MakeStreamWithPriority creates a stream with the given priority. The flags determines the behaviors of the stream.
@@ -35,52 +33,41 @@ func MakeStream(flags StreamFlags) (Stream, error) {
 // If the specified priority is outside the numerical range returned by `StreamPriorityRange`,
 // it will automatically be clamped to the lowest or the highest number in the range.
 func MakeStreamWithPriority(priority int, flags StreamFlags) (Stream, error) {
-	var s C.CUstream
-	if err := result(C.cuStreamCreateWithPriority(&s, C.uint(flags), C.int(priority))); err != nil {
-		return 0, err
-	}
-	return makeStream(s), nil
+	var s Stream
+	err := result(C.cuStreamCreateWithPriority(&s.s, C.uint(flags), C.int(priority)))
+	return s, err
 }
 
 // DestroyStream destroys the stream specified by hStream.
 //
 // In case the device is still doing work in the stream hStream when DestroyStrea() is called,
 // the function will return immediately and the resources associated with hStream will be released automatically once the device has completed all work in hStream.
-func DestroyStream(hStream *Stream) error {
-	stream := *hStream
-	s := stream.c()
-	*hStream = 0
-	return result(C.cuStreamDestroy(s))
+func (hStream *Stream) Destroy() error {
+	err := result(C.cuStreamDestroy(hStream.s))
+	*hStream = Stream{}
+	return err
 }
 
 func (ctx *Ctx) MakeStream(flags StreamFlags) (stream Stream, err error) {
-	var s C.CUstream
+	var s Stream
 
-	f := func() error { return result(C.cuStreamCreate(&s, C.uint(flags))) }
+	f := func() error { return result(C.cuStreamCreate(&s.s, C.uint(flags))) }
 	if err = ctx.Do(f); err != nil {
-		err = errors.Wrap(err, "MakeStream")
-		return
+		return s, errors.Wrap(err, "MakeStream")
 	}
-	stream = makeStream(s)
-	return
+	return s, nil
 }
 
-func (ctx *Ctx) MakeStreamWithPriority(priority int, flags StreamFlags) (stream Stream, err error) {
-	var s C.CUstream
-
-	f := func() error { return result(C.cuStreamCreateWithPriority(&s, C.uint(flags), C.int(priority))) }
-	if err = ctx.Do(f); err != nil {
-		err = errors.Wrap(err, "MakeStream With Priority")
-		return
+func (ctx *Ctx) MakeStreamWithPriority(priority int, flags StreamFlags) (Stream, error) {
+	var s Stream
+	f := func() error { return result(C.cuStreamCreateWithPriority(&s.s, C.uint(flags), C.int(priority))) }
+	if err := ctx.Do(f); err != nil {
+		return s, errors.Wrap(err, "MakeStream With Priority")
 	}
-	stream = makeStream(s)
-	return
+	return s, nil
 }
 
 func (ctx *Ctx) DestroyStream(hStream *Stream) {
-	stream := *hStream
-	s := stream.c()
-
-	f := func() error { return result(C.cuStreamDestroy(s)) }
+	f := func() error { return result(C.cuStreamDestroy(hStream.s)) }
 	ctx.err = ctx.Do(f)
 }