Merge pull request #294 from fixstars/fix/graph-cuda-interop

fix graph jit context
fixstars · Jul 8, 2024 · 9c57b7f · 9c57b7f
2 parents 18be51f + 3869542
commit 9c57b7f
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/include/ion/graph.h b/include/ion/graph.h
@@ -32,6 +32,11 @@ class Graph {
      * Run the pipeline immediately.
      */
     void run();
+    /**
+     * Set the user context which will be applied the pipeline built with this graph.
+     * @arg user_context_ptr: The pointer to the user context.
+     */
+    Graph& set_jit_context(Halide::JITUserContext *user_context_ptr);
 
     /**
      * Get the node list.

diff --git a/src/graph.cc b/src/graph.cc
@@ -54,6 +54,12 @@ Node Graph::add(const std::string& name)
     return n;
 }
 
+Graph& Graph::set_jit_context(Halide::JITUserContext *user_context_ptr) {
+    impl_->jit_ctx_ptr = user_context_ptr;
+    return *this;
+}
+
+
 void Graph::run()
 {
      if (!impl_->pipeline.defined()) {