Repository: numba/numba

Issue:
stencil kernel performance not scaling in njit(parallel=True)
<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

-->
```python
import numba as nb
import numpy as np

@nb.stencil
def kernel1(a):
    return 0.25 * (a[0, 1] + a[1, 0] + a[0, -1] + a[-1, 0])

@nb.njit
def calc1(in_arr):
    out_arr = kernel1(in_arr)
    return out_arr

@nb.njit(parallel= True) 
def calc2(in_arr):
    out_arr = kernel1(in_arr)
    return out_arr

@nb.njit
def calc3(in_arr):
    out_arr = np.zeros(in_arr.shape)
    for i in range(1, in_arr.shape[0]-1):
        for j in range(1, in_arr.shape[1]-1):
            out_arr[i,j] = 0.25 * (in_arr[i, j+1] + in_arr[i+1, j] + in_arr[i, j-1] + in_arr[i-1, j])
    return out_arr

@nb.njit(parallel=True)
def calc4(in_arr):
    out_arr = np.zeros(in_arr.shape)
    for i in range(1, in_arr.shape[0]-1):
        for j in range(1, in_arr.shape[1]-1):
            out_arr[i,j] = 0.25 * (in_arr[i, j+1] + in_arr[i+1, j] + in_arr[i, j-1] + in_arr[i-1, j])
    return out_arr


input_arr = np.linspace(1, 50000, 500000000).reshape((50000, 10000))
```

For a large matrix (`input_arr`), I am experiencing the following run times (avg. 5 runs):
`calc1(input_arr)` [njit + stencil] - `2.09s`
`calc2(input_arr)` [Parallel njit + stencil] - `2.35s`
`calc3(input_arr)` [njit + Loops] - `2.10s`
`calc4(input_arr)` [Parallel njit + Loops] - `1.08s`

As per the above timings, `njit + stencil` is performing similar to `njit + Loops`, as by-default the stencil kernels are not executed in parallel.
I am trying to understand why `Parallel njit + stencil` is not performing similar to `Parallel njit + Loops`?
How can I parallelize the stencil execution?
I have 16 threads available on my machine.

Any solution will be highly appreciated. Thank you!

The first time this issue was mentioned in Github - https://github.com/numba/numba/issues/2982

