token scratch - #1394
Conversation
| Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 256), | ||
| [&](int i) { shared_data(i) = ...; }); | ||
| }); | ||
| ``` |
There was a problem hiding this comment.
I think there’s a mismatch between the TeamPolicy example and what the current implementation guarantees.
As written, TokenScratchPool::acquire() just calls tokens_.acquire() and doesn’t know anything about the team (e.g. it doesn't pass the team_member). In a TeamPolicy kernel, the outer lambda runs once per team member, so multiple threads in the same team can call acquire() independently and potentially get different token IDs (i.e., different backing buffers) depending on the concurrency of the execution space. That would make the scratch effectively per-thread rather than per-team in the example above.
This also lines up with how the pool is sized (tokens_.size() * bytes_per_token), which reflects execution-space concurrency rather than number of teams. For true per-team scratch, I’d expect something along the lines of a PerTeam acquisition (or using AcquireTeamUniqueToken) so that all threads in a team share the same buffer.
Maybe the example is not really the intended usage though? It looks like there are no hierarchical tests in the unit test suite.
There was a problem hiding this comment.
Yeah---this doesn't work, and I didn't intend it to. I'm just removing it from the plan history doc. This was generated at the beginning of the effort to build the code and I made a bunch of changes after.
This raises a few related issues, which maybe are or are not best discussed here:
- SHOULD something akin to this be made to work? We could make it work by adding a constructor and token type for team scratch so we can use the same API for teams as for flat parallelism. IMO, given that Kokkos already supports something like this, the main benefit is we get to keep the same API.
- On the other hand, a simpler solution would maybe to just have our allocator object optionally be constructible from a team member, rather than our pool, and then it just is a thin wrapper around the equivalent kokkos API. Might be something we should think about in our loops?
- How should we handle things like this where the plan history generated by the robot isn't actually in sync with the final product? Updating it by hand as we have here is, I think, sensible... but it might also be useful, for transparency's sake, to show the original artifact for code archeology purposes if needed?
|
@adamdempsey90 @lroberts36 thanks for the comments! I think they're addressed now. Please re-review. |
adamdempsey90
left a comment
adamdempsey90
left a comment
There was a problem hiding this comment.
Hooking this up to a memory pool would be the next step to increase performance.
Agreed. Perhaps the object pools we already utilize could work for that. |
|
I've been doing some performance tuning on an AMD system. For flat loops, this mechanism is very slow due to an unavoidable atomic to synchronize the tokens. I don't know whether or not we want to keep this. Certainly I think the design is useful but may just not be something we want to encourage. Curious folks thoughts. @adamdempsey90 I think that may at least in part depend on whether or not you want this for artemis. |
We certainly need something like this for our dust coagulation routine in artemis. There's no other way to store potentially 100x100 matrices per thread. That being said, I can always use the ports of call memory pool + my own implementation of this in artemis. Have you verified that the slowness is due to this token acquisition and not the cost of allocating the view? And is this only slow on AMD systems? It might be worth having even if it's only performant on CPU and CUDA systems |
It's definitely the token acquisition. That said, I was doing something pretty dumb in my tests---I was basically always acquiring the token, whether or not I needed scratch. That was a pretty massive slowdown. A very easy/simple fix for that is to just only acquire a token if the number of requested bytes is > 0. I can push that change. In artemis, and in kernels that are very heavy the performance might not be noticable. It's probably most noticable where only a small amount of work is being done, in which case, the scratch should probably be replaced with another mechanism. |
|
@adamdempsey90 do you want to try this in Artemis? |
Pull request was converted to draft
|
@adamdempsey90 @lroberts36 do we still want this? Or should I close it? |
I don't feel strongly either way. It's not much to implement this in artemis if we ever get to it. And if/when we do and if it's helpful, we can just upstream here. But I do believe that there are monster kernels out there that would easily hide the token acquisition cost. |
Definitely agreed. So maybe it's worth merging. @lroberts36 what do you think? |
I think it makes sense to merge it. It works and we can see where it is a performance advantage and wire it into the loop abstraction |
|
OK sounds good. Then I'm removing WIP. If somebody can give it a quick second review. (Thanks @adamdempsey90 for the first one) we can get this in. @lroberts36 you already reviewed once, if you could give it a quick skim and click approve, we can merge. |
5 participants
PR Summary
This MR is motivated by a desire to pre-allocate scratch memory for flat par_for loops. It basically copies how level 1 scratch works in Kokkos with hierarchical loops, where you request a certain amount of memory per thread using Kokkos's unique token interface. Then in the loop body, you can request multiple unmanaged views based on the pre-allocated pool.
The advantage of this over, say, using a variable array as scratch is that it has a smaller memory footprint, one "scratch" per thread vs one per cell, which could be a big difference. We're thinking about changing how our riot loop patterns work (as we discussed in the last call) and this machinery is useful for facilitating cleaner loops for riot and hopefully other codes. I think @pdmullen also mentioned it'll be useful for artemis.
PR Checklist