np.arange (#191)

Author: Hardcode84

numba_dpcomp/numba_dpcomp/mlir/linalg_builder.py

@@ -143,7 +143,7 @@ def linalg_generic(self, inputs, outputs, iterators, maps, body):
     def linalg_index(self, dim):
         return self._linalg_index(self._context, dim)
 
-    def from_elements(self, values, dtype):
+    def from_elements(self, values, dtype=None):
         return self._from_elements(self._context, values, dtype)
 
     def extract(self, value, indices):

numba_dpcomp/numba_dpcomp/mlir/numpy/funcs.py

@@ -396,6 +396,36 @@ def ones_like_impl(builder, arr):
     return _init_impl(builder, arr.shape, arr.dtype, 1)
 
 
+@register_func("numpy.arange", numpy.arange)
+def arange_impl(builder, start, stop=None, step=None, dtype=None):
+    if stop is None:
+        stop = start
+        start = 0
+
+    if step is None:
+        step = 1
+
+    if dtype is None:
+        dtype = builder.int64
+
+    inc = builder.select(step < 0, 1, -1)
+    count = (stop - start + step + inc) // step
+    count = builder.select(count < 0, 0, count)
+
+    start = builder.from_elements(start)
+    step = builder.from_elements(step)
+    init = builder.init_tensor([count], dtype)
+
+    iterators = ["parallel"]
+    maps = ["(d0) -> (0)", "(d0) -> (0)", "(d0) -> (d0)"]
+
+    def body(a, b, c):
+        i = _linalg_index(0)
+        return a + b * i
+
+    return builder.linalg_generic((start, step), init, iterators, maps, body)
+
+
 @register_func("numpy.eye", numpy.eye)
 def eye_impl(builder, N, M=None, k=0, dtype=None):
     if M is None:

numba_dpcomp/numba_dpcomp/mlir/tests/test_numba_parfor.py

@@ -82,10 +82,10 @@ def _gen_tests():
         "test_fuse_argmin_argmax_max_min",  # numpy argmin, argmax
         "test_max",  # max reduction
         "test_min",  # min reduction
-        "test_arange",  # numpy.arange
+        "test_arange",  # select issue, complex
         "test_pi",  # np.random.ranf
         "test_simple20",  # AssertionError not raised
-        "test_simple24",  # numpy.arange
+        "test_simple24",  # getitem with array
         "test_0d_array",  # numpy prod
         "test_argmin",  # numpy.argmin
         "test_argmax",  # numpy.argmax
@@ -104,19 +104,18 @@ def _gen_tests():
         "test_simple19",  # np.dot unsupported args
         "test_no_hoisting_with_member_function_call",  # set support
         "test_parfor_dtype_type",  # dtype cast
-        "test_tuple3",  # numpy.arange
         "test_parfor_array_access3",  # TypeError: unsupported operand type(s) for -: 'NoneType' and 'NoneType'
         "test_preparfor_canonicalize_kws",  # array.argsort
         "test_parfor_array_access4",  # np.dot unsupported args
         "test_tuple_concat_with_reverse_slice",  # enumerate
         "test_reduce",  # functools.reduce
-        "test_two_d_array_reduction",  # np.arange
+        "test_two_d_array_reduction",  # 'memref<?x?xi64>' and result type 'memref<?x?xi32>' are cast incompatible
         "test_tuple_concat",  # tuple concat
         "test_two_d_array_reduction_with_float_sizes",  # np.array
-        "test_two_d_array_reduction_reuse",  # np.arange
-        "test_parfor_array_access_lower_slice",  # np.arange
+        "test_two_d_array_reduction_reuse",  # 'memref<?x?xi64>' and result type 'memref<?x?xi32>' are cast incompatible
+        "test_parfor_array_access_lower_slice",  # plier.getitem
         "test_size_assertion",  # AssertionError not raised
-        "test_parfor_slice18",  # np.arange
+        "test_parfor_slice18",  # cast types mismatch
         "test_simple12",  # complex128
         "test_parfor_slice2",  # AssertionError not raised
         "test_parfor_slice6",  # array.transpose

numba_dpcomp/numba_dpcomp/mlir/tests/test_numpy.py

@@ -838,6 +838,29 @@ def py_func(d):
     assert_equal(py_func(a), jit_func(a))
 
 
+@parametrize_function_variants(
+    "py_func",
+    [
+        "lambda : np.arange(0)",
+        "lambda : np.arange(1)",
+        "lambda : np.arange(7)",
+        "lambda : np.arange(-1)",
+        "lambda : np.arange(-1,6)",
+        "lambda : np.arange(-1,6,1)",
+        "lambda : np.arange(-1,6,2)",
+        "lambda : np.arange(-1,6,3)",
+        "lambda : np.arange(6,-1,-1)",
+        "lambda : np.arange(6,-1,-2)",
+        "lambda : np.arange(6,-1,-3)",
+        "lambda : np.arange(5,dtype=np.int32)",
+        "lambda : np.arange(5,dtype=np.float32)",
+    ],
+)
+def test_arange(py_func):
+    jit_func = njit(py_func)
+    assert_equal(py_func(), jit_func())
+
+
 @pytest.mark.parametrize("dtype", [np.int32, np.int64, np.float32, np.float64])
 def test_dtype_param(dtype):
     def py_func(dt):

numba_dpcomp/numba_dpcomp/mlir_compiler/lib/py_linalg_resolver.cpp

@@ -970,11 +970,21 @@ static py::object fromElementsImpl(py::capsule context, py::handle values,
   auto &ctx = getPyContext(context);
   auto &builder = ctx.builder;
   auto loc = ctx.loc;
-  auto type = unwrapType(dtype);
+  mlir::Type type;
+  if (!dtype.is_none())
+    type = unwrapType(dtype);
 
   llvm::SmallVector<mlir::Value> vals(containerSize(values));
   containerIterate(values, [&](auto index, py::handle obj) {
-    vals[index] = ctx.context.unwrapVal(loc, builder, obj, type);
+    auto val = [&]() -> mlir::Value {
+      if (type)
+        return ctx.context.unwrapVal(loc, builder, obj, type);
+
+      auto v = ctx.context.unwrapVal(loc, builder, obj);
+      type = v.getType();
+      return v;
+    }();
+    vals[index] = val;
   });
 
   if (vals.empty())