polish dataset use jittor with cuda

Jittor · Apr 23, 2022 · cc6cb28 · cc6cb28
1 parent 5e559cf
commit cc6cb28
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 1 deletion.
diff --git a/python/jittor/__init__.py b/python/jittor/__init__.py
@@ -9,7 +9,7 @@
 # file 'LICENSE.txt', which is part of this source code package.
 # ***************************************************************
 
-__version__ = '1.3.3.5'
+__version__ = '1.3.3.6'
 from jittor_utils import lock
 with lock.lock_scope():
     ori_int = int

diff --git a/python/jittor/dataset/dataset.py b/python/jittor/dataset/dataset.py
@@ -365,6 +365,7 @@ def _stop_all_workers(self):
         self.gid_obj.value = 0
 
     def _init_workers(self, index_list):
+        jt.migrate_all_to_cpu()
         jt.clean()
         jt.gc()
         self.index_list = mp.Array('i', self.real_len, lock=False)

diff --git a/python/jittor/src/var_holder.cc b/python/jittor/src/var_holder.cc
@@ -240,5 +240,16 @@ VarHolder* ternary_out_hint(VarHolder* cond, VarHolder* x, VarHolder* y) {
     return new VarHolder(make_ternary(cond->var, x->var, y->var));
 }
 
+void migrate_all_to_cpu() {
+    sync_all(true);
+#ifdef HAS_CUDA
+    for (auto vh : hold_vars) {
+        auto v = vh->var;
+        if (v->_outputs.size()) continue;
+        if (v->allocator && !v->allocator->is_cuda())
+            migrate_to_gpu(v, cpu_allocator);
+    }
+#endif
+}
 
 } // jittor
diff --git a/python/jittor/src/var_holder.h b/python/jittor/src/var_holder.h
@@ -359,4 +359,7 @@ inline vector<VarHolder*> make_vh_vector(vector<VarPtr>&& vps) {
 // @pyjt(ternary_out_hint)
 VarHolder* ternary_out_hint(VarHolder* cond, VarHolder* x, VarHolder* y);
 
+// @pyjt(migrate_all_to_cpu)
+void migrate_all_to_cpu();
+
 } // jittor