Env variable to copy literals with memory coloring

src/targets/gpu/hip.cpp

@@ -90,9 +90,15 @@ argument from_gpu(argument arg)
 
 void gpu_sync() { hipDeviceSynchronize(); }
 
-void copy_to_gpu(char* dst, const char* src, std::size_t size)
+void copy_to_gpu(argument src, argument dst)
 {
-    hipMemcpy(dst, src, size, hipMemcpyHostToDevice);
+    std::size_t src_size = src.get_shape().bytes();
+    std::size_t dst_size = dst.get_shape().bytes();
+    if(src_size > dst_size)
+        MIGRAPH_THROW("Not enough memory available in destination to do copy");
+    auto status = hipMemcpy(dst.data(), src.data(), src_size, hipMemcpyHostToDevice);
+    if(status != hipSuccess)
+        MIGRAPH_THROW("Copy to gpu failed: " + hip_error(status));
 }
 } // namespace gpu
 } // namespace migraph

src/targets/gpu/include/migraph/gpu/hip.hpp

@@ -15,7 +15,7 @@ migraph::argument from_gpu(migraph::argument arg);
 
 void gpu_sync();
 
-void copy_to_gpu(char* dst, const char* src, std::size_t size);
+void copy_to_gpu(argument src, argument dst);
 
 struct hip_allocate
 {
@@ -67,19 +67,19 @@ struct hip_write
     }
 };
 
-struct hip_memcpy
+struct hip_copy
 {
-    std::string name() const { return "hip_memcpy"; }
-    shape compute_shape(std::vector<shape> inputs) const { return inputs.at(1); }
-    argument compute(context&, shape output_shape, std::vector<argument> args) const
+    std::string name() const { return "hip_copy"; }
+    shape compute_shape(std::vector<shape> inputs) const 
+    { 
+        check_shapes{inputs}.has(2);
+        return inputs.at(1); 
+    }
+    argument compute(context&, const shape&, std::vector<argument> args) const
     {
-        char* dst        = args.at(0).data() + offset;
-        const char* src  = args.at(1).data();
-        std::size_t size = args.at(1).get_shape().bytes();
-        copy_to_gpu(dst, src, size);
-        return {std::move(output_shape), dst};
+        copy_to_gpu(args[0], args[1]);
+        return args[1];
     }
-    std::size_t offset = 0;
 };
 } // namespace gpu
 } // namespace migraph

src/targets/gpu/write_literals.cpp

@@ -2,11 +2,14 @@
 #include <migraph/iterator_for.hpp>
 #include <migraph/gpu/hip.hpp>
 #include <migraph/instruction.hpp>
+#include <migraph/env.hpp>
 
 namespace migraph {
 
 namespace gpu {
 
+MIGRAPH_DECLARE_ENV_VAR(MIGRAPH_COPY_LITERALS)
+
 struct hip_load_literal
 {
     shape s;
@@ -30,10 +33,21 @@ void write_literals::apply(program& p) const
     {
         if(ins->name() == "@literal")
         {
-            argument a    = to_gpu(ins->get_literal().get_argument());
-            std::size_t n = ctx->literals.size();
-            ctx->literals.push_back(a);
-            p.replace_instruction(ins, hip_load_literal{a.get_shape(), n});
+            if(enabled(MIGRAPH_COPY_LITERALS{}))
+            {
+                literal l = ins->get_literal();
+                auto pre  = p.add_literal(l);
+                auto s = p.add_outline(l.get_shape());
+                auto alloc = p.insert_instruction(std::next(pre), hip_allocate{}, s);
+                p.replace_instruction(ins, hip_copy{}, pre, alloc);
+            } 
+            else
+            {
+                argument a    = to_gpu(ins->get_literal().get_argument());
+                std::size_t n = ctx->literals.size();
+                ctx->literals.push_back(a);
+                p.replace_instruction(ins, hip_load_literal{a.get_shape(), n});
+            }
         }
     }
 }

test/gpu/literal.cpp

 #include <test.hpp>
 #include <basic_ops.hpp>
 #include <migraph/program.hpp>
+#include <migraph/instruction.hpp>
 #include <migraph/generate.hpp>
 #include <migraph/gpu/target.hpp>
 #include <migraph/gpu/hip.hpp>
@@ -11,8 +12,13 @@ void gpu_literal_test()
     auto lit = generate_literal(migraph::shape{migraph::shape::float_type, {4, 3, 3, 3}});
     p.add_literal(lit);
     p.compile(migraph::gpu::target{});
+    auto scratch = p.get_parameter("scratch");
+    if(scratch == p.end()) {
     auto result = p.eval({});
     EXPECT(lit == migraph::gpu::from_gpu(result));
+    } else {
+        EXPECT(scratch->get_shape().bytes() == lit.get_shape().bytes());
+    }
 }
 
 int main() { gpu_literal_test(); }