Function newDeviceAllocator

query/device_allocator.go:158–186

newDeviceAllocator returns a new device allocator instances.

()

// newDeviceAllocator returns a new device allocator instances.
func newDeviceAllocator() deviceAllocator {
	var da deviceAllocator
	deviceCount := cgoutils.GetDeviceCount()
	if cgoutils.IsPooledMemory() {
		utils.GetLogger().Info("Using pooled device memory manager")
		da = &pooledDeviceAllocatorImpl{}
	} else {
		utils.GetLogger().Info("Using memory tracking device memory manager")
		da = &memoryTrackingDeviceAllocatorImpl{
			memoryUsage: make([]int64, deviceCount),
		}
	}

	// Start memory usage reporting go routine.
	// Report the allocated memory of each device per memoryReportingInterval.
	timer := time.NewTimer(memoryReportingInterval)
	go func() {
		for {
			select {
			case <-timer.C:
				reportAllocatedMemory(deviceCount, da)
				// Since we already receive the event from channel,
				// there is no need to stop it and we can directly reset the timer.
				timer.Reset(memoryReportingInterval)
			}
		}
	}()
	return da
}

// memoryTrackingDeviceAllocatorImpl maintains the memory space for each device and reports the updated memory every time an
// allocation/free request is issued.