Add get_queue to context to get the current stream (#1097)

wrapped in a any_ptr class so the type can be checked at runtime for a mismatch.

src/include/migraphx/any_ptr.hpp

+#ifndef MIGRAPHX_GUARD_MIGRAPHX_ANY_PTR_HPP
+#define MIGRAPHX_GUARD_MIGRAPHX_ANY_PTR_HPP
+
+#include <migraphx/config.hpp>
+#include <migraphx/optional.hpp>
+#include <migraphx/errors.hpp>
+#include <migraphx/type_name.hpp>
+#include <cassert>
+#include <string_view>
+#include <typeindex>
+#include <type_traits>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+struct any_ptr
+{
+    any_ptr() = default;
+    template <class T>
+    any_ptr(T* p) : ptr(p), ti(typeid(T*)), name(get_name<T*>())
+    {
+    }
+
+    any_ptr(void* p, std::string_view pname) : ptr(p), name(pname) {}
+
+    void* get(std::string_view n) const
+    {
+        if(name != n)
+            MIGRAPHX_THROW("any_ptr: type mismatch: " + std::string{name} +
+                           " != " + std::string{n});
+        return ptr;
+    }
+
+    template <class T>
+    T get() const
+    {
+        static_assert(std::is_pointer<T>{}, "Must be a pointer");
+        assert(ptr != nullptr);
+        if(ti and std::type_index{typeid(T)} != *ti)
+            MIGRAPHX_THROW("any_ptr: type mismatch: " + std::string{name} + " != " + get_name<T>());
+        else if(name != get_name<T>())
+            MIGRAPHX_THROW("any_ptr: type mismatch: " + std::string{name} + " != " + get_name<T>());
+        return reinterpret_cast<T>(ptr);
+    }
+    void* unsafe_get() const { return ptr; }
+
+    private:
+    void* ptr                    = nullptr;
+    optional<std::type_index> ti = nullopt;
+    std::string_view name        = "";
+
+    template <class T>
+    static const std::string& get_name()
+    {
+        return get_type_name<std::remove_cv_t<std::remove_pointer_t<T>>>();
+    }
+};
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+#endif // MIGRAPHX_GUARD_MIGRAPHX_ANY_PTR_HPP

src/include/migraphx/context.hpp

@@ -9,6 +9,7 @@
 #include <utility>
 #include <migraphx/config.hpp>
 #include <migraphx/value.hpp>
+#include <migraphx/any_ptr.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -37,6 +38,12 @@ void from_value_context(T&, const value&)
 {
 }
 
+template <class T>
+any_ptr get_queue_context(T&)
+{
+    return {};
+}
+
 /*
  * Type-erased interface for:
  *
@@ -44,6 +51,7 @@ void from_value_context(T&, const value&)
  * {
  *      value to_value() const;
  *      void from_value(const value& v) ;
+ *      any_ptr get_queue() ;
  *      void finish() const;
  * };
  *
@@ -124,6 +132,12 @@ struct context
         (*this).private_detail_te_get_handle().from_value(v);
     }
 
+    any_ptr get_queue()
+    {
+        assert((*this).private_detail_te_handle_mem_var);
+        return (*this).private_detail_te_get_handle().get_queue();
+    }
+
     void finish() const
     {
         assert((*this).private_detail_te_handle_mem_var);
@@ -145,6 +159,7 @@ struct context
 
         virtual value to_value() const          = 0;
         virtual void from_value(const value& v) = 0;
+        virtual any_ptr get_queue()             = 0;
         virtual void finish() const             = 0;
     };
 
@@ -176,6 +191,19 @@ struct context
         from_value_context(private_detail_te_self, v);
     }
 
+    template <class T>
+    static auto private_detail_te_default_get_queue(char, T&& private_detail_te_self)
+        -> decltype(private_detail_te_self.get_queue())
+    {
+        return private_detail_te_self.get_queue();
+    }
+
+    template <class T>
+    static any_ptr private_detail_te_default_get_queue(float, T&& private_detail_te_self)
+    {
+        return get_queue_context(private_detail_te_self);
+    }
+
     template <typename PrivateDetailTypeErasedT>
     struct private_detail_te_handle_type : private_detail_te_handle_base_type
     {
@@ -216,6 +244,12 @@ struct context
             private_detail_te_default_from_value(char(0), private_detail_te_value, v);
         }
 
+        any_ptr get_queue() override
+        {
+
+            return private_detail_te_default_get_queue(char(0), private_detail_te_value);
+        }
+
         void finish() const override { private_detail_te_value.finish(); }
 
         PrivateDetailTypeErasedT private_detail_te_value;

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

@@ -235,6 +235,8 @@ struct context
         this->current_device = std::make_shared<hip_device>(0, n_streams);
     }
 
+    any_ptr get_queue() { return get_stream().get(); }
+
     private:
     // TODO: Make this a vector to support multiple devices
     std::shared_ptr<hip_device> current_device;

test/any_ptr.cpp

+#include <migraphx/any_ptr.hpp>
+#include <test.hpp>
+
+TEST_CASE(test_int_id)
+{
+    int i               = 1;
+    migraphx::any_ptr p = &i;
+    EXPECT(p.get<int*>() == &i);
+    EXPECT(p.get(migraphx::get_type_name(i)) == &i);
+    EXPECT(p.unsafe_get() == &i);
+    EXPECT(test::throws([&] { p.get<float*>(); }));
+    EXPECT(test::throws([&] { p.get(migraphx::get_type_name(&i)); }));
+}
+
+TEST_CASE(test_int_name)
+{
+    int i    = 1;
+    void* vp = &i;
+    migraphx::any_ptr p{vp, migraphx::get_type_name(i)};
+    EXPECT(p.get<int*>() == &i);
+    EXPECT(p.get(migraphx::get_type_name(i)) == &i);
+    EXPECT(p.unsafe_get() == &i);
+    EXPECT(test::throws([&] { p.get<float*>(); }));
+    EXPECT(test::throws([&] { p.get(migraphx::get_type_name(&i)); }));
+    EXPECT(test::throws([&] { p.get(migraphx::get_type_name(float{})); }));
+}
+
+int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/gpu/context_serialize.cpp

@@ -5,7 +5,7 @@
 #include <migraphx/context.hpp>
 #include "test.hpp"
 
-TEST_CASE(gpu_context)
+TEST_CASE(gpu_context_serialize)
 {
     migraphx::context ctx = migraphx::gpu::context{0, 3};
 
@@ -25,4 +25,10 @@ TEST_CASE(gpu_context)
     EXPECT(v == v1);
 }
 
+TEST_CASE(context_queue)
+{
+    migraphx::context ctx = migraphx::gpu::context{0, 3};
+    EXPECT(ctx.get_queue().get<hipStream_t>() != nullptr);
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }

tools/include/context.hpp

@@ -9,6 +9,7 @@
 #include <utility>
 #include <migraphx/config.hpp>
 #include <migraphx/value.hpp>
+#include <migraphx/any_ptr.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -33,12 +34,21 @@ value to_value_context(const T&)
 }
 
 template <class T>
-void from_value_context(T&, const value&){}
+void from_value_context(T&, const value&)
+{
+}
+
+template <class T>
+any_ptr get_queue_context(T&)
+{
+    return {};
+}
 
 <%
  interface('context',
            virtual('to_value', returns = 'value', const = True, default = 'to_value_context'),
            virtual('from_value', v = 'const value&', default = 'from_value_context'),
+           virtual('get_queue', returns = 'any_ptr', default = 'get_queue_context'),
            virtual('finish', returns = 'void', const = True)) %>
 
     inline void migraphx_to_value(value& v, const context& ctx)