Skip to content

Conversation

@abussy
Copy link

@abussy abussy commented Sep 30, 2025

This PR is part of an effort to port DFTK's stress and response calculations to the GPU. The main issue with the current implementation is that various types are not isbits. Any data passed to a GPU kernel must be isbits.

Changes:

  • Changed the type of functional parameters from ComponentArray to NamedTuple. Usage remains the same, while the new type is isbits.
  • Introduced a to_isbits() function that returns a isbits version of a functional. Issue is that the identifier field was restricted to be a Symbol, which is not isbits. The parametrization of the identifier is extended to accept Val, and the to_isbits() function returns a copy of the input functional with the identifier set to Val{identifier}()

With these changes (and a few more, for a later PR), it is possible to run DFTK stress calculations on the GPU.

Possible points of contention/discussion:

  • Some workaround is necessary to calculate the derivative of a functional wrt its parameters, since ForwardDiff expects functions to take Arrays as input (see runtests.jl and change_parameters dispatching in individual functional implementations)
  • Functional types are parametrized with plain NamedTuple, vs formerly ComponentArray{<:Number}. So the type of the individual parameters is not enforced. I am not 100% sure whether that is an issue.

Copy link
Member

@mfherbst mfherbst left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The thing with the identifer and Val identifier is very hackish. Is it not possible to directly make the struct always isbits. The notion to use symbols for this is not set in stone. This points to me to bigger problems in the design of DftFnuctionals. We should discuss this and then directly change whatever is needed to help the GPU implementation.

Some workaround is necessary to calculate the derivative of a functional wrt its parameters, since ForwardDiff expects functions to take Arrays as input (see runtests.jl and change_parameters dispatching in individual functional implementations)

If that can be hidden in a convenient way, I'm ok with the change from ComponentArrays to NamedTuple. The main reason for using ComponentArrays in the first place was that it makes such AD calls more convenient and intuitive. But again, this is not set in stone. Using NamedTuple should just not make AD very cumbersome. @niklasschmitz can probably comment on this.

PbeCorrelation(parameters, pbe.lda)
end
end
# Change functional parameters based on an array of values. Assumes consistent ordering.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What do you mean by assumes consistent ordering ? I would not assume that the keys are in the same order ... as this cannot be enforced in practice.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also can we not use a ComponentArray here ?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In any case I would settle on one way of getting and setting the parameters and avoid the duplication of the interface. Otherwise this may derail. DftFunctionals is anyway already too complicated for what it's doing, so if we need to refactor parts of it to make GPU work and reshuffle functionality, that's ok.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also can we not use a ComponentArray here ?

Yes, we can. We could have a single definition of change_parameters that accepts anything that seamlessly converts to a NamedTuple (including a ComponentArray):

function change_parameters(pbe::PbeCorrelation, parameters; keep_identifier=false)                                                    
    if keep_identifier                                                                               
        PbeCorrelation(NamedTuple(parameters), pbe.lda, pbe.identifier)                              
    else                                                                                             
        PbeCorrelation(NamedTuple(parameters), pbe.lda)                                              
    end                                                                                              
end

That would be safer, and make derivatives wrt functional parameters cleaner.

@abussy
Copy link
Author

abussy commented Oct 8, 2025

The thing with the identifer and Val identifier is very hackish.

It is. Actually, the problem does not show up for LDA functionals, as the identifier is not part of the struct. Arguably, the identifier field is not necessary, as it is only used at initialization, but never during calculations. How important is it to keep track of it, once the functional is created?

The range isbits types is quite limited, unfortunately. I guess we could also use an Int as identifier, and then refer to an array of symbols, but:

  1. It is very Fortran like
  2. It is not flexible

Finally, we could also split the functional struct into an isbits part and a non-isbits part, and only pass the former to GPU kernels. This would have to be explicit in functions such as potential_terms

@abussy
Copy link
Author

abussy commented Oct 28, 2025

Relaxed functional parameters type, such that constructor can take any parameter convertible to a NamedTuple (e.g. ComponentArray).

The debate on how to make a functional isbits is still on. While the current solution is hackish, it is probably the simplest. Still hackish, but possibly better, we could relax the typing in the functional struct:

struct PbeExchange{NT} <: Functional{:gga,:x} where {NT<:NamedTuple}
    parameters::NT
    identifier
end

This way, the to_isbits function could look like:

function to_isbits(pbe::PbeExchange)
    PbeExchange(pbe.parameters, nothing)
end

I would appreciate if we get this refactoring through. It is currently holding me from submitting further contributions.

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.

2 participants