nd4j-backends/nd4j-backend-impls/nd4j-cuda/src/main/java/org/nd4j/jita/memory/impl/CudaCachingZeroProvider.java
@Override
    public PointersPair malloc(AllocationShape shape, AllocationPoint point, AllocationStatus location) {
        long reqMemory = AllocationUtils.getRequiredMemory(shape);

        if (location == AllocationStatus.HOST && reqMemory < CudaEnvironment.getInstance().getConfiguration().getMaximumHostCacheableLength()) {

            CacheHolder cache = zeroCache.get(shape);
            if (cache != null) {
                Pointer pointer = cache.poll();
                if (pointer != null) {
                    cacheZeroHit.incrementAndGet();

                    // since this memory chunk is going to be used now, remove it's amount from
                    zeroCachedAmount.addAndGet(-1 * reqMemory);

                    PointersPair pair = new PointersPair();
                    pair.setDevicePointer(new CudaPointer(pointer.address()));
                    pair.setHostPointer(new CudaPointer(pointer.address()));

                    point.setAllocationStatus(AllocationStatus.HOST);
                    return pair;
                }
            }
            cacheZeroMiss.incrementAndGet();

            if (CudaEnvironment.getInstance().getConfiguration().isUsePreallocation() && zeroCachedAmount.get() < CudaEnvironment.getInstance().getConfiguration().getMaximumHostCache() / 10
                            && reqMemory < 16 * 1024 * 1024L) {
                CachePreallocator preallocator = new CachePreallocator(shape, location, CudaEnvironment.getInstance().getConfiguration().getPreallocationCalls());
                preallocator.start();
            }

            cacheZeroMiss.incrementAndGet();
            return super.malloc(shape, point, location);
        }

        return super.malloc(shape, point, location);
    }
@Override
    public PointersPair malloc(AllocationShape shape, AllocationPoint point, AllocationStatus location) {
        long reqMemory = AllocationUtils.getRequiredMemory(shape);

        if (location == AllocationStatus.HOST && reqMemory < CudaEnvironment.getInstance().getConfiguration().getMaximumHostCacheableLength()) {

            CacheHolder cache = zeroCache.get(shape);
            if (cache != null) {
                Pointer pointer = cache.poll();
                if (pointer != null) {
                    cacheZeroHit.incrementAndGet();

                    // since this memory chunk is going to be used now, remove it's amount from
                    zeroCachedAmount.addAndGet(-1 * reqMemory);

                    PointersPair pair = new PointersPair();
                    pair.setDevicePointer(new CudaPointer(pointer.address()));
                    pair.setHostPointer(new CudaPointer(pointer.address()));

                    point.setAllocationStatus(AllocationStatus.HOST);
                    return pair;
                }
            }
            cacheZeroMiss.incrementAndGet();

            if (CudaEnvironment.getInstance().getConfiguration().isUsePreallocation() && zeroCachedAmount.get() < CudaEnvironment.getInstance().getConfiguration().getMaximumHostCache() / 10
                            && reqMemory < 16 * 1024 * 1024L) {
                CachePreallocator preallocator = new CachePreallocator(shape, location, CudaEnvironment.getInstance().getConfiguration().getPreallocationCalls());
                preallocator.start();
            }

            cacheZeroMiss.incrementAndGet();
            return super.malloc(shape, point, location);
        }

        return super.malloc(shape, point, location);
    }
@ Override public PointersPair malloc ( AllocationShape shape , AllocationPoint point , AllocationStatus location ) { long reqMemory = AllocationUtils . getRequiredMemory ( shape ) ; if ( location == AllocationStatus . HOST && reqMemory < CudaEnvironment . getInstance ( ) . getConfiguration ( ) . getMaximumHostCacheableLength ( ) ) { CacheHolder cache = zeroCache . get ( shape ) ; if ( cache != null ) { Pointer pointer = cache . poll ( ) ; if ( pointer != null ) { cacheZeroHit . incrementAndGet ( ) ; // since this memory chunk is going to be used now, remove it's amount from zeroCachedAmount . addAndGet ( - 1 * reqMemory ) ; PointersPair pair = new PointersPair ( ) ; pair . setDevicePointer ( new CudaPointer ( pointer . address ( ) ) ) ; pair . setHostPointer ( new CudaPointer ( pointer . address ( ) ) ) ; point . setAllocationStatus ( AllocationStatus . HOST ) ; return pair ; } } cacheZeroMiss . incrementAndGet ( ) ; if ( CudaEnvironment . getInstance ( ) . getConfiguration ( ) . isUsePreallocation ( ) && zeroCachedAmount . get ( ) < CudaEnvironment . getInstance ( ) . getConfiguration ( ) . getMaximumHostCache ( ) / 10 && reqMemory < 16 * 1024 * 1024L ) { CachePreallocator preallocator = new CachePreallocator ( shape , location , CudaEnvironment . getInstance ( ) . getConfiguration ( ) . getPreallocationCalls ( ) ) ; preallocator . start ( ) ; } cacheZeroMiss . incrementAndGet ( ) ; return super . malloc ( shape , point , location ) ; } return super . malloc ( shape , point , location ) ; }
This method provides PointersPair to memory chunk specified by AllocationShape

PLEASE NOTE: This method can actually ignore malloc request, and give out previously cached free memory chunk with equal shape.

@param shape shape of desired memory chunk
@param point target AllocationPoint structure
@param location either HOST or DEVICE
@return
This method provides PointersPair to memory chunk specified by AllocationShape