Merge pull request #945 from IntelPython/fix-arange-issues

Fix arange issues

Author: oleksandr-pavlyk

dpctl/tensor/_ctors.py

@@ -483,21 +483,22 @@ def _coerce_and_infer_dt(*args, dt, sycl_queue, err_msg, allow_bool=False):
 
 def _round_for_arange(tmp):
     k = int(tmp)
-    if k > 0 and float(k) < tmp:
+    if k >= 0 and float(k) < tmp:
         tmp = tmp + 1
     return tmp
 
 
 def _get_arange_length(start, stop, step):
     "Compute length of arange sequence"
     span = stop - start
-    if type(step) in [int, float] and type(span) in [int, float]:
+    if hasattr(step, "__float__") and hasattr(span, "__float__"):
         return _round_for_arange(span / step)
     tmp = span / step
-    if type(tmp) is complex and tmp.imag == 0:
+    if hasattr(tmp, "__complex__"):
+        tmp = complex(tmp)
         tmp = tmp.real
     else:
-        return tmp
+        tmp = float(tmp)
     return _round_for_arange(tmp)
 
 
@@ -536,13 +537,18 @@ def arange(
     if stop is None:
         stop = start
         start = 0
+    if step is None:
+        step = 1
     dpctl.utils.validate_usm_type(usm_type, allow_none=False)
     sycl_queue = normalize_queue_device(sycl_queue=sycl_queue, device=device)
-    (start, stop, step,), dt = _coerce_and_infer_dt(
+    is_bool = False
+    if dtype:
+        is_bool = (dtype is bool) or (dpt.dtype(dtype) == dpt.bool)
+    (start_, stop_, step_), dt = _coerce_and_infer_dt(
         start,
         stop,
         step,
-        dt=dtype,
+        dt=dpt.int8 if is_bool else dtype,
         sycl_queue=sycl_queue,
         err_msg="start, stop, and step must be Python scalars",
         allow_bool=False,
@@ -554,18 +560,40 @@ def arange(
             sh = 0
     except TypeError:
         sh = 0
+    if is_bool and sh > 2:
+        raise ValueError("no fill-function for boolean data type")
     res = dpt.usm_ndarray(
         (sh,),
         dtype=dt,
         buffer=usm_type,
         order="C",
         buffer_ctor_kwargs={"queue": sycl_queue},
     )
-    _step = (start + step) - start
-    _step = dt.type(_step)
-    _start = dt.type(start)
+    sc_ty = dt.type
+    _first = sc_ty(start)
+    if sh > 1:
+        _second = sc_ty(start + step)
+        if dt in [dpt.uint8, dpt.uint16, dpt.uint32, dpt.uint64]:
+            int64_ty = dpt.int64.type
+            _step = int64_ty(_second) - int64_ty(_first)
+        else:
+            _step = _second - _first
+        _step = sc_ty(_step)
+    else:
+        _step = sc_ty(1)
+    _start = _first
     hev, _ = ti._linspace_step(_start, _step, res, sycl_queue)
     hev.wait()
+    if is_bool:
+        res_out = dpt.usm_ndarray(
+            (sh,),
+            dtype=dpt.bool,
+            buffer=usm_type,
+            order="C",
+            buffer_ctor_kwargs={"queue": sycl_queue},
+        )
+        res_out[:] = res
+        res = res_out
     return res
 
 

dpctl/tests/helper/__init__.py

@@ -19,14 +19,18 @@
 
 from ._helper import (
     create_invalid_capsule,
+    get_queue_or_skip,
     has_cpu,
     has_gpu,
     has_sycl_platforms,
+    skip_if_dtype_not_supported,
 )
 
 __all__ = [
     "create_invalid_capsule",
     "has_cpu",
     "has_gpu",
     "has_sycl_platforms",
+    "get_queue_or_skip",
+    "skip_if_dtype_not_supported",
 ]

dpctl/tests/helper/_helper.py

@@ -14,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
+
 import dpctl
 
 
@@ -39,3 +41,38 @@ def create_invalid_capsule():
     ctor.restype = ctypes.py_object
     ctor.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p]
     return ctor(id(ctor), b"invalid", 0)
+
+
+def get_queue_or_skip(args=tuple()):
+    try:
+        q = dpctl.SyclQueue(*args)
+    except dpctl.SyclQueueCreationError:
+        pytest.skip(f"Queue could not be created from {args}")
+    return q
+
+
+def skip_if_dtype_not_supported(dt, q_or_dev):
+    import dpctl.tensor as dpt
+
+    dt = dpt.dtype(dt)
+    if type(q_or_dev) is dpctl.SyclQueue:
+        dev = q_or_dev.sycl_device
+    elif type(q_or_dev) is dpctl.SyclDevice:
+        dev = q_or_dev
+    else:
+        raise TypeError(
+            "Expected dpctl.SyclQueue or dpctl.SyclDevice, "
+            f"got {type(q_or_dev)}"
+        )
+    dev_has_dp = dev.has_aspect_fp64
+    if dev_has_dp is False and dt in [dpt.float64, dpt.complex128]:
+        pytest.skip(
+            f"{dev.name} does not support double precision floating point types"
+        )
+    dev_has_hp = dev.has_aspect_fp16
+    if dev_has_hp is False and dt in [
+        dpt.float16,
+    ]:
+        pytest.skip(
+            f"{dev.name} does not support half precision floating point type"
+        )