perf(mge/utils): move astensor1d into C++

GitOrigin-RevId: e7c6659020d9db5a4b17f2a35a40ea2a99f7330d

imperative/python/megengine/core/_trace_option.py

@@ -9,6 +9,8 @@
 
 import os
 
+from ._imperative_rt.core2 import set_cpp_use_symbolic_shape
+
 _use_symbolic_shape = False
 if os.environ.get("MEGENGINE_USE_SYMBOLIC_SHAPE"):
     _use_symbolic_shape = True
@@ -25,3 +27,6 @@ def set_symbolic_shape(option: bool):
     _org = _use_symbolic_shape
     _use_symbolic_shape = option
     return _org
+
+
+set_cpp_use_symbolic_shape(use_symbolic_shape)

imperative/python/megengine/core/tensor/array_method.py

@@ -22,12 +22,12 @@ from .._imperative_rt.core2 import (
     astype_cpp,
     broadcast_cpp,
     dtype_promotion,
+    getitem_cpp,
 )
 from .._imperative_rt.core2 import reduce_to_scalar as _reduce_to_scalar
-from .._imperative_rt.core2 import reshape_cpp, squeeze_cpp, transpose_cpp
+from .._imperative_rt.core2 import reshape_cpp, setitem_cpp, squeeze_cpp, transpose_cpp
 from ..ops import builtin
 from . import amp
-from .indexing import getitem, setitem
 from .utils import _normalize_axis, astensor1d, cast_tensors, make_shape_tuple, subgraph
 
 _ElwMod = builtin.Elemwise.Mode
@@ -544,11 +544,11 @@ class ArrayMethodMixin(abc.ABC):
             yield self[i]
 
     def __getitem__(self, index):
-        return getitem(self, index)
+        return getitem_cpp(self, index)
 
     def __setitem__(self, index, value):
         if index is not Ellipsis:
-            value = setitem(self, index, value)
+            value = setitem_cpp(self, index, value)
         self._reset(value)
 
     __contains__ = _todo

imperative/python/megengine/core/tensor/indexing.py

-# -*- coding: utf-8 -*-
-# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
-#
-# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-from .._imperative_rt.core2 import (
-    getitem_cpp,
-    set_cpp_astensor1d,
-    set_cpp_use_symbolic_shape,
-    setitem_cpp,
-)
-from .._trace_option import use_symbolic_shape
-from .utils import astensor1d
-
-
-def getitem(tensor, index):
-    return getitem_cpp(tensor, index)
-
-
-def setitem(tensor, index, value):
-    return setitem_cpp(tensor, index, value)
-
-
-set_cpp_use_symbolic_shape(use_symbolic_shape)
-set_cpp_astensor1d(astensor1d)

imperative/python/megengine/core/tensor/utils.py

