Major refactoring inspired by Chad

Author: chewxy

array.go

@@ -2,7 +2,11 @@ package cu
 
 // #include <cuda.h>
 import "C"
-import "unsafe"
+import (
+	"unsafe"
+
+	"github.com/pkg/errors"
+)
 
 // Format is the type of array (think array types)
 type Format byte
@@ -18,10 +22,11 @@ const (
 	Float32 Format = C.CU_AD_FORMAT_FLOAT          // 32-bit floating point
 )
 
-// Array is the pointer to a CUDA array. The name is a bit of a misnomer, as it would lead one to imply that it's rangeable. It's not.
+// Array is the pointer to a CUDA array. The name is a bit of a misnomer,
+// as it would lead one to imply that it's rangeable. It's not.
 type Array uintptr
 
-func (arr Array) cuda() C.CUarray {
+func (arr Array) c() C.CUarray {
 	return *(*C.CUarray)(unsafe.Pointer(uintptr(arr)))
 }
 
@@ -129,3 +134,28 @@ func goArrayDesc(desc *C.CUDA_ARRAY_DESCRIPTOR) ArrayDesc {
 		NumChannels: uint(ad.NumChannels),
 	}
 }
+
+// Descriptor3 get a 3D CUDA array descriptor
+func (arr Array) Descriptor3() (Array3Desc, error) {
+	hArray := arr.c()
+	var desc C.CUDA_ARRAY3D_DESCRIPTOR
+	if err := result(C.cuArray3DGetDescriptor(&desc, hArray)); err != nil {
+		return Array3Desc{}, errors.Wrapf(err, "Array3DGetDescriptor")
+	}
+	return goArray3Desc(&desc), nil
+}
+
+// Descriptor gets a 1D or 2D CUDA array descriptor
+func (arr Array) Descriptor() (ArrayDesc, error) {
+	hArray := arr.c()
+	var desc C.CUDA_ARRAY_DESCRIPTOR
+	if err := result(C.cuArrayGetDescriptor(&desc, hArray)); err != nil {
+		return ArrayDesc{}, errors.Wrapf(err, "ArrayGetDescriptor")
+	}
+	return goArrayDesc(&desc), nil
+
+}
+
+func cuArrayToArray(arr *C.CUarray) Array {
+	return Array(uintptr(unsafe.Pointer(arr)))
+}

batch.c

