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

/**
 * Convert the underlying contents of this {@link Writable}
 * to an nd4j {@link DataBuffer}. Note that this is a *copy*
 * of the underlying buffer.
 * Also note that {@link java.nio.ByteBuffer#allocateDirect(int)}
 * is used for allocation.
 * This should be considered an expensive operation.
 *
 * This buffer should be cached when used. Once used, this can be
 * used in standard Nd4j operations.
 *
 * Beyond that, the reason we have to use allocateDirect
 * is due to nd4j data buffers being stored off heap (whether on cpu or gpu)
 * @param type the type of the data buffer
 * @param elementSize the size of each element in the buffer
 * @return the equivalent nd4j data buffer
 */
public DataBuffer asNd4jBuffer(DataType type, int elementSize) {
    int length = content.length / elementSize;
    DataBuffer ret = Nd4j.createBuffer(ByteBuffer.allocateDirect(content.length),type,length,0);
    for(int i = 0; i < length; i++) {
        switch(type) {
            case DOUBLE:
                ret.put(i,getDouble(i));
                break;
            case INT:
                ret.put(i,getInt(i));
                break;
            case FLOAT:
                ret.put(i,getFloat(i));
               break;
            case LONG:
                ret.put(i,getLong(i));
                break;
        }
    }
    return ret;
}
 

/**
 * Convert the underlying contents of this {@link Writable}
 * to an nd4j {@link DataBuffer}. Note that this is a *copy*
 * of the underlying buffer.
 * Also note that {@link java.nio.ByteBuffer#allocateDirect(int)}
 * is used for allocation.
 * This should be considered an expensive operation.
 *
 * This buffer should be cached when used. Once used, this can be
 * used in standard Nd4j operations.
 *
 * Beyond that, the reason we have to use allocateDirect
 * is due to nd4j data buffers being stored off heap (whether on cpu or gpu)
 * @param type the type of the data buffer
 * @param elementSize the size of each element in the buffer
 * @return the equivalent nd4j data buffer
 */
public DataBuffer asNd4jBuffer(DataBuffer.Type type,int elementSize) {
    int length = content.length / elementSize;
    DataBuffer ret = Nd4j.createBuffer(ByteBuffer.allocateDirect(content.length),type,length,0);
    for(int i = 0; i < length; i++) {
        switch(type) {
            case DOUBLE:
                ret.put(i,getDouble(i));
                break;
            case INT:
                ret.put(i,getInt(i));
                break;
            case FLOAT:
                ret.put(i,getFloat(i));
               break;
            case LONG:
                ret.put(i,getLong(i));
                break;
        }
    }
    return ret;
}
 

