[v1.7.x] Backport some numpy features + fixes

Makefile

@@ -460,9 +460,9 @@ endif
 
 all: lib/libmxnet.a lib/libmxnet.so $(BIN) extra-packages extension_libs
 
-SRC = $(wildcard src/*/*/*/*.cc src/*/*/*.cc src/*/*.cc src/*.cc)
+SRC = $(wildcard src/*/*/*/*/*.cc src/*/*/*/*.cc src/*/*/*.cc src/*/*.cc src/*.cc)
 OBJ = $(patsubst %.cc, build/%.o, $(SRC))
-CUSRC = $(wildcard src/*/*/*/*.cu src/*/*/*.cu src/*/*.cu src/*.cu)
+CUSRC = $(wildcard src/*/*/*/*.cu src/*/*/*/*.cu src/*/*/*.cu src/*/*.cu src/*.cu)
 CUOBJ = $(patsubst %.cu, build/%_gpu.o, $(CUSRC))
 
 ifeq ($(USE_TVM_OP), 1)

benchmark/python/ffi/benchmark_ffi.py

@@ -0,0 +1,231 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import timeit
+import itertools
+import argparse
+import os
+
+class OpArgMngr(object):
+    """Operator argument manager for storing operator workloads."""
+    args = {}
+
+    @staticmethod
+    def add_workload(funcname, *args, **kwargs):
+        if "_specifier" not in kwargs:
+            _specifier = funcname
+        else:
+            _specifier = kwargs["_specififer"]
+            del kwargs["_specififer"]
+        if _specifier in OpArgMngr.args:
+            raise ValueError("duplicate {}".format(_specifier))
+        OpArgMngr.args[_specifier] = {'args': args, 'kwargs': kwargs, 'funcname': funcname}
+
+
+def generate_workloads():
+    array_pool = {}
+    shapes = []
+    for ndim in range(4):
+        shapes.extend(list(itertools.product(range(4), repeat=ndim)))
+    for shape in shapes:
+        name = 'x'.join(str(i) for i in shape)
+        if name in array_pool:
+            raise ValueError("duplicate array {}".format(name))
+        array_pool[name] = dnp.ones(shape)
+    return array_pool
+
+
+def prepare_workloads():
+    pool = generate_workloads()
+    OpArgMngr.add_workload("zeros", (2, 2))
+    OpArgMngr.add_workload("einsum", "ii", pool['2x2'], optimize=False)
+    OpArgMngr.add_workload("unique", pool['1'], return_index=True, return_inverse=True, return_counts=True, axis=-1)
+    OpArgMngr.add_workload("dstack", (pool['2x1'], pool['2x1'], pool['2x1'], pool['2x1']))
+    OpArgMngr.add_workload("polyval", dnp.arange(10), pool['2x2'])
+    OpArgMngr.add_workload("ediff1d", pool['2x2'], pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("nan_to_num", pool['2x2'])
+    OpArgMngr.add_workload("tensordot", pool['2x2'], pool['2x2'], ((1, 0), (0, 1)))
+    OpArgMngr.add_workload("kron", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("cumsum", pool['3x2'], axis=0, out=pool['3x2'])
+    OpArgMngr.add_workload("sum", pool['2x2'], axis=0, keepdims=True, out=pool['1x2'])
+    OpArgMngr.add_workload("std", pool['2x2'], axis=0, ddof=0, keepdims=True, out=pool['1x2'])
+    OpArgMngr.add_workload("var", pool['2x2'], axis=0, ddof=1, keepdims=True, out=pool['1x2'])
+    OpArgMngr.add_workload("average", pool['2x2'], weights=pool['2'], axis=1, returned=True)
+    OpArgMngr.add_workload("histogram", pool['2x2'], bins=10, range=(0.0, 10.0))
+    OpArgMngr.add_workload("add", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("linalg.eig", pool['3x3'])
+    OpArgMngr.add_workload("linalg.eigh", pool['3x3'])
+    OpArgMngr.add_workload("linalg.det", pool['3x3'])
+    OpArgMngr.add_workload("linalg.slogdet", pool['3x3'])
+    OpArgMngr.add_workload("linalg.matrix_rank", pool['3x3'], pool['1'], hermitian=False)
+    OpArgMngr.add_workload("linalg.svd", pool['3x3'])
+    OpArgMngr.add_workload("linalg.cholesky", pool['1x1'])
+    OpArgMngr.add_workload("linalg.qr", pool['3x3'])
+    OpArgMngr.add_workload("linalg.lstsq", pool['2x1'], pool['2'], rcond=None)
+    OpArgMngr.add_workload("linalg.eigvals", pool['1x1'])
+    OpArgMngr.add_workload("linalg.eigvalsh", pool['1x1'], UPLO='L')
+    OpArgMngr.add_workload("linalg.inv", pool['1x1'])
+    OpArgMngr.add_workload("linalg.pinv", pool['2x3x3'], pool['1'], hermitian=False)
+    OpArgMngr.add_workload("linalg.solve", pool['1x1'], pool['1'])
+    OpArgMngr.add_workload("linalg.tensorinv", pool['1x1'], ind=2)
+    OpArgMngr.add_workload("linalg.tensorsolve", pool['1x1x1'], pool['1x1x1'], (2, 0, 1))
+    OpArgMngr.add_workload("linalg.svd", pool['3x3'])
+    OpArgMngr.add_workload("split", pool['3x3'], (0, 1, 2), axis=1)
+    OpArgMngr.add_workload("argmax", pool['3x2'], axis=-1)
+    OpArgMngr.add_workload("argmin", pool['3x2'], axis=-1)
+    OpArgMngr.add_workload("atleast_1d", pool['2'], pool['2x2'])
+    OpArgMngr.add_workload("atleast_2d", pool['2'], pool['2x2'])
+    OpArgMngr.add_workload("atleast_3d", pool['2'], pool['2x2'])
+    OpArgMngr.add_workload("indices", dimensions=(1, 2, 3))
+    OpArgMngr.add_workload("subtract", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("multiply", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("mod", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("remainder", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("divide", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("true_divide", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("power", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("lcm", pool['2x2'].astype('int32'), pool['2x2'].astype('int32'))
+    OpArgMngr.add_workload("diff", pool['2x2'], n=1, axis=-1)
+    OpArgMngr.add_workload("nonzero", pool['2x2'])
+    OpArgMngr.add_workload("tril", pool['2x2'], k=0)
+    OpArgMngr.add_workload("random.choice", pool['2'], size=(2, 2))
+    OpArgMngr.add_workload("take", pool['2'], dnp.array([1,0], dtype='int64'))
+    OpArgMngr.add_workload("clip", pool['2x2'], 0, 1)
+    OpArgMngr.add_workload("expand_dims", pool['2x2'], axis=0)
+    OpArgMngr.add_workload("broadcast_to", pool['2x2'], (2, 2, 2))
+    OpArgMngr.add_workload("full_like", pool['2x2'], 2)
+    OpArgMngr.add_workload("zeros_like", pool['2x2'])
+    OpArgMngr.add_workload("ones_like", pool['2x2'])
+    OpArgMngr.add_workload("bitwise_and", pool['2x2'].astype(int), pool['2x2'].astype(int))
+    OpArgMngr.add_workload("bitwise_xor", pool['2x2'].astype(int), pool['2x2'].astype(int))
+    OpArgMngr.add_workload("bitwise_or", pool['2x2'].astype(int), pool['2x2'].astype(int))
+    OpArgMngr.add_workload("copysign", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("arctan2", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("hypot", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("ldexp", pool['2x2'].astype(int), pool['2x2'].astype(int))
+    OpArgMngr.add_workload("random.uniform", low=0, high=1, size=1)
+    OpArgMngr.add_workload("random.exponential", scale=2, size=(2,2))
+    OpArgMngr.add_workload("random.rayleigh", scale=2, size=(2,2))
+    OpArgMngr.add_workload("random.weibull", a=2, size=(2,2))
+    OpArgMngr.add_workload("random.pareto", a=2, size=(2,2))
+    OpArgMngr.add_workload("random.power", a=2, size=(2,2))
+    OpArgMngr.add_workload("random.logistic", loc=2, scale=2, size=(2,2))
+    OpArgMngr.add_workload("random.gumbel", loc=2, scale=2, size=(2,2))
+    OpArgMngr.add_workload("where", pool['2x3'], pool['2x3'], pool['2x1'])
+    OpArgMngr.add_workload("fmax", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("fmin", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("fmod", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("may_share_memory", pool['2x3'][:0], pool['2x3'][:1])
+    OpArgMngr.add_workload("diag", pool['2x2'], k=1)
+    OpArgMngr.add_workload("diagonal", pool['2x2x2'], offset=-1, axis1=0, axis2=1)
+    OpArgMngr.add_workload("diag_indices_from", pool['2x2'])
+    OpArgMngr.add_workload("bincount", dnp.arange(3, dtype=int), pool['3'], minlength=4)
+    OpArgMngr.add_workload("percentile", pool['2x2x2'], 80, axis=0, out=pool['2x2'],\
+                           interpolation='midpoint')
+    OpArgMngr.add_workload("quantile", pool['2x2x2'], 0.8, axis=0, out=pool['2x2'],\
+                           interpolation='midpoint')
+    OpArgMngr.add_workload("all", pool['2x2x2'], axis=(0, 1),\
+                           out=dnp.array([False, False], dtype=bool), keepdims=False)
+    OpArgMngr.add_workload("any", pool['2x2x2'], axis=(0, 1),\
+                           out=dnp.array([False, False], dtype=bool), keepdims=False)
+    OpArgMngr.add_workload("roll", pool["2x2"], 1, axis=0)
+    OpArgMngr.add_workload("rot90", pool["2x2"], 2)
+    OpArgMngr.add_workload("triu", pool['3x3'])
+    OpArgMngr.add_workload("array_split", pool['2x2'], 2, axis=1)
+    OpArgMngr.add_workload("vsplit", pool['2x2'], 2)
+    OpArgMngr.add_workload("hsplit", pool['2x2'], 2)
+    OpArgMngr.add_workload("dsplit", pool['2x2x2'], 2)
+    OpArgMngr.add_workload("arange", 10)
+    OpArgMngr.add_workload("concatenate", (pool['1x2'], pool['1x2'], pool['1x2']), axis=0)
+    OpArgMngr.add_workload("append", pool['2x2'], pool['1x2'], axis=0)
+    OpArgMngr.add_workload("insert", pool['3x2'], 1, pool['1x1'], axis=0)
+    OpArgMngr.add_workload("delete", pool['3x2'], 1, axis=0)
+    OpArgMngr.add_workload("blackman", 12)
+    OpArgMngr.add_workload("eye", 5)
+    OpArgMngr.add_workload("hamming", 12)
+    OpArgMngr.add_workload("hanning", 12)
+    OpArgMngr.add_workload("linspace", 0, 10, 8, endpoint=False)
+    OpArgMngr.add_workload("logspace", 2.0, 3.0, num=4, base=2.0, dtype=onp.float32)
+    OpArgMngr.add_workload("matmul", pool['2x2'], pool['2x2'])
+    OpArgMngr.add_workload("mean", pool['2x2'], axis=0, keepdims=True)
+    OpArgMngr.add_workload("random.gamma", 1, size=(2, 3))
+    OpArgMngr.add_workload("random.normal", 1, size=(2, 3))
+
+
+def benchmark_helper(f, *args, **kwargs):
+    number = 10000
+    return timeit.timeit(lambda: f(*args, **kwargs), number=number) / number
+
+
+def get_op(module, funcname):
+    funcname = funcname.split(".")
+    for fname in funcname:
+        module = getattr(module, fname)
+    return module
+
+
+def run_benchmark(packages):
+    results = {}
+    for (k, v) in OpArgMngr.args.items():
+        result = {}
+        for (name, package) in packages.items():
+            print('{}.{} running...'.format(name, k))
+            op = get_op(package["module"], v["funcname"])
+            args = [package["data"](arg) for arg in v["args"]]
+            kwargs = {k: package["data"](v) for (k, v) in v["kwargs"].items()}
+            benchmark = benchmark_helper(op, *args, **kwargs)
+            result[name] = benchmark
+        results[k] = result
+    return results
+
+
+def show_results(results):
+    print("{:>24}{:>24}{:>24}".format("name", "package", "time(us)"))
+    for (specifier, d) in results.items():
+        for (k, v) in d.items():
+            print("{:>24}{:>24}{:>24}".format(specifier, k, v * 10 ** 6))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('ffi_type')
+    parsed = parser.parse_args()
+    if parsed.ffi_type == "cython":
+        os.environ['MXNET_ENABLE_CYTHON'] = '1'
+        os.environ['MXNET_ENFORCE_CYTHON'] = '1'
+    elif parsed.ffi_type == "ctypes":
+        os.environ['MXNET_ENABLE_CYTHON'] = '0'
+    else:
+        raise ValueError("unknown ffi_type {}",format(parsed.ffi_type))
+    os.environ["MXNET_ENGINE_TYPE"] = "NaiveEngine"
+    import mxnet as mx
+    import numpy as onp
+    from mxnet import np as dnp
+
+    mx.npx.set_np()
+    packages = {
+        "onp": {
+            "module": onp,
+            "data": lambda arr: arr.asnumpy() if isinstance(arr, dnp.ndarray) else arr
+        },
+        "dnp": {
+            "module": dnp,
+            "data": lambda arr: arr
+        }
+    }
+    prepare_workloads()
+    results = run_benchmark(packages)
+    show_results(results)

include/mxnet/ir/expr.h

@@ -141,7 +141,7 @@ class IntImmNode : public PrimExprNode {
   int64_t value;
 
   static constexpr const char* _type_key = "IntImm";
-  MXNET_DECLARE_FINAL_OBJECT_INFO(IntImmNode, PrimExprNode);
+  MXNET_DECLARE_FINAL_OBJECT_INFO(IntImmNode, PrimExprNode)
 };
 
 /*!
@@ -186,7 +186,7 @@ class FloatImmNode : public PrimExprNode {
   double value;
 
   static constexpr const char* _type_key = "FloatImm";
-  MXNET_DECLARE_FINAL_OBJECT_INFO(FloatImmNode, PrimExprNode);
+  MXNET_DECLARE_FINAL_OBJECT_INFO(FloatImmNode, PrimExprNode)
 };
 
 /*!

include/mxnet/node/container.h

@@ -42,7 +42,7 @@ class ArrayNode : public Object {
   std::vector<ObjectRef> data;
 
   static constexpr const char* _type_key = "Array";
-  MXNET_DECLARE_FINAL_OBJECT_INFO(ArrayNode, Object);
+  MXNET_DECLARE_FINAL_OBJECT_INFO(ArrayNode, Object)
 };
 
 /*!

include/mxnet/runtime/c_runtime_api.h

@@ -47,14 +47,11 @@ typedef enum {
   kNull = 4U,
   kMXNetType = 5U,
   kMXNetContext = 6U,
-  kArrayHandle = 7U,
-  kObjectHandle = 8U,
-  kModuleHandle = 9U,
-  kFuncHandle = 10U,
-  kStr = 11U,
-  kBytes = 12U,
-  kNDArrayContainer = 13U,
-  kNDArrayHandle = 14U,
+  kObjectHandle = 7U,
+  kStr = 8U,
+  kBytes = 9U,
+  kPyArg = 10U,
+  kNDArrayHandle = 11U,
   // Extension codes for other frameworks to integrate MXNet PackedFunc.
   // To make sure each framework's id do not conflict, use first and
   // last sections to mark ranges.
@@ -159,6 +156,24 @@ MXNET_DLL int MXNetFuncListGlobalNames(int* out_size,
  */
 MXNET_DLL int MXNetObjectFree(MXNetObjectHandle obj);
 
+
+/*!
+ * \brief Get the type_index from an object.
+ *
+ * \param obj The object handle.
+ * \param out_tindex the output type index.
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXNetObjectGetTypeIndex(MXNetObjectHandle obj, unsigned* out_tindex);
+
+/*!
+ * \brief Convert type key to type index.
+ * \param type_key The key of the type.
+ * \param out_tindex the corresponding type index.
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXNetObjectTypeKey2Index(const char* type_key, unsigned* out_tindex);
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif

include/mxnet/runtime/container.h

@@ -171,9 +171,9 @@ class ADTObj : public Object, public InplaceArrayBase<ADTObj, ObjectRef> {
   uint32_t size{0};
   // The fields of the structure follows directly in memory.
 
-  static constexpr const uint32_t _type_index = TypeIndex::kMXNetADT;
   static constexpr const char* _type_key = "MXNet.ADT";
-  MXNET_DECLARE_FINAL_OBJECT_INFO(ADTObj, Object);
+  static constexpr const uint32_t _type_index = TypeIndex::kMXNetADT;
+  MXNET_DECLARE_FINAL_OBJECT_INFO(ADTObj, Object)
 
  private:
   /*!
@@ -273,7 +273,7 @@ class ADT : public ObjectRef {
     return ADT(0, std::forward<Args>(args)...);
   }
 
-  MXNET_DEFINE_OBJECT_REF_METHODS(ADT, ObjectRef, ADTObj);
+  MXNET_DEFINE_OBJECT_REF_METHODS(ADT, ObjectRef, ADTObj)
 };
 
 }  // namespace runtime