Add an explicit graph node API

AbstractAPI.h

@@ -86,8 +86,46 @@ struct AbstractAPI {
 
   virtual bool isCapableOfGraphCapturing() = 0;
   virtual DeviceGraphHandle streamBeginCapture(std::vector<void*>& streamPtrs) = 0;
-  virtual void streamEndCapture(DeviceGraphHandle handle) = 0;
-  virtual void launchGraph(DeviceGraphHandle graphHandle, void* streamPtr) = 0;
+  virtual void streamEndCapture(const DeviceGraphHandle& handle) = 0;
+  virtual void launchGraph(const DeviceGraphHandle& graphHandle, void* streamPtr) = 0;
+
+  /**
+   * Explicit graph construction.
+   *
+   * Instead of recording a whole stream and letting the backend infer the dependency structure
+   * from events, the caller states the structure directly: every graphAddNode call contributes
+   * the work recorded by `recorder` and makes it depend on exactly `dependencies`. Fork/join is
+   * then a property of the graph rather than something that has to be expressed through streams
+   * and events.
+   *
+   * A single graph is built by one thread at a time. `recorder` receives a stream that is only a
+   * recording vehicle: the stream carries no ordering information beyond the extent of that one
+   * call, and the same stream may be reused for sibling nodes.
+   *
+   * If `recorder` enqueues nothing, the returned handle refers to `dependencies` themselves, so
+   * an empty recorder is a valid way to express a pure join node.
+   *
+   * graphBeginNode and graphEndNode are the same thing split in two, for callers that cannot
+   * wrap the recorded work in a callback and instead have to leave a node open across code they
+   * do not control. Only one node per stream may be open at a time.
+   */
+  virtual bool isCapableOfGraphNodes() = 0;
+  virtual DeviceGraphHandle graphCreate() = 0;
+  virtual void graphBeginNode(const DeviceGraphHandle& graphHandle,
+                              const std::vector<DeviceGraphNodeHandle>& dependencies,
+                              void* streamPtr) = 0;
+  virtual DeviceGraphNodeHandle graphEndNode(const DeviceGraphHandle& graphHandle,
+                                             void* streamPtr) = 0;
+  virtual void graphInstantiate(const DeviceGraphHandle& graphHandle) = 0;
+
+  DeviceGraphNodeHandle graphAddNode(const DeviceGraphHandle& graphHandle,
+                                     const std::vector<DeviceGraphNodeHandle>& dependencies,
+                                     void* streamPtr,
+                                     const std::function<void(void*)>& recorder) {
+    graphBeginNode(graphHandle, dependencies, streamPtr);
+    recorder(streamPtr);
+    return graphEndNode(graphHandle, streamPtr);
+  }
 
   virtual void* createStream(double priority = NAN) = 0;
   virtual void destroyGenericStream(void* streamPtr) = 0;

DataTypes.h

@@ -7,28 +7,64 @@
 
 #include <cstddef>
 #include <limits>
+#include <memory>
 
 namespace device {
-struct DeviceGraphHandle {
-  static const size_t invalidId{std::numeric_limits<size_t>::max()};
 
-  public:
-  explicit DeviceGraphHandle() : graphId(invalidId) {}
-  explicit DeviceGraphHandle(size_t id) : graphId(id) {}
+/**
+ * Backend-specific payload of a compute graph. Only the active interface implementation defines
+ * this type; every other translation unit sees an incomplete type and reaches the graph through
+ * DeviceGraphHandle.
+ */
+struct DeviceGraph;
 
-  DeviceGraphHandle(const DeviceGraphHandle& other) = default;
-  DeviceGraphHandle& operator=(const DeviceGraphHandle& other) = default;
+/**
+ * Owning handle to a compute graph.
+ *
+ * The backend resources (the graph and its executable instance) are released once the last handle
+ * pointing to them goes out of scope. A graph that is dropped from a cache therefore also frees
+ * its device-side resources.
+ */
+class DeviceGraphHandle {
+  public:
+  DeviceGraphHandle() = default;
+  explicit DeviceGraphHandle(std::shared_ptr<DeviceGraph> graphPtr) : graph(std::move(graphPtr)) {}
 
-  bool isInitialized() const { return graphId != invalidId; }
+  [[nodiscard]] bool isInitialized() const { return static_cast<bool>(graph); }
 
   operator bool() const { return isInitialized(); }
 
   bool operator!() const { return !isInitialized(); }
 
-  size_t getGraphId() { return graphId; }
+  [[nodiscard]] DeviceGraph* get() const { return graph.get(); }
+
+  void reset() { graph.reset(); }
+
+  private:
+  std::shared_ptr<DeviceGraph> graph;
+};
+
+/**
+ * Refers to the set of graph nodes produced by a single AbstractAPI::graphAddNode call.
+ *
+ * A node handle is an index into the graph that produced it and stays valid for that graph's
+ * lifetime. Passing it to a different graph is undefined.
+ */
+class DeviceGraphNodeHandle {
+  public:
+  static const size_t invalidId{std::numeric_limits<size_t>::max()};
+
+  DeviceGraphNodeHandle() = default;
+  explicit DeviceGraphNodeHandle(size_t id) : nodeId(id) {}
+
+  [[nodiscard]] bool isInitialized() const { return nodeId != invalidId; }
+
+  operator bool() const { return isInitialized(); }
+
+  [[nodiscard]] size_t getNodeId() const { return nodeId; }
 
   private:
-  size_t graphId{invalidId};
+  size_t nodeId{invalidId};
 };
 } // namespace device
 

interfaces/cuda/CudaWrappedAPI.h

@@ -83,8 +83,17 @@ class ConcreteAPI : public AbstractAPI {
 
   bool isCapableOfGraphCapturing() override;
   DeviceGraphHandle streamBeginCapture(std::vector<void*>& streamPtrs) override;
-  void streamEndCapture(DeviceGraphHandle handle) override;
-  void launchGraph(DeviceGraphHandle graphHandle, void* streamPtr) override;
+  void streamEndCapture(const DeviceGraphHandle& handle) override;
+  void launchGraph(const DeviceGraphHandle& graphHandle, void* streamPtr) override;
+
+  bool isCapableOfGraphNodes() override;
+  DeviceGraphHandle graphCreate() override;
+  void graphBeginNode(const DeviceGraphHandle& graphHandle,
+                      const std::vector<DeviceGraphNodeHandle>& dependencies,
+                      void* streamPtr) override;
+  DeviceGraphNodeHandle graphEndNode(const DeviceGraphHandle& graphHandle,
+                                     void* streamPtr) override;
+  void graphInstantiate(const DeviceGraphHandle& graphHandle) override;
 
   void* createStream(double priority) override;
   void destroyGenericStream(void* streamPtr) override;
@@ -127,14 +136,6 @@ class ConcreteAPI : public AbstractAPI {
 
   std::unordered_set<cudaStream_t> genericStreams{};
 
-  struct GraphDetails {
-    cudaGraph_t graph;
-    cudaGraphExec_t instance;
-    std::vector<void*> streamPtrs;
-    bool ready{false};
-  };
-  std::vector<GraphDetails> graphs;
-
   Statistics statistics{};
   std::unordered_map<void*, size_t> memToSizeMap{{nullptr, 0}};