Use argument queue_ref for empty array allocation for parfor

Author: ZzEeKkAa

numba_dpex/core/parfors/parfor_pass.py

@@ -17,8 +17,8 @@
 from numba.core import config, errors, ir, types, typing
 from numba.core.compiler_machinery import register_pass
 from numba.core.ir_utils import (
+    convert_size_to_var,
     dprint_func_ir,
-    mk_alloc,
     mk_unique_var,
     next_label,
 )
@@ -43,6 +43,7 @@
 )
 from numba.stencils.stencilparfor import StencilPass
 
+from numba_dpex.core.types.dpnp_ndarray_type import DpnpNdArray
 from numba_dpex.core.typing import dpnpdecl
 
 
@@ -58,6 +59,37 @@ class ConvertDPNPPass(ConvertNumpyPass):
     def __init__(self, pass_states):
         super().__init__(pass_states)
 
+    def _get_queue(self, queue_type, expr: tuple):
+        """
+        Extracts queue from the input arguments of the array operation.
+        """
+        pass_states = self.pass_states
+        typemap: map[str, any] = pass_states.typemap
+
+        var_with_queue = None
+
+        for var in expr[1]:
+            if isinstance(var, tuple):
+                res = self._get_queue(queue_type, var)
+                if res is not None:
+                    return res
+
+                continue
+
+            if not isinstance(var, ir.Var):
+                continue
+
+            _type = typemap[var.name]
+            if not isinstance(_type, DpnpNdArray):
+                continue
+            if queue_type != _type.queue:
+                continue
+
+            var_with_queue = var
+            break
+
+        return ir.Expr.getattr(var_with_queue, "sycl_queue", var_with_queue.loc)
+
     def _arrayexpr_to_parfor(self, equiv_set, lhs, arrayexpr, avail_vars):
         """generate parfor from arrayexpr node, which is essentially a
         map with recursive tree.
@@ -77,6 +109,10 @@ def _arrayexpr_to_parfor(self, equiv_set, lhs, arrayexpr, avail_vars):
             pass_states.typemap, size_vars, scope, loc
         )
 
+        # Expr is a tuple
+        ir_queue = self._get_queue(arr_typ.queue, expr)
+        assert ir_queue is not None
+
         # generate init block and body
         init_block = ir.Block(scope, loc)
         init_block.body = mk_alloc(
@@ -89,6 +125,7 @@ def _arrayexpr_to_parfor(self, equiv_set, lhs, arrayexpr, avail_vars):
             scope,
             loc,
             pass_states.typemap[lhs.name],
+            queue_ir_val=ir_queue,
         )
         body_label = next_label()
         body_block = ir.Block(scope, loc)
@@ -469,3 +506,66 @@ def _arrayexpr_tree_to_ir(
     typemap.pop(expr_out_var.name, None)
     typemap[expr_out_var.name] = el_typ
     return out_ir
+
+
+def mk_alloc(
+    typingctx,
+    typemap,
+    calltypes,
+    lhs,
+    size_var,
+    dtype,
+    scope,
+    loc,
+    lhs_typ,
+    **kws,
+):
+    """generate an array allocation with np.empty() and return list of nodes.
+    size_var can be an int variable or tuple of int variables.
+    lhs_typ is the type of the array being allocated.
+
+    Taken from numba, added kws argument to pass it to __allocate__
+    """
+    out = []
+    ndims = 1
+    size_typ = types.intp
+    if isinstance(size_var, tuple):
+        if len(size_var) == 1:
+            size_var = size_var[0]
+            size_var = convert_size_to_var(size_var, typemap, scope, loc, out)
+        else:
+            # tuple_var = build_tuple([size_var...])
+            ndims = len(size_var)
+            tuple_var = ir.Var(scope, mk_unique_var("$tuple_var"), loc)
+            if typemap:
+                typemap[tuple_var.name] = types.containers.UniTuple(
+                    types.intp, ndims
+                )
+            # constant sizes need to be assigned to vars
+            new_sizes = [
+                convert_size_to_var(s, typemap, scope, loc, out)
+                for s in size_var
+            ]
+            tuple_call = ir.Expr.build_tuple(new_sizes, loc)
+            tuple_assign = ir.Assign(tuple_call, tuple_var, loc)
+            out.append(tuple_assign)
+            size_var = tuple_var
+            size_typ = types.containers.UniTuple(types.intp, ndims)
+    if hasattr(lhs_typ, "__allocate__"):
+        return lhs_typ.__allocate__(
+            typingctx,
+            typemap,
+            calltypes,
+            lhs,
+            size_var,
+            dtype,
+            scope,
+            loc,
+            lhs_typ,
+            size_typ,
+            out,
+            **kws,
+        )
+
+    # Unused numba's code..
+    assert False

numba_dpex/core/types/dpnp_ndarray_type.py

@@ -2,14 +2,28 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from functools import partial
+
 import dpnp
 from numba.core import ir, types
 from numba.core.ir_utils import get_np_ufunc_typ, mk_unique_var
-from numba.core.pythonapi import NativeValue, PythonAPI, box, unbox
 
 from .usm_ndarray_type import USMNdArray
 
 
+def partialclass(cls, *args, **kwds):
+    """Creates fabric class of the original class with preset initialization
+    arguments."""
+    cls0 = partial(cls, *args, **kwds)
+    new_cls = type(
+        cls.__name__ + "Partial",
+        (cls,),
+        {"__new__": lambda cls, *args, **kwds: cls0(*args, **kwds)},
+    )
+
+    return new_cls
+
+
 class DpnpNdArray(USMNdArray):
     """
     The Numba type to represent an dpnp.ndarray. The type has the same
