Java Code Examples for org.nd4j.linalg.api.buffer.DataBuffer#length()

/**
 * Create a data buffer struct within
 * the passed in {@link FlatBufferBuilder}
 * @param bufferBuilder the existing flatbuffer
 *                      to use to serialize the {@link DataBuffer}
 * @param create the databuffer to serialize
 * @return an int representing the offset of the buffer
 */
public static int createDataBufferStruct(FlatBufferBuilder bufferBuilder,DataBuffer create) {
    bufferBuilder.prep(create.getElementSize(), (int) create.length() * create.getElementSize());
    for(int i = (int) (create.length() - 1); i >= 0; i--) {
        switch(create.dataType()) {
            case DOUBLE:
                double putDouble = create.getDouble(i);
                bufferBuilder.putDouble(putDouble);
                break;
            case FLOAT:
                float putFloat = create.getFloat(i);
                bufferBuilder.putFloat(putFloat);
                break;
            case INT:
                int putInt = create.getInt(i);
                bufferBuilder.putInt(putInt);
                break;
            case LONG:
                long putLong = create.getLong(i);
                bufferBuilder.putLong(putLong);
        }
    }

    return bufferBuilder.offset();

}
 

@Test
public void testCholeskyL() {
    INDArray A = Nd4j.create(new double[] {2, -1, 1, -1, 2, -1, 1, -1, 2,});
    A = A.reshape('c', 3, 3);
    INDArray O = Nd4j.create(A.dataType(), A.shape());
    Nd4j.copy(A, O);

    Nd4j.getBlasWrapper().lapack().potrf(A, true);

    A.mmuli(A.transpose());
    O.subi(A);
    DataBuffer db = O.data();
    for (int i = 0; i < db.length(); i++) {
        assertEquals(0, db.getFloat(i), 1e-5);
    }
}
 

@Override
public INDArray sortCooIndices(INDArray x) {

    if(x.getFormat() != SparseFormat.COO){
        throw new UnsupportedOperationException("Not a COO ndarray");
    }
    BaseSparseNDArrayCOO array = (BaseSparseNDArrayCOO) x;
    DataBuffer val = array.getValues();
    DataBuffer idx = array.getIndices();
    long length = val.length();
    int rank = array.underlyingRank();
    switch(val.dataType()){
        case FLOAT:
            NativeOpsHolder.getInstance().getDeviceNativeOps().sortCooIndicesFloat(null, (LongPointer) idx.addressPointer(), (FloatPointer) val.addressPointer(), length, rank);
            break;
        case DOUBLE:
            NativeOpsHolder.getInstance().getDeviceNativeOps().sortCooIndicesDouble(null, (LongPointer) idx.addressPointer(), (DoublePointer) val.addressPointer(), length, rank);
            break;
        default:
            throw new UnsupportedOperationException("Unknown datatype " + x.data().dataType());
    }

    return array;
}
 

@Test
public void testQRSquare() {
    INDArray A = Nd4j.create(new double[] {1, 2, 3, 4, 5, 6, 7, 8, 9});
    A = A.reshape('c', 3, 3);
    INDArray O = Nd4j.create(A.dataType(), A.shape());
    Nd4j.copy(A, O);
    INDArray R = Nd4j.create(A.dataType(), A.columns(), A.columns());

    Nd4j.getBlasWrapper().lapack().geqrf(A, R);

    A.mmuli(R);
    O.subi(A);
    DataBuffer db = O.data();
    for (int i = 0; i < db.length(); i++) {
        assertEquals(0, db.getFloat(i), 1e-5);
    }
}
 

@Test
public void testQRSquare() {
    INDArray A = Nd4j.create(new double[] {1, 2, 3, 4, 5, 6, 7, 8, 9});
    A = A.reshape('c', 3, 3);
    INDArray O = Nd4j.create(A.shape());
    Nd4j.copy(A, O);
    INDArray R = Nd4j.create(A.columns(), A.columns());

    Nd4j.getBlasWrapper().lapack().geqrf(A, R);

    A.mmuli(R);
    O.subi(A);
    DataBuffer db = O.data();
    for (int i = 0; i < db.length(); i++) {
        assertEquals(0, db.getFloat(i), 1e-5);
    }
}
 

