add slice optim

python/jittor/contrib.py

@@ -226,13 +226,11 @@ def concat(arr, dim):
     shape = list(a.shape)
     shape[dim] = total_dim
     s = jt.empty(shape, a.dtype)
-    slices = [slice(None)]*len(a.shape)
+    slices = [slice(None)]*(dim+1)
     for a in arr:
         if a.shape[dim] == 0:
             continue
         slices[dim] = slice(cdim, cdim+a.shape[dim])
-        # print(slices, type(a))
         s = s.setitem(tuple(slices), a)
-        # s = jt.setitem(s, tuple(slices), a)
         cdim += a.shape[dim]
     return s

python/jittor/test/test_setitem.py

@@ -8,80 +8,222 @@
 import unittest
 import jittor as jt
 import numpy as np
+from jittor.test.test_log import find_log_with_re
 skip_this_test = False
 
 @unittest.skipIf(skip_this_test, "No Torch found")
 class TestSetitem(unittest.TestCase):
-    def test_setitem(self):
-        arr0 = jt.random((4,2,2))
-        data0 = jt.ones((2,2))
-        arr0[1] = data0
-        arr0.sync()
-        data0.data[0,0] = 0
-        assert arr0[1,0,0] == 0
+    def test_getitem_grad_opt1(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            with jt.flag_scope(trace_py_var=2):
+                v = jt.random((1,2,3,4,5,6))
+                for i in range(6):
+                    loss = 0.
+                    ss = 1
+                    if v.shape[i] % 2 == 0: 
+                        ss = 2
+                    res = v.split(split_size=ss, dim=i)
+                    t_ = []
+                    for j in range(len(res)):
+                        t = np.random.random(res[j].shape).astype("float32")
+                        loss += res[j] * t
+                        t_.append(t)
+                    dv = jt.grad(loss, v)
+                    dv.sync()
+                    data = jt.dump_trace_data()
+                    jt.clear_trace_data()
+                    assert (dv.numpy() == np.concatenate(t_, i)).all()
+        logs = find_log_with_re(rep, "getitem_grad_opt happens")
+        assert len(logs) == 4
 
-        arr00 = jt.random((4,2,2))
-        data00 = jt.ones((2,2))
-        # share memory will fail if d has an edge to other nodes.
-        tmp = data00 + 1
-        arr00[1] = data00
-        arr00.sync()
-        data00.data[0,0] = 0
-        assert arr00[1,0,0] == 0
+    def test_setitem_grad_opt1(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            v = jt.random((4,10,10))
+            a = jt.ones((10,10))
+            b = jt.ones((10,10))
+            c = jt.ones((10,10))
+            d = jt.ones((10,10))
+            v[0] = a
+            v[1] = b
+            v[2] = c
+            v[3] = d
+            t = np.random.random((4,10,10)).astype("float32")
+            loss = v*t
+            da, db, dc, dd = jt.grad(loss, [a,b,c,d])
+            jt.sync([da, db, dc, dd])
+            assert (da.numpy() == t[0]).all()
+            assert (db.numpy() == t[1]).all()
+            assert (dc.numpy() == t[2]).all()
+            assert (dd.numpy() == t[3]).all()
+        logs = find_log_with_re(rep, "setitem_grad_opt happens")
+        logs1 = find_log_with_re(rep, "setitem_grad_opt set success")
+        assert len(logs) == 1
+        assert len(logs1) == 3
+
+    def test_setitem_grad_opt2(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            v = jt.random((4,10,10))
+            a = jt.ones((10,10))
+            b = jt.ones((10,10))
+            c = jt.ones((10,10))
+            d = jt.ones((10,10))
+            v[0] = a
+            v[2] = c
+            v[1] = b
+            v[3] = d
+            t = np.random.random((4,10,10)).astype("float32")
+            loss = v*t
+            da, db, dc, dd = jt.grad(loss, [a,b,c,d])
+            jt.sync([da, db, dc, dd])
+            assert (da.numpy() == t[0]).all()
+            assert (db.numpy() == t[1]).all()
+            assert (dc.numpy() == t[2]).all()
+            assert (dd.numpy() == t[3]).all()
+        logs = find_log_with_re(rep, "setitem_grad_opt happens")
+        logs1 = find_log_with_re(rep, "setitem_grad_opt set success")
+        assert len(logs) == 1
+        assert len(logs1) == 2
+
+    def test_setitem_grad_opt3(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            a = jt.ones((10,10,10))
+            b = jt.ones((10,10,10))
+            c = jt.ones((10,10,10))
+            d = jt.ones((10,10,10))
+            v = jt.contrib.concat([a,b,c,d], dim=1)
+            t = np.random.random((10,40,10)).astype("float32")
+            loss = v*t
+            da, db, dc, dd = jt.grad(loss, [a,b,c,d])
+            jt.sync([da, db, dc, dd])
+            assert (da.numpy() == t[:,:10]).all()
+            assert (db.numpy() == t[:,10:20]).all()
+            assert (dc.numpy() == t[:,20:30]).all()
+            assert (dd.numpy() == t[:,30:40]).all()
+        logs = find_log_with_re(rep, "setitem_grad_opt happens")
+        logs1 = find_log_with_re(rep, "setitem_grad_opt set success")
+        assert len(logs) == 1
+        assert len(logs1) == 3
 
-        arr1 = jt.random((4,2,2))
-        data1 = jt.zeros((2,2))
-        arr1[3,:,0:2] = data1
-        arr1.sync()
-        data1.data[0,0] = 1
-        assert arr1[3,0,0] == 1
+    def test_setitem1(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            v = jt.random((4,2,2))
+            d = jt.ones((2,2))
+            v[1] = d
+            del d
+            v.sync()
+        logs = find_log_with_re(rep, "setitem_inplace happens")
+        assert len(logs) == 1
 
-        arr21 = jt.ones((2,2))
-        arr22 = jt.ones((2,2)) * 2
-        arr2 = jt.contrib.concat([arr21, arr22], dim=0)
-        arr2.sync()
-        arr21.data[0,0] = 3
-        arr22.data[0,0] = 4
-        assert arr2[0,0] == 3
-        assert arr2[2,0] == 4
+    def test_setitem2(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            v = jt.random((4,2,2))
+            d = jt.ones((2,2))
+            v[1] = d
+            v.sync()
+        logs = find_log_with_re(rep, "setitem_inplace happens")
+        assert len(logs) == 0
 
-    def test_getitem(self):
-        # test for different slice type
-        arr0 = jt.random((4,3))
-        arr0_res = arr0[2,:]
-        arr0_res.data[1] = 1
-        assert arr0[2,1] == 1
+    def test_setitem3(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            v = jt.random((4,2,2))
+            v1 = v[1:3]
+            d = jt.ones((2,2))
+            v1[1] = d
+            v1.sync()
+        logs = find_log_with_re(rep, "setitem_inplace happens")
+        assert len(logs) == 0
+
+    def test_setitem4(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            v = jt.random((4,2,2))
+            v1 = v[1:3,0]
+            d = jt.ones((2,))
+            v1[1] = d
+            del d
+            v1.sync()
+        logs = find_log_with_re(rep, "setitem_inplace happens")
+        assert len(logs) == 1
+
+    def test_setitem5(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            v = jt.random((4,2,2))
+            v1 = v[1:3,0]
+            d = jt.ones((2,2))
+            d1 = d[0]
+            v1[1] = d1
+            v1.sync()
+        logs = find_log_with_re(rep, "setitem_inplace happens")
+        assert len(logs) == 0
+
+    def test_setitem6(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            v = jt.random((4,2,2))
+            v1 = v[1:3,0]
+            d = jt.ones((2,2,2))
+            d1 = d[0,0]
+            v1[1] = d1
+            del d1
+            v1.sync()
+        logs = find_log_with_re(rep, "setitem_inplace happens")
+        assert len(logs) == 1
 
-        arr1 = jt.array([1,2,3,4])
-        arr1_res = arr1[None]
-        arr1_res.data[0,2] = -1
-        assert arr1[2] == -1
+    def test_getitem1(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            v = jt.random((4,3))
+            v_res = v[2,:]
+            v_res.data[1] = 1
+        logs = find_log_with_re(rep, "getitem_inplace happens")
+        assert len(logs) == 1
 
-        arr2 = jt.array([1,2,3,4])
-        arr2_res = arr2[...]
-        arr2_res.data[2] = -1
-        assert arr2[2] == -1
+    def test_getitem2(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            v1 = jt.array([1,2,3,4])
+            v1_res = v1[None]
+            v1_res.data[0,2] = -1
+        logs = find_log_with_re(rep, "getitem_inplace happens")
+        assert len(logs) == 1
 
-        arr3 = jt.array([1,2,3,4])
-        arr3_res = arr3[3]
-        arr3_res.data[0] = -1
-        assert arr3[3] == -1
+    def test_getitem3(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            arr3 = jt.array([1,2,3,4])
+            arr3_res = arr3[3]
+            arr3_res.data[0] = -1
+        logs = find_log_with_re(rep, "getitem_inplace happens")
+        assert len(logs) == 1
 
-        arr4 = jt.random((4,2,3,3))
-        arr4_res = arr4[...,:,:]
-        arr4_res.data[0,0,1,1] = 1
-        assert arr4[0,0,1,1] == 1
+    def test_getitem4(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            arr4 = jt.random((4,2,3,3))
+            arr4_res = arr4[...,:,:]
+            arr4_res.data[0,0,1,1] = 1
+        logs = find_log_with_re(rep, "getitem_inplace happens")
+        assert len(logs) == 1
 
-        arr5 = jt.random((4,2,3,3))
-        arr5_res = arr5[1:3,:,:,:]
-        arr5_res.data[1,0,1,1] = 1
-        assert arr5[2,0,1,1] == 1
+    def test_getitem5(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            arr5 = jt.random((4,2,3,3))
+            arr5_res = arr5[1:3,:,:,:]
+            arr5_res.data[1,0,1,1] = 1
+        logs = find_log_with_re(rep, "getitem_inplace happens")
+        assert len(logs) == 1
 
-        arr6 = jt.random((4,2,3,3))
-        arr6_res = arr6[1]
-        arr6_res.data[0,1,1] = 1
-        assert arr6[1,0,1,1] == 1
+    def test_getitem6(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            arr6 = jt.random((4,2,3,3))
+            arr6_res = arr6[1]
+            arr6_res.data[0,1,1] = 1
+        logs = find_log_with_re(rep, "getitem_inplace happens")
+        assert len(logs) == 1
+
+    def test_getitem7(self):
+        with jt.log_capture_scope(log_vprefix="setitem_gopt=1000") as rep:
+            arr2 = jt.array([1,2,3,4])
+            arr2_res = arr2[...]
+            arr2_res.data[2] = -1
+        logs = find_log_with_re(rep, "getitem_inplace happens")
+        assert len(logs) == 1
 
+    def test_getitem8(self):
         # test for different data type (float32/float64/bool/int8/int32)
         arr_float32 = jt.random((4,2,3))
         arr_float32_res = arr_float32[1:3,:,:]

src/executor.cc

@@ -128,8 +128,10 @@ void Executor::run_sync(vector<Var*> vars, bool device_sync) {
             }
         }
         if (!need_opt) break;
+        toplogical_sort_forward_inplace(bfs_q, [&](Node* n) {});
+        SetupFreeBuffer setup_free_buffer;
         for (Node* n : bfs_q) {
-            if (n->flags.get(NodeFlags::_has_gopt)) {
+            if (!n->need_free() && n->flags.get(NodeFlags::_has_gopt)) {
                 n->op()->graph_optimize();
                 n->flags.set(NodeFlags::_has_gopt, 0);
             }

src/graph.h

@@ -111,6 +111,34 @@ void toplogical_sort_forward(vector<Node*>& nodes, vector<Node*>& sorted, Func&&
     ASSERTop(nodes.size(),==,sorted.size());
 }
 
+template <typename Func>
+void toplogical_sort_forward_inplace(vector<Node*>& nodes, Func&& func) {
+    auto t = ++Node::tflag_count;
+    int sorted_size = 0;
+    for (auto node : nodes) node->tflag = t;
+    for (auto node : nodes) {
+        auto& deps = node->custom_data;
+        deps = 0;
+        for (auto i : node->_inputs)
+            if (i.node->tflag == t)
+                deps++;
+        if (deps == 0)
+            nodes[sorted_size++] = node;
+    }
+    int i=0;
+    while (i < sorted_size) {
+        Node* node = nodes[i++];
+        for (auto o : node->_outputs)
+            if (o.node->tflag == t) {
+                o.node->custom_data--;
+                if (o.node->custom_data == 0)
+                    nodes[sorted_size++] = o.node;
+            }
+        func(node);
+    }
+    ASSERTop(nodes.size(),==,sorted_size);
+}
+
 
 template <typename Func>
 void toplogical_sort_backward(vector<Node*>& nodes, vector<Node*>& sorted, Func&& func) {

src/ops/getitem_op.cc

@@ -447,6 +447,9 @@ void GetitemOp::jit_run() {
     auto in = inputs().front();
     auto out = outputs().front();
     if (out->num == 0) return;
+    if (out->allocation == in->allocation && 
+        out->allocator == in->allocator)
+        return;
 
     @for(i, 0, ODIM, index_t oshape@i = o_shape[@i];)
     @if(ODIM>0,