cudnnOpTensor.go

package gocudnn

/*
#include <cudnn.h>
*/
import "C"
import (
	"fmt"
	"runtime"
	"unsafe"

	"github.com/dereklstinson/cutil"
)

//OPTensorD holds OP Tensor information
type OPTensorD struct {
	descriptor C.cudnnOpTensorDescriptor_t
	gogc       bool
}

//CreateOpTensorDescriptor creates and sets an OpTensor
func CreateOpTensorDescriptor() (*OPTensorD, error) {
	desc := new(OPTensorD)
	err := Status(C.cudnnCreateOpTensorDescriptor(&desc.descriptor)).error("CreateOpTensorDescriptor() ")
	if setfinalizer {
		desc.gogc = true
		runtime.SetFinalizer(desc, destroyopdesc)
	}

	return desc, err
}

//Set sets the OPTensorD.
func (t *OPTensorD) Set(op OpTensorOp, dtype DataType, nan NANProp) error {
	return Status(C.cudnnSetOpTensorDescriptor(t.descriptor, op.c(), dtype.c(), nan.c())).error("(t *OPTensorD) Set")

}

//Get returns the descriptor information with error
func (t *OPTensorD) Get() (op OpTensorOp, dtype DataType, nan NANProp, err error) {

	err = Status(C.cudnnGetOpTensorDescriptor(t.descriptor, op.cptr(), dtype.cptr(), nan.cptr())).error("(t *OPTensorD) Get()")

	return op, dtype, nan, err

}
func (t *OPTensorD) String() string {
	op, dtype, nan, err := t.Get()
	if err != nil {
		return fmt.Sprintf("OPTensorD{\nError: %v\n}\n", err)
	}
	return fmt.Sprintf("OPTensorD{\n%v,\n%v,\n%v\n}\n", op, dtype, nan)
}
func destroyopdesc(t *OPTensorD) error {
	return Status(C.cudnnDestroyOpTensorDescriptor(t.descriptor)).error(" destroyopdesc(t *OPTensorD)")
}

//Destroy destroys the descriptor
func (t *OPTensorD) Destroy() error {
	if setfinalizer || t.gogc {
		return nil
	}
	return destroyopdesc(t)
}

//OpTensor performs an operation on some tensors   C= operation( (alpha1 * A) , (alpha2 *B) ) + (beta *C)
func (t *OPTensorD) OpTensor(
	handle *Handle,
	alpha1 float64,
	aD *TensorD, A cutil.Mem,
	alpha2 float64,
	bD *TensorD, B cutil.Mem,
	beta float64,
	cD *TensorD, cmem cutil.Mem) error {
	a1 := cscalarbydatatype(aD.dtype, alpha1)
	a2 := cscalarbydatatype(bD.dtype, alpha2)
	b := cscalarbydatatype(cD.dtype, beta)
	if handle.w != nil {
		return handle.w.Work(func() error {
			return Status(C.cudnnOpTensor(
				handle.x,
				t.descriptor,
				a1.CPtr(),
				aD.descriptor, A.Ptr(),
				a2.CPtr(),
				bD.descriptor, B.Ptr(),
				b.CPtr(),
				cD.descriptor, cmem.Ptr())).error("(t *OPTensorD) OpTensor")
		})
	}
	return Status(C.cudnnOpTensor(
		handle.x,
		t.descriptor,
		a1.CPtr(),
		aD.descriptor, A.Ptr(),
		a2.CPtr(),
		bD.descriptor, B.Ptr(),
		b.CPtr(),
		cD.descriptor, cmem.Ptr())).error("(t *OPTensorD) OpTensor")

}

//OpTensorUS is like OpTensor but uses unsafe.Pointer instead of cutil.Mem
func (t *OPTensorD) OpTensorUS(
	handle *Handle,
	alpha1 float64,
	aD *TensorD, A unsafe.Pointer,
	alpha2 float64,
	bD *TensorD, B unsafe.Pointer,
	beta float64,
	cD *TensorD, cmem unsafe.Pointer) error {
	a1 := cscalarbydatatype(aD.dtype, alpha1)
	a2 := cscalarbydatatype(bD.dtype, alpha2)
	b := cscalarbydatatype(cD.dtype, beta)
	if handle.w != nil {
		return handle.w.Work(func() error {
			return Status(C.cudnnOpTensor(
				handle.x,
				t.descriptor,
				a1.CPtr(),
				aD.descriptor, A,
				a2.CPtr(),
				bD.descriptor, B,
				b.CPtr(),
				cD.descriptor, cmem)).error("(t *OPTensorD) OpTensorUS")
		})
	}
	return Status(C.cudnnOpTensor(
		handle.x,
		t.descriptor,
		a1.CPtr(),
		aD.descriptor, A,
		a2.CPtr(),
		bD.descriptor, B,
		b.CPtr(),
		cD.descriptor, cmem)).error("(t *OPTensorD) OpTensorUS")

}

//OpTensorOp is used for flags for the Optensor functions
type OpTensorOp C.cudnnOpTensorOp_t

//Add sets o to OpTensorOp(C.CUDNN_OP_TENSOR_ADD) and returns the new value
func (o *OpTensorOp) Add() OpTensorOp { *o = OpTensorOp(C.CUDNN_OP_TENSOR_ADD); return *o }

//Mul sets o to OpTensorOp(C.CUDNN_OP_TENSOR_MUL) and returns the new value
func (o *OpTensorOp) Mul() OpTensorOp { *o = OpTensorOp(C.CUDNN_OP_TENSOR_MUL); return *o }

//Min sets o to OpTensorOp(C.CUDNN_OP_TENSOR_MIN)  and returns the new value
func (o *OpTensorOp) Min() OpTensorOp { *o = OpTensorOp(C.CUDNN_OP_TENSOR_MIN); return *o }

//Max sets o to OpTensorOp(C.CUDNN_OP_TENSOR_MAX) and returns the new value
func (o *OpTensorOp) Max() OpTensorOp { *o = OpTensorOp(C.CUDNN_OP_TENSOR_MAX); return *o }

//Sqrt sets o to OpTensorOp(C.CUDNN_OP_TENSOR_SQRT) and returns the new value
func (o *OpTensorOp) Sqrt() OpTensorOp { *o = OpTensorOp(C.CUDNN_OP_TENSOR_SQRT); return *o }

//Not returns OpTensorOp(C.CUDNN_OP_TENSOR_NOT) and returns the new value
func (o *OpTensorOp) Not() OpTensorOp { *o = OpTensorOp(C.CUDNN_OP_TENSOR_NOT); return *o }
func (o OpTensorOp) String() string {
	var x string
	oflg := o
	switch o {
	case oflg.Add():
		x = "Add"
	case oflg.Mul():
		x = "Mul"
	case oflg.Min():
		x = "Min"
	case oflg.Max():
		x = "Max"
	case oflg.Sqrt():
		x = "Sqrt"
	case oflg.Not():
		x = "Not"
	default:
		x = "Unsupported Flag"
	}
	return "OpTensorOp: " + x

}
func (o OpTensorOp) c() C.cudnnOpTensorOp_t      { return C.cudnnOpTensorOp_t(o) }
func (o *OpTensorOp) cptr() *C.cudnnOpTensorOp_t { return (*C.cudnnOpTensorOp_t)(o) }