gorgonia.org/gorgonia@v0.9.17/perf.go (about)

     1  package gorgonia
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/chewxy/hm"
     7  	"gorgonia.org/tensor"
     8  )
     9  
    10  var nodePool = &sync.Pool{
    11  	New: func() interface{} { return new(Node) },
    12  }
    13  
    14  func borrowNode() *Node { return nodePool.Get().(*Node) }
    15  
    16  func returnNode(n *Node) {
    17  	// if the node is being returned to the pool then it should be removed from the graph that it is linked too as well
    18  	if n.g != nil {
    19  		n.g.RemoveNode(n)
    20  	}
    21  
    22  	// zero out any data in the node
    23  	ReturnType(n.t)
    24  	tensor.ReturnInts(n.shape)
    25  
    26  	n.t = nil
    27  	n.shape = nil
    28  	n.op = nil
    29  	n.children = nil
    30  	n.name = ""
    31  	n.group = ""
    32  	n.groups = nil
    33  	n.g = nil
    34  	n.boundTo = nil
    35  	n.derivOf = nil
    36  	n.deriv = nil
    37  	n.hash = 0
    38  	n.hashed = false
    39  	n.inferredShape = false
    40  	n.unchanged = false
    41  	n.isStmt = false
    42  	n.ofInterest = false
    43  
    44  	nodePool.Put(n)
    45  }
    46  
    47  // ReturnNode returns a node to the pool. It does not check that the *Node has been removed from the graph. USE WITH CAUTION.
    48  func ReturnNode(n *Node) {
    49  	n.g = nil
    50  	returnNode(n)
    51  }
    52  
    53  // handles Returning of Values
    54  
    55  var dvpool = &sync.Pool{
    56  	New: func() interface{} { return new(dualValue) },
    57  }
    58  
    59  func borrowDV() *dualValue { return dvpool.Get().(*dualValue) }
    60  
    61  func returnDV(dv *dualValue) {
    62  	returnValue(dv.d)
    63  	returnValue(dv.Value)
    64  	// if dvdT, ok := dv.d.(tensor.Tensor); ok {
    65  	// 	returnTensor(dvdT)
    66  	// }
    67  	// if dvvT, ok := dv.Value.(tensor.Tensor); ok {
    68  	// 	returnTensor(dvvT)
    69  	// }
    70  
    71  	dv.d = nil
    72  	dv.Value = nil
    73  	dvpool.Put(dv)
    74  }
    75  
    76  func returnTensor(t tensor.Tensor) {
    77  	tensor.ReturnTensor(t)
    78  }
    79  
    80  func returnValue(v Value) {
    81  	if t, ok := v.(tensor.Tensor); ok {
    82  		returnTensor(t)
    83  	}
    84  }
    85  
    86  var dimSizerPool = new(sync.Map)
    87  
    88  func borrowDimSizers(size int) []DimSizer {
    89  	var pool *sync.Pool
    90  	p, ok := dimSizerPool.Load(size)
    91  
    92  	if !ok {
    93  		s := size
    94  		pool = &sync.Pool{
    95  			New: func() interface{} { return make([]DimSizer, s, s) },
    96  		}
    97  		dimSizerPool.Store(size, pool)
    98  	} else {
    99  		pool = p.(*sync.Pool)
   100  	}
   101  	return pool.Get().([]DimSizer)
   102  }
   103  
   104  func returnDimSizers(ds []DimSizer) {
   105  	p, ok := dimSizerPool.Load(cap(ds))
   106  	if !ok {
   107  		return
   108  	}
   109  	pool := p.(*sync.Pool)
   110  	for i := range ds {
   111  		ds[i] = nil
   112  	}
   113  	pool.Put(ds)
   114  }
   115  
   116  var tensorTypePool = &sync.Pool{
   117  	New: func() interface{} { return new(TensorType) },
   118  }
   119  
   120  func borrowTensorType() *TensorType {
   121  	return tensorTypePool.Get().(*TensorType)
   122  }
   123  
   124  func returnTensorType(t *TensorType) {
   125  	switch t {
   126  	case vecF64, vecF32:
   127  		return
   128  	case matF64, matF32:
   129  		return
   130  	case ten3F64, ten3F32:
   131  		return
   132  	}
   133  	t.Of = nil
   134  	t.Dims = 0
   135  	tensorTypePool.Put(t)
   136  }
   137  
   138  // ReturnType ...
   139  func ReturnType(t hm.Type) {
   140  	switch tt := t.(type) {
   141  	case *TensorType:
   142  		returnTensorType(tt)
   143  	case TensorType:
   144  		// do nothing
   145  	case tensor.Dtype:
   146  		// do nothing
   147  	case *hm.FunctionType:
   148  		hm.ReturnFnType(tt)
   149  	}
   150  }