Remove get_argument_gpu function

src/targets/miopen/include/rtg/miopen/hip.hpp

@@ -18,15 +18,20 @@ inline hip_ptr gpu_allocate(std::size_t sz)
     return hip_ptr{result};
 }
 
+inline hip_ptr write_to_gpu(const void* x, std::size_t sz)
+{
+    auto result = gpu_allocate(sz);
+    // TODO: Check status
+    hipMemcpy(result.get(), x, sz, hipMemcpyHostToDevice);
+    return result;
+}
+
 template <class T>
 hip_ptr write_to_gpu(const T& x)
 {
     using type  = typename T::value_type;
     auto size   = x.size() * sizeof(type);
-    auto result = gpu_allocate(size);
-    // TODO: Check status
-    hipMemcpy(result.get(), x.data(), size, hipMemcpyHostToDevice);
-    return result;
+    return write_to_gpu(x.data(), size);
 }
 
 template <class T>
@@ -38,6 +43,23 @@ std::vector<T> read_from_gpu(const void* x, std::size_t sz)
     return result;
 }
 
+inline rtg::argument to_gpu(rtg::argument arg)
+{
+    auto p = share(write_to_gpu(arg.data(), arg.get_shape().bytes()));
+    return {arg.get_shape(), [p]() mutable { return reinterpret_cast<char*>(p.get()); }};
+}
+
+inline rtg::argument from_gpu(rtg::argument arg)
+{
+    rtg::argument result;
+    arg.visit([&](auto x) {
+        using type = typename decltype(x)::value_type;
+        auto v = read_from_gpu<type>(arg.data(), x.get_shape().bytes() / sizeof(type));
+        result     = {x.get_shape(), [v]() mutable { return reinterpret_cast<char*>(v.data()); }};
+    });
+    return result;
+}
+
 struct hip_allocate
 {
     std::string name() const { return "hip::allocate"; }

test/miopen/miopen.cpp

@@ -23,14 +23,6 @@ rtg::program create_program()
     return p;
 }
 
-// TODO: Move to header
-rtg::argument get_tensor_argument_gpu(rtg::shape s)
-{
-    auto v = rtg::generate_tensor_data<float>(s);
-    auto p = rtg::share(rtg::miopen::write_to_gpu(v));
-    return {s, [p]() mutable { return reinterpret_cast<char*>(p.get()); }};
-}
-
 std::vector<float> cpu()
 {
     std::vector<float> result;
@@ -48,10 +40,10 @@ std::vector<float> gpu()
 {
     std::vector<float> result;
     auto p = create_program();
-    auto x = get_tensor_argument_gpu({rtg::shape::float_type, {4, 3, 3, 3}});
-    auto w = get_tensor_argument_gpu({rtg::shape::float_type, {4, 3, 3, 3}});
+    auto x = rtg::miopen::to_gpu(rtg::generate_argument({rtg::shape::float_type, {4, 3, 3, 3}}));
+    auto w = rtg::miopen::to_gpu(rtg::generate_argument({rtg::shape::float_type, {4, 3, 3, 3}}));
     p.compile(rtg::miopen::miopen_target{});
-    auto y      = get_tensor_argument_gpu(p.get_parameter_shape("output"));
+    auto y      = rtg::miopen::to_gpu(rtg::generate_argument(p.get_parameter_shape("output")));
     auto handle = rtg::miopen::make_obj<rtg::miopen::miopen_handle>(&miopenCreate);
     auto r      = p.eval(
         {{"x", x}, {"w", w}, {"output", y}, {"handle", {rtg::shape::any_type, handle.get()}}});