gorgonia.org/gorgonia@v0.9.17/cuda/arena.go (about)

     1  package cuda
     2  
     3  import (
     4  	"gorgonia.org/cu"
     5  	"gorgonia.org/tensor"
     6  )
     7  
     8  // this file implements the arena
     9  
    10  var _ Arena = &Engine{}
    11  
    12  // Get allocates memory of certain size and returns a pointer to it
    13  func (e *Engine) Get(size int64) (tensor.Memory, error) {
    14  	ptr, err := e.a.alloc(size)
    15  	return cu.DevicePtr(ptr), err
    16  }
    17  
    18  // Put releases a chunk of memory of certain size
    19  func (e *Engine) Put(mem tensor.Memory, size int64) {
    20  	addr := uintptr(mem.Uintptr())
    21  	e.a.free(addr)
    22  }
    23  
    24  // ResetAllocator releases used memory of the engine
    25  func (e *Engine) ResetAllocator() {
    26  	used := make([]uintptr, 0, len(e.a.used))
    27  	for k := range e.a.used {
    28  		used = append(used, k)
    29  	}
    30  
    31  	for _, ptr := range used {
    32  		e.a.free(ptr + e.a.start)
    33  	}
    34  	e.a.coalesce()
    35  	e.a.reset() // reset statistics
    36  }