feat: Add support for CUDA streams in GNN plugin #4012
Conversation
WalkthroughIntroduce a new Changes
Suggested Labels
Suggested Reviewers
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (2)
⏰ Context from checks skipped due to timeout of 90000ms (25)
🔇 Additional comments (6)
Finishing Touches
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🔭 Outside diff range comments (1)
Plugins/ExaTrkX/include/Acts/Plugins/ExaTrkX/TorchEdgeClassifier.hpp (1)
Line range hint
31-39: In Config struct, deviceID parameter remains, yet ExecutionContext we now embrace.Confusing this dual configuration is. Choose one path we must:
- Remove deviceID from Config, letting ExecutionContext handle all device management
- Document clearly how deviceID and ExecutionContext interact, we should
The Force suggests option 1, for clarity it brings.
struct Config { std::string modelPath; std::vector<int> selectedFeatures = {}; float cut = 0.21; int nChunks = 1; // NOTE for GNN use 1 bool undirected = false; - int deviceID = 0; bool useEdgeFeatures = false; };Also applies to: 45-45
🧹 Nitpick comments (5)
Plugins/ExaTrkX/include/Acts/Plugins/ExaTrkX/TorchMetricLearning.hpp (1)
Line range hint
35-36: Update the documentation for device handling, we must.The comment "default is the first GPU if available" outdated it now is. With ExecutionContext we work, yet old GPU assumptions in comments remain. Update documentation to reflect new CUDA streams approach, we should.
int knnVal = 500; bool shuffleDirections = false; - int deviceID = 0; // default is the first GPU if available + // Device and stream configuration now handled through ExecutionContextAlso applies to: 47-47
Plugins/ExaTrkX/include/Acts/Plugins/ExaTrkX/OnnxMetricLearning.hpp (1)
Line range hint
38-38: A unified approach to CUDA streams, we need.Across all files, scattered device handling exists. To bring balance to the Force:
- Create a central ExecutionContext management utility
- Remove all direct device handling (device(), deviceID)
- Document CUDA stream best practices
- Add examples of ExecutionContext usage
Help with implementation of these suggestions, would you like?
Also applies to: 45-45, 47-47, 42-42
Plugins/ExaTrkX/include/Acts/Plugins/ExaTrkX/Stages.hpp (1)
24-28: Hmmmm, good the structure is, but documentation it needs!Well-designed the ExecutionContext structure is, with CPU as default device and optional CUDA stream. But document its purpose and usage, you should.
Add this documentation, you must:
/// Capture the context of the execution +/// @brief Encapsulates device and stream information for GPU-accelerated operations +/// @details The device defaults to CPU, and stream is optional for CUDA devices struct ExecutionContext { torch::Device device{torch::kCPU}; std::optional<c10::cuda::CUDAStream> stream; };Plugins/ExaTrkX/src/BoostTrackBuilding.cpp (1)
51-51: Use the Force of CUDA streams, you do not!ExecutionContext you receive, but its power you do not harness. Opportunity for parallel execution with CUDA streams, you miss.
Consider using CUDA streams for graph operations when available:
+ if (execContext.stream) { + // Use CUDA-aware boost graph operations with the stream + // TODO: Implement CUDA-accelerated graph components + }Plugins/ExaTrkX/src/TorchEdgeClassifier.cpp (1)
81-85: Reorganize the guards, we should.Though functional your implementation is, improve readability we can. Move device_guard declaration closer to its usage, we should.
- std::optional<c10::cuda::CUDAGuard> device_guard; std::optional<c10::cuda::CUDAStreamGuard> streamGuard; if (device.is_cuda()) { + std::optional<c10::cuda::CUDAGuard> device_guard{device.index()}; - device_guard.emplace(device.index()); streamGuard.emplace(execContext.stream.value()); }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
Plugins/ExaTrkX/include/Acts/Plugins/ExaTrkX/OnnxEdgeClassifier.hpp(1 hunks)Plugins/ExaTrkX/include/Acts/Plugins/ExaTrkX/OnnxMetricLearning.hpp(1 hunks)Plugins/ExaTrkX/include/Acts/Plugins/ExaTrkX/Stages.hpp(4 hunks)Plugins/ExaTrkX/include/Acts/Plugins/ExaTrkX/TorchEdgeClassifier.hpp(1 hunks)Plugins/ExaTrkX/include/Acts/Plugins/ExaTrkX/TorchMetricLearning.hpp(1 hunks)Plugins/ExaTrkX/src/BoostTrackBuilding.cpp(1 hunks)Plugins/ExaTrkX/src/ExaTrkXPipeline.cpp(3 hunks)Plugins/ExaTrkX/src/OnnxEdgeClassifier.cpp(1 hunks)Plugins/ExaTrkX/src/TorchEdgeClassifier.cpp(1 hunks)Plugins/ExaTrkX/src/TorchMetricLearning.cpp(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (5)
- GitHub Check: merge-sentinel
- GitHub Check: linux_ubuntu_extra (ubuntu2204_clang, 20)
- GitHub Check: linux_ubuntu_extra (ubuntu2204, 20)
- GitHub Check: linux_ubuntu
- GitHub Check: build_debug
🔇 Additional comments (3)
Plugins/ExaTrkX/include/Acts/Plugins/ExaTrkX/OnnxEdgeClassifier.hpp (1)
38-38: Hmmmm, inconsistent these methods are.Device accessor method you keep, while ExecutionContext you now use. Confusion this may bring, when streams and device information from different sources flow. Consider deprecating the device() method, you should.
Run this command to check usage of device() method, we must:
Also applies to: 40-41
Plugins/ExaTrkX/include/Acts/Plugins/ExaTrkX/Stages.hpp (1)
44-49: Approve this change, I do!Wise decision it is, to use ExecutionContext with default value. Consistent with the Force, this change is.
Plugins/ExaTrkX/src/TorchEdgeClassifier.cpp (1)
70-72: Wise implementation of CUDA streams, this is.Properly managed, the CUDA streams are. Good practices you follow, with ExecutionContext and CUDAStreamGuard.
Also applies to: 81-81, 84-84
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
Plugins/ExaTrkX/include/Acts/Plugins/ExaTrkX/BoostTrackBuilding.hpp(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (18)
- GitHub Check: merge-sentinel
- GitHub Check: CI Bridge / lcg_106a: [alma9, gcc13]
- GitHub Check: CI Bridge / lcg_105: [alma9, clang16]
- GitHub Check: unused_files
- GitHub Check: CI Bridge / lcg_105: [alma9, gcc13]
- GitHub Check: CI Bridge / clang_tidy
- GitHub Check: CI Bridge / linux_ubuntu_2204
- GitHub Check: CI Bridge / build_exatrkx
- GitHub Check: CI Bridge / linux_ubuntu_2204_clang
- GitHub Check: CI Bridge / clang_tidy
- GitHub Check: CI Bridge / build_exatrkx_cpu
- GitHub Check: missing_includes
- GitHub Check: CI Bridge / build_linux_ubuntu
- GitHub Check: CI Bridge / build_linux_ubuntu
- GitHub Check: CI Bridge / build_exatrkx
- GitHub Check: linux_ubuntu_extra (ubuntu2204, 20)
- GitHub Check: linux_ubuntu
- GitHub Check: docs
🔇 Additional comments (1)
Plugins/ExaTrkX/include/Acts/Plugins/ExaTrkX/BoostTrackBuilding.hpp (1)
28-28: Verify the impact of ExecutionContext changes across the codebase, we must.Throughout the ExaTrkX plugin, ripples of this change may spread. Ensure consistent usage of ExecutionContext in derived classes and calling code, we should.
Run this script to verify the changes:
✅ Verification successful
Verified, the ExecutionContext changes are. Consistent with the Force, they remain.
Throughout the ExaTrkX plugin, a harmonious pattern we observe. In all components - TorchMetricLearning, TorchEdgeClassifier, OnnxMetricLearning, OnnxEdgeClassifier, and BoostTrackBuilding - the same ExecutionContext usage flows. The change in BoostTrackBuilding.hpp, with the codebase's established patterns, aligns perfectly.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Search for any remaining direct torch::Device usage in the ExaTrkX plugin echo "Searching for torch::Device usage in ExaTrkX plugin..." rg "torch::Device" -A 5 "Plugins/ExaTrkX" # Search for ExecutionContext usage echo "Verifying ExecutionContext adoption..." rg "ExecutionContext" -A 5 "Plugins/ExaTrkX"Length of output: 22218
|
Uses the torch framework (as it is required anyways currently) as a source of cuda streams. Extends the interface of the some components to use the streams.
--- END COMMIT MESSAGE ---
Any further description goes here, @-mentions are ok here!
feat,fix,refactor,docs,choreandbuildtypes.Summary by CodeRabbit
New Features
ExecutionContextstructure to manage device and stream information across ExaTrkX plugin components.Refactor
ExecutionContextinstead of directtorch::Deviceparameters.Chores