Skip to content
Permalink
Browse files Browse the repository at this point in the history
Merge pull request #50508 from wamuir:fix-tstring-dealloc
PiperOrigin-RevId: 384557722
Change-Id: I72858edf72952fd4e7e0a1d9776c9408a7081d42
  • Loading branch information
tensorflower-gardener committed Jul 13, 2021
2 parents 5a14b2e + 49499c1 commit 8721ba9
Showing 1 changed file with 17 additions and 13 deletions.
30 changes: 17 additions & 13 deletions tensorflow/go/tensor.go
Expand Up @@ -98,9 +98,9 @@ func NewTensor(value interface{}) (*Tensor, error) {

raw := tensorData(t.c)

runtime.SetFinalizer(t, func(t *Tensor) {
defer runtime.SetFinalizer(t, func(t *Tensor) {
if dataType == String {
t.clearTStrings(raw, nflattened)
t.clearTStrings(raw, int64(nbytes/C.sizeof_TF_TString))
}

t.finalize()
Expand All @@ -111,15 +111,18 @@ func NewTensor(value interface{}) (*Tensor, error) {
if isAllArray(val.Type()) {
// We have arrays all the way down, or just primitive types. We can
// just copy the memory in as it is all contiguous.
if err := copyPtr(buf, unpackEFace(value).data, int(val.Type().Size())); err != nil {
if _, err := copyPtr(buf, unpackEFace(value).data, int(val.Type().Size())); err != nil {
return nil, err
}
} else {
// When there are slices involved the memory for each leaf slice may
// not be contiguous with the others or in the order we might
// expect, so we need to work our way down to each slice of
// primitives and copy them individually
if err := encodeTensorWithSlices(buf, val, shape); err != nil {
if n, err := encodeTensorWithSlices(buf, val, shape); err != nil {
// Set nbytes to count of bytes written for deferred call to
// runtime.SetFinalizer
nbytes = uintptr(n)
return nil, err
}
}
Expand Down Expand Up @@ -486,13 +489,13 @@ func sizeVarUint(v uint64) int {

// encodeTensorWithSlices writes v to the specified buffer using the format specified in
// c_api.h. Use stringEncoder for String tensors.
func encodeTensorWithSlices(w *bytes.Buffer, v reflect.Value, shape []int64) error {
func encodeTensorWithSlices(w *bytes.Buffer, v reflect.Value, shape []int64) (int, error) {
// If current dimension is a slice, verify that it has the expected size
// Go's type system makes that guarantee for arrays.
if v.Kind() == reflect.Slice {
expected := int(shape[0])
if v.Len() != expected {
return fmt.Errorf("mismatched slice lengths: %d and %d", v.Len(), expected)
return 0, fmt.Errorf("mismatched slice lengths: %d and %d", v.Len(), expected)
}
} else if v.Kind() == reflect.String {
s := v.Interface().(string)
Expand All @@ -501,7 +504,7 @@ func encodeTensorWithSlices(w *bytes.Buffer, v reflect.Value, shape []int64) err
ptr := unsafe.Pointer(&tstr)
return copyPtr(w, ptr, C.sizeof_TF_TString)
} else if v.Kind() != reflect.Array {
return fmt.Errorf("unsupported type %v", v.Type())
return 0, fmt.Errorf("unsupported type %v", v.Type())
}

// Once we have just a single dimension we can just copy the data
Expand All @@ -514,15 +517,17 @@ func encodeTensorWithSlices(w *bytes.Buffer, v reflect.Value, shape []int64) err
return copyPtr(w, ptr, v.Len()*int(elt.Type().Size()))
}

n := 0
subShape := shape[1:]
for i := 0; i < v.Len(); i++ {
err := encodeTensorWithSlices(w, v.Index(i), subShape)
j, err := encodeTensorWithSlices(w, v.Index(i), subShape)
if err != nil {
return err
return n + j, err
}
n += j
}

return nil
return n, nil
}

// It isn't safe to use reflect.SliceHeader as it uses a uintptr for Data and
Expand All @@ -536,15 +541,14 @@ type sliceHeader struct {
// copyPtr copies the backing data for a slice or array directly into w. Note
// we don't need to worry about byte ordering because we want the natural byte
// order for the machine we're running on.
func copyPtr(w *bytes.Buffer, ptr unsafe.Pointer, l int) error {
func copyPtr(w *bytes.Buffer, ptr unsafe.Pointer, l int) (int, error) {
// Convert our slice header into a []byte so we can call w.Write
b := *(*[]byte)(unsafe.Pointer(&sliceHeader{
Data: ptr,
Len: l,
Cap: l,
}))
_, err := w.Write(b)
return err
return w.Write(b)
}

func bug(format string, args ...interface{}) error {
Expand Down

0 comments on commit 8721ba9

Please sign in to comment.