[Feature Request] Memory allocation watcher

[sa-]

Review Mojo's priorities

• I have read the roadmap and priorities and I believe this request falls within the priorities.

What is your request?

Description

As a developer using Mojo, I would like to be able to unit test the number of allocations made by a function, maybe similar to how Julia's @time macro allows for measuring allocations.

• It would allow for more robust performance testing and ensure that functions are not making unnecessary copies or allocations
• In a language where the compiler will sometimes decide to make a copy under the hood, this is a good guardrail. Esp in loops where the cost of a copy can multiply
• I can see this being particularly useful when someone accidentally introduces a performance regression but the unit test is able to catch it before the code is ever merged or published

Proposed Implementation

One possible implementation of this feature could be the addition of a new @allocations macro, which would allow developers to assert the number of allocations made by a function in a unit test. Below are some possible ways to express this. I'm not sure if either of these are the best, and readers are encouraged to contribute ideas of their own!
Edit: Option 4 is my favourite so far

Option 1: if Mojo has macros

def test_my_function():
    # asserting less than 64B have been allocated
    @assert_max_allocations("64B") my_function()

Option 2: A pythonic take using with

def test_my_function():
    with assert_max_allocations("64B"):  # asserting less than 64B have been allocated
        my_function()

[edit from feedback] Option 3: Same as option 2 but without strings

def test_my_function():
    with assert_max_allocations(64): # asserting less than 64B have been allocated
        my_function()

[Another edit 5/22/24] Option 4: Most flexibility with allocation values

def test_my_function():
    # the allocation count and the bytes allocated are stored in 2 variables
    with allocation_watcher() as alloc_count, alloc_bytes:
        my_function()
    assert(alloc_count < 20)  # asserting less than 20 allocations have been made
    assert(alloc_bytes < 64)  # asserting less than 64B have been allocated
    assert(alloc_bytes < sys.info.sizeof[SomeRelevantType]() * 64)  # credits to @lsh for this example

What is your motivation for this change?

Currently, when writing performance-critical code, developers often have to rely on other tools (e.g. xctrace on macos) to measure allocations and ensure that they are minimized. But these tools are not standard across platforms. Having the ability to assert the number of allocations in unit tests would

1. Make it easier to catch performance issues early on in the development process
2. Create a standard way of doing this across platforms

Any other details?

Julia's @time macro in action

julia> @time h()
  1.151921 seconds (915.60 k allocations: 48.166 MiB, 1.64% gc time)
(0.5, 0.0)
 
julia> @time h()
  0.000013 seconds (21 allocations: 720 bytes)
(0.5, 0.0)

[sa-]

@mzaks @bgreni thanks for the interest!

I've just added option 4 which I believe is the best version so far

[bgreni]

@mzaks @bgreni thanks for the interest!

I've just added option 4 which I believe is the best version so far

I agree, the added flexibility could be useful.