From 955fc2d55385f953b28e02e9fc1617ff054b2ebe Mon Sep 17 00:00:00 2001 From: Jan Pfeifer Date: Tue, 11 Mar 2025 12:14:48 +0100 Subject: [PATCH] Fixed small memory leak of a VectorData wrapper when converting to HLO/StableHLO (#31).Fixed small memory leak of a VectorData wrapper when converting to HLO/StableHLO (#31). --- cbuffer/cbuffer.go | 2 +- docs/CHANGELOG.md | 4 ++++ xlabuilder/xlacomputation.go | 10 ++++++++-- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/cbuffer/cbuffer.go b/cbuffer/cbuffer.go index a9faf32..20e02f1 100644 --- a/cbuffer/cbuffer.go +++ b/cbuffer/cbuffer.go @@ -68,7 +68,7 @@ func (b *CBuffer) Free() { // Bytes returns the buffer as a byte slice. // -// Ownership is not transferred: remember to free CBuffer afterwards. +// Ownership is not transferred: remember to free CBuffer afterward. func (b *CBuffer) Bytes() []byte { if b.data == nil { return nil diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 557b7fc..7569ea2 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -1,5 +1,9 @@ # Gopjrt Changelog +# Next + +* Fixed small memory leak of a VectorData wrapper when converting to HLO/StableHLO (#31). + # v0.6.2 - 2025/02/26 * Fixed C/C++ wrapper version. diff --git a/xlabuilder/xlacomputation.go b/xlabuilder/xlacomputation.go index bc35a6d..a9d63f2 100644 --- a/xlabuilder/xlacomputation.go +++ b/xlabuilder/xlacomputation.go @@ -77,7 +77,9 @@ func (comp *XlaComputation) SerializedHLO() *cbuffer.CBuffer { } var vectorData *C.VectorData vectorData = (*C.VectorData)(C.XlaComputationSerializedHLO(unsafe.Pointer(comp.cXlaComputation))) - return cbuffer.New(unsafe.Pointer(vectorData.data), int(vectorData.count), true) + cBuf := cbuffer.New(unsafe.Pointer(vectorData.data), int(vectorData.count), true) + C.free(unsafe.Pointer(vectorData)) + return cBuf } // HasStableHLO returns whether StableHLO support was included in the build -- it's very large, so by default it is not. @@ -111,7 +113,9 @@ func (comp *XlaComputation) SerializedStableHLO() (*cbuffer.CBuffer, error) { if err != nil { return nil, errors.Wrapf(err, "while converting XlaComputation to StableHLO") } - return cbuffer.New(unsafe.Pointer(vectorData.data), int(vectorData.count), true), nil + cBuf := cbuffer.New(unsafe.Pointer(vectorData.data), int(vectorData.count), true) + C.free(unsafe.Pointer(vectorData)) + return cBuf, nil } // TextHLO generates the HLO program as a and returns its text representation. @@ -130,6 +134,8 @@ func (comp *XlaComputation) TextHLO() string { } // TextStableHLO generates the StableHLO program. +// +// It returns an error if StableHLO code was not linked in (it's large). This can be checked with HasStableHLO. func (comp *XlaComputation) TextStableHLO() (string, error) { if comp.IsNil() { return "", errors.New("XlaComputation is nil, maybe it has already been destroyed?")