/**
 * 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;
}
 

@Override
public INDArray nextFloat(char order, long[] shape) {
    long length = ArrayUtil.prodLong(shape);
    INDArray ret = Nd4j.create(shape, order);

    DataBuffer data = ret.data();
    for (long i = 0; i < length; i++) {
        data.put(i, nextFloat());
    }

    return ret;
}
 

@Override
public INDArray nextDouble(char order, long[] shape) {
    long length = ArrayUtil.prodLong(shape);
    INDArray ret = Nd4j.create(shape, order);

    DataBuffer data = ret.data();
    for (long i = 0; i < length; i++) {
        data.put(i, nextDouble());
    }

    return ret;
}
 

@Override
public INDArray nextGaussian(char order, long[] shape) {
    long length = ArrayUtil.prodLong(shape);
    INDArray ret = Nd4j.create(shape, order);

    DataBuffer data = ret.data();
    for (long i = 0; i < length; i++) {
        data.put(i, nextGaussian());
    }

    return ret;
}
 

public static DataBuffer createFromByteBuffer(ByteBuffer bb,int bb_pos,DataBuffer.Type type,int length,int elementSize) {
    bb.order(ByteOrder.LITTLE_ENDIAN);
    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;
}
 

@Override
public INDArray deserialize(JsonParser jp, DeserializationContext deserializationContext) throws IOException {
    JsonNode node = jp.getCodec().readTree(jp);
    JsonNode arr = node.get("dataBuffer");
    int rank = node.get("rankField").asInt();
    int numElements = node.get("numElements").asInt();
    int offset = node.get("offsetField").asInt();
    JsonNode shape = node.get("shapeField");
    JsonNode stride = node.get("strideField");
    String type = node.get("typeField").asText();
    int[] realShape = new int[rank];
    int[] realStride = new int[rank];
    DataBuffer buff = Nd4j.createBuffer(numElements);
    for (int i = 0; i < numElements; i++) {
        buff.put(i, arr.get(i).asDouble());
    }

    String ordering = node.get("orderingField").asText();
    for (int i = 0; i < rank; i++) {
        realShape[i] = shape.get(i).asInt();
        realStride[i] = stride.get(i).asInt();
    }

    INDArray ret = type.equals("real") ? Nd4j.create(buff, realShape, realStride, offset, ordering.charAt(0))
                    : Nd4j.createComplex(buff, realShape, realStride, offset, ordering.charAt(0));
    return ret;
}
 

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 INDArray nextInt(long[] shape) {
    int length = ArrayUtil.prod(shape);
    INDArray ret = Nd4j.create(shape);

    DataBuffer data = ret.data();
    for (int i = 0; i < length; i++) {
        data.put(i, nextInt());
    }

    return ret;
}
 

@Override
public INDArray nextInt(long[] shape) {
    int length = ArrayUtil.prod(shape);
    INDArray ret = Nd4j.create(shape);

    DataBuffer data = ret.data();
    for (int i = 0; i < length; i++) {
        data.put(i, nextInt());
    }

    return ret;
}
 

@Override
public INDArray nextDouble(char order, long[] shape) {
    long length = ArrayUtil.prodLong(shape);
    INDArray ret = Nd4j.create(shape, order);

    DataBuffer data = ret.data();
    for (long i = 0; i < length; i++) {
        data.put(i, nextDouble());
    }

    return ret;
}
 

@Override
public INDArray deserialize(JsonParser jp, DeserializationContext deserializationContext) throws IOException, JsonProcessingException {
    JsonNode node = jp.getCodec().readTree(jp);
    if(node.has("dataBuffer")){
        //Must be legacy format serialization
        JsonNode arr = node.get("dataBuffer");
        int rank = node.get("rankField").asInt();
        int numElements = node.get("numElements").asInt();
        int offset = node.get("offsetField").asInt();
        JsonNode shape = node.get("shapeField");
        JsonNode stride = node.get("strideField");
        int[] shapeArr = new int[rank];
        int[] strideArr = new int[rank];
        DataBuffer buff = Nd4j.createBuffer(numElements);
        for (int i = 0; i < numElements; i++) {
            buff.put(i, arr.get(i).asDouble());
        }

        String ordering = node.get("orderingField").asText();
        for (int i = 0; i < rank; i++) {
            shapeArr[i] = shape.get(i).asInt();
            strideArr[i] = stride.get(i).asInt();
        }

        return Nd4j.create(buff, shapeArr, strideArr, offset, ordering.charAt(0));
    }
    //Standard/new format
    return new NDArrayTextDeSerializer().deserialize(node);
}
 

@Test
public void testPut() throws Exception {
    DataBuffer buffer = Nd4j.createBuffer(new float[]{1f,2f,3f,4f});
    buffer.put(2, 16f);

    assertEquals(16.0f, buffer.getFloat(2), 0.001f);

    System.out.println("Data: " + buffer);
}
 

public DataBuffer shiftLeft(DataBuffer buffer, int from, int offset, long datalength) {
    for (int i = from; i < datalength; i++) {
        buffer.put(i - offset, buffer.getDouble(i));
    }
    return buffer;
}
 

@Override
public INDArray thresholdEncode(INDArray input, double threshold, Integer boundary) {

    MatchCondition condition = new MatchCondition(input, Conditions.absGreaterThanOrEqual(threshold));
    int cntAbs = Nd4j.getExecutioner().exec(condition, Integer.MAX_VALUE).getInt(0);

    if (cntAbs < 2)
        return null;

    if (boundary != null)
        cntAbs = Math.min(cntAbs, boundary);

    DataBuffer buffer = input.data();

    long originalLength = buffer.length() * Nd4j.sizeOfDataType(buffer.dataType());
    int compressedLength = cntAbs + 4;
    // first 3 elements contain header

    DataBuffer encodedBuffer = Nd4j.getMemoryManager().getCurrentWorkspace() == null ? Nd4j.getDataBufferFactory().createInt(4+cntAbs, false) : Nd4j.getDataBufferFactory().createInt(4+cntAbs, false, Nd4j.getMemoryManager().getCurrentWorkspace());

    encodedBuffer.put(0, cntAbs);
    encodedBuffer.put(1, (int) buffer.length());
    encodedBuffer.put(2, Float.floatToIntBits((float) threshold));

    // format id
    encodedBuffer.put(3, ThresholdCompression.FLEXIBLE_ENCODING);

    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("THRESHOLD");
    descriptor.setCompressionType(CompressionType.LOSSLESS);

    //CompressedDataBuffer cbuff = new CompressedDataBuffer(pointer, descriptor);

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

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

    return Nd4j.createArrayFromShapeBuffer(encodedBuffer, input.shapeInfoDataBuffer());
}
 

@Override
public Pair<DataBuffer, DataBuffer> getTADOnlyShapeInfo(INDArray array, int[] dimension) {
    if (dimension != null && dimension.length > 1)
        Arrays.sort(dimension);

    if (dimension == null)
        dimension = new int[] {Integer.MAX_VALUE};

    boolean isScalar = dimension == null || (dimension.length == 1 && dimension[0] == Integer.MAX_VALUE);

    // FIXME: this is fast triage, remove it later
    int targetRank = isScalar ? 2 : array.rank(); //dimensionLength <= 1 ? 2 : dimensionLength;
    long offsetLength = 0;
    long tadLength = 1;

    if(!isScalar)
        for (int i = 0; i < dimension.length; i++) {
            tadLength *= array.shape()[dimension[i]];
        }

    if(!isScalar)
        offsetLength = array.lengthLong() / tadLength;
    else
        offsetLength = 1;
    //     logger.info("Original shape info before TAD: {}", array.shapeInfoDataBuffer());
    //    logger.info("dimension: {}, tadLength: {}, offsetLength for TAD: {}", Arrays.toString(dimension),tadLength, offsetLength);

    DataBuffer outputBuffer = new CudaLongDataBuffer(targetRank * 2 + 4);
    DataBuffer offsetsBuffer = new CudaLongDataBuffer(offsetLength);

    AtomicAllocator.getInstance().getAllocationPoint(outputBuffer).tickHostWrite();
    AtomicAllocator.getInstance().getAllocationPoint(offsetsBuffer).tickHostWrite();

    DataBuffer dimensionBuffer = AtomicAllocator.getInstance().getConstantBuffer(dimension);
    Pointer dimensionPointer = AtomicAllocator.getInstance().getHostPointer(dimensionBuffer);

    Pointer xShapeInfo = AddressRetriever.retrieveHostPointer(array.shapeInfoDataBuffer());
    Pointer targetPointer = AddressRetriever.retrieveHostPointer(outputBuffer);
    Pointer offsetsPointer = AddressRetriever.retrieveHostPointer(offsetsBuffer);
    if(!isScalar)
        nativeOps.tadOnlyShapeInfo((LongPointer) xShapeInfo, (IntPointer) dimensionPointer, dimension.length,
                (LongPointer) targetPointer, new LongPointerWrapper(offsetsPointer));

    else  {
        outputBuffer.put(0,2);
        outputBuffer.put(1,1);
        outputBuffer.put(2,1);
        outputBuffer.put(3,1);
        outputBuffer.put(4,1);
        outputBuffer.put(5,0);
        outputBuffer.put(6,0);
        outputBuffer.put(7,99);

    }

    AtomicAllocator.getInstance().getAllocationPoint(outputBuffer).tickHostWrite();
    AtomicAllocator.getInstance().getAllocationPoint(offsetsBuffer).tickHostWrite();

    //   logger.info("TAD shapeInfo after construction: {}", Arrays.toString(TadDescriptor.dataBufferToArray(outputBuffer)));
    // now we need to copy this buffer to either device global memory or device cache

    return new Pair<>(outputBuffer, offsetsBuffer);

}
 

@Override
public long bitmapEncode(INDArray indArray, INDArray target, double threshold) {
    long length = indArray.lengthLong();
    long tLen = target.data().length();

    if (tLen != (length / 16 + 5))
        throw new ND4JIllegalStateException("Length of target array should be " + (length / 16 + 5));

    if (target.data().dataType() != DataBuffer.Type.INT)
        throw new ND4JIllegalStateException("Target array should have INT dataType");

    DataBuffer buffer = target.data();
    buffer.put(0, (int) length);
    buffer.put(1, (int) length);
    buffer.put(2, Float.floatToIntBits((float) threshold));

    // format id
    buffer.put(3, ThresholdCompression.BITMAP_ENCODING);

    CudaContext context = AtomicAllocator.getInstance().getFlowController().prepareAction(indArray);

    if (extraz.get() == null)
        extraz.set(new PointerPointer(32));


    PointerPointer extras = extraz.get().put(
            AtomicAllocator.getInstance().getHostPointer(indArray),
            context.getOldStream(),
            context.getBufferScalar(),
            context.getBufferReduction()
    );

    long val = 0;
    if (indArray.data().dataType() == DataBuffer.Type.FLOAT) {
        val = nativeOps.encodeBitmapFloat(extras,
                (FloatPointer) AtomicAllocator.getInstance().getPointer(indArray, context),
                length,
                (IntPointer) AtomicAllocator.getInstance().getPointer(buffer, context),
                (float) threshold
        );
    } else if (indArray.data().dataType() == DataBuffer.Type.DOUBLE) {
        val = nativeOps.encodeBitmapDouble(extras,
                (DoublePointer) AtomicAllocator.getInstance().getPointer(indArray, context),
                length,
                (IntPointer) AtomicAllocator.getInstance().getPointer(buffer, context),
                (float) threshold
        );
    } else if (indArray.data().dataType() == DataBuffer.Type.HALF) {
        val = nativeOps.encodeBitmapHalf(extras,
                (ShortPointer) AtomicAllocator.getInstance().getPointer(indArray, context),
                length,
                (IntPointer) AtomicAllocator.getInstance().getPointer(buffer, context),
                (float) threshold
        );
    } else
        throw new ND4JIllegalStateException("Unknown dataType " + indArray.data().dataType());


    AtomicAllocator.getInstance().getFlowController().registerAction(context, indArray);

    AtomicAllocator.getInstance().getAllocationPoint(buffer).tickDeviceWrite();

    return val;
}
 

/**
 * Get the element wise stride for the
 * shape info buffer
 * @param buffer the buffer to get the element
 *               wise stride from
 * @return the element wise stride for the buffer
 */
public static void setElementWiseStride(DataBuffer buffer, int elementWiseStride) {
    int length2 = shapeInfoLength(Shape.rank(buffer));
    //if (1 > 0) throw new RuntimeException("setElementWiseStride called: [" + elementWiseStride + "], buffer: " + buffer);
    buffer.put(length2 - 2, elementWiseStride);
}
 

/**
 * Get the element wise stride for the
 * shape info buffer
 * @param buffer the buffer to get the element
 *               wise stride from
 * @return the element wise stride for the buffer
 */
public static void setElementWiseStride(DataBuffer buffer, int elementWiseStride) {
    int length2 = shapeInfoLength(Shape.rank(buffer));
    //if (1 > 0) throw new RuntimeException("setElementWiseStride called: [" + elementWiseStride + "], buffer: " + buffer);
    buffer.put(length2 - 2, elementWiseStride);
}