nd4j-backends/nd4j-backend-impls/nd4j-cuda/src/main/java/org/nd4j/linalg/jcublas/ops/executioner/CudaExecutioner.java
public void exec(CustomOp op) {

        Nd4j.getExecutioner().commit();

        if (op.opName().equalsIgnoreCase("im2col")) {
            val dtype = Nd4j.dataType();

            val xArr = op.inputArguments()[0];
            val zArr = op.outputArguments()[0];

            CudaContext context = AtomicAllocator.getInstance().getFlowController().prepareAction(zArr, xArr);

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

            PointerPointer xShapeHost =
                    extraz.get().put(AddressRetriever.retrieveHostPointer(xArr.shapeInfoDataBuffer()), // 0
                            context.getOldStream(), // 1
                            AtomicAllocator.getInstance().getDeviceIdPointer(), // 2
                            context.getBufferAllocation(), // 3
                            context.getBufferReduction(), // 4
                            context.getBufferScalar(), // 5
                            context.getBufferSpecial(),
                            null,
                            AddressRetriever.retrieveHostPointer(zArr.shapeInfoDataBuffer())
                    );


            val x = AtomicAllocator.getInstance().getPointer(xArr, context);
            val z = AtomicAllocator.getInstance().getPointer(zArr, context);

            val xShape = AtomicAllocator.getInstance().getPointer(xArr.shapeInfoDataBuffer(), context);
            val zShape = AtomicAllocator.getInstance().getPointer(zArr.shapeInfoDataBuffer(), context);

            double zeroPad = 0.0;
            if(op.tArgs() != null && op.tArgs().length > 0){
                zeroPad = op.tArgs()[0];
            }
            val extrass = new double[]{op.iArgs()[0], op.iArgs()[1], op.iArgs()[2], op.iArgs()[3], op.iArgs()[4], op.iArgs()[5], op.iArgs()[6], op.iArgs()[7], op.iArgs()[8], zeroPad};
            val extraArgsBuff = Nd4j.getConstantHandler().getConstantBuffer(extrass);
            val extraArgs = AtomicAllocator.getInstance().getPointer(extraArgsBuff, context);


            if (dtype == DataBuffer.Type.DOUBLE) {
                nativeOps.execTransformDouble(xShapeHost, 37, (DoublePointer) x, (LongPointer) xShape, (DoublePointer) z, (LongPointer) zShape, (DoublePointer) extraArgs);
            } else if (dtype == DataBuffer.Type.FLOAT) {
                nativeOps.execTransformFloat(xShapeHost, 37, (FloatPointer) x, (LongPointer) xShape, (FloatPointer) z, (LongPointer) zShape, (FloatPointer) extraArgs);
            } else if (dtype == DataBuffer.Type.HALF) {
                nativeOps.execTransformHalf(xShapeHost, 37, (ShortPointer) x, (LongPointer) xShape, (ShortPointer) z, (LongPointer) zShape, (ShortPointer) extraArgs);
            }

            //AtomicAllocator.getInstance().getAllocationPoint(zArr).tickDeviceWrite();
            AtomicAllocator.getInstance().getFlowController().registerAction(context, zArr, xArr);

            //Nd4j.getExecutioner().commit();

            return;
        } else if (op.opName().equalsIgnoreCase("col2im")) {
            val dtype = Nd4j.dataType();

            val xArr = op.inputArguments()[0];
            val zArr = op.outputArguments()[0];

            CudaContext context = AtomicAllocator.getInstance().getFlowController().prepareAction(zArr, xArr);

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

            PointerPointer xShapeHost =
                    extraz.get().put(AddressRetriever.retrieveHostPointer(xArr.shapeInfoDataBuffer()), // 0
                            context.getOldStream(), // 1
                            AtomicAllocator.getInstance().getDeviceIdPointer(), // 2
                            context.getBufferAllocation(), // 3
                            context.getBufferReduction(), // 4
                            context.getBufferScalar(), // 5
                            context.getBufferSpecial(),
                            null,
                            AddressRetriever.retrieveHostPointer(zArr.shapeInfoDataBuffer())
                    );


            val x = AtomicAllocator.getInstance().getPointer(xArr, context);
            val z = AtomicAllocator.getInstance().getPointer(zArr, context);

            val xShape = AtomicAllocator.getInstance().getPointer(xArr.shapeInfoDataBuffer(), context);
            val zShape = AtomicAllocator.getInstance().getPointer(zArr.shapeInfoDataBuffer(), context);

            val extrass = new double[]{op.iArgs()[0], op.iArgs()[1], op.iArgs()[2], op.iArgs()[3], op.iArgs()[4], op.iArgs()[5], op.iArgs()[6], op.iArgs()[7]};
            val extraArgsBuff = Nd4j.getConstantHandler().getConstantBuffer(extrass);
            val extraArgs = AtomicAllocator.getInstance().getPointer(extraArgsBuff, context);

            if (dtype == DataBuffer.Type.DOUBLE) {
                nativeOps.execTransformDouble(xShapeHost, 36, (DoublePointer) x, (LongPointer) xShape, (DoublePointer) z, (LongPointer) zShape, (DoublePointer) extraArgs);
            } else if (dtype == DataBuffer.Type.FLOAT) {
                nativeOps.execTransformFloat(xShapeHost, 36, (FloatPointer) x, (LongPointer) xShape, (FloatPointer) z, (LongPointer) zShape, (FloatPointer) extraArgs);
            } else if (dtype == DataBuffer.Type.HALF) {
                nativeOps.execTransformHalf(xShapeHost, 36, (ShortPointer) x, (LongPointer) xShape, (ShortPointer) z, (LongPointer) zShape, (ShortPointer) extraArgs);
            }

            //AtomicAllocator.getInstance().getAllocationPoint(zArr).tickDeviceWrite();
            AtomicAllocator.getInstance().getFlowController().registerAction(context, zArr, xArr);

            //Nd4j.getExecutioner().commit();

            return;
        } else if (op.opName().equalsIgnoreCase("pooling2d")) {
            val dtype = Nd4j.dataType();

            val xArr = op.inputArguments()[0];
            val zArr = op.outputArguments()[0];

            CudaContext context = AtomicAllocator.getInstance().getFlowController().prepareAction(zArr, xArr);

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

            PointerPointer xShapeHost =
                    extraz.get().put(AddressRetriever.retrieveHostPointer(xArr.shapeInfoDataBuffer()), // 0
                            context.getOldStream(), // 1
                            AtomicAllocator.getInstance().getDeviceIdPointer(), // 2
                            context.getBufferAllocation(), // 3
                            context.getBufferReduction(), // 4
                            context.getBufferScalar(), // 5
                            context.getBufferSpecial(),
                            null,
                            AddressRetriever.retrieveHostPointer(zArr.shapeInfoDataBuffer())
                    );


            val x = AtomicAllocator.getInstance().getPointer(xArr, context);
            val z = AtomicAllocator.getInstance().getPointer(zArr, context);

            val xShape = AtomicAllocator.getInstance().getPointer(xArr.shapeInfoDataBuffer(), context);
            val zShape = AtomicAllocator.getInstance().getPointer(zArr.shapeInfoDataBuffer(), context);

            val extrass = new double[]{op.iArgs()[0], op.iArgs()[1], op.iArgs()[2], op.iArgs()[3], op.iArgs()[4], op.iArgs()[5], op.iArgs()[6], op.iArgs()[7], op.iArgs()[8]};
            val extraArgsBuff = Nd4j.getConstantHandler().getConstantBuffer(extrass);
            val extraArgs = AtomicAllocator.getInstance().getPointer(extraArgsBuff, context);

            if (dtype == DataBuffer.Type.DOUBLE) {
                nativeOps.execTransformDouble(xShapeHost, 71, (DoublePointer) x, (LongPointer) xShape, (DoublePointer) z, (LongPointer) zShape, (DoublePointer) extraArgs);
            } else if (dtype == DataBuffer.Type.FLOAT) {
                nativeOps.execTransformFloat(xShapeHost, 71, (FloatPointer) x, (LongPointer) xShape, (FloatPointer) z, (LongPointer) zShape, (FloatPointer) extraArgs);
            } else if (dtype == DataBuffer.Type.HALF) {
                nativeOps.execTransformHalf(xShapeHost, 71, (ShortPointer) x, (LongPointer) xShape, (ShortPointer) z, (LongPointer) zShape, (ShortPointer) extraArgs);
            }

            // AtomicAllocator.getInstance().getAllocationPoint(zArr).tickDeviceWrite();
            AtomicAllocator.getInstance().getFlowController().registerAction(context, zArr, xArr);

            //Nd4j.getExecutioner().commit();

            return;
        }

        Nd4j.getExecutioner().commit();

        long st = profilingHookIn(op);

        CudaContext context =(CudaContext) AtomicAllocator.getInstance().getDeviceContext().getContext();
        //AtomicAllocator.getInstance().getFlowController().prepareActionAllWrite(op.outputArguments());

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


        PointerPointer extras = extraz.get().put(
                new CudaPointer(1),
                context.getOldStream(),
                context.getBufferScalar(),
                context.getBufferReduction());

        val outputArgs = op.outputArguments();
        val inputArgs = op.inputArguments();

        if (outputArgs.length == 0 && !op.isInplaceCall())
            throw new ND4JIllegalStateException("You can't execute non-inplace CustomOp without outputs being specified");

        val lc = op.opName().toLowerCase();
        val hash = op.opHash();


        val inputShapes = new PointerPointer<>(inputArgs.length * 2);
        val inputBuffers = new PointerPointer<>(inputArgs.length * 2);

        int cnt= 0;
        for (val in: inputArgs) {
            val hp = AtomicAllocator.getInstance().getHostPointer(in.shapeInfoDataBuffer());
            inputBuffers.put(cnt,  AtomicAllocator.getInstance().getHostPointer(in));
            inputShapes.put(cnt, hp);


            val dp = AtomicAllocator.getInstance().getPointer(in.shapeInfoDataBuffer(), context);

            inputBuffers.put(cnt + inputArgs.length, AtomicAllocator.getInstance().getPointer(in, context));
            inputShapes.put(cnt+ inputArgs.length, dp);

            if (op.isInplaceCall())
                AtomicAllocator.getInstance().getAllocationPoint(in).tickHostWrite();

            cnt++;
        }


        val outputShapes = new PointerPointer<>(outputArgs.length * 2);
        val outputBuffers = new PointerPointer<>(outputArgs.length * 2);

        cnt= 0;
        for (val out: outputArgs) {
            outputBuffers.put(cnt,  AtomicAllocator.getInstance().getHostPointer(out));
            outputShapes.put(cnt,  AtomicAllocator.getInstance().getHostPointer(out.shapeInfoDataBuffer()));

            outputBuffers.put(cnt + outputArgs.length,  AtomicAllocator.getInstance().getPointer(out, context));
            outputShapes.put(cnt + outputArgs.length,  AtomicAllocator.getInstance().getPointer(out.shapeInfoDataBuffer(), context));

            AtomicAllocator.getInstance().getAllocationPoint(out).tickHostWrite();

            cnt++;
        }

        if (Nd4j.dataType() == DataBuffer.Type.FLOAT) {
            val tArgs = op.tArgs().length > 0 ? new FloatPointer(op.tArgs().length) : null;
            val iArgs = op.iArgs().length > 0 ? new LongPointer(op.iArgs().length) : null;

            cnt = 0;
            for (val t: op.tArgs())
                tArgs.put(cnt++, (float) t);

            cnt = 0;
            for (val i: op.iArgs())
                iArgs.put(cnt++, i);

            val status = OpStatus.byNumber(nativeOps.execCustomOpFloat(extras, hash, inputBuffers, inputShapes, inputArgs.length, outputBuffers, outputShapes, outputArgs.length, tArgs, op.tArgs().length, iArgs, op.iArgs().length, op.isInplaceCall()));
            if (status != OpStatus.ND4J_STATUS_OK)
                throw new ND4JIllegalStateException("Op execution failed: " + status);
        } else if (Nd4j.dataType() == DataBuffer.Type.DOUBLE) {
            val tArgs = op.tArgs().length > 0 ? new DoublePointer(op.tArgs().length) : null;
            val iArgs = op.iArgs().length > 0 ? new LongPointer(op.iArgs().length) : null;

            cnt = 0;
            for (val t: op.tArgs())
                tArgs.put(cnt++, t);

            for (val i: op.iArgs())
                iArgs.put(cnt++, i);

            val status = OpStatus.byNumber(nativeOps.execCustomOpDouble(extras, hash, inputBuffers, inputShapes, inputArgs.length, outputBuffers, outputShapes, outputArgs.length, tArgs, op.tArgs().length, iArgs, op.iArgs().length, op.isInplaceCall()));
            if (status != OpStatus.ND4J_STATUS_OK)
                throw new ND4JIllegalStateException("Op execution failed: " + status);
        } else if (Nd4j.dataType() == DataBuffer.Type.HALF) {
            val tArgs = op.tArgs().length > 0 ? new ShortPointer(op.tArgs().length) : null;
            val iArgs = op.iArgs().length > 0 ? new LongPointer(op.iArgs().length) : null;

            cnt = 0;
            for (val t: op.tArgs())
                tArgs.put(cnt++, ArrayUtil.toHalf((float) t));

            cnt = 0;
            for (val i: op.iArgs())
                iArgs.put(cnt++, i);

            val status = OpStatus.byNumber(nativeOps.execCustomOpHalf(extras, hash, inputBuffers, inputShapes, inputArgs.length, outputBuffers, outputShapes, outputArgs.length, tArgs, op.tArgs().length, iArgs, op.iArgs().length, op.isInplaceCall()));
            if (status != OpStatus.ND4J_STATUS_OK)
                throw new ND4JIllegalStateException("Op execution failed: " + status);
        }

        //AtomicAllocator.getInstance().getFlowController().prepareActionAllWrite(op.outputArguments());

        profilingHookOut(op, st);
    }
