optimize ring buffer and copy free array op

Author: cjld

python/jittor/__init__.py

@@ -7,9 +7,12 @@
 # This file is subject to the terms and conditions defined in
 # file 'LICENSE.txt', which is part of this source code package.
 # ***************************************************************
-__version__ = '1.2.0.7'
+__version__ = '1.2.0.8'
 from . import lock
 with lock.lock_scope():
+    ori_int = int
+    ori_float = float
+    ori_bool = bool
     from . import compiler
     from .compiler import LOG, has_cuda
     from .compiler import compile_custom_ops, compile_custom_op
@@ -874,16 +877,13 @@ def to_float(v):
 def to_bool(v):
     dtype = str(v.dtype)
     assert dtype.startswith("int") or dtype=="bool"
-    return bool(v.item())
+    return ori_bool(v.item())
 
 Var.item = item
 Var.__int__ = to_int
 Var.__float__ = to_float
 Var.__bool__ = to_bool
 
-ori_int = int
-ori_float = float
-
 int = int32
 Var.int = Var.int32
 float = float32

python/jittor/dataset/dataset.py

@@ -15,8 +15,6 @@
 from collections.abc import Sequence, Mapping
 import pathlib
 from PIL import Image
-from jittor_utils import ring_buffer
-from jittor_utils.ring_buffer import RingBuffer
 import multiprocessing as mp
 import signal
 from jittor_utils import LOG
@@ -30,8 +28,8 @@
 
 class Worker:
     def __init__(self, target, args, buffer_size):
-        buffer = mp.Array('c', buffer_size, lock=False)
-        self.buffer = RingBuffer(buffer)
+        self.buffer = jt.RingBuffer(buffer_size)
+
         self.status = mp.Array('f', 5, lock=False)
         self.p = mp.Process(target=target, args=args+(self.buffer,self.status))
         self.p.daemon = True
@@ -253,13 +251,12 @@ class YourDataset(Dataset):
         msg.append(f"progress:{self.last_id}/{self.batch_len}")
         msg.append(f"batch(s): {self.batch_time:.3f}\twait(s):{self.wait_time:.3f}")
         msg.append(f"recv(s): {self.recv_time:.3f}\tto_jittor(s):{self.to_jittor_time:.3f}")
-        msg.append(f"recv_raw_call: {ring_buffer.recv_raw_call}")
         msg.append(f"last 10 workers: {self.idmap[max(0, self.last_id-9):self.last_id+1]}")
         msg.append(f"ID\twait(s)\topen(s)\tload(s)\tsend(s)\ttotal(s)")
         for i in range(self.num_workers):
             w = self.workers[i]
             s = w.status
-            msg.append(f"#{i}\t{s[0]:.3f}\t{s[4]:.3f}\t{s[1]:.3f}\t{s[2]:.3f}\t{s[3]:.3f}\t{w.buffer.allocator}")
+            msg.append(f"#{i}\t{s[0]:.3f}\t{s[4]:.3f}\t{s[1]:.3f}\t{s[2]:.3f}\t{s[3]:.3f}\t{w.buffer}")
         LOG.i('\n'.join(msg))
 
     def _stop_all_workers(self):

python/jittor/pyjt_compiler.py

@@ -31,8 +31,10 @@ def parse_attrs(s):
     "uint": ["PyLong_AsUnsignedLong", "PyLong_FromUnsignedLong", "PyLong_CheckExact"],
     "uint64": ["PyLong_AsUnsignedLongLong", "PyLong_FromUnsignedLongLong", "PyLong_CheckExact"],
     "void": ["...", "GET_PY_NONE", "..."],
+    "PyObject*": ["","",""],
 }
 def get_pytype_map(T, i):
+    assert T != ""
     if T in pytype_map:
         return pytype_map[T][i]
     return ["from_py_object", "to_py_object", "is_type"][i]+"<"+T+">"
@@ -204,7 +206,7 @@ def get_def_code(df, scope_name, pyname, self_as_arg0=False):
             func_call = f"(GET_RAW_PTR({scope_name},self))->" + func_call
     if pyname == "__init__":
         # XXX->xxx(...) ---> new XXX xxx(...)
-        assert "->" in func_call
+        assert "->" in func_call, func_call
         func_call = "new " + func_call.replace("->", " ")
     if no_need_convert:
         func_quick_check_runable = ""

python/jittor/test/test_reduce_op.py

@@ -39,7 +39,7 @@ def check(a, op, dims):
         idims = [(), (0,), (1,), (2,), (3,), (0, 2), (1,3), (1,2,3), 2, 3]
 
         iop = [ op[7:] for op in dir(jt) if op.startswith("reduce_")]
-        assert len(iop) >= 10
+        assert len(iop) >= 10, iop
         for a in ia:
             check(a, iop[0], idims[0])
         for op in iop:

python/jittor/test/test_relu.py