@@ -57,7 +57,7 @@ CUresult processFn(fnargs_t* args){
 		abort();
 		break;
 	case fn_mallocManaged:
-		abort();
+		ret = cuMemAllocManaged(&args->devPtr0, args->size, CU_MEM_ATTACH_GLOBAL);
 		break;
 	case fn_memfreeD:
 		// fprintf(stderr, "memfree %p\n", args->devPtr0);

batch.go

@@ -133,7 +133,7 @@ func (fn *fnargs) c() C.uintptr_t {
 // For the moment, BatchedContext only supports a limited number of CUDA Runtime APIs.
 // Feel free to send a pull request with more APIs.
 type BatchedContext struct {
-	Context
+	CUContext
 	Device
 
 	workAvailable chan struct{} // an empty struct is sent down workAvailable when there is work
@@ -149,10 +149,10 @@ type BatchedContext struct {
 }
 
 // NewBatchedContext creates a batched CUDA context.
-func NewBatchedContext(c Context, d Device) *BatchedContext {
+func NewBatchedContext(c CUContext, d Device) *BatchedContext {
 	return &BatchedContext{
-		Context: c,
-		Device:  d,
+		CUContext: c,
+		Device:    d,
 
 		workAvailable: make(chan struct{}, 1),
 		work:          make(chan call, workBufLen),
@@ -232,7 +232,7 @@ func (ctx *BatchedContext) DoWork() {
 		addQueueLength(len(ctx.queue))
 		addBlockingCallers()
 
-		cctx := C.CUcontext(unsafe.Pointer(uintptr(ctx.Context)))
+		cctx := C.CUcontext(unsafe.Pointer(uintptr(ctx.CUContext)))
 		ctx.results = ctx.results[:cap(ctx.results)]                         // make sure of the maximum availability for ctx.results
 		C.process(cctx, &ctx.fns[0], &ctx.results[0], C.int(len(ctx.queue))) // process the queue
 		ctx.results = ctx.results[:len(ctx.queue)]                           // then  truncate it to the len of queue for reporting purposes
@@ -252,6 +252,8 @@ func (ctx *BatchedContext) DoWork() {
 				ctx.retVal <- DevicePtr(retVal.devptr0)
 			case C.fn_mallocH:
 			case C.fn_mallocManaged:
+				retVal = (*fnargs)(unsafe.Pointer(uintptr(ctx.fns[len(ctx.fns)-1])))
+				ctx.retVal <- DevicePtr(retVal.devptr0)
 			case C.fn_allocAndCopy:
 				retVal = (*fnargs)(unsafe.Pointer(uintptr(ctx.fns[len(ctx.fns)-1])))
 				ctx.retVal <- DevicePtr(retVal.devptr0)
@@ -295,7 +297,7 @@ func (ctx *BatchedContext) FirstError() error {
 func (ctx *BatchedContext) SetCurrent() {
 	fn := &fnargs{
 		fn:  C.fn_setCurrent,
-		ctx: C.CUcontext(unsafe.Pointer(uintptr(ctx.Context))),
+		ctx: C.CUcontext(unsafe.Pointer(uintptr(ctx.CUContext))),
 	}
 	c := call{fn, false}
 	ctx.enqueue(c)
@@ -311,6 +313,15 @@ func (ctx *BatchedContext) MemAlloc(bytesize int64) (retVal DevicePtr, err error
 	return ctx.enqueue(c)
 }
 
+func (ctx *BatchedContext) MemAllocManaged(bytesize int64, flags MemAttachFlags) (retVal DevicePtr, err error) {
+	fn := &fnargs{
+		fn:   C.fn_mallocManaged,
+		size: C.size_t(bytesize),
+	}
+	c := call{fn, true}
+	return ctx.enqueue(c)
+}
+
 func (ctx *BatchedContext) Memcpy(dst, src DevicePtr, byteCount int64) {
 	// pc, _, _, _ := runtime.Caller(1)
 	// logf("Memcpy %v %v| called by %v", dst, src, runtime.FuncForPC(pc).Name())

batch_test.go

@@ -9,7 +9,7 @@ import (
 func TestBatchContext(t *testing.T) {
 	var err error
 	var dev Device
-	var ctx Context
+	var ctx CUContext
 	var mod Module
 	var fn Function
 
@@ -97,7 +97,7 @@ loop:
 func TestLargeBatch(t *testing.T) {
 	var err error
 	var dev Device
-	var ctx Context
+	var ctx CUContext
 	var mod Module
 	var fn Function
 
@@ -210,7 +210,7 @@ func BenchmarkNoBatching(bench *testing.B) {
 	defer runtime.UnlockOSThread()
 
 	var err error
-	var ctx Context
+	var ctx CUContext
 	var mod Module
 	var fn Function
 
@@ -289,7 +289,7 @@ func BenchmarkBatching(bench *testing.B) {
 
 	var err error
 	var dev Device
-	var ctx Context
+	var ctx CUContext
 	var mod Module
 	var fn Function
 

batchedPatterns_test.go

@@ -8,7 +8,7 @@ import (
 
 func TestAttributes(t *testing.T) {
 	var dev Device
-	var ctx Context
+	var ctx CUContext
 	var err error
 
 	if dev, ctx, err = testSetup(); err != nil {
@@ -51,7 +51,7 @@ func TestAttributes(t *testing.T) {
 
 func TestLaunchAndSync(t *testing.T) {
 	var err error
-	var ctx Context
+	var ctx CUContext
 	var mod Module
 	var fn Function
 
@@ -128,7 +128,7 @@ func TestLaunchAndSync(t *testing.T) {
 
 func TestAllocAndCopy(t *testing.T) {
 	var err error
-	var ctx Context
+	var ctx CUContext
 	var mem DevicePtr
 
 	if _, ctx, err = testSetup(); err != nil {

cublas/cmd/cublasgen/binding.go cmd/gencublas/binding.go

@@ -28,8 +28,8 @@ var (
 )
 
 const (
+	typ     = "impl *Standalone"
 	header  = "cublasgen.h"
-	typ     = "impl *Implementation"
 	prefix  = "cublas"
 	warning = "Float32 implementations are autogenerated and not directly tested."
 )

cublas/cmd/cublasgen/main.go cmd/gencublas/main.go

@@ -37,11 +37,11 @@ type drotmParams struct {
 	h    [4]float64
 }
 
-func (impl *Implementation) Srotg(a float32, b float32) (c float32, s float32, r float32, z float32) {
+func (impl *Standalone) Srotg(a float32, b float32) (c float32, s float32, r float32, z float32) {
 	impl.e = status(C.cublasSrotg(C.cublasHandle_t(impl.h), (*C.float)(&a), (*C.float)(&b), (*C.float)(&c), (*C.float)(&s)))
 	return c, s, a, b
 }
-func (impl *Implementation) Srotmg(d1 float32, d2 float32, b1 float32, b2 float32) (p blas.SrotmParams, rd1 float32, rd2 float32, rb1 float32) {
+func (impl *Standalone) Srotmg(d1 float32, d2 float32, b1 float32, b2 float32) (p blas.SrotmParams, rd1 float32, rd2 float32, rb1 float32) {
 	if impl.e != nil {
 			return
 	}
@@ -50,7 +50,7 @@ func (impl *Implementation) Srotmg(d1 float32, d2 float32, b1 float32, b2 float3
 	return blas.SrotmParams{Flag: blas.Flag(pi.flag), H: pi.h}, d1, d2, b1
 }
 
-func (impl *Implementation) Srotm(n int, x []float32, incX int, y []float32, incY int, p blas.SrotmParams) {
+func (impl *Standalone) Srotm(n int, x []float32, incX int, y []float32, incY int, p blas.SrotmParams) {
 	if impl.e != nil {
 			return
 	}
@@ -83,15 +83,15 @@ func (impl *Implementation) Srotm(n int, x []float32, incX int, y []float32, inc
 	impl.e = status(C.cublasSrotm(C.cublasHandle_t(impl.h), C.int(n), (*C.float)(&x[0]), C.int(incX), (*C.float)(&y[0]), C.int(incY), (*C.float)(unsafe.Pointer(&pi))))
 }
 
-func (impl *Implementation) Drotg(a float64, b float64) (c float64, s float64, r float64, z float64) {
+func (impl *Standalone) Drotg(a float64, b float64) (c float64, s float64, r float64, z float64) {
 	if impl.e != nil {
 			return
 	}
 	impl.e = status(C.cublasDrotg(C.cublasHandle_t(impl.h), (*C.double)(&a), (*C.double)(&b), (*C.double)(&c), (*C.double)(&s)))
 	return c, s, a, b
 }
 
-func (impl *Implementation) Drotmg(d1 float64, d2 float64, b1 float64, b2 float64) (p blas.DrotmParams, rd1 float64, rd2 float64, rb1 float64) {
+func (impl *Standalone) Drotmg(d1 float64, d2 float64, b1 float64, b2 float64) (p blas.DrotmParams, rd1 float64, rd2 float64, rb1 float64) {
 	if impl.e != nil {
 			return
 	}
@@ -100,7 +100,7 @@ func (impl *Implementation) Drotmg(d1 float64, d2 float64, b1 float64, b2 float6
 	return blas.DrotmParams{Flag: blas.Flag(pi.flag), H: pi.h}, d1, d2, b1
 }
 
-func (impl *Implementation) Drotm(n int, x []float64, incX int, y []float64, incY int, p blas.DrotmParams) {
+func (impl *Standalone) Drotm(n int, x []float64, incX int, y []float64, incY int, p blas.DrotmParams) {
 	if impl.e != nil {
 			return
 	}
@@ -132,7 +132,7 @@ func (impl *Implementation) Drotm(n int, x []float64, incX int, y []float64, inc
 	impl.e = status(C.cublasDrotm(C.cublasHandle_t(impl.h), C.int(n), (*C.double)(&x[0]), C.int(incX), (*C.double)(&y[0]), C.int(incY), (*C.double)(unsafe.Pointer(&pi))))
 }
 
-func (impl *Implementation) Cdotu(n int, x []complex64, incX int, y []complex64, incY int) (dotu complex64) {
+func (impl *Standalone) Cdotu(n int, x []complex64, incX int, y []complex64, incY int) (dotu complex64) {
 	if impl.e != nil {
 			return
 	}
@@ -157,7 +157,7 @@ func (impl *Implementation) Cdotu(n int, x []complex64, incX int, y []complex64,
 	impl.e = status(C.cublasCdotu(C.cublasHandle_t(impl.h), C.int(n), (*C.cuComplex)(unsafe.Pointer(&x[0])), C.int(incX), (*C.cuComplex)(unsafe.Pointer(&y[0])), C.int(incY), (*C.cuComplex)(unsafe.Pointer(&dotu))))
 	return dotu
 }
-func (impl *Implementation) Cdotc(n int, x []complex64, incX int, y []complex64, incY int) (dotc complex64) {
+func (impl *Standalone) Cdotc(n int, x []complex64, incX int, y []complex64, incY int) (dotc complex64) {
 	if impl.e != nil {
 			return
 	}
@@ -183,7 +183,7 @@ func (impl *Implementation) Cdotc(n int, x []complex64, incX int, y []complex64,
 	impl.e = status(C.cublasCdotc(C.cublasHandle_t(impl.h), C.int(n), (*C.cuComplex)(unsafe.Pointer(&x[0])), C.int(incX), (*C.cuComplex)(unsafe.Pointer(&y[0])), C.int(incY), (*C.cuComplex)(unsafe.Pointer(&dotc))))
 	return dotc
 }
-func (impl *Implementation) Zdotu(n int, x []complex128, incX int, y []complex128, incY int) (dotu complex128) {
+func (impl *Standalone) Zdotu(n int, x []complex128, incX int, y []complex128, incY int) (dotu complex128) {
 	if impl.e != nil {
 			return
 	}
@@ -209,7 +209,7 @@ func (impl *Implementation) Zdotu(n int, x []complex128, incX int, y []complex12
 	impl.e = status(C.cublasZdotu(C.cublasHandle_t(impl.h), C.int(n), (*C.cuDoubleComplex)(unsafe.Pointer(&x[0])), C.int(incX), (*C.cuDoubleComplex)(unsafe.Pointer(&y[0])), C.int(incY), (*C.cuDoubleComplex)(unsafe.Pointer(&dotu))))
 	return dotu
 }
-func (impl *Implementation) Zdotc(n int, x []complex128, incX int, y []complex128, incY int) (dotc complex128) {
+func (impl *Standalone) Zdotc(n int, x []complex128, incX int, y []complex128, incY int) (dotc complex128) {
 	if impl.e != nil {
 			return
 	}
@@ -235,27 +235,27 @@ func (impl *Implementation) Zdotc(n int, x []complex128, incX int, y []complex12
 	return dotc
 }
 
-func (impl *Implementation) Sdsdot(n int, alpha float32, x []float32, incX int, y []float32, incY int) float32 {
+func (impl *Standalone) Sdsdot(n int, alpha float32, x []float32, incX int, y []float32, incY int) float32 {
 	panic("Unimplemented in cuBLAS. Please contact nvidia.")
 }
 
-func (impl *Implementation) Dsdot(n int, x []float32, incX int, y []float32, incY int) float64 {
+func (impl *Standalone) Dsdot(n int, x []float32, incX int, y []float32, incY int) float64 {
 	panic("Unimplemented in cuBLAS. Please contact nvidia.")
 }
 
-func (impl *Implementation) Strmm(s blas.Side, ul blas.Uplo, tA blas.Transpose, d blas.Diag, m, n int, alpha float32, a []float32, lda int, b []float32, ldb int){
+func (impl *Standalone) Strmm(s blas.Side, ul blas.Uplo, tA blas.Transpose, d blas.Diag, m, n int, alpha float32, a []float32, lda int, b []float32, ldb int){
 	panic("Unimplemented in cuBLAS. Please contact nvidia.")	
 }
 
-func (impl *Implementation) Dtrmm(s blas.Side, ul blas.Uplo, tA blas.Transpose, d blas.Diag, m, n int, alpha float64, a []float64, lda int, b []float64, ldb int){
+func (impl *Standalone) Dtrmm(s blas.Side, ul blas.Uplo, tA blas.Transpose, d blas.Diag, m, n int, alpha float64, a []float64, lda int, b []float64, ldb int){
 	panic("Unimplemented in cuBLAS. Please contact nvidia.")	
 }
 
-func (impl *Implementation) Ctrmm(s blas.Side, ul blas.Uplo, tA blas.Transpose, d blas.Diag, m, n int, alpha complex64, a []complex64, lda int, b []complex64, ldb int){
+func (impl *Standalone) Ctrmm(s blas.Side, ul blas.Uplo, tA blas.Transpose, d blas.Diag, m, n int, alpha complex64, a []complex64, lda int, b []complex64, ldb int){
 	panic("Unimplemented in cuBLAS. Please contact nvidia.")	
 }
 
-func (impl *Implementation) Ztrmm(s blas.Side, ul blas.Uplo, tA blas.Transpose, d blas.Diag, m, n int, alpha complex128, a []complex128, lda int, b []complex128, ldb int){
+func (impl *Standalone) Ztrmm(s blas.Side, ul blas.Uplo, tA blas.Transpose, d blas.Diag, m, n int, alpha complex128, a []complex128, lda int, b []complex128, ldb int){
 	panic("Unimplemented in cuBLAS. Please contact nvidia.")	
 }
 
@@ -265,10 +265,10 @@ func (impl *Implementation) Ztrmm(s blas.Side, ul blas.Uplo, tA blas.Transpose,
 
 // TODO: complex scale
 const complexScales = `
-func (impl *Implementation) Cscal(n int, alpha complex64, x []complex64, incX int) {}
-func (impl *Implementation) Zscal(n int, alpha complex64, x []complex128, incX int){}
-func (impl *Implementation) Csscal(n int, alpha float32, x []complex64, incX int) {}
-func (impl *Implementation) Zsscal(n int, alpha float64, x []complex128, incX int){}
+func (impl *Standalone) Cscal(n int, alpha complex64, x []complex64, incX int) {}
+func (impl *Standalone) Zscal(n int, alpha complex64, x []complex128, incX int){}
+func (impl *Standalone) Csscal(n int, alpha float32, x []complex64, incX int) {}
+func (impl *Standalone) Zsscal(n int, alpha float64, x []complex128, incX int){}
 `
 
 const amaxRaw = `