@@ -40,15 +54,22 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
         Returns: The DpnpNdArray class.
         """
         if method == "__call__":
-            if not all(
-                (
-                    isinstance(inp, DpnpNdArray)
-                    or isinstance(inp, types.abstract.Number)
-                )
-                for inp in inputs
-            ):
+            dpnp_type = None
+
+            for inp in inputs:
+                if isinstance(inp, DpnpNdArray):
+                    dpnp_type = inp
+                    continue
+                if isinstance(inp, types.abstract.Number):
+                    continue
+
                 return NotImplemented
-            return DpnpNdArray
+
+            assert dpnp_type is not None
+
+            return partialclass(
+                DpnpNdArray, queue=dpnp_type.queue, usm_type=dpnp_type.usm_type
+            )
         else:
             return
 
@@ -71,6 +92,8 @@ def __allocate__(
         lhs_typ,
         size_typ,
         out,
+        # dpex specific argument:
+        queue_ir_val=None,
     ):
         """Generates the Numba typed IR representing the allocation of a new
         DpnpNdArray using the dpnp.ndarray overload.
@@ -94,6 +117,10 @@ def __allocate__(
 
         Returns: The IR Value for the allocated array
         """
+        # TODO: it looks like it is being called only for parfor allocations,
+        # so can we rely on it? We can grab information from input arguments
+        # from rhs, but doc does not set any restriction on parfor use only.
+        assert queue_ir_val is not None
         g_np_var = ir.Var(scope, mk_unique_var("$np_g_var"), loc)
         if typemap:
             typemap[g_np_var.name] = types.misc.Module(dpnp)
@@ -132,11 +159,13 @@ def __allocate__(
         usm_typ_var = ir.Var(scope, mk_unique_var("$np_usm_type_var"), loc)
         # A default device string arg added as a placeholder
         device_typ_var = ir.Var(scope, mk_unique_var("$np_device_var"), loc)
+        queue_typ_var = ir.Var(scope, mk_unique_var("$np_queue_var"), loc)
 
         if typemap:
             typemap[layout_var.name] = types.literal(lhs_typ.layout)
             typemap[usm_typ_var.name] = types.literal(lhs_typ.usm_type)
-            typemap[device_typ_var.name] = types.literal(lhs_typ.device)
+            typemap[device_typ_var.name] = types.none
+            typemap[queue_typ_var.name] = lhs_typ.queue
 
         layout_var_assign = ir.Assign(
             ir.Const(lhs_typ.layout, loc), layout_var, loc
@@ -145,16 +174,29 @@ def __allocate__(
             ir.Const(lhs_typ.usm_type, loc), usm_typ_var, loc
         )
         device_typ_var_assign = ir.Assign(
-            ir.Const(lhs_typ.device, loc), device_typ_var, loc
+            ir.Const(None, loc), device_typ_var, loc
         )
+        queue_typ_var_assign = ir.Assign(queue_ir_val, queue_typ_var, loc)
 
         out.extend(
-            [layout_var_assign, usm_typ_var_assign, device_typ_var_assign]
+            [
+                layout_var_assign,
+                usm_typ_var_assign,
+                device_typ_var_assign,
+                queue_typ_var_assign,
+            ]
         )
 
         alloc_call = ir.Expr.call(
             attr_var,
-            [size_var, typ_var, layout_var, device_typ_var, usm_typ_var],
+            [
+                size_var,
+                typ_var,
+                layout_var,
+                device_typ_var,
+                usm_typ_var,
+                queue_typ_var,
+            ],
             (),
             loc,
         )
@@ -170,6 +212,7 @@ def __allocate__(
                         layout_var,
                         device_typ_var,
                         usm_typ_var,
+                        queue_typ_var,
                     ]
                 ],
                 {},