[AIMIGRAPHX-885][AIMIGRAPHX-987] Use External Stream Contexts

src/include/migraphx/context.hpp

@@ -67,6 +67,16 @@ any_ptr get_queue_context(T&)
     return {};
 }
 
+template <class T>
+void set_queue_context(T&, any_ptr)
+{
+}
+
+template <class T>
+void restore_queue_context(T&)
+{
+}
+
 template <class T>
 void wait_for_context(T&, any_ptr)
 {
@@ -89,6 +99,10 @@ struct MIGRAPHX_EXPORT context
     // (optional)
     any_ptr get_queue();
     // (optional)
+    void set_queue(any_ptr queue);
+    // (optional)
+    void restore_queue();
+    // (optional)
     void wait_for(any_ptr queue);
     // (optional)
     void finish_on(any_ptr queue);
@@ -142,6 +156,33 @@ struct context
         return get_queue_context(private_detail_te_self);
     }
 
+    template <class T>
+    static auto private_detail_te_default_set_queue(char, T&& private_detail_te_self, any_ptr queue)
+        -> decltype(private_detail_te_self.set_queue(queue))
+    {
+        private_detail_te_self.set_queue(queue);
+    }
+
+    template <class T>
+    static void
+    private_detail_te_default_set_queue(float, T&& private_detail_te_self, any_ptr queue)
+    {
+        set_queue_context(private_detail_te_self, queue);
+    }
+
+    template <class T>
+    static auto private_detail_te_default_restore_queue(char, T&& private_detail_te_self)
+        -> decltype(private_detail_te_self.restore_queue())
+    {
+        private_detail_te_self.restore_queue();
+    }
+
+    template <class T>
+    static void private_detail_te_default_restore_queue(float, T&& private_detail_te_self)
+    {
+        restore_queue_context(private_detail_te_self);
+    }
+
     template <class T>
     static auto private_detail_te_default_wait_for(char, T&& private_detail_te_self, any_ptr queue)
         -> decltype(private_detail_te_self.wait_for(queue))
@@ -192,6 +233,10 @@ struct context
                                                       std::declval<const value&>()),
                  private_detail_te_default_get_queue(char(0),
                                                      std::declval<PrivateDetailTypeErasedT>()),
+                 private_detail_te_default_set_queue(
+                     char(0), std::declval<PrivateDetailTypeErasedT>(), std::declval<any_ptr>()),
+                 private_detail_te_default_restore_queue(char(0),
+                                                         std::declval<PrivateDetailTypeErasedT>()),
                  private_detail_te_default_wait_for(
                      char(0), std::declval<PrivateDetailTypeErasedT>(), std::declval<any_ptr>()),
                  private_detail_te_default_finish_on(
@@ -289,6 +334,18 @@ struct context
         return (*this).private_detail_te_get_handle().get_queue();
     }
 
+    void set_queue(any_ptr queue)
+    {
+        assert((*this).private_detail_te_handle_mem_var);
+        (*this).private_detail_te_get_handle().set_queue(queue);
+    }
+
+    void restore_queue()
+    {
+        assert((*this).private_detail_te_handle_mem_var);
+        (*this).private_detail_te_get_handle().restore_queue();
+    }
+
     void wait_for(any_ptr queue)
     {
         assert((*this).private_detail_te_handle_mem_var);
@@ -323,6 +380,8 @@ struct context
         virtual value to_value() const          = 0;
         virtual void from_value(const value& v) = 0;
         virtual any_ptr get_queue()             = 0;
+        virtual void set_queue(any_ptr queue)   = 0;
+        virtual void restore_queue()            = 0;
         virtual void wait_for(any_ptr queue)    = 0;
         virtual void finish_on(any_ptr queue)   = 0;
         virtual void finish() const             = 0;
@@ -373,6 +432,18 @@ struct context
             return private_detail_te_default_get_queue(char(0), private_detail_te_value);
         }
 