public void exec(CustomOp op) {

        Nd4j.getExecutioner().commit();

        if (op.opName().equalsIgnoreCase("im2col")) {
            val dtype = Nd4j.dataType();

            val xArr = op.inputArguments()[0];
            val zArr = op.outputArguments()[0];

            CudaContext context = AtomicAllocator.getInstance().getFlowController().prepareAction(zArr, xArr);

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

            PointerPointer xShapeHost =
                    extraz.get().put(AddressRetriever.retrieveHostPointer(xArr.shapeInfoDataBuffer()), // 0
                            context.getOldStream(), // 1
                            AtomicAllocator.getInstance().getDeviceIdPointer(), // 2
                            context.getBufferAllocation(), // 3
                            context.getBufferReduction(), // 4
                            context.getBufferScalar(), // 5
                            context.getBufferSpecial(),
                            null,
                            AddressRetriever.retrieveHostPointer(zArr.shapeInfoDataBuffer())
                    );


            val x = AtomicAllocator.getInstance().getPointer(xArr, context);
            val z = AtomicAllocator.getInstance().getPointer(zArr, context);

            val xShape = AtomicAllocator.getInstance().getPointer(xArr.shapeInfoDataBuffer(), context);
            val zShape = AtomicAllocator.getInstance().getPointer(zArr.shapeInfoDataBuffer(), context);

            double zeroPad = 0.0;
            if(op.tArgs() != null && op.tArgs().length > 0){
                zeroPad = op.tArgs()[0];
            }
            val extrass = new double[]{op.iArgs()[0], op.iArgs()[1], op.iArgs()[2], op.iArgs()[3], op.iArgs()[4], op.iArgs()[5], op.iArgs()[6], op.iArgs()[7], op.iArgs()[8], zeroPad};
            val extraArgsBuff = Nd4j.getConstantHandler().getConstantBuffer(extrass);
            val extraArgs = AtomicAllocator.getInstance().getPointer(extraArgsBuff, context);


            if (dtype == DataBuffer.Type.DOUBLE) {
                nativeOps.execTransformDouble(xShapeHost, 37, (DoublePointer) x, (LongPointer) xShape, (DoublePointer) z, (LongPointer) zShape, (DoublePointer) extraArgs);
            } else if (dtype == DataBuffer.Type.FLOAT) {
                nativeOps.execTransformFloat(xShapeHost, 37, (FloatPointer) x, (LongPointer) xShape, (FloatPointer) z, (LongPointer) zShape, (FloatPointer) extraArgs);
            } else if (dtype == DataBuffer.Type.HALF) {
                nativeOps.execTransformHalf(xShapeHost, 37, (ShortPointer) x, (LongPointer) xShape, (ShortPointer) z, (LongPointer) zShape, (ShortPointer) extraArgs);
            }

            //AtomicAllocator.getInstance().getAllocationPoint(zArr).tickDeviceWrite();
            AtomicAllocator.getInstance().getFlowController().registerAction(context, zArr, xArr);

            //Nd4j.getExecutioner().commit();

            return;
        } else if (op.opName().equalsIgnoreCase("col2im")) {
            val dtype = Nd4j.dataType();

            val xArr = op.inputArguments()[0];
            val zArr = op.outputArguments()[0];

            CudaContext context = AtomicAllocator.getInstance().getFlowController().prepareAction(zArr, xArr);

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

            PointerPointer xShapeHost =
                    extraz.get().put(AddressRetriever.retrieveHostPointer(xArr.shapeInfoDataBuffer()), // 0
                            context.getOldStream(), // 1
                            AtomicAllocator.getInstance().getDeviceIdPointer(), // 2
                            context.getBufferAllocation(), // 3
                            context.getBufferReduction(), // 4
                            context.getBufferScalar(), // 5
                            context.getBufferSpecial(),
                            null,
                            AddressRetriever.retrieveHostPointer(zArr.shapeInfoDataBuffer())
                    );


            val x = AtomicAllocator.getInstance().getPointer(xArr, context);
            val z = AtomicAllocator.getInstance().getPointer(zArr, context);

            val xShape = AtomicAllocator.getInstance().getPointer(xArr.shapeInfoDataBuffer(), context);
            val zShape = AtomicAllocator.getInstance().getPointer(zArr.shapeInfoDataBuffer(), context);

            val extrass = new double[]{op.iArgs()[0], op.iArgs()[1], op.iArgs()[2], op.iArgs()[3], op.iArgs()[4], op.iArgs()[5], op.iArgs()[6], op.iArgs()[7]};
            val extraArgsBuff = Nd4j.getConstantHandler().getConstantBuffer(extrass);
            val extraArgs = AtomicAllocator.getInstance().getPointer(extraArgsBuff, context);

            if (dtype == DataBuffer.Type.DOUBLE) {
                nativeOps.execTransformDouble(xShapeHost, 36, (DoublePointer) x, (LongPointer) xShape, (DoublePointer) z, (LongPointer) zShape, (DoublePointer) extraArgs);
            } else if (dtype == DataBuffer.Type.FLOAT) {
                nativeOps.execTransformFloat(xShapeHost, 36, (FloatPointer) x, (LongPointer) xShape, (FloatPointer) z, (LongPointer) zShape, (FloatPointer) extraArgs);
            } else if (dtype == DataBuffer.Type.HALF) {
                nativeOps.execTransformHalf(xShapeHost, 36, (ShortPointer) x, (LongPointer) xShape, (ShortPointer) z, (LongPointer) zShape, (ShortPointer) extraArgs);
            }

            //AtomicAllocator.getInstance().getAllocationPoint(zArr).tickDeviceWrite();
            AtomicAllocator.getInstance().getFlowController().registerAction(context, zArr, xArr);

            //Nd4j.getExecutioner().commit();

            return;
        } else if (op.opName().equalsIgnoreCase("pooling2d")) {
            val dtype = Nd4j.dataType();

            val xArr = op.inputArguments()[0];
            val zArr = op.outputArguments()[0];

            CudaContext context = AtomicAllocator.getInstance().getFlowController().prepareAction(zArr, xArr);

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

            PointerPointer xShapeHost =
                    extraz.get().put(AddressRetriever.retrieveHostPointer(xArr.shapeInfoDataBuffer()), // 0
                            context.getOldStream(), // 1
                            AtomicAllocator.getInstance().getDeviceIdPointer(), // 2
                            context.getBufferAllocation(), // 3
                            context.getBufferReduction(), // 4
                            context.getBufferScalar(), // 5
                            context.getBufferSpecial(),
                            null,
                            AddressRetriever.retrieveHostPointer(zArr.shapeInfoDataBuffer())
                    );


            val x = AtomicAllocator.getInstance().getPointer(xArr, context);
            val z = AtomicAllocator.getInstance().getPointer(zArr, context);

            val xShape = AtomicAllocator.getInstance().getPointer(xArr.shapeInfoDataBuffer(), context);
            val zShape = AtomicAllocator.getInstance().getPointer(zArr.shapeInfoDataBuffer(), context);

            val extrass = new double[]{op.iArgs()[0], op.iArgs()[1], op.iArgs()[2], op.iArgs()[3], op.iArgs()[4], op.iArgs()[5], op.iArgs()[6], op.iArgs()[7], op.iArgs()[8]};
            val extraArgsBuff = Nd4j.getConstantHandler().getConstantBuffer(extrass);
            val extraArgs = AtomicAllocator.getInstance().getPointer(extraArgsBuff, context);

            if (dtype == DataBuffer.Type.DOUBLE) {
                nativeOps.execTransformDouble(xShapeHost, 71, (DoublePointer) x, (LongPointer) xShape, (DoublePointer) z, (LongPointer) zShape, (DoublePointer) extraArgs);
            } else if (dtype == DataBuffer.Type.FLOAT) {
                nativeOps.execTransformFloat(xShapeHost, 71, (FloatPointer) x, (LongPointer) xShape, (FloatPointer) z, (LongPointer) zShape, (FloatPointer) extraArgs);
            } else if (dtype == DataBuffer.Type.HALF) {
                nativeOps.execTransformHalf(xShapeHost, 71, (ShortPointer) x, (LongPointer) xShape, (ShortPointer) z, (LongPointer) zShape, (ShortPointer) extraArgs);
            }

            // AtomicAllocator.getInstance().getAllocationPoint(zArr).tickDeviceWrite();
            AtomicAllocator.getInstance().getFlowController().registerAction(context, zArr, xArr);

            //Nd4j.getExecutioner().commit();

            return;
        }

        Nd4j.getExecutioner().commit();

        long st = profilingHookIn(op);

        CudaContext context =(CudaContext) AtomicAllocator.getInstance().getDeviceContext().getContext();
        //AtomicAllocator.getInstance().getFlowController().prepareActionAllWrite(op.outputArguments());

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


        PointerPointer extras = extraz.get().put(
                new CudaPointer(1),
                context.getOldStream(),
                context.getBufferScalar(),
                context.getBufferReduction());

        val outputArgs = op.outputArguments();
        val inputArgs = op.inputArguments();

        if (outputArgs.length == 0 && !op.isInplaceCall())
            throw new ND4JIllegalStateException("You can't execute non-inplace CustomOp without outputs being specified");

        val lc = op.opName().toLowerCase();
        val hash = op.opHash();


        val inputShapes = new PointerPointer<>(inputArgs.length * 2);
        val inputBuffers = new PointerPointer<>(inputArgs.length * 2);

        int cnt= 0;
        for (val in: inputArgs) {
            val hp = AtomicAllocator.getInstance().getHostPointer(in.shapeInfoDataBuffer());
            inputBuffers.put(cnt,  AtomicAllocator.getInstance().getHostPointer(in));
            inputShapes.put(cnt, hp);


            val dp = AtomicAllocator.getInstance().getPointer(in.shapeInfoDataBuffer(), context);

            inputBuffers.put(cnt + inputArgs.length, AtomicAllocator.getInstance().getPointer(in, context));
            inputShapes.put(cnt+ inputArgs.length, dp);

            if (op.isInplaceCall())
                AtomicAllocator.getInstance().getAllocationPoint(in).tickHostWrite();

            cnt++;
        }


        val outputShapes = new PointerPointer<>(outputArgs.length * 2);
        val outputBuffers = new PointerPointer<>(outputArgs.length * 2);

        cnt= 0;
        for (val out: outputArgs) {
            outputBuffers.put(cnt,  AtomicAllocator.getInstance().getHostPointer(out));
            outputShapes.put(cnt,  AtomicAllocator.getInstance().getHostPointer(out.shapeInfoDataBuffer()));

            outputBuffers.put(cnt + outputArgs.length,  AtomicAllocator.getInstance().getPointer(out, context));
            outputShapes.put(cnt + outputArgs.length,  AtomicAllocator.getInstance().getPointer(out.shapeInfoDataBuffer(), context));

            AtomicAllocator.getInstance().getAllocationPoint(out).tickHostWrite();

            cnt++;
        }

        if (Nd4j.dataType() == DataBuffer.Type.FLOAT) {
            val tArgs = op.tArgs().length > 0 ? new FloatPointer(op.tArgs().length) : null;
            val iArgs = op.iArgs().length > 0 ? new LongPointer(op.iArgs().length) : null;

            cnt = 0;
            for (val t: op.tArgs())
                tArgs.put(cnt++, (float) t);

            cnt = 0;
            for (val i: op.iArgs())
                iArgs.put(cnt++, i);

            val status = OpStatus.byNumber(nativeOps.execCustomOpFloat(extras, hash, inputBuffers, inputShapes, inputArgs.length, outputBuffers, outputShapes, outputArgs.length, tArgs, op.tArgs().length, iArgs, op.iArgs().length, op.isInplaceCall()));
            if (status != OpStatus.ND4J_STATUS_OK)
                throw new ND4JIllegalStateException("Op execution failed: " + status);
        } else if (Nd4j.dataType() == DataBuffer.Type.DOUBLE) {
            val tArgs = op.tArgs().length > 0 ? new DoublePointer(op.tArgs().length) : null;
            val iArgs = op.iArgs().length > 0 ? new LongPointer(op.iArgs().length) : null;

            cnt = 0;
            for (val t: op.tArgs())
                tArgs.put(cnt++, t);

            for (val i: op.iArgs())
                iArgs.put(cnt++, i);

            val status = OpStatus.byNumber(nativeOps.execCustomOpDouble(extras, hash, inputBuffers, inputShapes, inputArgs.length, outputBuffers, outputShapes, outputArgs.length, tArgs, op.tArgs().length, iArgs, op.iArgs().length, op.isInplaceCall()));
            if (status != OpStatus.ND4J_STATUS_OK)
                throw new ND4JIllegalStateException("Op execution failed: " + status);
        } else if (Nd4j.dataType() == DataBuffer.Type.HALF) {
            val tArgs = op.tArgs().length > 0 ? new ShortPointer(op.tArgs().length) : null;
            val iArgs = op.iArgs().length > 0 ? new LongPointer(op.iArgs().length) : null;

            cnt = 0;
            for (val t: op.tArgs())
                tArgs.put(cnt++, ArrayUtil.toHalf((float) t));

            cnt = 0;
            for (val i: op.iArgs())
                iArgs.put(cnt++, i);

            val status = OpStatus.byNumber(nativeOps.execCustomOpHalf(extras, hash, inputBuffers, inputShapes, inputArgs.length, outputBuffers, outputShapes, outputArgs.length, tArgs, op.tArgs().length, iArgs, op.iArgs().length, op.isInplaceCall()));
            if (status != OpStatus.ND4J_STATUS_OK)
                throw new ND4JIllegalStateException("Op execution failed: " + status);
        }

        //AtomicAllocator.getInstance().getFlowController().prepareActionAllWrite(op.outputArguments());

        profilingHookOut(op, st);
    }
