Java JCublas类代码示例

import jcuda.jcublas.JCublas; //导入依赖的package包/类
/**
 * Retrieve the data from the gpu
 * @param arr the array to getFloat the data for
 * @param from the origin pointer
 * @param to the to pointer (DO NOT AUGMENT WITH OFFSET: THIS IS TAKEN CARE OF FOR YOU)
 */
public static void getData(JCublasNDArray arr,Pointer from,Pointer to) {

    //p is typically the data vector which is strided access
    if(arr.length() == arr.data().length())
        JCublas.cublasGetVector(
                arr.length(),
                size(arr),
                from,
                1,
                to.withByteOffset(arr.offset() * size(arr)),
                1);
    else
        JCublas.cublasGetVector(
                arr.length(),
                size(arr),
                from,
                1,
                to.withByteOffset(arr.offset() * size(arr)),
                arr.majorStride());




}
 

import jcuda.jcublas.JCublas; //导入依赖的package包/类
/**
 * Returns the index of the max element
 * in the given ndarray
 * @param x
 * @return
 */
public static int iamax(INDArray x) {
    JCublas.cublasInit();

    JCublasNDArray xC = (JCublasNDArray) x;
    Pointer xCPointer = alloc(xC);


    int max = JCublas.cublasIsamax(
            x.length(),
            xCPointer,
            1);
    free(xCPointer);
    return max - 1;

}
 

import jcuda.jcublas.JCublas; //导入依赖的package包/类
/**
 * Multiply the given ndarray
 *  by alpha
 * @param alpha
 * @param x
 * @return
 */
public static INDArray scal(double alpha, INDArray x) {
    DataTypeValidation.assertDouble(x);

    JCublas.cublasInit();
    JCublasNDArray xC = (JCublasNDArray) x;

    Pointer xCPointer = alloc(xC);
    JCublas.cublasDscal(
            xC.length(),
            alpha,
            xCPointer,
            1);
    getData(xC, xCPointer, Pointer.to(xC.data().asDouble()));
    free(xCPointer);
    return x;

}
 

import jcuda.jcublas.JCublas; //导入依赖的package包/类
public static IComplexDouble dot(IComplexNDArray x, IComplexNDArray y) {
    DataTypeValidation.assertSameDataType(x,y);
    JCublas.cublasInit();

    JCublasComplexNDArray aC = (JCublasComplexNDArray) x;
    JCublasComplexNDArray bC = (JCublasComplexNDArray) y;

    Pointer aCPointer = alloc(aC);
    Pointer bCPointer = alloc(bC);


    jcuda.cuDoubleComplex dott = JCublas.cublasZdotc(
            x.length(),
            aCPointer,
            1,
            bCPointer,
            1);

    IComplexDouble ret =   Nd4j.createDouble(dott.x, dott.y);
    free(aCPointer,bCPointer);
    return ret;
}
 

import jcuda.jcublas.JCublas; //导入依赖的package包/类
/**
 * Simpler version of saxpy
 * taking in to account the parameters of the ndarray
 * @param alpha the alpha to scale by
 * @param x the x
 * @param y the y
 */
public static void axpy(double alpha, INDArray x, INDArray y) {
    DataTypeValidation.assertDouble(x,y);
    JCublas.cublasInit();
    JCublasNDArray xC = (JCublasNDArray) x;
    JCublasNDArray yC = (JCublasNDArray) y;

    Pointer xCPointer = alloc(xC);
    Pointer yCPointer = alloc(yC);

    JCublas.cublasDaxpy(x.length(), alpha, xCPointer, 1, yCPointer, 1);


    getData(yC,yCPointer,Pointer.to(yC.data().asDouble()));
    free(xCPointer,yCPointer);

}
 

import jcuda.jcublas.JCublas; //导入依赖的package包/类
/**
 * Simpler version of saxpy
 * taking in to account the parameters of the ndarray
 * @param alpha the alpha to scale by
 * @param x the x
 * @param y the y
 */
public static void saxpy(float alpha, INDArray x, INDArray y) {
    DataTypeValidation.assertFloat(x,y);
    JCublas.cublasInit();
    JCublasNDArray xC = (JCublasNDArray) x;
    JCublasNDArray yC = (JCublasNDArray) y;

    Pointer xCPointer = alloc(xC);
    Pointer yCPointer = alloc(yC);

    JCublas.cublasSaxpy(x.length(),alpha,xCPointer,1,yCPointer,1);


    getData(yC,yCPointer,Pointer.to(yC.data().asFloat()));
    free(xCPointer,yCPointer);

}
 