@Test
public void testQRRect() {
    INDArray A = Nd4j.create(new double[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
    A = A.reshape('f', 4, 3);
    INDArray O = Nd4j.create(A.shape());
    Nd4j.copy(A, O);

    INDArray R = Nd4j.create(A.columns(), A.columns());
    Nd4j.getBlasWrapper().lapack().geqrf(A, R);

    A.mmuli(R);
    O.subi(A);
    DataBuffer db = O.data();
    for (int i = 0; i < db.length(); i++) {
        assertEquals(0, db.getFloat(i), 1e-5);
    }
}
 

private DataBuffer addAtPosition(DataBuffer buf, long dataSize, int pos, double value) {

        DataBuffer buffer = (buf.length() == dataSize) ? reallocate(buf) : buf;
        double[] tail = buffer.getDoublesAt(pos, (int) dataSize - pos);

        buffer.put(pos, value);
        for (int i = 0; i < tail.length; i++) {
            buffer.put(i + pos + 1, tail[i]);
        }
        return buffer;
    }
 

@Override
public DataBuffer compress(DataBuffer buffer) {
    INDArray temp = Nd4j.createArrayFromShapeBuffer(buffer, Nd4j.getShapeInfoProvider().createShapeInformation(new int[]{1, (int) buffer.length()}).getFirst());
    double max = temp.amaxNumber().doubleValue();

    int cntAbs = temp.scan(Conditions.absGreaterThanOrEqual(max - (max * threshold))).intValue();

    long originalLength = buffer.length() * Nd4j.sizeOfDataType(buffer.dataType());
    int compressedLength = cntAbs + 4;
    // first 3 elements contain header
    IntPointer pointer = new IntPointer(compressedLength);
    pointer.put(0, cntAbs);
    pointer.put(1, (int) buffer.length());
    pointer.put(2, Float.floatToIntBits(threshold)); // please note, this value will be ovewritten anyway
    pointer.put(3, 0);

    CompressionDescriptor descriptor = new CompressionDescriptor();
    descriptor.setCompressedLength(compressedLength * 4); // sizeOf(INT)
    descriptor.setOriginalLength(originalLength);
    descriptor.setOriginalElementSize(Nd4j.sizeOfDataType(buffer.dataType()));
    descriptor.setNumberOfElements(buffer.length());

    descriptor.setCompressionAlgorithm(getDescriptor());
    descriptor.setCompressionType(getCompressionType());

    CompressedDataBuffer cbuff = new CompressedDataBuffer(pointer, descriptor);

    Nd4j.getNDArrayFactory().convertDataEx(getBufferTypeEx(buffer), buffer.addressPointer(), DataBuffer.TypeEx.FTHRESHOLD, pointer, buffer.length());

    Nd4j.getAffinityManager().tagLocation(buffer, AffinityManager.Location.HOST);

    return cbuff;
}
 

@Override
public DataBuffer compress(DataBuffer buffer) {
    INDArray temp = Nd4j.createArrayFromShapeBuffer(buffer, Nd4j.getShapeInfoProvider().createShapeInformation(new int[]{1, (int) buffer.length()}).getFirst());
    MatchCondition condition = new MatchCondition(temp, Conditions.absGreaterThanOrEqual(threshold));
    int cntAbs = Nd4j.getExecutioner().exec(condition, Integer.MAX_VALUE).getInt(0);


    //log.info("density ratio: {}", String.format("%.2f", cntAbs * 100.0f / buffer.length()));

    if (cntAbs < 2)
        return null;

    long originalLength = buffer.length() * Nd4j.sizeOfDataType(buffer.dataType());
    int compressedLength = cntAbs + 4;
    // first 3 elements contain header
    IntPointer pointer = new IntPointer(compressedLength);
    pointer.put(0, cntAbs);
    pointer.put(1, (int) buffer.length());
    pointer.put(2, Float.floatToIntBits(threshold));
    pointer.put(3, 0);

    CompressionDescriptor descriptor = new CompressionDescriptor();
    descriptor.setCompressedLength(compressedLength * 4); // sizeOf(INT)
    descriptor.setOriginalLength(originalLength);
    descriptor.setOriginalElementSize(Nd4j.sizeOfDataType(buffer.dataType()));
    descriptor.setNumberOfElements(buffer.length());

    descriptor.setCompressionAlgorithm(getDescriptor());
    descriptor.setCompressionType(getCompressionType());



    CompressedDataBuffer cbuff = new CompressedDataBuffer(pointer, descriptor);

    Nd4j.getNDArrayFactory().convertDataEx(getBufferTypeEx(buffer), buffer.addressPointer(), DataBuffer.TypeEx.THRESHOLD, pointer, buffer.length());

    Nd4j.getAffinityManager().tagLocation(buffer, AffinityManager.Location.HOST);

    return cbuff;
}
 

@Override
public DataBuffer compress(DataBuffer buffer) {
    INDArray temp = Nd4j.createArrayFromShapeBuffer(buffer, Nd4j.getShapeInfoProvider().createShapeInformation(new long[]{1, buffer.length()}, buffer.dataType()).getFirst());
    MatchCondition condition = new MatchCondition(temp, Conditions.absGreaterThanOrEqual(threshold));
    int cntAbs = Nd4j.getExecutioner().exec(condition).getInt(0);


    //log.info("density ratio: {}", String.format("%.2f", cntAbs * 100.0f / buffer.length()));

    if (cntAbs < 2)
        return null;

    long originalLength = buffer.length() * Nd4j.sizeOfDataType(buffer.dataType());
    int compressedLength = cntAbs + 4;
    // first 3 elements contain header
    IntPointer pointer = new IntPointer(compressedLength);
    pointer.put(0, cntAbs);
    pointer.put(1, (int) buffer.length());
    pointer.put(2, Float.floatToIntBits(threshold));
    pointer.put(3, 0);

    CompressionDescriptor descriptor = new CompressionDescriptor();
    descriptor.setCompressedLength(compressedLength * 4); // sizeOf(INT)
    descriptor.setOriginalLength(originalLength);
    descriptor.setOriginalElementSize(Nd4j.sizeOfDataType(buffer.dataType()));
    descriptor.setNumberOfElements(buffer.length());

    descriptor.setCompressionAlgorithm(getDescriptor());
    descriptor.setCompressionType(getCompressionType());



    CompressedDataBuffer cbuff = new CompressedDataBuffer(pointer, descriptor);

    Nd4j.getNDArrayFactory().convertDataEx(getBufferTypeEx(buffer), buffer.addressPointer(), DataTypeEx.THRESHOLD, pointer, buffer.length());

    Nd4j.getAffinityManager().tagLocation(buffer, AffinityManager.Location.HOST);

    return cbuff;
}
 

@Test
public void testCholeskyU() {
    INDArray A = Nd4j.create(new double[] {2, -1, 2, -1, 2, -1, 2, -1, 2,});
    A = A.reshape('f', 3, 3);
    INDArray O = Nd4j.create(A.shape());
    Nd4j.copy(A, O);

    Nd4j.getBlasWrapper().lapack().potrf(A, false);
    A = A.transpose().mmul(A);
    O.subi(A);
    DataBuffer db = O.data();
    for (int i = 0; i < db.length(); i++) {
        assertEquals(0, db.getFloat(i), 1e-5);
    }
}
 

/**
 * Create a  compression descriptor from the given
 * data buffer elements
 * @param buffer the databuffer to create
 *               the compression descriptor from
 */
public CompressionDescriptor(DataBuffer buffer) {
    this.originalLength = buffer.length() * buffer.getElementSize();
    this.numberOfElements = buffer.length();
    this.originalElementSize = buffer.getElementSize();
    this.originalDataType = buffer.dataType();
}
 

/**
 * Create a {@link DataBuffer} from a
 * byte buffer. This is meant to be used with flatbuffers
 * @param bb the flat buffers buffer
 * @param bb_pos the position to start from
 * @param type the type of buffer to create
 * @param length the length of the buffer to create
 * @return the created databuffer
 */
public static DataBuffer createFromByteBuffer(ByteBuffer bb, int bb_pos, DataType type, int length) {
    bb.order(ByteOrder.LITTLE_ENDIAN);
    int elementSize = DataTypeUtil.lengthForDtype(type);
    DataBuffer ret = Nd4j.createBuffer(ByteBuffer.allocateDirect(length *   elementSize),type,length,0);

    switch(type) {
        case DOUBLE:
            for(int i = 0; i < ret.length(); i++) {
                double doubleGet = bb.getDouble(bb.capacity() - bb_pos + (i * elementSize));
                ret.put(i,doubleGet);
            }
            break;
        case FLOAT:
            for(int i = 0; i < ret.length(); i++) {
                float floatGet = bb.getFloat(bb.capacity() - bb_pos + (i * elementSize));
                ret.put(i,floatGet);
            }
            break;
        case INT:
            for(int i = 0; i < ret.length(); i++) {
                int intGet = bb.getInt(bb.capacity() - bb_pos  + (i * elementSize));
                ret.put(i,intGet);
            }
            break;
        case LONG:
            for(int i = 0; i < ret.length(); i++) {
                long longGet = bb.getLong(bb.capacity() - bb_pos  + (i * elementSize));
                ret.put(i,longGet);
            }
            break;
    }

    return ret;
}
 

/**
 * PLEASE NOTE: Specific implementation, on systems without special devices can return HostPointer here
 *
 * @param buffer
 * @return
 */
@Override
public org.bytedeco.javacpp.Pointer getDevicePointer(DataBuffer buffer, CudaContext context) {
    // TODO: It would be awesome to get rid of typecasting here
    //getCudaContext().syncOldStream();
    AllocationPoint dstPoint = ((BaseCudaDataBuffer) buffer).getAllocationPoint();

    //log.info("getDevicePointer called");
    /*
    if (configuration.getMemoryModel() == Configuration.MemoryModel.DELAYED && dstPoint.getAllocationStatus() == AllocationStatus.HOST) {
    
        // if we have constant buffer (aka shapeInfo or other constant stuff)
        if (buffer.isConstant()) {
            Nd4j.getConstantHandler().moveToConstantSpace(buffer);
        } else {
            PointersPair pair = memoryProvider.malloc(dstPoint.getShape(), dstPoint, AllocationStatus.DEVICE);
    
            if (pair != null) {
                Integer deviceId = getDeviceId();
    
                dstPoint.getPointers().setDevicePointer(pair.getDevicePointer());
                dstPoint.setAllocationStatus(AllocationStatus.DEVICE);
    
                deviceAllocations.get(deviceId).put(dstPoint.getObjectId(), dstPoint.getObjectId());
    
                zeroAllocations.get(dstPoint.getBucketId()).remove(dstPoint.getObjectId());
                deviceMemoryTracker.addToAllocation(Thread.currentThread().getId(), deviceId, AllocationUtils.getRequiredMemory(dstPoint.getShape()));
    
    
                dstPoint.tickHostWrite();
            }
        }
    }
    */
    // here's the place, where we do care about promotion. but we only care about promotion of original  buffers
    if (dstPoint.getAllocationStatus() == AllocationStatus.HOST && buffer.offset() == 0 && 1 < 0) {
        if (dstPoint.getDeviceTicks() > configuration.getMinimumRelocationThreshold()) {
            // at this point we know, that this request is done withing some existent context
            long requiredMemory = AllocationUtils.getRequiredMemory(dstPoint.getShape());
            if (deviceMemoryTracker.reserveAllocationIfPossible(Thread.currentThread().getId(), getDeviceId(),
                            requiredMemory) && pingDeviceForFreeMemory(getDeviceId(), requiredMemory)) {
                // so, memory is reserved
                promoteObject(buffer);
            }
        }
    }


    // if that's device state, we probably might want to update device memory state
    if (dstPoint.getAllocationStatus() == AllocationStatus.DEVICE) {
        if (!dstPoint.isActualOnDeviceSide()) {
            //                log.info("Relocating to GPU");
            relocate(AllocationStatus.HOST, AllocationStatus.DEVICE, dstPoint, dstPoint.getShape(), context);
        } else {
            //  log.info("Buffer is actual on device side: " + dstPoint.getShape());
        }
    } //else log.info("Not on [DEVICE]");


    //  we update memory use counter, to announce that it's somehow used on device
    dstPoint.tickDeviceRead();

    // return pointer with offset if needed. length is specified for constructor compatibility purposes
    CudaPointer p = new CudaPointer(dstPoint.getPointers().getDevicePointer(), buffer.length(),
                    (buffer.offset() * buffer.getElementSize()));
    switch (buffer.dataType()) {
        case DOUBLE:
            return p.asDoublePointer();
        case FLOAT:
            return p.asFloatPointer();
        case INT:
            return p.asIntPointer();
        case HALF:
            return p.asShortPointer();
        case LONG:
            return p.asLongPointer();
        default:
            return p;
    }
}
 

/**
 * Symmetric in place shuffle of an ndarray
 * along a specified set of dimensions. Each array in list should have it's own dimension at the same index of dimensions array
 *
 * @param arrays      the ndarrays to shuffle
 * @param dimensions the dimensions to do the shuffle
 * @return
 */
@Override
public void shuffle(List<INDArray> arrays, Random rnd, List<int[]> dimensions) {
    if (dimensions == null || dimensions.size() == 0)
        throw new RuntimeException("Dimension can't be null or 0-length");

    if (arrays == null || arrays.size() == 0)
        throw new RuntimeException("No input arrays provided");

    if (dimensions.size() > 1 && arrays.size() != dimensions.size())
        throw new IllegalStateException("Number of dimensions do not match number of arrays to shuffle");

    int tadLength = 1;
    for (int i = 0; i < dimensions.get(0).length; i++) {
        tadLength *= arrays.get(0).shape()[dimensions.get(0)[i]];
    }

    long numTads = arrays.get(0).length() / tadLength;

    val map = ArrayUtil.buildInterleavedVector(rnd, (int) numTads);

    PointerPointer dataPointers = new PointerPointer(arrays.size());
    PointerPointer shapePointers = new PointerPointer(arrays.size());
    PointerPointer tadPointers = new PointerPointer(arrays.size());
    PointerPointer offsetPointers = new PointerPointer(arrays.size());

    PointerPointer dummy = new PointerPointer(new Pointer[] {null});

    List<Pair<DataBuffer, DataBuffer>> list = new ArrayList<>();

    TADManager tadManager = Nd4j.getExecutioner().getTADManager();

    val ptrMap = new IntPointer(map);

    long[] ptrs = new long[arrays.size()];


    for (int i = 0; i < arrays.size(); i++) {
        INDArray array = arrays.get(i);

        Nd4j.getCompressor().autoDecompress(array);


        int[] dimension = dimensions.size() > 1 ? dimensions.get(i) : dimensions.get(0);

        Pair<DataBuffer, DataBuffer> tadBuffers = tadManager.getTADOnlyShapeInfo(array, dimension);
        list.add(tadBuffers);

        Pointer hostTadShapeInfo = tadBuffers.getFirst().addressPointer();

        DataBuffer offsets = tadBuffers.getSecond();

        if (offsets.length() != numTads)
            throw new ND4JIllegalStateException("Can't symmetrically shuffle arrays with non-equal number of TADs");

        if (offsets == null)
            throw new ND4JIllegalStateException("Offsets for shuffle can't be null");


        dataPointers.put(i, array.data().addressPointer());
        shapePointers.put(i, array.shapeInfoDataBuffer().addressPointer());
        offsetPointers.put(i, offsets.addressPointer());
        tadPointers.put(i, tadBuffers.getFirst().addressPointer());
    }

    if (Nd4j.dataType() == DataBuffer.Type.DOUBLE) {
        nativeOps.shuffleDouble(dummy, dataPointers, shapePointers, dataPointers, shapePointers, arrays.size(),
                ptrMap, tadPointers, offsetPointers);
    } else if (Nd4j.dataType() == DataBuffer.Type.FLOAT) {
        nativeOps.shuffleFloat(dummy, dataPointers, shapePointers, dataPointers, shapePointers, arrays.size(),
                ptrMap, tadPointers, offsetPointers);
    } else {
        // HALFs
    }

    dataPointers.address();
    shapePointers.address();
    tadPointers.address();
    offsetPointers.address();
}
 

/**
 * PLEASE NOTE: Specific implementation, on systems without special devices can return HostPointer here
 *
 * @param buffer
 * @return
 */
@Override
public org.bytedeco.javacpp.Pointer getDevicePointer(DataBuffer buffer, CudaContext context) {
    // TODO: It would be awesome to get rid of typecasting here
    AllocationPoint dstPoint = ((BaseCudaDataBuffer) buffer).getAllocationPoint();

    // if that's device state, we probably might want to update device memory state
    if (dstPoint.getAllocationStatus() == AllocationStatus.DEVICE) {
        if (!dstPoint.isActualOnDeviceSide()) {
            //relocate(AllocationStatus.HOST, AllocationStatus.DEVICE, dstPoint, dstPoint.getShape(), context);
            throw new UnsupportedOperationException("Pew-pew");
        }
    }

    if (dstPoint.getDevicePointer() == null)
        return null;


    // return pointer. length is specified for constructor compatibility purposes. Offset is accounted at C++ side
    val p = new CudaPointer(dstPoint.getDevicePointer(), buffer.length(), 0);

    if (OpProfiler.getInstance().getConfig().isCheckLocality())
         NativeOpsHolder.getInstance().getDeviceNativeOps().tryPointer(context.getOldStream(), p, 1);

    switch (buffer.dataType()) {
        case DOUBLE:
            return p.asDoublePointer();
        case FLOAT:
            return p.asFloatPointer();
        case UINT32:
        case INT:
            return p.asIntPointer();
        case SHORT:
        case UINT16:
        case HALF:
        case BFLOAT16:
            return p.asShortPointer();
        case UINT64:
        case LONG:
            return p.asLongPointer();
        case UTF8:
        case UBYTE:
        case BYTE:
            return p.asBytePointer();
        case BOOL:
            return p.asBooleanPointer();
        default:
            return p;
    }
}
 

@Override
public DataBuffer compress(DataBuffer buffer) {

    CompressionDescriptor descriptor = new CompressionDescriptor(buffer, this);

    BytePointer ptr = new BytePointer(buffer.length() * buffer.getElementSize());
    CompressedDataBuffer result = new CompressedDataBuffer(ptr, descriptor);

    Nd4j.getMemoryManager().memcpy(result, buffer);

    return result;
}
 

@Override
public DataBuffer compress(DataBuffer buffer) {

    CompressionDescriptor descriptor = new CompressionDescriptor(buffer, this);

    BytePointer ptr = new BytePointer(buffer.length() * buffer.getElementSize());
    CompressedDataBuffer result = new CompressedDataBuffer(ptr, descriptor);

    Nd4j.getMemoryManager().memcpy(result, buffer);

    return result;
}
 

/**
 * Return if there is enough allocated memory space to add data of a given length in the databuffer
 * @param buffer a databuffer in which we want to add data
 * @param length the length of the data
 * @return a boolean if the insertion is possible
 * */
public boolean canInsert(DataBuffer buffer, int length) {
    return buffer.capacity() - buffer.length() >= length;
}
 

/**
 * Create a  compression descriptor from the given
 * data buffer elements
 * @param buffer the databuffer to create
 *               the compression descriptor from
 */
public CompressionDescriptor(DataBuffer buffer) {
    this.originalLength = buffer.length() * buffer.getElementSize();
    this.numberOfElements = buffer.length();
    this.originalElementSize = buffer.getElementSize();
}