@@ -65,8 +65,9 @@ def test_relu(self):
         # ***************************************************************
         # Test GELU Layer
         # ***************************************************************
-        arr = np.random.randn(16,10,224,224)
-        check_equal(arr, jnn.GELU(), tnn.GELU())
+        if hasattr(tnn, "GELU"):
+            arr = np.random.randn(16,10,224,224)
+            check_equal(arr, jnn.GELU(), tnn.GELU())
 
         # ***************************************************************
         # Test Softplus  Layer

python/jittor/test/test_ring_buffer2.py

@@ -0,0 +1,86 @@
+# ***************************************************************
+# Copyright (c) 2020 Jittor. Authors: 
+#     Dun Liang <randonlang@gmail.com>. 
+# All Rights Reserved.
+# This file is subject to the terms and conditions defined in
+# file 'LICENSE.txt', which is part of this source code package.
+# ***************************************************************
+import jittor as jt
+import unittest
+import numpy as np
+import random
+from .test_core import expect_error
+from jittor.dataset.mnist import MNIST
+import jittor.transform as trans
+from tqdm import tqdm
+
+def test_ring_buffer():
+    buffer = jt.RingBuffer(1000)
+    def test_send_recv(data):
+        print("test send recv", type(data))
+        buffer.push(data)
+        recv = buffer.pop()
+        if isinstance(data, np.ndarray):
+            assert (recv == data).all()
+        else:
+            assert data == recv
+
+    n_byte = 0
+    test_send_recv(1)
+    n_byte += 1 + 8
+    assert n_byte == buffer.total_pop() and n_byte == buffer.total_push()
+    test_send_recv(100000000000)
+    n_byte += 1 + 8
+    assert n_byte == buffer.total_pop() and n_byte == buffer.total_push()
+
+    test_send_recv(1e-5)
+    n_byte += 1 + 8
+    assert n_byte == buffer.total_pop() and n_byte == buffer.total_push()
+    test_send_recv(100000000000.0)
+    n_byte += 1 + 8
+    assert n_byte == buffer.total_pop() and n_byte == buffer.total_push()
+
+    test_send_recv("float32")
+    n_byte += 1 + 8 + 7
+    assert n_byte == buffer.total_pop() and n_byte == buffer.total_push()
+    test_send_recv("")
+    n_byte += 1 + 8 + 0
+    assert n_byte == buffer.total_pop() and n_byte == buffer.total_push()
+    test_send_recv("xxxxxxxxxx")
+    n_byte += 1 + 8 + 10
+    assert n_byte == buffer.total_pop() and n_byte == buffer.total_push()
+
+    test_send_recv([1,0.2])
+    n_byte += 1 + 8 + 1 + 8 + 1 + 8
+    assert n_byte == buffer.total_pop() and n_byte == buffer.total_push()
+    test_send_recv({'asd':1})
+    n_byte += 1 + 8 + 1 + 8 + 3 + 1 + 8
+    assert n_byte == buffer.total_pop() and n_byte == buffer.total_push()
+
+    test_send_recv(np.random.rand(10,10))
+    n_byte += 1 + 16 + 2 + 10*10*8
+    assert n_byte == buffer.total_pop() and n_byte == buffer.total_push()
+    test_send_recv(test_ring_buffer)
+
+    expect_error(lambda: test_send_recv(np.random.rand(10,1000)))
+
+
+class TestRingBuffer(unittest.TestCase):
+
+    def test_ring_buffer(self):
+        test_ring_buffer()
+
+    def test_dataset(self):
+        return
+        self.train_loader = MNIST(train=True, transform=trans.Resize(224)) \
+            .set_attrs(batch_size=300, shuffle=True)
+        self.train_loader.num_workers = 1
+        for batch_idx, (data, target) in tqdm(enumerate(self.train_loader)):
+            # self.train_loader.display_worker_status()
+            if batch_idx > 30:
+                break
+            pass
+        
+
+if __name__ == "__main__":
+    unittest.main()

python/jittor/transform/__init__.py

