Skip to content

Specialise on the parameter values#326

Draft
Mikolaj-A-Kowalski wants to merge 2 commits intomainfrom
mak/valued-params
Draft

Specialise on the parameter values#326
Mikolaj-A-Kowalski wants to merge 2 commits intomainfrom
mak/valued-params

Conversation

@Mikolaj-A-Kowalski
Copy link
Collaborator

Related to #320.

It is marked as draft since it is intended to be a starting point for the discussion on whether something like that should be merged into the code, and to show off some proposal on how it can be done. Hence, I have only applied it to the SugarKelp

In terms of performance, as discussed in the issue, it doesn't seem to give much. As expected, the register usage is significantly reduced, but it is not enough to increase occupancy. Small saving related to the fact that some computation can be compile-time evaluated is minimal (<2%) and on the verge of measurable. Based on that alone I would say it is probably not worth the effort... but on the other hand, for some kernels, in some circumstances, it may be enough to push a kernel to the higher occupancy band, and the performance improvements when that happens can be significant :-/

In terms of implementation I went with what I thought to be least invasive, which is to move numeric parameters to separate @kwdef struct and embed it as a parameter of the model struct (as suggested by @glwagner). Then, to 'unwrap' the value from Val{V} when using . (getproperty) access automatically, the model struct can subtype UnwrapValueFields. I wish it could be done with a 'Holy Trait', but it would result horrible type piracy (specialising Base.getproperty(::Any, ::Symbol) 🏴‍☠️ 😬 ).

The alternative would be to use a custom 'property-getter' function. I went with what I went, but I don't have any preference for either.

For the callable types, I have added a custom Val parametric type so they will be automatically unwrapped when called. As far as I am able to tell (looking into the generated code + poking it with @code_... macros) it does specialise the call for an instance. But I am still a bit uneasy if there isn't some subtle trap/edge case lurking around.

By introducing an internal CallableVal parametric type it is possible
to store callable objects as type parameters and automatically unwrap
them at the call site.
@Mikolaj-A-Kowalski
Copy link
Collaborator Author

Mikolaj-A-Kowalski commented Dec 17, 2025

Some notes:

  1. Tests are failing but it is only due to the regression in the summary of the SugarKelp. I will fix it later on.
  2. Being unsure if embedding the values inside the type is a good idea (given non-obvious performance benefit). It should be possible to make it user input. That is refactor along the lines:
struct SugarKelp{TL, SO, PS} <: UnwrapValueFields
                    temperature_limit :: TL
                    solver :: SO
                    params :: PS

    function SugarKelp(FT = Float64, valued=true; 
                       temperature_limit::TL = LinearOptimalTemperatureRange{FT}(),
                       solver::SO = NewtonRaphsonSolver{FT, Int}(; atol = eps(FT(1e-9))),
                       kwargs...
                       ) where {TL, SO}
        params = SugarKelpParameters{FT}(; kwargs...)
        if valued
            return new{CallableVal{temperature_limit}, CallableVal{solver}, Val{params}}(CallableVal(temperature_limit), CallableVal(solver), Val(params))
        else
            return new{TL, SO, typeof(params)}(temperature_limit, solver, params)
        end
    end
end 

i.e. remove the Value annotations and just, if requested, embed Val{...} as the type parameter. Perhaps it is something to consider.

@jagoosw
Copy link
Collaborator

jagoosw commented Dec 27, 2025

I'm not sure about the value of the trade off between the added source code complexity vs the potential peformance benefits. It would be interesting to see if we can manufacture a case where this gives tangible peformance improvements.

I don't know if some automation and abstraction could perhaps off set some of the added complexity making it worth while?

@glwagner
Copy link
Collaborator

glwagner commented Dec 27, 2025

This could possibly be developed as an extension to KernelAbstractions syntax (because it requires both Val(params) as well as a new function to unpack the params), eg

@kernel function func_with_params(a, b, @Val c, d)
    # etc
 end

under the hood this would generate something

func_with_params(a, b, ::Val{c}, d) where c = func_with_params(a, b, c, d)

then on the user side, you would have the option of Val-wrapping or not when you build the kernel.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants