InfrastructureSystems.jl icon indicating copy to clipboard operation
InfrastructureSystems.jl copied to clipboard

Implement non-allocating `@scoped_enum`

Open GabrielKS opened this issue 8 months ago • 4 comments

The fundamental problem behind https://github.com/NREL-Sienna/PowerSystems.jl/pull/1216 is that @scoped_enum is backed by a Dict, which is mutable, so we can't compile away the lookup even in cases where it looks like a constant (e.g., my_units == UnitSystem.NATURAL_UNITS). There is no notion of a "constant dictionary" in Julia. However, a dictionary is just a function with a very small domain, and Val-based function application is more amenable to compile-time optimization. Here, I use that idea to implement a non-allocating @scoped_enum. This should eliminate the need for @scoped_enum users to implement their own patchwork fixes like https://github.com/NREL-Sienna/PowerSystems.jl/pull/1216 when they need performance.

Draft because I'm not confident that this is the right approach: I've tested that it doesn't allocate, but I still need to test other aspects of performance. Also, this is not how @enum is implemented in @Base, maybe an implementation along those lines would be even better.

GabrielKS avatar Mar 28 '25 21:03 GabrielKS

Codecov Report

Attention: Patch coverage is 83.33333% with 3 lines in your changes missing coverage. Please review.

Project coverage is 79.13%. Comparing base (711caa4) to head (8bee413). Report is 4 commits behind head on main.

Files with missing lines Patch % Lines
src/utils/utils.jl 83.33% 3 Missing :warning:
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main     #458      +/-   ##
==========================================
+ Coverage   77.79%   79.13%   +1.33%     
==========================================
  Files          71       71              
  Lines        5818     5736      -82     
==========================================
+ Hits         4526     4539      +13     
+ Misses       1292     1197      -95     
Flag Coverage Δ
unittests 79.13% <83.33%> (+1.33%) :arrow_up:

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
src/utils/utils.jl 65.93% <83.33%> (+1.24%) :arrow_up:

... and 4 files with indirect coverage changes

:rocket: New features to boost your workflow:
  • :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

codecov[bot] avatar Mar 28 '25 21:03 codecov[bot]

@daniel-thom do you remember why we went with @scoped_enum? I don't recall why we went that way

jd-lara avatar Apr 02 '25 17:04 jd-lara

Okay, this is ready for review. It doesn't solve all the issues with, for instance, rapidly getting a bunch of values that need to be unit converted from PowerSystems — there's some type instability there that will need to be fixed elsewhere — but as far as I can tell it will not impede such a fix and it does fix anything with ACBusTypes we might be doing in PowerFlows. If you run the included performance tests with the old version of @scoped_enum, you can see how much allocation this saves in certain cases.

GabrielKS avatar Apr 03 '25 21:04 GabrielKS

So EnumX.jl will kind of work, as long as we accept bitcasting internally, and we figure out the de/serialization

@enumx Fruit APPLE = 1 ORANGE = 2

@testset "Test scoped_enum performance" begin
    is_apple(x::Fruit.T) = x == Fruit.APPLE
    function f()
        @test (@allocated Fruit.APPLE) == 0
        @test (@allocated instances(Fruit.T)) == 0

        rng = Random.Xoshiro(47)
        my_fruits = [rand(rng, instances(Fruit.T)) for _ in 1:1_000]
        my_results = Vector{Bool}(undef, length(my_fruits))
        # Ref(x) rather than [x] is necessary to avoid spurious allocation
        @test (@allocated my_results .= (my_fruits .== Ref(Fruit.APPLE))) == 0
        # After compilation, here is observed the most drastic difference between the Dict-based implementation and the multiple dispach-based one:
        @test (@allocated my_results .= is_apple.(my_fruits)) == 0
        @test (@allocated my_results .= (my_fruits .< Ref(Fruit.ORANGE))) == 0
        @test (@allocated my_fruits .= Fruit.T.(Int32(3) .- Base.bitcast.(Int32, my_fruits))) == 0
        @test (@allocated(
            my_fruits .= Base.bitcast.(Fruit.T, Int32(3) .- Base.bitcast.(Int32, my_fruits)))) == 0
    end
    f()
    f()
end

josephmckinsey avatar Apr 03 '25 22:04 josephmckinsey

Let's merge this for now. It doesn't commit us to anything — we can easily switch to something else if we find something better.

GabrielKS avatar Apr 08 '25 16:04 GabrielKS