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 }