Support CPU-only, task-only subgraphs with compiled execution - #417
Support CPU-only, task-only subgraphs with compiled execution#417rohany wants to merge 13 commits into
Conversation
|
@rohany I assume this needs to be reviewed? |
|
yes! |
| // expensive, we can pivot to a work-counter based implementation | ||
| // that skips this entire check if it is known that no subgraphs | ||
| // have been enqueued since the last check. | ||
| RWLock::AutoReaderLock al(pending_subgraphs_lock); |
There was a problem hiding this comment.
I don't understand why reader lock is the correct semantics here given the subgraph queue is being modified here.
There was a problem hiding this comment.
Yes, good catch, this is a piece of logic that didn't make it over from the port. I will update it to upgrade the lock to a writer lock before popping.
There was a problem hiding this comment.
[A note for me later]: When re-acquiring the writer lock, we shouldn't need to check again that the queue of subgraphs contains an element, as there is only one subgraph executor per ThreadedTaskScheduler. However, writing the code to just check again after acquiring the writer lock will probably be more future proof if that assertion changes later (i.e. when we add support for waiting on events in the subgraph).
| // SubgraphResourceReaper is a background work item that processes | ||
| // SubgraphInstantiationCleanup items asynchronously, freeing resources | ||
| // allocated for subgraph instantiation after the subgraph has finished. | ||
| class SubgraphResourceReaper : public BackgroundWorkItem { |
There was a problem hiding this comment.
nit: Perhaps SubgraphResourceDeleter, SubgraphResourceCleanup?
|
|
||
| // ProcSubgraphExecutor manages the logic of what a Processor should | ||
| // actually do when executing components of a compiled subgraph. | ||
| class ProcSubgraphExecutor { |
There was a problem hiding this comment.
I am wondering if there is an opportunity here to:
- Move subgraph under a separate sub-graph folder
- Start making changes with one-class-one-file type of granulairty
This would otherwise have to re-factored on a longer run. We made this "class-per-file" decision a while ago and obviously attempting to follow that for new abstractions.
There was a problem hiding this comment.
sure, i can move some new parts and classes declared in these files into a sub-folder. I don't want to (as part of this PR) to untangle the existing subgraph code into separate files.
| thread->start_subgraph_task_execution(); | ||
|
|
||
| // We can't hold the lock while executing tasks. | ||
| scheduler->lock.unlock(); |
There was a problem hiding this comment.
I kind of don't like this design of subgraph executor reaching out back to the scheduler for unlock/lock semantics. I think the subgraph executor should not really be aware of any higher level locking semantics. That would be cleaner. One way to fix this is to unlock at the higher level and then move the per-subgraph-executor synchronization into local mutex. Alternative option will be going lock-free on the executor.
There was a problem hiding this comment.
I wrote it this way originally at your suggestion because you wanted the ThreadScheduler itself to be unaware of what was happening inside the subgraph executor. I don't think it is possible to write this in a lock-free way without a significant rewrite of the ThreadScheduler itself which is not worth it in my opinion.
There was a problem hiding this comment.
I also want to maintain a handle on the lock here as there are other cases in this function where we may want to release the scheduler lock to trigger an event later. Separating the two logical components too much might make for headaches later.
|
|
||
| // Execution mode, which controls the optimization strategy used | ||
| // to execute the subgraph. | ||
| enum ExecutionMode |
There was a problem hiding this comment.
I don't understand why do we need this execution modes here. The obvious motivation for using the sub-graph here would for the application to get it compiled and gain the lowest possible overhead. Now, I have a feeling that the interpreted mode here is being added as as implementation artifact due to the artificial restrictions on only doing a compiled version for single-node use cases?
I think we need to compile everything and work-out the design strategy for optimized inter-node signaling here bypassing the host mediation and/or the active message protocol business. So for example (device-side), if we are running two GPU kernels inside two different tasks across nodes over IB and we know that consumer should start the device-side execution as soon as the producer device work has completed, it's possible to signal with RDMA put into the pre-registered device-side slot on the consumer. The consumer device would be waiting for this location with cuStreamWaitValue etc .
This PR obviously adds a CPU-only path - is this something is planned to be added as a follow up here?
There was a problem hiding this comment.
I don't understand why do we need this execution modes here. The obvious motivation for using the sub-graph here would for the application to get it compiled and gain the lowest possible overhead. Now, I have a feeling that the interpreted mode here is being added as as implementation artifact due to the artificial restrictions on only doing a compiled version for single-node use cases?
No, the interpreted mode is to handle/represent the fallback case into the current implementation of subgraphs, which essentially is a loop of launches through all tasks/copies in the subgraph. Once this full implementation is complete, we should deprecate and then remove the existing subgraph implementation, but we're not at that stage yet.
Yes, I agree that we need to compile the graph and do a variety of optimizations, which are all present in the full prototype implementation in our paper. As written in the description of this PR, I will add support for more features/optimizations incrementally instead of doing them all in here.
So for example (device-side), if we are running two GPU kernels inside two different tasks across nodes over IB and we know that consumer should start the device-side execution as soon as the producer device work has completed, it's possible to signal with RDMA put into the pre-registered device-side slot on the consumer.
My initial implementation of subgraphs actually did optimizations like this only within a node, and leveraged Realm's normal event/barrier infrastructure for cross-node dependencies. There were a few reasons for this:
- it was easier and has less likelihood of making a large difference because the dominating factor is a network message anyway. There's a few percent to be gained from doing this though at small task sizes (as shown in our paper), but it's more complicated and would need to be thought about a little carefully.
- Legion doesn't generate graphs that this optimization is applicable for. Legion control replicates and shards traces that it lowers to Realm graphs, and those sharded graphs only represent work done on the shard, and the graphs on each shard communicate through barriers. Legion does this to avoid representing the graph on every node, enabling scalability of the compilation process. We can imagine doing something like this in follow-on work where Realm perhaps has an option to control replicate graphs itself, in which case it can use a more direct communication mechanism for the cross-node operations of the graph.
2 participants
This PR adds initial support for a limited case of compiled Realm subgraphs to have a solid base of adding more features. It currently supports:
All of these will get resolved in follow-up work, but this PR adds the initial layer that later pieces can be built on. I'm in the process of adding more tests and will continue to do so, but this is ready to start getting some looks. I did a bunch of simplification from the prototype implementation which shrunk a good amount of this code.