Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Completely Refactor #179

Draft
wants to merge 53 commits into
base: main
Choose a base branch
from
Draft

Completely Refactor #179

wants to merge 53 commits into from

Conversation

willtebbutt
Copy link
Member

@willtebbutt willtebbutt commented Feb 27, 2025

Per #177 this is an attempt to significantly refactor the internals for robustness.

Todo:

  • understand what the desired semantics are regarding copying, per the discussion in Questions About Semantics #178
  • determine whether the task field of a TapedTask ought actually to be part of the public interface -- issue to discussion open at Accessing task field of Libtask TapedTask AdvancedPS.jl#113
  • implement an opt-in mechanism for nested calls containing produce statements
  • tidy up main body of code
  • document how PhiNodes are handled more thoroughly (include a worked example)
  • document the intent of the code + the high-level structure. In particular, re-work the README to explain what the public interface is, and what it does.
  • get all CI passing
  • check that we can make AdvancedPS and Turing work using this PR
  • kwarg support
  • improve nesting mechanism to ensure that we always catch everything
  • optimise nested calls (currently throwing away all return type information to get prototype running)
  • benchmark to ensure no regressions
  • add BBCode section to Mooncake docs so that this implementation is intelligible to everyone
  • add more high-level documentation

Linked Issues:

Closes #165
Closes #167
Closes #171
Closes #176

@willtebbutt
Copy link
Member Author

willtebbutt commented Mar 4, 2025

I've been thinking a bit further about this, and I wonder whether we ought just to use the new ScopedValues feature in Base? We could use this, and then just move all trace-specific items from task-local storage to a trace-local.

e.g. if we add a variable to Libtask

const task_cache = ScopedValue{Any}(nothing) # maybe change types for stability

then the implementation of consume can be something like

function consume(t::TapedTask)
    with(task_cache => t.cache) do
        return t.mc(t.args...)
    end
end

Then if you want to use something stored in the cache field of a TapedTask, you just refer to Libtask.task_cache in the function you've got produce statements in, rather than current_task().storage.

edit: this feature only officially landed in Julia 1.11, but this is based on an implementation which supports 1.8+, so we should be able to use one or the other selectively.

edit2: alternatively, we only apply the upgrade in this PR to 1.11 and up.

@willtebbutt
Copy link
Member Author

Status Update: nested calls now basically work, I just need to tidy up the implementation a bit. I've also tidied up the whole implementation and documented it much more thoroughly. There's more that needs doing, but it's almost at he point that we can try and get it to work with AdvancedPS + Turing, to make sure that we do indeed have all of the features that we think we need.

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.

Use MistyClosures for Libtask 2.0?
2 participants