public void exec ( CustomOp op ) { Nd4j . getExecutioner ( ) . commit ( ) ; if ( op . opName ( ) . equalsIgnoreCase ( "im2col" ) ) { val dtype = Nd4j . dataType ( ) ; val xArr = op . inputArguments ( ) [ 0 ] ; val zArr = op . outputArguments ( ) [ 0 ] ; CudaContext context = AtomicAllocator . getInstance ( ) . getFlowController ( ) . prepareAction ( zArr , xArr ) ; if ( extraz . get ( ) == null ) extraz . set ( new PointerPointer ( 32 ) ) ; PointerPointer xShapeHost = extraz . get ( ) . put ( AddressRetriever . retrieveHostPointer ( xArr . shapeInfoDataBuffer ( ) ) , // 0 context . getOldStream ( ) , // 1 AtomicAllocator . getInstance ( ) . getDeviceIdPointer ( ) , // 2 context . getBufferAllocation ( ) , // 3 context . getBufferReduction ( ) , // 4 context . getBufferScalar ( ) , // 5 context . getBufferSpecial ( ) , null , AddressRetriever . retrieveHostPointer ( zArr . shapeInfoDataBuffer ( ) ) ) ; val x = AtomicAllocator . getInstance ( ) . getPointer ( xArr , context ) ; val z = AtomicAllocator . getInstance ( ) . getPointer ( zArr , context ) ; val xShape = AtomicAllocator . getInstance ( ) . getPointer ( xArr . shapeInfoDataBuffer ( ) , context ) ; val zShape = AtomicAllocator . getInstance ( ) . getPointer ( zArr . shapeInfoDataBuffer ( ) , context ) ; double zeroPad = 0.0 ; if ( op . tArgs ( ) != null && op . tArgs ( ) . length > 0 ) { zeroPad = op . tArgs ( ) [ 0 ] ; } val extrass = new double [ ] { op . iArgs ( ) [ 0 ] , op . iArgs ( ) [ 1 ] , op . iArgs ( ) [ 2 ] , op . iArgs ( ) [ 3 ] , op . iArgs ( ) [ 4 ] , op . iArgs ( ) [ 5 ] , op . iArgs ( ) [ 6 ] , op . iArgs ( ) [ 7 ] , op . iArgs ( ) [ 8 ] , zeroPad } ; val extraArgsBuff = Nd4j . getConstantHandler ( ) . getConstantBuffer ( extrass ) ; val extraArgs = AtomicAllocator . getInstance ( ) . getPointer ( extraArgsBuff , context ) ; if ( dtype == DataBuffer . Type . DOUBLE ) { nativeOps . execTransformDouble ( xShapeHost , 37 , ( DoublePointer ) x , ( LongPointer ) xShape , ( DoublePointer ) z , ( LongPointer ) zShape , ( DoublePointer ) extraArgs ) ; } else if ( dtype == DataBuffer . Type . FLOAT ) { nativeOps . execTransformFloat ( xShapeHost , 37 , ( FloatPointer ) x , ( LongPointer ) xShape , ( FloatPointer ) z , ( LongPointer ) zShape , ( FloatPointer ) extraArgs ) ; } else if ( dtype == DataBuffer . Type . HALF ) { nativeOps . execTransformHalf ( xShapeHost , 37 , ( ShortPointer ) x , ( LongPointer ) xShape , ( ShortPointer ) z , ( LongPointer ) zShape , ( ShortPointer ) extraArgs ) ; } //AtomicAllocator.getInstance().getAllocationPoint(zArr).tickDeviceWrite(); AtomicAllocator . getInstance ( ) . getFlowController ( ) . registerAction ( context , zArr , xArr ) ; //Nd4j.getExecutioner().commit(); return ; } else if ( op . opName ( ) . equalsIgnoreCase ( "col2im" ) ) { val dtype = Nd4j . dataType ( ) ; val xArr = op . inputArguments ( ) [ 0 ] ; val zArr = op . outputArguments ( ) [ 0 ] ; CudaContext context = AtomicAllocator . getInstance ( ) . getFlowController ( ) . prepareAction ( zArr , xArr ) ; if ( extraz . get ( ) == null ) extraz . set ( new PointerPointer ( 32 ) ) ; PointerPointer xShapeHost = extraz . get ( ) . put ( AddressRetriever . retrieveHostPointer ( xArr . shapeInfoDataBuffer ( ) ) , // 0 context . getOldStream ( ) , // 1 AtomicAllocator . getInstance ( ) . getDeviceIdPointer ( ) , // 2 context . getBufferAllocation ( ) , // 3 context . getBufferReduction ( ) , // 4 context . getBufferScalar ( ) , // 5 context . getBufferSpecial ( ) , null , AddressRetriever . retrieveHostPointer ( zArr . shapeInfoDataBuffer ( ) ) ) ; val x = AtomicAllocator . getInstance ( ) . getPointer ( xArr , context ) ; val z = AtomicAllocator . getInstance ( ) . getPointer ( zArr , context ) ; val xShape = AtomicAllocator . getInstance ( ) . getPointer ( xArr . shapeInfoDataBuffer ( ) , context ) ; val zShape = AtomicAllocator . getInstance ( ) . getPointer ( zArr . shapeInfoDataBuffer ( ) , context ) ; val extrass = new double [ ] { op . iArgs ( ) [ 0 ] , op . iArgs ( ) [ 1 ] , op . iArgs ( ) [ 2 ] , op . iArgs ( ) [ 3 ] , op . iArgs ( ) [ 4 ] , op . iArgs ( ) [ 5 ] , op . iArgs ( ) [ 6 ] , op . iArgs ( ) [ 7 ] } ; val extraArgsBuff = Nd4j . getConstantHandler ( ) . getConstantBuffer ( extrass ) ; val extraArgs = AtomicAllocator . getInstance ( ) . getPointer ( extraArgsBuff , context ) ; if ( dtype == DataBuffer . Type . DOUBLE ) { nativeOps . execTransformDouble ( xShapeHost , 36 , ( DoublePointer ) x , ( LongPointer ) xShape , ( DoublePointer ) z , ( LongPointer ) zShape , ( DoublePointer ) extraArgs ) ; } else if ( dtype == DataBuffer . Type . FLOAT ) { nativeOps . execTransformFloat ( xShapeHost , 36 , ( FloatPointer ) x , ( LongPointer ) xShape , ( FloatPointer ) z , ( LongPointer ) zShape , ( FloatPointer ) extraArgs ) ; } else if ( dtype == DataBuffer . Type . HALF ) { nativeOps . execTransformHalf ( xShapeHost , 36 , ( ShortPointer ) x , ( LongPointer ) xShape , ( ShortPointer ) z , ( LongPointer ) zShape , ( ShortPointer ) extraArgs ) ; } //AtomicAllocator.getInstance().getAllocationPoint(zArr).tickDeviceWrite(); AtomicAllocator . getInstance ( ) . getFlowController ( ) . registerAction ( context , zArr , xArr ) ; //Nd4j.getExecutioner().commit(); return ; } else if ( op . opName ( ) . equalsIgnoreCase ( "pooling2d" ) ) { val dtype = Nd4j . dataType ( ) ; val xArr = op . inputArguments ( ) [ 0 ] ; val zArr = op . outputArguments ( ) [ 0 ] ; CudaContext context = AtomicAllocator . getInstance ( ) . getFlowController ( ) . prepareAction ( zArr , xArr ) ; if ( extraz . get ( ) == null ) extraz . set ( new PointerPointer ( 32 ) ) ; PointerPointer xShapeHost = extraz . get ( ) . put ( AddressRetriever . retrieveHostPointer ( xArr . shapeInfoDataBuffer ( ) ) , // 0 context . getOldStream ( ) , // 1 AtomicAllocator . getInstance ( ) . getDeviceIdPointer ( ) , // 2 context . getBufferAllocation ( ) , // 3 context . getBufferReduction ( ) , // 4 context . getBufferScalar ( ) , // 5 context . getBufferSpecial ( ) , null , AddressRetriever . retrieveHostPointer ( zArr . shapeInfoDataBuffer ( ) ) ) ; val x = AtomicAllocator . getInstance ( ) . getPointer ( xArr , context ) ; val z = AtomicAllocator . getInstance ( ) . getPointer ( zArr , context ) ; val xShape = AtomicAllocator . getInstance ( ) . getPointer ( xArr . shapeInfoDataBuffer ( ) , context ) ; val zShape = AtomicAllocator . getInstance ( ) . getPointer ( zArr . shapeInfoDataBuffer ( ) , context ) ; val extrass = new double [ ] { op . iArgs ( ) [ 0 ] , op . iArgs ( ) [ 1 ] , op . iArgs ( ) [ 2 ] , op . iArgs ( ) [ 3 ] , op . iArgs ( ) [ 4 ] , op . iArgs ( ) [ 5 ] , op . iArgs ( ) [ 6 ] , op . iArgs ( ) [ 7 ] , op . iArgs ( ) [ 8 ] } ; val extraArgsBuff = Nd4j . getConstantHandler ( ) . getConstantBuffer ( extrass ) ; val extraArgs = AtomicAllocator . getInstance ( ) . getPointer ( extraArgsBuff , context ) ; if ( dtype == DataBuffer . Type . DOUBLE ) { nativeOps . execTransformDouble ( xShapeHost , 71 , ( DoublePointer ) x , ( LongPointer ) xShape , ( DoublePointer ) z , ( LongPointer ) zShape , ( DoublePointer ) extraArgs ) ; } else if ( dtype == DataBuffer . Type . FLOAT ) { nativeOps . execTransformFloat ( xShapeHost , 71 , ( FloatPointer ) x , ( LongPointer ) xShape , ( FloatPointer ) z , ( LongPointer ) zShape , ( FloatPointer ) extraArgs ) ; } else if ( dtype == DataBuffer . Type . HALF ) { nativeOps . execTransformHalf ( xShapeHost , 71 , ( ShortPointer ) x , ( LongPointer ) xShape , ( ShortPointer ) z , ( LongPointer ) zShape , ( ShortPointer ) extraArgs ) ; } // AtomicAllocator.getInstance().getAllocationPoint(zArr).tickDeviceWrite(); AtomicAllocator . getInstance ( ) . getFlowController ( ) . registerAction ( context , zArr , xArr ) ; //Nd4j.getExecutioner().commit(); return ; } Nd4j . getExecutioner ( ) . commit ( ) ; long st = profilingHookIn ( op ) ; CudaContext context = ( CudaContext ) AtomicAllocator . getInstance ( ) . getDeviceContext ( ) . getContext ( ) ; //AtomicAllocator.getInstance().getFlowController().prepareActionAllWrite(op.outputArguments()); if ( extraz . get ( ) == null ) extraz . set ( new PointerPointer ( 32 ) ) ; PointerPointer extras = extraz . get ( ) . put ( new CudaPointer ( 1 ) , context . getOldStream ( ) , context . getBufferScalar ( ) , context . getBufferReduction ( ) ) ; val outputArgs = op . outputArguments ( ) ; val inputArgs = op . inputArguments ( ) ; if ( outputArgs . length == 0 && ! op . isInplaceCall ( ) ) throw new ND4JIllegalStateException ( "You can't execute non-inplace CustomOp without outputs being specified" ) ; val lc = op . opName ( ) . toLowerCase ( ) ; val hash = op . opHash ( ) ; val inputShapes = new PointerPointer <> ( inputArgs . length * 2 ) ; val inputBuffers = new PointerPointer <> ( inputArgs . length * 2 ) ; int cnt = 0 ; for ( val in : inputArgs ) { val hp = AtomicAllocator . getInstance ( ) . getHostPointer ( in . shapeInfoDataBuffer ( ) ) ; inputBuffers . put ( cnt , AtomicAllocator . getInstance ( ) . getHostPointer ( in ) ) ; inputShapes . put ( cnt , hp ) ; val dp = AtomicAllocator . getInstance ( ) . getPointer ( in . shapeInfoDataBuffer ( ) , context ) ; inputBuffers . put ( cnt + inputArgs . length , AtomicAllocator . getInstance ( ) . getPointer ( in , context ) ) ; inputShapes . put ( cnt + inputArgs . length , dp ) ; if ( op . isInplaceCall ( ) ) AtomicAllocator . getInstance ( ) . getAllocationPoint ( in ) . tickHostWrite ( ) ; cnt ++ ; } val outputShapes = new PointerPointer <> ( outputArgs . length * 2 ) ; val outputBuffers = new PointerPointer <> ( outputArgs . length * 2 ) ; cnt = 0 ; for ( val out : outputArgs ) { outputBuffers . put ( cnt , AtomicAllocator . getInstance ( ) . getHostPointer ( out ) ) ; outputShapes . put ( cnt , AtomicAllocator . getInstance ( ) . getHostPointer ( out . shapeInfoDataBuffer ( ) ) ) ; outputBuffers . put ( cnt + outputArgs . length , AtomicAllocator . getInstance ( ) . getPointer ( out , context ) ) ; outputShapes . put ( cnt + outputArgs . length , AtomicAllocator . getInstance ( ) . getPointer ( out . shapeInfoDataBuffer ( ) , context ) ) ; AtomicAllocator . getInstance ( ) . getAllocationPoint ( out ) . tickHostWrite ( ) ; cnt ++ ; } if ( Nd4j . dataType ( ) == DataBuffer . Type . FLOAT ) { val tArgs = op . tArgs ( ) . length > 0 ? new FloatPointer ( op . tArgs ( ) . length ) : null ; val iArgs = op . iArgs ( ) . length > 0 ? new LongPointer ( op . iArgs ( ) . length ) : null ; cnt = 0 ; for ( val t : op . tArgs ( ) ) tArgs . put ( cnt ++ , ( float ) t ) ; cnt = 0 ; for ( val i : op . iArgs ( ) ) iArgs . put ( cnt ++ , i ) ; val status = OpStatus . byNumber ( nativeOps . execCustomOpFloat ( extras , hash , inputBuffers , inputShapes , inputArgs . length , outputBuffers , outputShapes , outputArgs . length , tArgs , op . tArgs ( ) . length , iArgs , op . iArgs ( ) . length , op . isInplaceCall ( ) ) ) ; if ( status != OpStatus . ND4J_STATUS_OK ) throw new ND4JIllegalStateException ( "Op execution failed: " + status ) ; } else if ( Nd4j . dataType ( ) == DataBuffer . Type . DOUBLE ) { val tArgs = op . tArgs ( ) . length > 0 ? new DoublePointer ( op . tArgs ( ) . length ) : null ; val iArgs = op . iArgs ( ) . length > 0 ? new LongPointer ( op . iArgs ( ) . length ) : null ; cnt = 0 ; for ( val t : op . tArgs ( ) ) tArgs . put ( cnt ++ , t ) ; for ( val i : op . iArgs ( ) ) iArgs . put ( cnt ++ , i ) ; val status = OpStatus . byNumber ( nativeOps . execCustomOpDouble ( extras , hash , inputBuffers , inputShapes , inputArgs . length , outputBuffers , outputShapes , outputArgs . length , tArgs , op . tArgs ( ) . length , iArgs , op . iArgs ( ) . length , op . isInplaceCall ( ) ) ) ; if ( status != OpStatus . ND4J_STATUS_OK ) throw new ND4JIllegalStateException ( "Op execution failed: " + status ) ; } else if ( Nd4j . dataType ( ) == DataBuffer . Type . HALF ) { val tArgs = op . tArgs ( ) . length > 0 ? new ShortPointer ( op . tArgs ( ) . length ) : null ; val iArgs = op . iArgs ( ) . length > 0 ? new LongPointer ( op . iArgs ( ) . length ) : null ; cnt = 0 ; for ( val t : op . tArgs ( ) ) tArgs . put ( cnt ++ , ArrayUtil . toHalf ( ( float ) t ) ) ; cnt = 0 ; for ( val i : op . iArgs ( ) ) iArgs . put ( cnt ++ , i ) ; val status = OpStatus . byNumber ( nativeOps . execCustomOpHalf ( extras , hash , inputBuffers , inputShapes , inputArgs . length , outputBuffers , outputShapes , outputArgs . length , tArgs , op . tArgs ( ) . length , iArgs , op . iArgs ( ) . length , op . isInplaceCall ( ) ) ) ; if ( status != OpStatus . ND4J_STATUS_OK ) throw new ND4JIllegalStateException ( "Op execution failed: " + status ) ; } //AtomicAllocator.getInstance().getFlowController().prepareActionAllWrite(op.outputArguments()); profilingHookOut ( op , st ) ; }
This method executes given CustomOp

PLEASE NOTE: You're responsible for input/output validation
PLEASE NOTE: right now this operations are executing on CPU
@param op
This method executes given CustomOp