[IPython] [opengl] Make OpenGL kernel print work in IPython & IDLE (t…

…aichi-dev#1303) * [skip ci] add ti.imdisplay for IPython users * merge master * not print until ti.sync(); fix imcook bug * [skip ci] revert unrelated * [skip ci] Update docs/syntax.rst Co-authored-by: Yuanming Hu <[email protected]> Co-authored-by: Yuanming Hu <[email protected]>
archibate · Jun 27, 2020 · 22dc529 · 22dc529
1 parent 7862c09
commit 22dc529
Show file tree

Hide file tree

Showing 6 changed files with 45 additions and 9 deletions.
diff --git a/docs/syntax.rst b/docs/syntax.rst
@@ -226,16 +226,14 @@ Debug your program with ``print()`` in Taichi-scope. For example:
         print('v is', v)
         #=> v is [3, 4]
 
-.. note::
-
-    For now, print is only supported on CPU, CUDA and OpenGL backends.
+.. warning::
 
-    For the CUDA backend, the printed result won't shows up until ``ti.sync()``:
+    General speaking, the printed result won't shows up until ``ti.sync()``:
 
     .. code-block:: python
 
         import taichi as ti
-        ti.init(arch=ti.cuda)
+        ti.init(arch=ti.gpu)
 
         @ti.kernel
         def kern():
@@ -254,6 +252,9 @@ Debug your program with ``print()`` in Taichi-scope. For example:
         before kernel
         after kernel
         inside kernel
-        after
+        after sync
+
+    This is because GPU memory is only copied when necessary, for performance considerations.
 
-    Also note that host access or program end will also implicitly invoke for ``ti.sync()``.
+    Also note that host access or program end will also implicitly invoke for
+    ``ti.sync()``.
diff --git a/python/taichi/__init__.py b/python/taichi/__init__.py
@@ -7,7 +7,7 @@
 from taichi.misc import *
 from taichi.misc.gui import GUI
 from taichi.misc.np2ply import PLYWriter
-from taichi.misc.image import imread, imwrite, imshow
+from taichi.misc.image import *
 from taichi.misc.task import Task
 from taichi.misc.test import *
 from taichi.misc import settings as settings

diff --git a/python/taichi/lang/impl.py b/python/taichi/lang/impl.py
@@ -198,6 +198,9 @@ def get_tape(self, loss=None):
     def sync(self):
         self.materialize()
         self.prog.synchronize()
+        # print's in kernel won't take effect until ti.sync(), discussion:
+        # https://github.com/taichi-dev/taichi/pull/1303#discussion_r444897102
+        print(taichi_lang_core.pop_python_print_buffer(), end='')
 
 
 pytaichi = PyTaichi()

diff --git a/taichi/backends/opengl/opengl_api.cpp b/taichi/backends/opengl/opengl_api.cpp
@@ -5,6 +5,7 @@
 #include "taichi/program/kernel.h"
 #include "taichi/program/program.h"
 #include "taichi/util/environ_config.h"
+#include "taichi/python/print_buffer.h"
 #include "taichi/backends/opengl/runtime.h"
 
 #ifdef TI_WITH_OPENGL
@@ -487,7 +488,7 @@ struct CompiledProgram::Impl {
             TI_WARN("[glsl] Unexpected serialization type: {}, ignoring", type);
             break;
         };
-        std::cout << str;
+        py_cout << str;
       }
     }
     rt_buf->msg_count = 0;

diff --git a/taichi/python/export_misc.cpp b/taichi/python/export_misc.cpp
@@ -7,6 +7,7 @@
 #include "taichi/common/task.h"
 #include "taichi/math/math.h"
 #include "taichi/python/exception.h"
+#include "taichi/python/print_buffer.h"
 #include "taichi/python/export.h"
 #include "taichi/system/benchmark.h"
 #include "taichi/system/profiler.h"
@@ -20,6 +21,8 @@
 
 TI_NAMESPACE_BEGIN
 
+PythonPrintBuffer py_cout;
+
 Config config_from_py_dict(py::dict &c) {
   Config config;
   for (auto item : c) {
@@ -153,6 +156,9 @@ void export_misc(py::module &m) {
     }
     printf("test was successful.\n");
   });
+  m.def("pop_python_print_buffer", []() {
+      return py_cout.pop_content();
+  });
   m.def("with_cuda", is_cuda_api_available);
   m.def("with_metal", taichi::lang::metal::is_metal_api_available);
   m.def("with_opengl", taichi::lang::opengl::is_opengl_api_available);

diff --git a/taichi/python/print_buffer.h b/taichi/python/print_buffer.h
@@ -0,0 +1,25 @@
+#pragma once
+
+#include <sstream>
+
+TI_NAMESPACE_BEGIN
+
+struct PythonPrintBuffer {
+  /* holds kernel print result before switching back to python */
+  std::stringstream ss;
+
+  template <typename T>
+  PythonPrintBuffer &operator<<(const T &t) {
+    ss << t;
+    return *this;
+  }
+  std::string pop_content() {
+    auto ret = ss.str();
+    ss = std::stringstream();
+    return ret;
+  }
+};
+
+extern PythonPrintBuffer py_cout;
+
+TI_NAMESPACE_END