import jcuda.jcublas.JCublas; //导入依赖的package包/类
private void updateParam(float learning_rate)
{	
	float alpha = -learning_rate;
	for(int i=0; i<model.getParamSize().size(); ++i)
	{
		JCublas.cublasSaxpy(model.getParamSize().get(i), alpha, model.getGradDev().get(i), 1, model.getParamDev().get(i), 1);
	}
}
 

import jcuda.jcublas.JCublas; //导入依赖的package包/类
/**
 * Allocate and return a pointer
 * based on the length of the ndarray
 * @param ndarray the ndarray to allocate
 * @return the allocated pointer
 */
public static Pointer alloc(JCublasComplexNDArray ndarray) {
    Pointer ret = new Pointer();
    //allocate memory for the pointer
    JCublas.cublasAlloc(
            ndarray.length() * 2,
            size(ndarray)
            , ret);

    /* Copy from data to pointer at majorStride() (you want to stride through the data properly) incrementing by 1 for the pointer on the GPU.
    * This allows us to copy only what we need. */
    Pointer toData =null;
    if(ndarray.data().dataType().equals(DataBuffer.FLOAT))
        toData = Pointer.to(ndarray.data().asFloat()).withByteOffset(ndarray.offset() * size(ndarray));
    else
        toData =  Pointer.to(ndarray.data().asDouble()).withByteOffset(ndarray.offset() * size(ndarray));

    if(ndarray.length() == ndarray.data().length())
        JCublas.cublasSetVector(
                ndarray.length() * 2,
                size(ndarray),
                toData,
                1,
                ret,
                1);
    else
        JCublas.cublasSetVector(
                ndarray.length() * 2,
                size(ndarray),
                toData,
                1,
                ret,
                1);

    return ret;

}
 

import jcuda.jcublas.JCublas; //导入依赖的package包/类
/**
 * General matrix vector multiplication
 * @param A
 * @param B
 * @param C
 * @param alpha
 * @param beta
 * @return
 */
public static INDArray gemv(INDArray A, INDArray B, INDArray C, double alpha, double beta) {

    DataTypeValidation.assertDouble(A,B,C);
    JCublas.cublasInit();

    JCublasNDArray cA = (JCublasNDArray) A;
    JCublasNDArray cB = (JCublasNDArray) B;
    JCublasNDArray cC = (JCublasNDArray) C;

    Pointer cAPointer = alloc(cA);
    Pointer cBPointer = alloc(cB);
    Pointer cCPointer = alloc(cC);


    JCublas.cublasDgemv(
            'N',
            A.rows(),
            A.columns(),
            alpha,
            cAPointer,
            A.rows(),
            cBPointer,
            1,
            beta,
            cCPointer,
            1);

    getData(cC,cCPointer,Pointer.to(cC.data().asDouble()));
    free(cAPointer,cBPointer,cCPointer);


    return C;
}
 

import jcuda.jcublas.JCublas; //导入依赖的package包/类
/**
 * General matrix multiply
 * @param A
 * @param B
 * @param a
 * @param C
 * @param b
 * @return
 */
public static IComplexNDArray gemm(IComplexNDArray A, IComplexNDArray B, IComplexFloat a,IComplexNDArray C
        , IComplexFloat b) {
    DataTypeValidation.assertFloat(A,B,C);
    JCublas.cublasInit();

    JCublasComplexNDArray cA = (JCublasComplexNDArray) A;
    JCublasComplexNDArray cB = (JCublasComplexNDArray) B;
    JCublasComplexNDArray cC = (JCublasComplexNDArray) C;

    Pointer cAPointer = alloc(cA);
    Pointer cBPointer = alloc(cB);
    Pointer cCPointer = alloc(cC);


    cuComplex alpha = cuComplex.cuCmplx(a.realComponent().floatValue(),b.imaginaryComponent().floatValue());
    cuComplex beta = cuComplex.cuCmplx(b.realComponent().floatValue(),b.imaginaryComponent().floatValue());

    JCublas.cublasCgemm(
            'n', //trans
            'n',
            cC.rows(),  // m
            cC.columns(), // n
            cA.columns(), //k,
            alpha,
            cAPointer, // A
            A.rows(),  // lda
            cBPointer, // x
            B.rows(), // ldb
            beta,  // beta
            cCPointer, // y
            C.rows()); // ldc


    getData(cC,cCPointer,Pointer.to(cC.data().asFloat()));
    free(cAPointer,cBPointer,cCPointer);

    return C;

}
 