- [x] I have tried using the latest released version of Numba (most recent is
 visible in the change log (https://github.com/numba/numba/blob/master/CHANGE_LOG).
- [x] I have included below a minimal working reproducer (if you are unsure how
 to write one see http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports).




Discussion:
I can confirm the performance seen here. I saw multiple threads running and the process taking 200% cpu time in `parallel=True` versions. 

Ping @DrTodd13.
@animator I believe the problem is that when stencil parfor calls a kernel it allocates the result with np.full with the stencil's cval.  I guess np.full with 0 is much slower than np.zeros.  If you pre-allocate the output using np.zeros and then pass it to the kernel as calc2(in_arr, out=out_arr) then I think you'll see that that is equivalent to calc4 (my local testing shows this to be true).  This is a workaround but I am working on a permanent fix to allocate the output with np.empty and to only fill in the borders with cval.  Hopefully that is much faster than np.full.
Thank you for investigating @DrTodd13 . Eagerly waiting for the fix.

Regards.

Resolving patch:
diff --git a/numba/stencils/stencil.py b/numba/stencils/stencil.py
--- a/numba/stencils/stencil.py
+++ b/numba/stencils/stencil.py
@@ -602,18 +602,25 @@ def cval_as_str(cval):
         if result is None:
             return_type_name = numpy_support.as_dtype(
                                return_type.dtype).type.__name__
+            out_init ="{} = np.empty({}, dtype=np.{})\n".format(
+                        out_name, shape_name, return_type_name)
+
             if "cval" in self.options:
                 cval = self.options["cval"]
-                if return_type.dtype != typing.typeof.typeof(cval):
+                cval_ty = typing.typeof.typeof(cval)
+                if not self._typingctx.can_convert(cval_ty, return_type.dtype):
                     msg = "cval type does not match stencil return type."
                     raise NumbaValueError(msg)
-                out_init ="{} = np.full({}, {}, dtype=np.{})\n".format(
-                            out_name, shape_name, cval_as_str(cval),
-                            return_type_name)
             else:
-                out_init ="{} = np.zeros({}, dtype=np.{})\n".format(
-                            out_name, shape_name, return_type_name)
+                 cval = 0
             func_text += "    " + out_init
+            for dim in range(the_array.ndim):
+                start_items = [":"] * the_array.ndim
+                end_items = [":"] * the_array.ndim
+                start_items[dim] = ":-{}".format(self.neighborhood[dim][0])
+                end_items[dim] = "-{}:".format(self.neighborhood[dim][1])
+                func_text += "    " + "{}[{}] = {}\n".format(out_name, ",".join(start_items), cval_as_str(cval))
+                func_text += "    " + "{}[{}] = {}\n".format(out_name, ",".join(end_items), cval_as_str(cval))
         else: # result is present, if cval is set then use it
             if "cval" in self.options:
                 cval = self.options["cval"]
diff --git a/numba/stencils/stencilparfor.py b/numba/stencils/stencilparfor.py
--- a/numba/stencils/stencilparfor.py
+++ b/numba/stencils/stencilparfor.py
@@ -176,6 +176,8 @@ def _mk_stencil_parfor(self, label, in_args, out_arr, stencil_ir,
 
         if config.DEBUG_ARRAY_OPT >= 1:
             print("stencil_blocks after replace stencil accesses")
+            print("start_lengths:", start_lengths)
+            print("end_lengths:", end_lengths)
             ir_utils.dump_blocks(stencil_blocks)
 
         # create parfor loop nests
@@ -184,11 +186,15 @@ def _mk_stencil_parfor(self, label, in_args, out_arr, stencil_ir,
         in_arr_dim_sizes = equiv_set.get_shape(in_arr)
 
         assert ndims == len(in_arr_dim_sizes)
+        start_inds = []
+        last_inds = []
         for i in range(ndims):
             last_ind = self._get_stencil_last_ind(in_arr_dim_sizes[i],
                                         end_lengths[i], gen_nodes, scope, loc)
             start_ind = self._get_stencil_start_ind(
                                         start_lengths[i], gen_nodes, scope, loc)
+            start_inds.append(start_ind)
+            last_inds.append(last_ind)
             # start from stencil size to avoid invalid array access
             loopnests.append(numba.parfors.parfor.LoopNest(parfor_vars[i],
                                 start_ind, last_ind, 1))
@@ -235,7 +241,8 @@ def _mk_stencil_parfor(self, label, in_args, out_arr, stencil_ir,
             if "cval" in stencil_func.options:
                 cval = stencil_func.options["cval"]
                 # TODO: Loosen this restriction to adhere to casting rules.
-                if return_type.dtype != typing.typeof.typeof(cval):
+                cval_ty = typing.typeof.typeof(cval)
+                if not self.typingctx.can_convert(cval_ty, return_type.dtype):
                     raise ValueError("cval type does not match stencil return type.")
 
                 temp2 = return_type.dtype(cval)
@@ -263,13 +270,133 @@ def _mk_stencil_parfor(self, label, in_args, out_arr, stencil_ir,
             dtype_attr_assign = ir.Assign(dtype_np_attr_call, dtype_attr_var, loc)
             init_block.body.append(dtype_attr_assign)
 
-            stmts = ir_utils.gen_np_call("full",
-                                       np.full,
+            stmts = ir_utils.gen_np_call("empty",
+                                       np.empty,
                                        out_arr,
-                                       [shape_var, zero_var, dtype_attr_var],
+                                       [shape_var, dtype_attr_var],
                                        self.typingctx,
                                        self.typemap,
                                        self.calltypes)
+            # ------------------
+            # Generate the code to fill just the border with zero_var.
+
+            # Generate a none var to use in slicing.
+            none_var = ir.Var(scope, mk_unique_var("$none_var"), loc)
+            none_assign = ir.Assign(ir.Const(None, loc), none_var, loc)
+            stmts.append(none_assign)
+            self.typemap[none_var.name] = types.none
+            # Generate a zero var to use in slicing.
+            zero_index_var = ir.Var(scope, mk_unique_var("$zero_index_var"), loc)
+            zero_index_assign = ir.Assign(ir.Const(0, loc), zero_index_var, loc)
+            stmts.append(zero_index_assign)
+            self.typemap[zero_index_var.name] = types.intp
+            # Generate generic ":" slice.
+            # ---- Generate var to hold slice func var.
+            slice_func_var = ir.Var(scope, mk_unique_var("$slice_func_var"), loc)
+            slice_fn_ty = self.typingctx.resolve_value_type(slice)
+            self.typemap[slice_func_var.name] = slice_fn_ty
+            slice_g = ir.Global('slice', slice, loc)
+            slice_assign = ir.Assign(slice_g, slice_func_var, loc)
+            stmts.append(slice_assign)
+            # ---- Generate call to slice func.
+            sig = self.typingctx.resolve_function_type(slice_fn_ty,
+                                                       (types.none,) * 2,
+                                                       {})
+            slice_callexpr = ir.Expr.call(func=slice_func_var,
+                                          args=(none_var, none_var),
+                                          kws=(),
+                                          loc=loc)
+            self.calltypes[slice_callexpr] = sig
+            # ---- Generate slice var
+            slice_var = ir.Var(scope, mk_unique_var("$slice"), loc)
+            self.typemap[slice_var.name] = types.slice2_type
+            slice_assign = ir.Assign(slice_callexpr, slice_var, loc)
+            stmts.append(slice_assign)
+
+            def handle_border(slice_fn_ty,
+                              dim,
+                              scope,
+                              loc,
+                              slice_func_var,
+                              stmts,
+                              border_inds,
+                              border_tuple_items,
+                              other_arg,
+                              other_first):
+                # Handle the border for start or end of the index range.
+                # ---- Generate call to slice func.
+                sig = self.typingctx.resolve_function_type(
+                    slice_fn_ty,
+                    (types.intp,) * 2,
+                    {})
+                si = border_inds[dim]
+                assert(isinstance(si, (int, ir.Var)))
+                si_var = ir.Var(scope, mk_unique_var("$border_ind"), loc)
+                self.typemap[si_var.name] = types.intp
+                if isinstance(si, int):
+                    si_assign = ir.Assign(ir.Const(si, loc), si_var, loc)
+                else:
+                    si_assign = ir.Assign(si, si_var, loc)
+                stmts.append(si_assign)
+
+                slice_callexpr = ir.Expr.call(
+                    func=slice_func_var,
+                    args=(other_arg, si_var) if other_first else (si_var, other_arg),
+                    kws=(),
+                    loc=loc)
+                self.calltypes[slice_callexpr] = sig
+                # ---- Generate slice var
+                border_slice_var = ir.Var(scope, mk_unique_var("$slice"), loc)
+                self.typemap[border_slice_var.name] = types.slice2_type
+                slice_assign = ir.Assign(slice_callexpr, border_slice_var, loc)
+                stmts.append(slice_assign)
+
+                border_tuple_items[dim] = border_slice_var
+                border_ind_var = ir.Var(scope, mk_unique_var(
+                    "$border_index_tuple_var"), loc)
+                self.typemap[border_ind_var.name] = types.containers.UniTuple(
+                    types.slice2_type, ndims)
+                tuple_call = ir.Expr.build_tuple(border_tuple_items, loc)
+                tuple_assign = ir.Assign(tuple_call, border_ind_var, loc)
+                stmts.append(tuple_assign)
+
+                setitem_call = ir.SetItem(out_arr, border_ind_var, zero_var, loc)
+                self.calltypes[setitem_call] = signature(
+                                                types.none, self.typemap[out_arr.name],
+                                                self.typemap[border_ind_var.name],
+                                                self.typemap[out_arr.name].dtype
+                                                )
+                stmts.append(setitem_call)
+
+            # For each dimension, add setitem to set border values.
+            for dim in range(in_arr_typ.ndim):
+                # First, fill all entries with ":".
+                start_tuple_items = [slice_var] * in_arr_typ.ndim
+                last_tuple_items = [slice_var] * in_arr_typ.ndim
+
+                handle_border(slice_fn_ty,
+                              dim,
+                              scope,
+                              loc,
+                              slice_func_var,
+                              stmts,
+                              start_inds,
+                              start_tuple_items,
+                              zero_index_var,
+                              True)
+                handle_border(slice_fn_ty,
+                              dim,
+                              scope,
+                              loc,
+                              slice_func_var,
+                              stmts,
+                              last_inds,
+                              last_tuple_items,
+                              in_arr_dim_sizes[dim],
+                              False)
+
+            # ------------------
+
             equiv_set.insert_equiv(out_arr, in_arr_dim_sizes)
             init_block.body.extend(stmts)
         else: # out is present


Test patch:
diff --git a/numba/tests/test_stencils.py b/numba/tests/test_stencils.py
--- a/numba/tests/test_stencils.py
+++ b/numba/tests/test_stencils.py
@@ -2066,15 +2066,11 @@ def kernel(a):
         self.check(kernel, a)
 
     def test_basic35(self):
-        """simple cval """
+        """simple cval where cval is int but castable to dtype of float"""
         def kernel(a):
             return a[0, 1]
         a = np.arange(12.).reshape(3, 4)
-        ex = self.exception_dict(
-            stencil=NumbaValueError,
-            parfor=ValueError,
-            njit=NumbaValueError)
-        self.check(kernel, a, options={'cval': 5}, expected_exception=ex)
+        self.check(kernel, a, options={'cval': 5})
 
     def test_basic36(self):
         """more complex with cval"""