nd4j-backends/nd4j-backend-impls/nd4j-cuda/src/main/java/org/nd4j/jita/memory/CudaMemoryManager.java
@Override
    public Pointer allocate(long bytes, MemoryKind kind, boolean initialize) {
        AtomicAllocator allocator = AtomicAllocator.getInstance();

        //log.info("Allocating {} bytes in {} memory...", bytes, kind);

        if (kind == MemoryKind.HOST) {
            Pointer ptr = NativeOpsHolder.getInstance().getDeviceNativeOps().mallocHost(bytes, 0);

            if (ptr == null)
                throw new RuntimeException("Failed to allocate " + bytes + " bytes from HOST memory");

            if (initialize)
                Pointer.memset(ptr, 0, bytes);

            return ptr;//allocator.getMemoryHandler().alloc(AllocationStatus.HOST, null, null, initialize).getHostPointer();
        } else if (kind == MemoryKind.DEVICE) {
            Pointer ptr = NativeOpsHolder.getInstance().getDeviceNativeOps().mallocDevice(bytes, null, 0);


            //log.info("Allocating {} bytes for device_{}", bytes, Nd4j.getAffinityManager().getDeviceForCurrentThread());

            if (ptr == null)
                throw new RuntimeException("Failed to allocate " + bytes + " bytes from DEVICE [" + Nd4j.getAffinityManager().getDeviceForCurrentThread() + "] memory");

            if (initialize) {
                CudaContext context = (CudaContext) AtomicAllocator.getInstance().getDeviceContext().getContext();

                int i = NativeOpsHolder.getInstance().getDeviceNativeOps().memsetAsync(ptr, 0, bytes, 0, context.getSpecialStream());
                if (i == 0)
                    throw new ND4JIllegalStateException("memset failed on device_" + Nd4j.getAffinityManager().getDeviceForCurrentThread());

                context.getSpecialStream().synchronize();
            }


            return ptr; //allocator.getMemoryHandler().alloc(AllocationStatus.HOST, null, null, initialize).getDevicePointer();
        } else
            throw new RuntimeException("Unknown MemoryKind requested: " + kind);
    }
@Override
    public Pointer allocate(long bytes, MemoryKind kind, boolean initialize) {
        AtomicAllocator allocator = AtomicAllocator.getInstance();

        //log.info("Allocating {} bytes in {} memory...", bytes, kind);

        if (kind == MemoryKind.HOST) {
            Pointer ptr = NativeOpsHolder.getInstance().getDeviceNativeOps().mallocHost(bytes, 0);

            if (ptr == null)
                throw new RuntimeException("Failed to allocate " + bytes + " bytes from HOST memory");

            if (initialize)
                Pointer.memset(ptr, 0, bytes);

            return ptr;//allocator.getMemoryHandler().alloc(AllocationStatus.HOST, null, null, initialize).getHostPointer();
        } else if (kind == MemoryKind.DEVICE) {
            Pointer ptr = NativeOpsHolder.getInstance().getDeviceNativeOps().mallocDevice(bytes, null, 0);


            //log.info("Allocating {} bytes for device_{}", bytes, Nd4j.getAffinityManager().getDeviceForCurrentThread());

            if (ptr == null)
                throw new RuntimeException("Failed to allocate " + bytes + " bytes from DEVICE [" + Nd4j.getAffinityManager().getDeviceForCurrentThread() + "] memory");

            if (initialize) {
                CudaContext context = (CudaContext) AtomicAllocator.getInstance().getDeviceContext().getContext();

                int i = NativeOpsHolder.getInstance().getDeviceNativeOps().memsetAsync(ptr, 0, bytes, 0, context.getSpecialStream());
                if (i == 0)
                    throw new ND4JIllegalStateException("memset failed on device_" + Nd4j.getAffinityManager().getDeviceForCurrentThread());

                context.getSpecialStream().synchronize();
            }


            return ptr; //allocator.getMemoryHandler().alloc(AllocationStatus.HOST, null, null, initialize).getDevicePointer();
        } else
            throw new RuntimeException("Unknown MemoryKind requested: " + kind);
    }
@ Override public Pointer allocate ( long bytes , MemoryKind kind , boolean initialize ) { AtomicAllocator allocator = AtomicAllocator . getInstance ( ) ; //log.info("Allocating {} bytes in {} memory...", bytes, kind); if ( kind == MemoryKind . HOST ) { Pointer ptr = NativeOpsHolder . getInstance ( ) . getDeviceNativeOps ( ) . mallocHost ( bytes , 0 ) ; if ( ptr == null ) throw new RuntimeException ( "Failed to allocate " + bytes + " bytes from HOST memory" ) ; if ( initialize ) Pointer . memset ( ptr , 0 , bytes ) ; return ptr ; //allocator.getMemoryHandler().alloc(AllocationStatus.HOST, null, null, initialize).getHostPointer(); } else if ( kind == MemoryKind . DEVICE ) { Pointer ptr = NativeOpsHolder . getInstance ( ) . getDeviceNativeOps ( ) . mallocDevice ( bytes , null , 0 ) ; //log.info("Allocating {} bytes for device_{}", bytes, Nd4j.getAffinityManager().getDeviceForCurrentThread()); if ( ptr == null ) throw new RuntimeException ( "Failed to allocate " + bytes + " bytes from DEVICE [" + Nd4j . getAffinityManager ( ) . getDeviceForCurrentThread ( ) + "] memory" ) ; if ( initialize ) { CudaContext context = ( CudaContext ) AtomicAllocator . getInstance ( ) . getDeviceContext ( ) . getContext ( ) ; int i = NativeOpsHolder . getInstance ( ) . getDeviceNativeOps ( ) . memsetAsync ( ptr , 0 , bytes , 0 , context . getSpecialStream ( ) ) ; if ( i == 0 ) throw new ND4JIllegalStateException ( "memset failed on device_" + Nd4j . getAffinityManager ( ) . getDeviceForCurrentThread ( ) ) ; context . getSpecialStream ( ) . synchronize ( ) ; } return ptr ; //allocator.getMemoryHandler().alloc(AllocationStatus.HOST, null, null, initialize).getDevicePointer(); } else throw new RuntimeException ( "Unknown MemoryKind requested: " + kind ) ; }
This method returns Pointer to allocated memory chunk

PLEASE NOTE: Cache options depend on specific implementations

@param bytes
@param kind
@param initialize
This method returns Pointer to allocated memory chunk