+        void set_queue(any_ptr queue) override
+        {
+
+            private_detail_te_default_set_queue(char(0), private_detail_te_value, queue);
+        }
+
+        void restore_queue() override
+        {
+
+            private_detail_te_default_restore_queue(char(0), private_detail_te_value);
+        }
+
         void wait_for(any_ptr queue) override
         {
 

src/program.cpp

@@ -644,7 +644,7 @@ std::vector<argument> program::eval(const parameter_map& params,
     if(exec_env.async)
     {
         assert(contexts.size() == 1);
-        contexts.front().wait_for(exec_env.queue);
+        contexts.front().set_queue(exec_env.queue);
     }
 
     // When MIGRAPHX_TRACE_EVAL is set, overwrite any user-provided trace callback with our trace
@@ -697,7 +697,7 @@ std::vector<argument> program::eval(const parameter_map& params,
     if(exec_env.async)
     {
         assert(contexts.size() == 1);
-        contexts.front().finish_on(exec_env.queue);
+        contexts.front().restore_queue();
     }
 
     return ret;

src/targets/gpu/include/migraphx/gpu/context.hpp

@@ -42,6 +42,7 @@
 #include <migraphx/gpu/hsa_chiplet.hpp>
 #include <unordered_map>
 #include <memory>
+#include <optional>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -90,6 +91,8 @@ struct hip_device
 
         hipStream_t get()
         {
+            if(external_stream.has_value())
+                return external_stream.value();
             if(not enabled(MIGRAPHX_ENABLE_NULL_STREAM{}))
             {
                 setup();
@@ -144,12 +147,39 @@ struct hip_device
         }
 #endif
 
+        void set_raw_stream(hipStream_t raw_stream)
+        {
+#if MIGRAPHX_USE_MIOPEN
+            if(mihandle != nullptr)
+                miopenSetStream(mihandle.get(), raw_stream);
+#endif
+#if MIGRAPHX_USE_ROCBLAS
+            if(rbhandle != nullptr)
+                rocblas_set_stream(rbhandle.get(), raw_stream);
+#endif
+        }
+
+        bool has_external_stream() const { return external_stream.has_value(); }
+
+        void set_queue(hipStream_t q)
+        {
+            external_stream = q;
+            set_raw_stream(q);
+        }
+
+        void restore_queue()
+        {
+            external_stream.reset();
+            set_raw_stream(s.get());
+        }
+
         void wait() const
         {
-            if(s == nullptr)
+            hipStream_t cur = external_stream.value_or(s.get());
+            if(cur == nullptr)
                 return;
             setup();
-            auto status = hipStreamSynchronize(s.get());
+            auto status = hipStreamSynchronize(cur);
             if(status != hipSuccess)
                 MIGRAPHX_THROW("Failed to wait: " + hip_error(status));
         }
@@ -173,6 +203,8 @@ struct hip_device
         private:
         std::size_t id           = 0;
         shared<hip_stream_ptr> s = nullptr;
+        std::optional<hipStream_t> external_stream{};
+
 #if MIGRAPHX_USE_MIOPEN
         shared<miopen_handle> mihandle = nullptr;
 #endif
@@ -334,25 +366,45 @@ struct context
         this->current_device = std::make_shared<hip_device>(device, n_streams);
     }
 
+    // Pure event-based synchronization point.  Records an event on the
+    // caller's queue and makes the context's current stream wait on it.
     void wait_for(any_ptr queue)
     {
         auto status = hipEventRecord(begin_event.get(), queue.get<hipStream_t>());
         if(status != hipSuccess)
             MIGRAPHX_THROW("Failed to record: " + hip_error(status));
-
         get_stream().wait(begin_event.get());
     }
 
+    // Symmetric counterpart of wait_for().  Records an event on the context's
+    // current stream and makes the caller's queue wait on it.
     void finish_on(any_ptr queue)
     {
         get_stream().record(finish_event.get());
-
         auto status = hipStreamWaitEvent(queue.get<hipStream_t>(), finish_event.get(), 0);
         if(status != hipSuccess)
             MIGRAPHX_THROW("Failed to wait on event: " + hip_error(status));
     }
 
-    any_ptr get_queue() { return get_stream().get(); }
+    any_ptr get_queue()
+    {
+        auto* s = get_stream().get();
+        return s == nullptr ? any_ptr{} : any_ptr{s};
+    }
+
+    // Bind a caller-provided queue for subsequent submissions.
+    // Passing an empty / null any_ptr is equivalent to binding the HIP
+    // default stream (nullptr), which is a distinct, valid operation from
+    // restore_queue() -- the two must not be conflated.  We bypass the typed
+    // any_ptr accessor when the pointer is null because a default-constructed
+    // any_ptr carries no type name and would otherwise throw on get<>().
+    void set_queue(any_ptr queue)
+    {
+        hipStream_t s = queue.unsafe_get() == nullptr ? nullptr : queue.get<hipStream_t>();
+        get_stream().set_queue(s);
+    }
+
+    void restore_queue() { get_stream().restore_queue(); }
 
     std::pair<hipEvent_t, hipEvent_t> get_perf_events() const
     {

test/eval_test.cpp

@@ -28,6 +28,8 @@
 #include <migraphx/stringutils.hpp>
 #include <migraphx/compile_options.hpp>
 #include <migraphx/make_op.hpp>
+#include <migraphx/execution_environment.hpp>
+#include <migraphx/any_ptr.hpp>
 #include <sstream>
 #include "test.hpp"
 #include <basic_ops.hpp>
@@ -151,6 +153,24 @@ struct double_invert_target
     migraphx::context get_context() const { return {}; }
 };
 
+// Minimal context that implements the optional set_queue/restore_queue members
+// of the context concept
+// Each call bumps a counter so a test can verify the dispatch routed through.
+struct tracked_ctx
+{
+    int set_calls     = 0;
+    int restore_calls = 0;
+    migraphx::any_ptr last_queue{};
+
+    void finish() const {}
+    void set_queue(migraphx::any_ptr q)
+    {
+        ++set_calls;
+        last_queue = q;
+    }
+    void restore_queue() { ++restore_calls; }
+};
+
 TEST_CASE(literal_test1)
 {
     migraphx::program p;
@@ -645,4 +665,57 @@ TEST_CASE(eval_trace_with_target_test)
     EXPECT(not fired_ops.empty());
 }
 
+TEST_CASE(async_eval_on_cpu_target_invokes_set_and_restore_queue)
+{
+    // The async branches of program::eval() call contexts.front().set_queue()
+    // before generic_eval and contexts.front().restore_queue() after.  id_target
+    // wraps an id_target::context that has no set_queue/restore_queue members,
+    // so this exercises the type-erased facade end-to-end on a non-GPU build:
+    //   - program::eval async prologue + epilogue (set_queue / restore_queue)
+    //   - context::set_queue / context::restore_queue facade bodies
+    //   - the float-overload (no-member) dispatchers
+    //   - the default set_queue_context / restore_queue_context free-function
+    //     fallbacks
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    auto one = mm->add_literal(1);
+    auto two = mm->add_literal(2);
+    mm->add_instruction(migraphx::make_op("add"), one, two);
+    p.compile(id_target{});
+
+    int dummy = 0;
+    migraphx::execution_environment exec_env;
+    exec_env.queue = migraphx::any_ptr{&dummy};
+    exec_env.async = true;
+
+    auto result = p.eval({}, exec_env).back();
+    EXPECT(result == migraphx::literal{3});
+
+    // A default-constructed any_ptr is a legal exec_env.queue and must also
+    // round-trip through set_queue / restore_queue without throwing.
+    migraphx::execution_environment exec_env_null;
+    exec_env_null.async = true;
+    auto result2        = p.eval({}, exec_env_null).back();
+    EXPECT(result2 == migraphx::literal{3});
+}
+
+TEST_CASE(context_facade_dispatches_to_member_set_and_restore_queue)
+{
+    // Sister test of the one above: tracked_ctx *does* implement set_queue and
+    // restore_queue.
+    migraphx::context ctx{tracked_ctx{}};
+
+    int dummy = 0;
+    migraphx::any_ptr q{&dummy};
+    ctx.set_queue(q);
+    ctx.set_queue(q);
+    ctx.restore_queue();
+
+    auto* held = ctx.any_cast<tracked_ctx>();
+    EXPECT(held != nullptr);
+    EXPECT(held->set_calls == 2);
+    EXPECT(held->restore_calls == 1);
+    EXPECT(held->last_queue.unsafe_get() == &dummy);
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/gpu/context_serialize.cpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2025 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -51,7 +51,10 @@ TEST_CASE(gpu_context_serialize)
 TEST_CASE(context_queue)
 {
     migraphx::context ctx = migraphx::gpu::context{0, 3};
-    EXPECT(ctx.get_queue().get<hipStream_t>() != nullptr);
+    // unsafe_get() avoids a type-mismatch throw if MIGRAPHX_ENABLE_NULL_STREAM
+    // is set: context::get_queue() returns an untyped any_ptr when the bound
+    // stream is nullptr, so get<hipStream_t>() would not yield "false".
+    EXPECT(ctx.get_queue().unsafe_get() != nullptr);
 }
 
 int main(int argc, const char* argv[]) { test::run(argc, argv); }