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

Accessing `task` field of Libtask TapedTask

Open willtebbutt opened this issue 9 months ago • 1 comments

In the extension for Libtask, AdvancedPS often accesses the task field of a Libtask TapedTask. I would like to know why this is done, and whether it is strictly necessary? See e.g. https://github.com/TuringLang/AdvancedPS.jl/blob/8b4b558b9c008f923f11d2af243e26bfd4b3d914/ext/AdvancedPSLibtaskExt.jl#L86 and https://github.com/TuringLang/AdvancedPS.jl/blob/8b4b558b9c008f923f11d2af243e26bfd4b3d914/ext/AdvancedPSLibtaskExt.jl#L179 .

For context, I'm looking at reimplementing the TapedTask -- see https://github.com/TuringLang/Libtask.jl/pull/179 -- and my proposal doesn't actually contain a Julia Task object anymore. I obviously need it to be compatible with this library, so if for some reason it's actually necessary to have a Julia Task, I'll have to rethink my design, and officially make the task field part of the public interface (I don't believe that it currently is).

willtebbutt avatar Feb 28 '25 15:02 willtebbutt

That was not a good design, nor is it necessary. IIRC, back in the earlier version of Libtask using C code, I couldn't find a way to enumerate all variables currently existing in the task's function. That is because Libtask only have the compiled machine code for each Task, so finding all the variables involves scanning the current stack and heap, which is challenging. However, we must duplicate VarInfo for copied particles to implement particle filters correctly. That was impossible without knowing the location of VarInfo in the C-level code. So, instead, I tweaked the @model macro to save a pointer to the current particle's (or TapedTask's) VarInfo. Then, a flag was attached to each particle to indicate whether it had been recently copied but had not duplicated its VarInfo.

Later, we added particle local rng so both rnd and varinfo are stored locally in each task's local storage.

Here is the code for assessing these particle-local (or task-local) rng and varinfo objects:

https://github.com/TuringLang/Turing.jl/blob/ef6c257fe8c027afd38e9db7217be8c611f31681/src/mcmc/particle_mcmc.jl#L349-L375

Here is the code for deep copying these local rng and varinfo objects:

https://github.com/TuringLang/AdvancedPS.jl/blob/8b4b558b9c008f923f11d2af243e26bfd4b3d914/ext/AdvancedPSLibtaskExt.jl#L140

yebai avatar Feb 28 '25 15:02 yebai