Remove hip_free operator

src/targets/miopen/hip.cpp

 
 #include <rtg/miopen/hip.hpp>
 
+#include <rtg/manage_ptr.hpp>
+#include <miopen/miopen.h>
+
+#include <vector>
+
 namespace rtg {
 namespace miopen {
 
-hip_ptr gpu_allocate(std::size_t sz)
+using hip_ptr = RTG_MANAGE_PTR(void, hipFree);
+
+hip_ptr allocate_gpu(std::size_t sz)
 {
     void* result;
     // TODO: Check status
@@ -12,14 +19,37 @@ hip_ptr gpu_allocate(std::size_t sz)
     return hip_ptr{result};
 }
 
+template <class T>
+hip_ptr write_to_gpu(const T& x)
+{
+    using type = typename T::value_type;
+    auto size  = x.size() * sizeof(type);
+    return write_to_gpu(x.data(), size);
+}
+
+template <class T>
+std::vector<T> read_from_gpu(const void* x, std::size_t sz)
+{
+    std::vector<T> result(sz);
+    // TODO: Check status
+    hipMemcpy(result.data(), x, sz * sizeof(T), hipMemcpyDeviceToHost);
+    return result;
+}
+
 hip_ptr write_to_gpu(const void* x, std::size_t sz)
 {
-    auto result = gpu_allocate(sz);
+    auto result = allocate_gpu(sz);
     // TODO: Check status
     hipMemcpy(result.get(), x, sz, hipMemcpyHostToDevice);
     return result;
 }
 
+rtg::argument allocate_gpu(rtg::shape s)
+{
+    auto p = share(allocate_gpu(s.bytes()));
+    return {s, [p]() mutable { return reinterpret_cast<char*>(p.get()); }};
+}
+
 rtg::argument to_gpu(rtg::argument arg)
 {
     auto p = share(write_to_gpu(arg.data(), arg.get_shape().bytes()));

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

@@ -11,33 +11,12 @@
 namespace rtg {
 namespace miopen {
 
-using hip_ptr = RTG_MANAGE_PTR(void, hipFree);
-
-hip_ptr gpu_allocate(std::size_t sz);
-
-hip_ptr write_to_gpu(const void* x, std::size_t sz);
+rtg::argument allocate_gpu(rtg::shape s);
 
 rtg::argument to_gpu(rtg::argument arg);
 
 rtg::argument from_gpu(rtg::argument arg);
 
-template <class T>
-hip_ptr write_to_gpu(const T& x)
-{
-    using type = typename T::value_type;
-    auto size  = x.size() * sizeof(type);
-    return write_to_gpu(x.data(), size);
-}
-
-template <class T>
-std::vector<T> read_from_gpu(const void* x, std::size_t sz)
-{
-    std::vector<T> result(sz);
-    // TODO: Check status
-    hipMemcpy(result.data(), x, sz * sizeof(T), hipMemcpyDeviceToHost);
-    return result;
-}
-
 struct hip_allocate
 {
     std::string name() const { return "hip::allocate"; }
@@ -48,26 +27,7 @@ struct hip_allocate
     }
     argument compute(shape output_shape, std::vector<argument>) const
     {
-        char* data = nullptr;
-        // TODO: Check return status
-        hipMalloc(&data, output_shape.bytes());
-        return {output_shape, data};
-    }
-};
-
-struct hip_free
-{
-    std::string name() const { return "hip::free"; }
-    shape compute_shape(std::vector<shape> inputs) const
-    {
-        check_shapes{inputs}.has(1);
-        return {};
-    }
-    argument compute(shape, std::vector<argument> args) const
-    {
-        // TODO: Check return status
-        hipFree(args.front().data());
-        return {};
+        return allocate_gpu(output_shape);
     }
 };
 

src/targets/miopen/miopen_target.cpp

@@ -118,7 +118,6 @@ struct miopen_apply
         {
             auto is     = prog->add_outline(s);
             auto result = prog->insert_instruction(ins, hip_allocate{}, is);
-            prog->insert_instruction(++ins, hip_free{}, result);
             return result;
         }
     }