@@ -20,6 +20,7 @@ from .._imperative_rt.core2 import (
     _get_convert_inputs,
     _set_convert_inputs,
     apply,
+    astensor1d_cpp,
     astype_cpp,
     convert_inputs_cpp,
     convert_single_value_cpp,
@@ -50,14 +51,6 @@ def set_convert_inputs(flag):
     return _set_convert_inputs(flag)
 
 
-def concatenate(inputs, axis=0, *, device=None):
-    inputs = convert_inputs(*inputs)
-    if device is None:
-        device = get_device(inputs)
-    (result,) = apply(builtin.Concat(axis=axis, comp_node=device), *inputs)
-    return result
-
-
 def convert_single_value(v, *, dtype=None, device=None):
     return convert_single_value_cpp(v, dtype, device)
 
@@ -104,34 +97,7 @@ def astensor1d(x, *reference, dtype=None, device=None):
       * numpy array
       * tensor (returned as is, regardless of dtype and device)
     """
-    try:
-        ndim = x.ndim
-    except AttributeError:
-        pass
-    except ValueError:
-        if dtype is not None and dtype != x.dtype:
-            x = astype_cpp(x, dtype)
-        if device is not None:
-            cn = as_device(device).to_c()
-            (x,) = apply(builtin.Copy(comp_node=cn), x)
-        return x
-    else:
-        if ndim != 0 and ndim != 1:
-            raise ValueError("ndim != 1 or 0, get : %d" % ndim)
-        if not isinstance(x, (Tensor, SymbolVar)):
-            x = Const(x, dtype, device, reference)
-        return x
-
-    if not isinstance(x, collections.abc.Sequence):
-        raise TypeError
-
-    if any(isinstance(i, (Tensor, SymbolVar)) for i in x):
-        x = concatenate(x, device=device) if len(x) > 1 else x[0]
-        if dtype is not None:
-            x = astype_cpp(x, dtype)
-        return x
-    x = Const(x, dtype, device, reference)
-    return x
+    return astensor1d_cpp(x, dtype, device, reference)
 
 
 def _normalize_axis(

imperative/python/src/tensor.cpp

@@ -104,13 +104,12 @@ struct SymbolVarContext {
 
 interpreter::Interpreter::Channel* interpreter_for_py = nullptr;
 PyTypeObject* py_tensor_type = nullptr;
-PyObject *cpp_use_symbolic_shape, *cpp_astensor1d;
+PyObject* cpp_use_symbolic_shape;
 
 #define REGISTE_APPLY_FUNC(mode) \
     void set_##mode(py::object pyf) { mode = pyf.ptr(); }
 
 REGISTE_APPLY_FUNC(cpp_use_symbolic_shape)
-REGISTE_APPLY_FUNC(cpp_astensor1d)
 
 #undef REGISTE_APPLY_FUNC
 
@@ -426,6 +425,7 @@ WRAP_FUNC_PY35(Const);
 WRAP_FUNC_PY35(astype_cpp);
 WRAP_FUNC_PY35(convert_single_value_cpp);
 WRAP_FUNC_PY35(convert_inputs_cpp);
+WRAP_FUNC_PY35(astensor1d_cpp);
 #undef WRAP_FUNC_PY35
 #define MGE_PY_INTERFACE(NAME, FUNC) \
     { #NAME, (PyCFunction)py35_##FUNC, METH_VARARGS, nullptr }
@@ -568,6 +568,7 @@ void init_tensor(py::module m) {
             MGE_PY_INTERFACE(astype_cpp, astype_cpp),
             MGE_PY_INTERFACE(convert_single_value_cpp, convert_single_value_cpp),
             MGE_PY_INTERFACE(convert_inputs_cpp, convert_inputs_cpp),
+            MGE_PY_INTERFACE(astensor1d_cpp, astensor1d_cpp),
             {nullptr, nullptr, 0, nullptr}};
     for (auto&& def : method_defs) {
         if (def.ml_meth != nullptr) {
@@ -957,8 +958,6 @@ void init_tensor(py::module m) {
 
     m.def("set_cpp_use_symbolic_shape", &set_cpp_use_symbolic_shape);
 
-    m.def("set_cpp_astensor1d", &set_cpp_astensor1d);
-
     m.def("set_module_tracing", [=] { get_module_trace()->enable(); });
 
     m.def("unset_module_tracing", [=] { get_module_trace()->disable(); });

imperative/python/src/tensor_utils.h

@@ -32,4 +32,6 @@ PyObject* convert_single_value_cpp(PyObject* self, PyObject* const* args, size_t
 
 PyObject* convert_inputs_cpp(PyObject* self, PyObject* const* args, size_t nargs);
 
+PyObject* astensor1d_cpp(PyObject* self, PyObject* const* args, size_t nargs);
+
 }  // namespace mgb::imperative::python
\ No newline at end of file

imperative/python/test/unit/core/test_indexing_op.py

@@ -511,6 +511,20 @@ def test_advance_indexing_with_bool(test_varnode):
         network = Network()
     else:
         network = None
+
+    a = np.array([[True, False], [False, True]])
+    b = np.array([1])
+    aa = make_tensor(a, network)
+    bb = make_tensor(b, network)
+    np.testing.assert_equal(a[b], get_value(aa[bb]))
+    b = np.array([[True, True], [False, True]])
+    bb = make_tensor(b, network)
+    np.testing.assert_equal(a[b], get_value(aa[bb]))
+    if not test_varnode:
+        a[b] = False
+        aa[bb] = False
+        np.testing.assert_equal(a, get_value(aa))
+
     a = np.arange(9).reshape(3, 3).astype(np.float32)
     b = np.array([1, 2, 3])
     c = np.array([1, 2, 3])
@@ -525,67 +539,68 @@ def test_advance_indexing_with_bool(test_varnode):
     a = np.arange(9).reshape(3, 3).astype(np.float32)
     b = np.array([False, True, True])
     c = np.array([2, 0]).astype(np.int32)
-    aa = Tensor(a)
-    bb = Tensor(b)
-    cc = Tensor(c)
-    np.testing.assert_equal(a[b, c], aa[bb, cc].numpy())
+    aa = make_tensor(a, network)
+    bb = make_tensor(b, network)
+    cc = make_tensor(c, network)
+    np.testing.assert_equal(a[b, c], get_value(aa[bb, cc]))
     a[b, c] = -1.0
     aa[bb, cc] = -1.0
-    np.testing.assert_equal(a, aa.numpy())
+    np.testing.assert_equal(a, get_value(aa))
     d = np.array([-1, -2], dtype=np.float32)
-    dd = Tensor(d)
+    dd = make_tensor(d, network)
     a[b, c] = d
     aa[bb, cc] = dd
-    np.testing.assert_equal(a, aa.numpy())
+    np.testing.assert_equal(a, get_value(aa))
 
     a = np.ones((2, 2))
     b = np.array([[True, False], [False, True]])
-    aa = Tensor(a)
-    bb = Tensor(b)
-    np.testing.assert_equal(a[b], aa[bb].numpy())
+    aa = make_tensor(a, network)
+    bb = make_tensor(b, network)
+    np.testing.assert_equal(a[b], get_value(aa[bb]))
     b[:] = True
     bb[:] = True
-    np.testing.assert_equal(a[b], aa[bb].numpy())
-    np.testing.assert_equal(a[:, [True, False]], aa[:, [True, False]].numpy())
+    np.testing.assert_equal(a[b], get_value(aa[bb]))
+    np.testing.assert_equal(a[:, [True, False]], get_value(aa[:, [True, False]]))
 
     a = np.array([[True, False], [False, True]])
     b = np.array([1])
-    aa = Tensor(a)
-    bb = Tensor(b)
-    np.testing.assert_equal(a[b], aa[bb].numpy())
+    aa = make_tensor(a, network)
+    bb = make_tensor(b, network)
+    np.testing.assert_equal(a[b], get_value(aa[bb]))
     b = np.array([[True, True], [False, True]])
-    bb = Tensor(b)
-    np.testing.assert_equal(a[b], aa[bb].numpy())
-    a[b] = False
-    aa[bb] = False
-    np.testing.assert_equal(a, aa.numpy())
+    bb = make_tensor(b, network)
+    np.testing.assert_equal(a[b], get_value(aa[bb]))
+    if not test_varnode:
+        a[b] = False
+        aa[bb] = False
+        np.testing.assert_equal(a, get_value(aa))
 
     a = np.ones((2, 2), dtype=np.int32)
     b = np.array([[False, False], [False, False]])
-    aa = Tensor(a)
-    bb = Tensor(b)
-    np.testing.assert_equal(a[b], aa[b].numpy())
-    np.testing.assert_equal(a[b], aa[bb].numpy())
+    aa = make_tensor(a, network)
+    bb = make_tensor(b, network)
+    np.testing.assert_equal(a[b], get_value(aa[b]))
+    np.testing.assert_equal(a[b], get_value(aa[bb]))
 
     b = np.array([False, False])
-    bb = Tensor(b)
-    np.testing.assert_equal(a[b], aa[bb].numpy().reshape(a[b].shape))
+    bb = make_tensor(b, network)
+    np.testing.assert_equal(a[b], get_value(aa[bb]).reshape(a[b].shape))
 
     a = np.arange(576).reshape(2, 3, 4, 3, 4, 2).astype("int32")
-    aa = Tensor(a)
+    aa = make_tensor(a, network)
 
     b = (np.random.sample((2, 3, 4)) > 0.5).astype("bool")
-    bb = Tensor(b)
-    np.testing.assert_equal(a[b, :, 0:4:2], aa[bb, :, 0:4:2].numpy())
+    bb = make_tensor(b, network)
+    np.testing.assert_equal(a[b, :, 0:4:2], get_value(aa[bb, :, 0:4:2]))
 
     b = (np.random.sample((4, 3, 4)) > 0.5).astype("bool")
-    bb = Tensor(b)
-    np.testing.assert_equal(a[..., b, 0:2], aa[..., bb, 0:2].numpy())
+    bb = make_tensor(b, network)
+    np.testing.assert_equal(a[..., b, 0:2], get_value(aa[..., bb, 0:2]))
 
     b = (np.random.sample((3, 4, 3)) > 0.5).astype("bool")
-    bb = Tensor(b)
+    bb = make_tensor(b, network)
     np.testing.assert_equal(
-        a[:, b, 0:2, [True, False]], aa[:, bb, 0:2, [True, False]].numpy()
+        a[:, b, 0:2, [True, False]], get_value(aa[:, bb, 0:2, [True, False]])
     )