MINOR: Move runtime.SetFinalizer into freeXXX

This commit is contained in:
pjongy 2023-10-19 03:31:22 +00:00
parent cd0565716f
commit 8031bb85f5
2 changed files with 16 additions and 10 deletions

View File

@ -17,8 +17,10 @@ type Scalar struct {
// free releases C allocated memory.
func freeCScalar(x *Scalar) error {
nbytes := x.nbytes()
atomic.AddInt64(&AllocatedMem, -nbytes)
if gotch.Debug {
nbytes := x.nbytes()
atomic.AddInt64(&AllocatedMem, -nbytes)
}
lock.Lock()
delete(ExistingScalars, x.name)
lock.Unlock()
@ -53,8 +55,10 @@ func newScalar(cscalar lib.Cscalar, nameOpt ...string) *Scalar {
}
atomic.AddInt64(&ScalarCount, 1)
nbytes := x.nbytes()
atomic.AddInt64(&AllocatedMem, nbytes)
if gotch.Debug {
nbytes := x.nbytes()
atomic.AddInt64(&AllocatedMem, nbytes)
}
lock.Lock()
ExistingScalars[x.name] = struct{}{}
lock.Unlock()

View File

@ -67,8 +67,10 @@ func newTensor(ctensor lib.Ctensor, nameOpt ...string) *Tensor {
x.d = new(bigStruct)
atomic.AddInt64(&TensorCount, 1)
nbytes := x.nbytes()
atomic.AddInt64(&AllocatedMem, nbytes)
if gotch.Debug {
nbytes := x.nbytes()
atomic.AddInt64(&AllocatedMem, nbytes)
}
lock.Lock()
if _, ok := ExistingTensors[name]; ok {
name = fmt.Sprintf("%s_%09d", name, TensorCount)
@ -170,6 +172,10 @@ func freeCTensor(ts *Tensor) error {
// IMPORTANT. make it nil so won't double free.
ts.ctensor = nil
// Clear SetFinalizer on ts so no double free tensor.
// Ref. https://pkg.go.dev/runtime#SetFinalizer
runtime.SetFinalizer(ts, nil)
return nil
}
@ -1190,10 +1196,6 @@ func (ts *Tensor) Drop() error {
return nil
}
// Clear SetFinalizer on ts so no double free tensor.
// Ref. https://pkg.go.dev/runtime#SetFinalizer
runtime.SetFinalizer(ts, nil)
ts.calledFrom = "ts.Drop()"
return freeCTensor(ts)
}