import jcuda.jcublas.JCublas; //导入依赖的package包/类
/**
 * Copy x to y
 * @param x the origin
 * @param y the destination
 */
public static void copy(IComplexNDArray x, IComplexNDArray y) {
    DataTypeValidation.assertSameDataType(x,y);
    JCublas.cublasInit();

    JCublasComplexNDArray xC = (JCublasComplexNDArray) x;
    JCublasComplexNDArray yC = (JCublasComplexNDArray) y;

    Pointer xCPointer = alloc(xC);
    Pointer yCPointer = alloc(yC);

    if(xC.data().dataType().equals(DataBuffer.FLOAT)) {
        JCublas.cublasScopy(
                x.length(),
                xCPointer,
                1,
                yCPointer,
                1);


        getData(yC,yCPointer,Pointer.to(yC.data().asFloat()));

    }

    else {
        JCublas.cublasDcopy(
                x.length(),
                xCPointer,
                1,
                yCPointer,
                1);


        getData(yC,yCPointer,Pointer.to(yC.data().asDouble()));

    }


    free(xCPointer,yCPointer);
}
 

import jcuda.jcublas.JCublas; //导入依赖的package包/类
/**
 *
 * @param x
 * @return
 */
public static float asum(IComplexNDArray x) {

    JCublas.cublasInit();

    JCublasComplexNDArray xC = (JCublasComplexNDArray) x;

    Pointer xCPointer = alloc(xC);
    float sum = JCublas.cublasScasum(x.length(), xCPointer, 1);
    free(xCPointer);
    return sum;
}
 

import jcuda.jcublas.JCublas; //导入依赖的package包/类
/**
 * Swap the elements in each ndarray
 * @param x
 * @param y
 */
public static void swap(INDArray x, INDArray y) {

    DataTypeValidation.assertSameDataType(x,y);
    JCublas.cublasInit();

    JCublasNDArray xC = (JCublasNDArray) x;
    JCublasNDArray yC = (JCublasNDArray) y;
    Pointer xCPointer = alloc(xC);
    Pointer yCPointer = alloc(yC);


    if(xC.data().dataType().equals(DataBuffer.FLOAT)) {
        JCublas.cublasSswap(
                xC.length(),
                xCPointer,
                1,
                yCPointer,
                1);

        getData(yC,yCPointer,Pointer.to(yC.data().asFloat()));
    }

    else {
        JCublas.cublasDswap(
                xC.length(),
                xCPointer,
                1,
                yCPointer,
                1);

        getData(yC,yCPointer,Pointer.to(yC.data().asDouble()));
    }

    free(xCPointer,yCPointer);

}
 

import jcuda.jcublas.JCublas; //导入依赖的package包/类
/**
 * Returns the norm2 of the given ndarray
 * @param x
 * @return
 */
public static float nrm2(INDArray x) {
    JCublas.cublasInit();
    JCublasNDArray xC = (JCublasNDArray) x;
    Pointer xCPointer = alloc(xC);


    float normal2 = JCublas.cublasSnrm2(x.length(), xCPointer, 1);
    JCublas.cublasFree(xCPointer);
    return normal2;
}
 

import jcuda.jcublas.JCublas; //导入依赖的package包/类
public static INDArray ger(INDArray A, INDArray B, INDArray C, double alpha) {
    DataTypeValidation.assertDouble(A,B,C);
    JCublas.cublasInit();
    // = alpha * A * transpose(B) + C
    JCublasNDArray aC = (JCublasNDArray) A;
    JCublasNDArray bC = (JCublasNDArray) B;
    JCublasNDArray cC = (JCublasNDArray) C;

    Pointer aCPointer = alloc(aC);
    Pointer bCPointer = alloc(bC);
    Pointer cCPointer = alloc(cC);


    JCublas.cublasDger(
            A.rows(),   // m
            A.columns(),// n
            alpha,      // alpha
            aCPointer,        // d_A or x
            A.rows(),   // incx
            bCPointer,        // dB or y
            B.rows(),   // incy
            cCPointer,        // dC or A
            C.rows()    // lda
    );

    getData(cC,cCPointer,Pointer.to(cC.data().asDouble()));
    free(aCPointer, bCPointer, cCPointer);

    return C;
}