@@ -213,7 +213,7 @@ def to_tensor(img):
         img_ = transform.to_tensor(img)
     """
     if isinstance(img, Image.Image):
-        return np.array(img).transpose((2,0,1)) / np.float32(255)
+        return np.array(img).transpose((2,0,1)) * np.float32(1.0/255.0)
     return img
 
 
@@ -323,7 +323,7 @@ def __call__(self, img):
         if isinstance(img, Image.Image):
             img = (np.array(img).transpose((2,0,1)) \
                 - self.mean*np.float32(255.)) \
-                / (self.std*np.float32(255.))
+                * (np.float32(1./255.)/self.std)
         else:
             img = (img - self.mean) / self.std
         return img

python/jittor_utils/ring_buffer.py

@@ -127,7 +127,7 @@ def str_to_char_array(s, array_len):
     return a
 
 def char_array_to_str(a):
-    return str(a.tostring(), 'ascii').strip()
+    return str(a.tobytes(), 'ascii').strip()
 
 class RingBuffer:
     def __init__(self, buffer):

src/misc/nano_vector.h

@@ -64,7 +64,7 @@ struct NanoVector {
     // @pyjt(__init__)
     inline NanoVector(const NanoVector& nv) : data(nv.data), offset(nv.offset) {}
 
-    void clear() { data = offset = 0; }
+    inline void clear() { data = offset = 0; }
 
     // @pyjt(__len__, __map_len__)
     inline int size() const {
@@ -158,10 +158,22 @@ struct NanoVector {
         for (auto a : v) push_back_check_overflow(a);
     }
 
+    inline static NanoVector make(const int64* v, int n) {
+        NanoVector nv;
+        for (int i=0; i<n; i++) nv.push_back_check_overflow(v[i]);
+        return nv;
+    }
+
+    inline static NanoVector make(const int32* v, int n) {
+        NanoVector nv;
+        for (int i=0; i<n; i++) nv.push_back_check_overflow(v[i]);
+        return nv;
+    }
+
     inline NanoVector(int64 x) { push_back(x); }
 
     // @pyjt(__repr__)
-    string to_string() {
+    inline string to_string() {
         string s="[";
         for (int i=0; i<size(); i++) {
             s += S(at(i));

src/misc/ring_buffer.cc

@@ -0,0 +1,91 @@
+// ***************************************************************
+// Copyright (c) 2020 Jittor. All Rights Reserved.
+// Authors: Dun Liang <randonlang@gmail.com>.
+// This file is subject to the terms and conditions defined in
+// file 'LICENSE.txt', which is part of this source code package.
+// ***************************************************************
+#include <chrono>
+#include <thread>
+#include <sys/mman.h>
+#include "common.h"
+#include "misc/ring_buffer.h"
+
+namespace jittor {
+
+RingBuffer::RingBuffer(uint64 size, bool multiprocess) : m(multiprocess), cv(multiprocess) {
+    int i=0;
+    for (;(1ll<<i)<size;i++);
+    size_mask = (1ll<<i)-1;
+    this->size = size_mask+1;
+    size_bit = i;
+    l = r = is_wait = 0;
+    is_multiprocess = multiprocess;
+}
+
+RingBuffer::~RingBuffer() {
+}
+
+
+RingBuffer* RingBuffer::make_ring_buffer(uint64 size, bool multiprocess) {
+    int i=0;
+    for (;(1ll<<i)<size;i++);
+    uint64 size_mask = (1ll<<i)-1;
+    size = size_mask+1;
+    uint64 total_size = sizeof(RingBuffer) + size;
+    void* ptr = multiprocess ?
+        // mmap(NULL, total_size, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS | MAP_HUGETLB, -1, 0) :
+        mmap(NULL, total_size, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, -1, 0) :
+        // mmap(NULL, total_size, PROT_READ | PROT_WRITE, MAP_SHARED, -1, 0) :
+        (void*)malloc(total_size);
+    std::memset(ptr, 0, total_size);
+    auto rb = (RingBuffer*)ptr;
+    new (rb) RingBuffer(size, multiprocess);
+    return rb;
+}
+
+void RingBuffer::free_ring_buffer(RingBuffer* rb) {
+    uint64 total_size = sizeof(RingBuffer) + rb->size;
+    auto is_multiprocess = rb->is_multiprocess;
+    rb->~RingBuffer();
+    if (is_multiprocess) {
+        munmap(rb, total_size);
+    } else {
+        rb->~RingBuffer();
+        free((void*)rb);
+    }
+}
+
+// test
+
+JIT_TEST(ring_buffer_benchmark) {
+    size_t n = 1ll << 20;
+    size_t size = 1<<15;
+    // size_t n = 1ll << 30;
+    // size_t size = 1<<20;
+    // size_t n = 1ll << 10;
+    // size_t size = 1<<5;
+    RingBuffer* rb = RingBuffer::make_ring_buffer(size, 0);
+    std::thread p([&]() {
+        for (size_t i=0; i<n; i++) {
+            rb->push_t<int>(i);
+        }
+    });
+    auto start = std::chrono::high_resolution_clock::now();
+    size_t s = 0;
+    for (size_t i=0; i<n; i++) {
+        auto x = rb->pop_t<int>();
+        s += x;
+    }
+    auto finish = std::chrono::high_resolution_clock::now();
+    auto tt =  std::chrono::duration_cast<std::chrono::nanoseconds>(finish-start).count();
+    p.join();
+    expect_error([&]() { rb->push(size+1); });
+    RingBuffer::free_ring_buffer(rb);
+
+    LOGi << tt << tt*1.0/n;
+    LOGi << s << (n*(n-1)/2); 
+    ASSERTop(s,==,(n*(n-1)/2));
+    ASSERTop(tt*1.0/n,<=,50);
+}
+
+}