Fix test breakage when address/memory/thread sanitizer was enabled.

Update the production code to not set the `*_allocated_bytes` fields
when using address/memory/thread sanitizer, since those sanitizers don't
support `mallinfo()` -- `mallinfo()` may return invalid results when sanitizers
are enabled .

Update test to not expect the `*_allocated_bytes` fields to be valid when using
address/memory/thread sanitizer.

Also update docs in `memory_info.h` to make it clearer that even if `isSupported()`
returns true, it is possible that only _some_ of the fields are supported.

Also document that `memory_info.h` supports Windows; this was already clear from
other parts of the header file that described the behaviour on Windows.

PiperOrigin-RevId: 825505016
This commit is contained in:
Fergus Henderson 2025-10-29 05:40:25 -07:00 committed by TensorFlower Gardener
parent 684717efe0
commit 98a24eb07e
3 changed files with 29 additions and 5 deletions

View File

@ -92,7 +92,9 @@ MemoryUsage GetMemoryUsage() {
result.private_footprint_bytes = (vm_swap_kb + res.ru_maxrss) * 1024;
}
}
#if defined(__NO_MALLINFO__) || !defined(__GLIBC__)
#if defined(__NO_MALLINFO__) || !defined(__GLIBC__) || \
defined(ADDRESS_SANITIZER) || defined(MEMORY_SANITIZER) || \
defined(THREAD_SANITIZER)
result.total_allocated_bytes = -1;
result.in_use_allocated_bytes = -1;
#elif __GLIBC_MINOR__ >= 33
@ -103,7 +105,9 @@ MemoryUsage GetMemoryUsage() {
const auto mem = mallinfo();
result.total_allocated_bytes = mem.arena;
result.in_use_allocated_bytes = mem.uordblks;
#endif // defined(__NO_MALLINFO__)
#endif // defined(__NO_MALLINFO__) || !defined(__GLIBC__) || \
// defined(ADDRESS_SANITIZER) || defined(MEMORY_SANITIZER) ||
// defined(THREAD_SANITIZER)
#elif defined(__APPLE__)
struct task_vm_info vm_info;
mach_msg_type_number_t count = TASK_VM_INFO_COUNT;

View File

@ -30,6 +30,9 @@ struct MemoryUsage {
// Indicates whether obtaining memory usage is supported on the platform, thus
// indicating whether the values defined in this struct make sense or not.
// Note that even if this returns true, some of the fields in the struct may
// not be supported by GetMemoryUsage(); in such cases, unsupported fields
// will be set to kValueNotSet (zero) or -1.
static bool IsSupported();
MemoryUsage()
@ -127,7 +130,7 @@ struct MemoryUsage {
};
// Return the memory usage from the system.
// Note: this currently only works on Linux-based and Apple systems.
// Note: this currently only works on Linux-based, Apple, and Windows systems.
MemoryUsage GetMemoryUsage();
} // namespace memory

View File

@ -15,7 +15,6 @@ limitations under the License.
#include "tensorflow/lite/profiling/memory_info.h"
#include <memory>
#include <new>
#include <sstream>
#include <string>
@ -75,10 +74,28 @@ TEST(MemoryUsage, GetMemoryUsage) {
}
EXPECT_GE(result.mem_footprint_kb, size / 1024);
#if !defined(__linux__) || \
(!defined(ADDRESS_SANITIZER) && !defined(MEMORY_SANITIZER) && \
!defined(THREAD_SANITIZER))
EXPECT_GE(result.total_allocated_bytes, size);
EXPECT_GE(result.in_use_allocated_bytes, size);
#else
// The mallinfo() function, which is used on Linux, returns invalid
// results when address/memory/thread sanitizer is enabled, e.g.
// <https://github.com/google/sanitizers/issues/1845>, so the
// *_allocated_bytes fields are not supported in those cases,
// and should be set to either -1 or kValueNotSet(0).
if (result.total_allocated_bytes != -1) {
EXPECT_EQ(result.total_allocated_bytes, MemoryUsage::kValueNotSet);
}
if (result.in_use_allocated_bytes != -1) {
EXPECT_EQ(result.in_use_allocated_bytes, MemoryUsage::kValueNotSet);
}
#endif // !defined(__linux__) || \
// (!defined(ADDRESS_SANITIZER) && !defined(MEMORY_SANITIZER) && \
// !defined(THREAD_SANITIZER))
EXPECT_GE(result.private_footprint_bytes, size);
#endif
#endif // defined(__linux__) || defined(__APPLE__) || defined(_WIN32)
}
// The main aim of this test is just to exercise the code for