Skip to content
Permalink
master
Go to file
 
 
Cannot retrieve contributors at this time
4580 lines (4086 sloc) 190 KB
#
# Copyright (c) 2017 Intel Corporation
# SPDX-License-Identifier: BSD-2-Clause
#
"""
This module transforms data-parallel operations such as Numpy calls into
'Parfor' nodes, which are nested loops that can be parallelized.
It also implements optimizations such as loop fusion, and extends the rest of
compiler analysis and optimizations to support Parfors.
This is similar to ParallelAccelerator package in Julia:
https://github.com/IntelLabs/ParallelAccelerator.jl
'Parallelizing Julia with a Non-invasive DSL', T. Anderson et al., ECOOP'17.
"""
import types as pytypes # avoid confusion with numba.types
import sys, math
import os
import textwrap
import copy
import inspect
import linecache
from functools import reduce
from collections import defaultdict, OrderedDict, namedtuple
from contextlib import contextmanager
import operator
import numba.core.ir
from numba.core import types, typing, utils, errors, ir, analysis, postproc, rewrites, typeinfer, config, ir_utils
from numba import prange, pndindex
from numba.np.numpy_support import as_dtype
from numba.core.typing.templates import infer_global, AbstractTemplate
from numba.stencils.stencilparfor import StencilPass
from numba.core.extending import register_jitable
from numba.core.ir_utils import (
mk_unique_var,
next_label,
mk_alloc,
get_np_ufunc_typ,
mk_range_block,
mk_loop_header,
get_name_var_table,
replace_vars,
replace_vars_inner,
visit_vars,
visit_vars_inner,
remove_dead,
copy_propagate,
get_block_copies,
apply_copy_propagate,
dprint_func_ir,
find_topo_order,
get_stmt_writes,
rename_labels,
get_call_table,
simplify,
simplify_CFG,
has_no_side_effect,
canonicalize_array_math,
add_offset_to_labels,
find_callname,
find_build_sequence,
guard,
require,
GuardException,
compile_to_numba_ir,
get_definition,
build_definitions,
replace_arg_nodes,
replace_returns,
is_getitem,
is_setitem,
is_get_setitem,
index_var_of_get_setitem,
set_index_var_of_get_setitem,
find_potential_aliases,
replace_var_names)
from numba.core.analysis import (compute_use_defs, compute_live_map,
compute_dead_maps, compute_cfg_from_blocks)
from numba.core.controlflow import CFGraph
from numba.core.typing import npydecl, signature
from numba.core.types.functions import Function
from numba.parfors.array_analysis import (random_int_args, random_1arg_size,
random_2arg_sizelast, random_3arg_sizelast,
random_calls, assert_equiv)
from numba.core.extending import overload
import copy
import numpy
import numpy as np
from numba.parfors import array_analysis
import numba.cpython.builtins
from numba.stencils import stencilparfor
# circular dependency: import numba.npyufunc.dufunc.DUFunc
# wrapped pretty print
_termwidth = 80
_txtwrapper = textwrap.TextWrapper(width=_termwidth, drop_whitespace=False)
def print_wrapped(x):
for l in x.splitlines():
[print(y) for y in _txtwrapper.wrap(l)]
sequential_parfor_lowering = False
# init_prange is a sentinel call that specifies the start of the initialization
# code for the computation in the upcoming prange call
# This lets the prange pass to put the code in the generated parfor's init_block
def init_prange():
return
@overload(init_prange)
def init_prange_overload():
def no_op():
return
return no_op
class internal_prange(object):
def __new__(cls, *args):
return range(*args)
def min_parallel_impl(return_type, arg):
# XXX: use prange for 1D arrays since pndindex returns a 1-tuple instead of
# integer. This causes type and fusion issues.
if arg.ndim == 0:
def min_1(in_arr):
return in_arr[()]
elif arg.ndim == 1:
def min_1(in_arr):
numba.parfors.parfor.init_prange()
min_checker(len(in_arr))
val = numba.cpython.builtins.get_type_max_value(in_arr.dtype)
for i in numba.parfors.parfor.internal_prange(len(in_arr)):
val = min(val, in_arr[i])
return val
else:
def min_1(in_arr):
numba.parfors.parfor.init_prange()
min_checker(len(in_arr))
val = numba.cpython.builtins.get_type_max_value(in_arr.dtype)
for i in numba.pndindex(in_arr.shape):
val = min(val, in_arr[i])
return val
return min_1
def max_parallel_impl(return_type, arg):
if arg.ndim == 0:
def max_1(in_arr):
return in_arr[()]
elif arg.ndim == 1:
def max_1(in_arr):
numba.parfors.parfor.init_prange()
max_checker(len(in_arr))
val = numba.cpython.builtins.get_type_min_value(in_arr.dtype)
for i in numba.parfors.parfor.internal_prange(len(in_arr)):
val = max(val, in_arr[i])
return val
else:
def max_1(in_arr):
numba.parfors.parfor.init_prange()
max_checker(len(in_arr))
val = numba.cpython.builtins.get_type_min_value(in_arr.dtype)
for i in numba.pndindex(in_arr.shape):
val = max(val, in_arr[i])
return val
return max_1
def argmin_parallel_impl(in_arr):
numba.parfors.parfor.init_prange()
argmin_checker(len(in_arr))
A = in_arr.ravel()
init_val = numba.cpython.builtins.get_type_max_value(A.dtype)
ival = typing.builtins.IndexValue(0, init_val)
for i in numba.parfors.parfor.internal_prange(len(A)):
curr_ival = typing.builtins.IndexValue(i, A[i])
ival = min(ival, curr_ival)
return ival.index
def argmax_parallel_impl(in_arr):
numba.parfors.parfor.init_prange()
argmax_checker(len(in_arr))
A = in_arr.ravel()
init_val = numba.cpython.builtins.get_type_min_value(A.dtype)
ival = typing.builtins.IndexValue(0, init_val)
for i in numba.parfors.parfor.internal_prange(len(A)):
curr_ival = typing.builtins.IndexValue(i, A[i])
ival = max(ival, curr_ival)
return ival.index
def dotvv_parallel_impl(a, b):
numba.parfors.parfor.init_prange()
l = a.shape[0]
m = b.shape[0]
# TODO: investigate assert_equiv
#assert_equiv("sizes of l, m do not match", l, m)
s = 0
for i in numba.parfors.parfor.internal_prange(l):
s += a[i] * b[i]
return s
def dotvm_parallel_impl(a, b):
numba.parfors.parfor.init_prange()
l = a.shape
m, n = b.shape
# TODO: investigate assert_equiv
#assert_equiv("Sizes of l, m do not match", l, m)
c = np.zeros(n, a.dtype)
# TODO: evaluate dotvm implementation options
#for i in prange(n):
# s = 0
# for j in range(m):
# s += a[j] * b[j, i]
# c[i] = s
for i in numba.parfors.parfor.internal_prange(m):
c += a[i] * b[i, :]
return c
def dotmv_parallel_impl(a, b):
numba.parfors.parfor.init_prange()
m, n = a.shape
l = b.shape
# TODO: investigate assert_equiv
#assert_equiv("sizes of n, l do not match", n, l)
c = np.empty(m, a.dtype)
for i in numba.parfors.parfor.internal_prange(m):
s = 0
for j in range(n):
s += a[i, j] * b[j]
c[i] = s
return c
def dot_parallel_impl(return_type, atyp, btyp):
# Note that matrix matrix multiply is not translated.
if (isinstance(atyp, types.npytypes.Array) and
isinstance(btyp, types.npytypes.Array)):
if atyp.ndim == btyp.ndim == 1:
return dotvv_parallel_impl
# TODO: evaluate support for dotvm and enable
#elif atyp.ndim == 1 and btyp.ndim == 2:
# return dotvm_parallel_impl
elif atyp.ndim == 2 and btyp.ndim == 1:
return dotmv_parallel_impl
def sum_parallel_impl(return_type, arg):
zero = return_type(0)
if arg.ndim == 0:
def sum_1(in_arr):
return in_arr[()]
elif arg.ndim == 1:
def sum_1(in_arr):
numba.parfors.parfor.init_prange()
val = zero
for i in numba.parfors.parfor.internal_prange(len(in_arr)):
val += in_arr[i]
return val
else:
def sum_1(in_arr):
numba.parfors.parfor.init_prange()
val = zero
for i in numba.pndindex(in_arr.shape):
val += in_arr[i]
return val
return sum_1
def prod_parallel_impl(return_type, arg):
one = return_type(1)
if arg.ndim == 0:
def prod_1(in_arr):
return in_arr[()]
elif arg.ndim == 1:
def prod_1(in_arr):
numba.parfors.parfor.init_prange()
val = one
for i in numba.parfors.parfor.internal_prange(len(in_arr)):
val *= in_arr[i]
return val
else:
def prod_1(in_arr):
numba.parfors.parfor.init_prange()
val = one
for i in numba.pndindex(in_arr.shape):
val *= in_arr[i]
return val
return prod_1
def mean_parallel_impl(return_type, arg):
# can't reuse sum since output type is different
zero = return_type(0)
if arg.ndim == 0:
def mean_1(in_arr):
return in_arr[()]
elif arg.ndim == 1:
def mean_1(in_arr):
numba.parfors.parfor.init_prange()
val = zero
for i in numba.parfors.parfor.internal_prange(len(in_arr)):
val += in_arr[i]
return val/len(in_arr)
else:
def mean_1(in_arr):
numba.parfors.parfor.init_prange()
val = zero
for i in numba.pndindex(in_arr.shape):
val += in_arr[i]
return val/in_arr.size
return mean_1
def var_parallel_impl(return_type, arg):
if arg.ndim == 0:
def var_1(in_arr):
return 0
elif arg.ndim == 1:
def var_1(in_arr):
# Compute the mean
m = in_arr.mean()
# Compute the sum of square diffs
numba.parfors.parfor.init_prange()
ssd = 0
for i in numba.parfors.parfor.internal_prange(len(in_arr)):
val = in_arr[i] - m
ssd += np.real(val * np.conj(val))
return ssd / len(in_arr)
else:
def var_1(in_arr):
# Compute the mean
m = in_arr.mean()
# Compute the sum of square diffs
numba.parfors.parfor.init_prange()
ssd = 0
for i in numba.pndindex(in_arr.shape):
val = in_arr[i] - m
ssd += np.real(val * np.conj(val))
return ssd / in_arr.size
return var_1
def std_parallel_impl(return_type, arg):
def std_1(in_arr):
return in_arr.var() ** 0.5
return std_1
def arange_parallel_impl(return_type, *args):
dtype = as_dtype(return_type.dtype)
def arange_1(stop):
return np.arange(0, stop, 1, dtype)
def arange_2(start, stop):
return np.arange(start, stop, 1, dtype)
def arange_3(start, stop, step):
return np.arange(start, stop, step, dtype)
if any(isinstance(a, types.Complex) for a in args):
def arange_4(start, stop, step, dtype):
numba.parfors.parfor.init_prange()
nitems_c = (stop - start) / step
nitems_r = math.ceil(nitems_c.real)
nitems_i = math.ceil(nitems_c.imag)
nitems = int(max(min(nitems_i, nitems_r), 0))
arr = np.empty(nitems, dtype)
for i in numba.parfors.parfor.internal_prange(nitems):
arr[i] = start + i * step
return arr
else:
def arange_4(start, stop, step, dtype):
numba.parfors.parfor.init_prange()
nitems_r = math.ceil((stop - start) / step)
nitems = int(max(nitems_r, 0))
arr = np.empty(nitems, dtype)
val = start
for i in numba.parfors.parfor.internal_prange(nitems):
arr[i] = start + i * step
return arr
if len(args) == 1:
return arange_1
elif len(args) == 2:
return arange_2
elif len(args) == 3:
return arange_3
elif len(args) == 4:
return arange_4
else:
raise ValueError("parallel arange with types {}".format(args))
def linspace_parallel_impl(return_type, *args):
dtype = as_dtype(return_type.dtype)
def linspace_2(start, stop):
return np.linspace(start, stop, 50)
def linspace_3(start, stop, num):
numba.parfors.parfor.init_prange()
arr = np.empty(num, dtype)
div = num - 1
delta = stop - start
arr[0] = start
for i in numba.parfors.parfor.internal_prange(num):
arr[i] = start + delta * (i / div)
return arr
if len(args) == 2:
return linspace_2
elif len(args) == 3:
return linspace_3
else:
raise ValueError("parallel linspace with types {}".format(args))
replace_functions_map = {
('argmin', 'numpy'): lambda r,a: argmin_parallel_impl,
('argmax', 'numpy'): lambda r,a: argmax_parallel_impl,
('min', 'numpy'): min_parallel_impl,
('max', 'numpy'): max_parallel_impl,
('amin', 'numpy'): min_parallel_impl,
('amax', 'numpy'): max_parallel_impl,
('sum', 'numpy'): sum_parallel_impl,
('prod', 'numpy'): prod_parallel_impl,
('mean', 'numpy'): mean_parallel_impl,
('var', 'numpy'): var_parallel_impl,
('std', 'numpy'): std_parallel_impl,
('dot', 'numpy'): dot_parallel_impl,
('arange', 'numpy'): arange_parallel_impl,
('linspace', 'numpy'): linspace_parallel_impl,
}
def fill_parallel_impl(return_type, arr, val):
"""Parallel implemention of ndarray.fill. The array on
which to operate is retrieved from get_call_name and
is passed along with the value to fill.
"""
if arr.ndim == 1:
def fill_1(in_arr, val):
numba.parfors.parfor.init_prange()
for i in numba.parfors.parfor.internal_prange(len(in_arr)):
in_arr[i] = val
return None
else:
def fill_1(in_arr, val):
numba.parfors.parfor.init_prange()
for i in numba.pndindex(in_arr.shape):
in_arr[i] = val
return None
return fill_1
replace_functions_ndarray = {
'fill': fill_parallel_impl,
}
@register_jitable
def max_checker(arr_size):
if arr_size == 0:
raise ValueError(("zero-size array to reduction operation "
"maximum which has no identity"))
@register_jitable
def min_checker(arr_size):
if arr_size == 0:
raise ValueError(("zero-size array to reduction operation "
"minimum which has no identity"))
@register_jitable
def argmin_checker(arr_size):
if arr_size == 0:
raise ValueError("attempt to get argmin of an empty sequence")
@register_jitable
def argmax_checker(arr_size):
if arr_size == 0:
raise ValueError("attempt to get argmax of an empty sequence")
checker_impl = namedtuple('checker_impl', ['name', 'func'])
replace_functions_checkers_map = {
('argmin', 'numpy') : checker_impl('argmin_checker', argmin_checker),
('argmax', 'numpy') : checker_impl('argmax_checker', argmax_checker),
('min', 'numpy') : checker_impl('min_checker', min_checker),
('max', 'numpy') : checker_impl('max_checker', max_checker),
('amin', 'numpy') : checker_impl('min_checker', min_checker),
('amax', 'numpy') : checker_impl('max_checker', max_checker),
}
class LoopNest(object):
'''The LoopNest class holds information of a single loop including
the index variable (of a non-negative integer value), and the
range variable, e.g. range(r) is 0 to r-1 with step size 1.
'''
def __init__(self, index_variable, start, stop, step):
self.index_variable = index_variable
self.start = start
self.stop = stop
self.step = step
def __repr__(self):
return ("LoopNest(index_variable = {}, range = ({}, {}, {}))".
format(self.index_variable, self.start, self.stop, self.step))
def list_vars(self):
all_uses = []
all_uses.append(self.index_variable)
if isinstance(self.start, ir.Var):
all_uses.append(self.start)
if isinstance(self.stop, ir.Var):
all_uses.append(self.stop)
if isinstance(self.step, ir.Var):
all_uses.append(self.step)
return all_uses
class Parfor(ir.Expr, ir.Stmt):
id_counter = 0
def __init__(
self,
loop_nests,
init_block,
loop_body,
loc,
index_var,
equiv_set,
pattern,
flags,
no_sequential_lowering=False,
races=set()):
super(Parfor, self).__init__(
op='parfor',
loc=loc
)
self.id = type(self).id_counter
type(self).id_counter += 1
#self.input_info = input_info
#self.output_info = output_info
self.loop_nests = loop_nests
self.init_block = init_block
self.loop_body = loop_body
self.index_var = index_var
self.params = None # filled right before parallel lowering
self.equiv_set = equiv_set
# The parallel patterns this parfor was generated from and their options
# for example, a parfor could be from the stencil pattern with
# the neighborhood option
assert len(pattern) > 1
self.patterns = [pattern]
self.flags = flags
# if True, this parfor shouldn't be lowered sequentially even with the
# sequential lowering option
self.no_sequential_lowering = no_sequential_lowering
self.races = races
if config.DEBUG_ARRAY_OPT_STATS:
fmt = 'Parallel for-loop #{} is produced from pattern \'{}\' at {}'
print(fmt.format(
self.id, pattern, loc))
def __repr__(self):
return "id=" + str(self.id) + repr(self.loop_nests) + \
repr(self.loop_body) + repr(self.index_var)
def list_vars(self):
"""list variables used (read/written) in this parfor by
traversing the body and combining block uses.
"""
all_uses = []
for l, b in self.loop_body.items():
for stmt in b.body:
all_uses += stmt.list_vars()
for loop in self.loop_nests:
all_uses += loop.list_vars()
for stmt in self.init_block.body:
all_uses += stmt.list_vars()
return all_uses
def get_shape_classes(self, var, typemap=None):
"""get the shape classes for a given variable.
If a typemap is specified then use it for type resolution
"""
# We get shape classes from the equivalence set but that
# keeps its own typemap at a time prior to lowering. So
# if something is added during lowering then we can pass
# in a type map to use. We temporarily replace the
# equivalence set typemap, do the work and then restore
# the original on the way out.
if typemap is not None:
save_typemap = self.equiv_set.typemap
self.equiv_set.typemap = typemap
res = self.equiv_set.get_shape_classes(var)
if typemap is not None:
self.equiv_set.typemap = save_typemap
return res
def dump(self, file=None):
file = file or sys.stdout
print(("begin parfor {}".format(self.id)).center(20, '-'), file=file)
print("index_var = ", self.index_var, file=file)
print("params = ", self.params, file=file)
print("races = ", self.races, file=file)
for loopnest in self.loop_nests:
print(loopnest, file=file)
print("init block:", file=file)
self.init_block.dump(file)
for offset, block in sorted(self.loop_body.items()):
print('label %s:' % (offset,), file=file)
block.dump(file)
print(("end parfor {}".format(self.id)).center(20, '-'), file=file)
def validate_params(self, typemap):
"""
Check that Parfors params are of valid types.
"""
if self.params is None:
msg = ("Cannot run parameter validation on a Parfor with params "
"not set")
raise ValueError(msg)
for p in self.params:
ty = typemap.get(p)
if ty is None:
msg = ("Cannot validate parameter %s, there is no type "
"information available")
raise ValueError(msg)
if isinstance(ty, types.BaseTuple):
if ty.count > config.PARFOR_MAX_TUPLE_SIZE:
msg = ("Use of a tuple (%s) of length %d in a parallel region "
"exceeds the maximum supported tuple size. Since "
"Generalized Universal Functions back parallel regions "
"and those do not support tuples, tuples passed to "
"parallel regions are unpacked if their size is below "
"a certain threshold, currently configured to be %d. "
"This threshold can be modified using the Numba "
"environment variable NUMBA_PARFOR_MAX_TUPLE_SIZE.")
raise errors.UnsupportedParforsError(msg %
(p, ty.count, config.PARFOR_MAX_TUPLE_SIZE),
self.loc)
def _analyze_parfor(parfor, equiv_set, typemap, array_analysis):
"""Recursive array analysis for parfor nodes.
"""
func_ir = array_analysis.func_ir
parfor_blocks = wrap_parfor_blocks(parfor)
# Since init_block get label 0 after wrap, we need to save
# the equivset for the real block label 0.
backup_equivset = array_analysis.equiv_sets.get(0, None)
array_analysis.run(parfor_blocks, equiv_set)
unwrap_parfor_blocks(parfor, parfor_blocks)
parfor.equiv_set = array_analysis.equiv_sets[0]
# Restore equivset for block 0 after parfor is unwrapped
if backup_equivset:
array_analysis.equiv_sets[0] = backup_equivset
return [], []
array_analysis.array_analysis_extensions[Parfor] = _analyze_parfor
class ParforDiagnostics(object):
"""Holds parfor diagnostic info, this is accumulated throughout the
PreParforPass and ParforPass, also in the closure inlining!
"""
def __init__(self):
# holds ref to the function for which this is providing diagnostics
self.func = None
# holds a map of the replaced functions
self.replaced_fns = dict()
# used to identify "internal" parfor functions
self.internal_name = '__numba_parfor_gufunc'
self.fusion_info = defaultdict(list)
self.nested_fusion_info = defaultdict(list)
self.fusion_reports = []
self.hoist_info = {}
self.has_setup = False
def setup(self, func_ir, fusion_enabled):
self.func_ir = func_ir
self.name = self.func_ir.func_id.func_qualname
self.line = self.func_ir.loc
self.fusion_enabled = fusion_enabled
if self.internal_name in self.name:
self.purpose = 'Internal parallel function'
else:
self.purpose = 'Function %s, %s' % (self.name, self.line)
# we store a reference to the parfors prior to fusion etc, the parfors
# do get mangled in the fusion process but in a predetermined manner
# and by holding a reference here the "before" state can be printed
self.initial_parfors = self.get_parfors()
self.has_setup = True
@property
def has_setup(self):
return self._has_setup
@has_setup.setter
def has_setup(self, state):
self._has_setup = state
def count_parfors(self, blocks=None):
return len(self.get_parfors())
def _get_nested_parfors(self, parfor, parfors_list):
blocks = wrap_parfor_blocks(parfor)
self._get_parfors(blocks, parfors_list)
unwrap_parfor_blocks(parfor)
def _get_parfors(self, blocks, parfors_list):
for label, blk in blocks.items():
for stmt in blk.body:
if isinstance(stmt, Parfor):
parfors_list.append(stmt)
self._get_nested_parfors(stmt, parfors_list)
def get_parfors(self):
parfors_list = []
self._get_parfors(self.func_ir.blocks, parfors_list)
return parfors_list
def hoisted_allocations(self):
allocs = []
for pf_id, data in self.hoist_info.items():
stmt = data.get('hoisted', [])
for inst in stmt:
if isinstance(inst.value, ir.Expr):
if inst.value.op == 'call':
call = guard(find_callname, self.func_ir, inst.value)
if call is not None and call == ('empty', 'numpy'):
allocs.append(inst)
return allocs
def compute_graph_info(self, _a):
"""
compute adjacency list of the fused loops
and find the roots in of the lists
"""
a = copy.deepcopy(_a)
if a == {}:
return [], set()
vtx = set()
for v in a.values():
for x in v:
vtx.add(x)
# find roots
potential_roots = set(a.keys())
roots = potential_roots - vtx
if roots is None:
roots = set()
# populate rest of adjacency list
not_roots = set()
for x in range(max(set(a.keys()).union(vtx)) + 1):
val = a.get(x)
if val is not None:
a[x] = val
elif val is []:
not_roots.add(x) # debug only
else:
a[x] = []
# fold adjacency list into an actual list ordered
# by vtx
l = []
for x in sorted(a.keys()):
l.append(a[x])
return l, roots #, not_roots
def get_stats(self, fadj, nadj, root):
"""
Computes the number of fused and serialized loops
based on a fusion adjacency list `fadj` and a nested
parfors adjacency list `nadj` for the root, `root`
"""
def count_root(fadj, nadj, root, nfused, nserial):
for k in nadj[root]:
nserial += 1
if nadj[k] == []:
nfused += len(fadj[k])
else:
nf, ns = count_root(fadj, nadj, k, nfused, nserial)
nfused += nf
nserial = ns
return nfused, nserial
nfused, nserial = count_root(fadj, nadj, root, 0, 0)
return nfused, nserial
def reachable_nodes(self, adj, root):
"""
returns a list of nodes reachable in an adjacency list from a
specified root
"""
fusers = []
fusers.extend(adj[root])
for k in adj[root]:
if adj[k] != []:
fusers.extend(self.reachable_nodes(adj, k))
return fusers
def sort_pf_by_line(self, pf_id, parfors_simple):
"""
pd_id - the parfors id
parfors_simple - the simple parfors map
"""
# this sorts parfors by source line number
pf = parfors_simple[pf_id][0]
pattern = pf.patterns[0]
line = max(0, pf.loc.line - 1) # why are these out by 1 ?!
filename = self.func_ir.loc.filename
nadj, nroots = self.compute_graph_info(self.nested_fusion_info)
fadj, froots = self.compute_graph_info(self.fusion_info)
graphs = [nadj, fadj]
# If the parfor is internal, like internal prange, then the
# default line number is from its location in the numba source
# To get a more accurate line number, this first checks the
# adjacency graph for fused parfors that might not be internal
# and uses the minimum line number from there. If that fails
# (case where there's just a single internal parfor) the IR
# is walked backwards from the parfor location and the first non
# parfor statement line number is used.
if isinstance(pattern, tuple):
if pattern[1] == 'internal':
reported_loc = pattern[2][1]
if reported_loc.filename == filename:
return max(0, reported_loc.line - 1)
else:
# first recurse and check the adjacency list for
# something that is not an in internal parfor
tmp = []
for adj in graphs:
if adj: # graph may be empty, e.g. no nesting
for k in adj[pf_id]:
tmp.append(self.sort_pf_by_line(k, parfors_simple))
if tmp:
return max(0, min(tmp) - 1)
# second run through the parfor block to see if there's
# and reference to a line number in the user source
for blk in pf.loop_body.values():
for stmt in blk.body:
if stmt.loc.filename == filename:
return max(0, stmt.loc.line - 1)
# finally run through the func_ir and look for the
# first non-parfor statement prior to this one and
# grab the line from that
for blk in self.func_ir.blocks.values():
try:
idx = blk.body.index(pf)
for i in range(idx - 1, 0, -1):
stmt = blk.body[i]
if not isinstance(stmt, Parfor):
line = max(0, stmt.loc.line - 1)
break
except ValueError:
pass
return line
def dump(self, level=1):
if not self.has_setup:
raise RuntimeError("self.setup has not been called")
name = self.func_ir.func_id.func_qualname
line = self.func_ir.loc
if self.internal_name in name:
purpose_str = 'Internal parallel functions '
purpose = 'internal'
else:
purpose_str = ' Function %s, %s ' % (name, line)
purpose = 'user'
print_loop_search = False
print_source_listing = False
print_fusion_search = False
print_fusion_summary = False
print_loopnest_rewrite = False
print_pre_optimised = False
print_post_optimised = False
print_allocation_hoist = False
print_instruction_hoist = False
print_internal = False
# each level switches on progressively more output
if level in (1, 2, 3, 4):
print_source_listing = True
print_post_optimised = True
else:
raise ValueError("Report level unknown, should be one of 1, 2, 3, 4")
if level in (2, 3, 4):
print_pre_optimised = True
if level in (3, 4):
print_allocation_hoist = True
if level == 3:
print_fusion_summary = True
print_loopnest_rewrite = True
if level == 4:
print_fusion_search = True
print_instruction_hoist = True
print_internal = True
if purpose == 'internal' and not print_internal:
return
print_wrapped('\n ')
print_wrapped(_termwidth * "=")
print_wrapped((" Parallel Accelerator Optimizing: %s " % purpose_str).center(_termwidth, '='))
print_wrapped(_termwidth * "=")
print_wrapped("")
#----------- search section
if print_loop_search:
print_wrapped('Looking for parallel loops'.center(_termwidth, '-'))
parfors_simple = dict()
# print in line order, parfors loop id is based on discovery order
for pf in sorted(self.initial_parfors, key=lambda x: x.loc.line):
# use 0 here, the parfors are mutated by the time this routine
# is called, however, fusion appends the patterns so we can just
# pull in the first as a "before fusion" emulation
r_pattern = pf.patterns[0]
pattern = pf.patterns[0]
loc = pf.loc
if isinstance(pattern, tuple):
if pattern[0] == 'prange':
if pattern[1] == 'internal':
replfn = '.'.join(reversed(list(pattern[2][0])))
loc = pattern[2][1]
r_pattern = '%s %s' % (replfn, '(internal parallel version)')
elif pattern[1] == 'user':
r_pattern = "user defined prange"
elif pattern[1] == 'pndindex':
r_pattern = "internal pndindex" #FIXME: trace this!
else:
assert 0
fmt = 'Parallel for-loop #%s: is produced from %s:\n %s\n \n'
if print_loop_search:
print_wrapped(fmt % (pf.id, loc, r_pattern))
parfors_simple[pf.id] = (pf, loc, r_pattern)
count = self.count_parfors()
if print_loop_search:
print_wrapped("\nFound %s parallel loops." % count)
print_wrapped('-' * _termwidth)
#----------- augmented source section
filename = self.func_ir.loc.filename
try:
# Try to get a relative path
# ipython/jupyter input just returns as filename
path = os.path.relpath(filename)
except ValueError:
# Fallback to absolute path if error occurred in getting the
# relative path.
# This may happen on windows if the drive is different
path = os.path.abspath(filename)
if print_source_listing:
func_name = self.func_ir.func_id.func
try:
lines = inspect.getsource(func_name).splitlines()
except OSError: # generated function
lines = None
if lines:
src_width = max([len(x) for x in lines])
map_line_to_pf = defaultdict(list) # parfors can alias lines
for k, v in parfors_simple.items():
# TODO: do a better job of tracking parfors that are not in
# this file but are referred to, e.g. np.arange()
if parfors_simple[k][1].filename == filename:
match_line = self.sort_pf_by_line(k, parfors_simple)
map_line_to_pf[match_line].append(str(k))
max_pf_per_line = max([1] + [len(x) for x in map_line_to_pf.values()])
width = src_width + (1 + max_pf_per_line * (len(str(count)) + 2))
newlines = []
newlines.append('\n')
newlines.append('Parallel loop listing for %s' % purpose_str)
newlines.append(width * '-' + '|loop #ID')
fmt = '{0:{1}}| {2}'
# why are these off by 1?
lstart = max(0, self.func_ir.loc.line - 1)
for no, line in enumerate(lines, lstart):
pf_ids = map_line_to_pf.get(no, None)
if pf_ids is not None:
pfstr = '#' + ', '.join(pf_ids)
else:
pfstr = ''
stripped = line.strip('\n')
srclen = len(stripped)
if pf_ids:
l = fmt.format(width * '-', width, pfstr)
else:
l = fmt.format(width * ' ', width, pfstr)
newlines.append(stripped + l[srclen:])
print('\n'.join(newlines))
else:
print("No source available")
#---------- these are used a lot here on in
sword = '+--'
parfors = self.get_parfors() # this is the mutated parfors
parfor_ids = [x.id for x in parfors]
n_parfors = len(parfor_ids)
#----------- loop fusion section
if print_fusion_search or print_fusion_summary:
if not sequential_parfor_lowering:
print_wrapped(' Fusing loops '.center(_termwidth, '-'))
msg = ("Attempting fusion of parallel loops (combines loops "
"with similar properties)...\n")
print_wrapped(msg)
else:
msg = "Performing sequential lowering of loops...\n"
print_wrapped(msg)
print_wrapped(_termwidth * '-')
# if there are some parfors, print information about them!
if n_parfors > -1:
def dump_graph_indented(a, root_msg, node_msg):
fac = len(sword)
def print_graph(adj, roots):
def print_g(adj, root, depth):
for k in adj[root]:
print_wrapped(fac * depth * ' ' + '%s%s %s' % (sword, k, node_msg))
if adj[k] != []:
print_g(adj, k, depth + 1)
for r in roots:
print_wrapped('%s%s %s' % (sword, r, root_msg))
print_g(l, r, 1)
print_wrapped("")
l, roots = self.compute_graph_info(a)
print_graph(l, roots)
if print_fusion_search:
for report in self.fusion_reports:
l1, l2, msg = report
print_wrapped(" Trying to fuse loops #%s and #%s:" % (l1, l2))
print_wrapped(" %s" % msg)
if self.fusion_info != {}:
if print_fusion_summary:
print_wrapped("\n \nFused loop summary:\n")
dump_graph_indented(self.fusion_info, 'has the following loops fused into it:', '(fused)')
if print_fusion_summary:
if self.fusion_enabled:
after_fusion = "Following the attempted fusion of parallel for-loops"
else:
after_fusion = "With fusion disabled"
print_wrapped(('\n{} there are {} parallel for-loop(s) (originating from loops labelled: {}).').format(
after_fusion, n_parfors, ', '.join(['#%s' % x for x in parfor_ids])))
print_wrapped(_termwidth * '-')
print_wrapped("")
#----------- loop nest section
if print_loopnest_rewrite:
if self.nested_fusion_info != {}:
print_wrapped((" Optimising loop nests ").center(_termwidth, '-'))
print_wrapped("Attempting loop nest rewrites (optimising for the largest parallel loops)...\n ")
root_msg = 'is a parallel loop'
node_msg = '--> rewritten as a serial loop'
dump_graph_indented(self.nested_fusion_info, root_msg, node_msg)
print_wrapped(_termwidth * '-')
print_wrapped("")
#---------- compute various properties and orderings in the data for subsequent use
# ensure adjacency lists are the same size for both sets of info
# (nests and fusion may not traverse the same space, for
# convenience [] is used as a condition to halt recursion)
fadj, froots = self.compute_graph_info(self.fusion_info)
nadj, _nroots = self.compute_graph_info(self.nested_fusion_info)
if len(fadj) > len(nadj):
lim = len(fadj)
tmp = nadj
else:
lim = len(nadj)
tmp = fadj
for x in range(len(tmp), lim):
tmp.append([])
# This computes the roots of true loop nests (i.e. loops containing
# loops opposed to just a loop that's a root).
nroots = set()
if _nroots:
for r in _nroots:
if nadj[r] != []:
nroots.add(r)
all_roots = froots ^ nroots
# This computes all the parfors at the top level that are either:
# - roots of loop fusion
# - roots of true loop nests
# it then combines these based on source line number for ease of
# producing output ordered in a manner similar to the code structure
froots_lines = {}
for x in froots:
line = self.sort_pf_by_line(x, parfors_simple)
froots_lines[line] = 'fuse', x, fadj
nroots_lines = {}
for x in nroots:
line = self.sort_pf_by_line(x, parfors_simple)
nroots_lines[line] = 'nest', x, nadj
all_lines = froots_lines.copy()
all_lines.update(nroots_lines)
# nroots, froots, nadj and fadj are all set up correctly
# define some print functions
def print_unoptimised(lines):
# This prints the unoptimised parfors state
fac = len(sword)
def print_nest(fadj_, nadj_, theroot, reported, region_id):
def print_g(fadj_, nadj_, nroot, depth):
print_wrapped(fac * depth * ' ' + '%s%s %s' % (sword, nroot, '(parallel)'))
for k in nadj_[nroot]:
if nadj_[k] == []:
msg = []
msg.append(fac * (depth + 1) * ' ' + '%s%s %s' % (sword, k, '(parallel)'))
if fadj_[k] != [] and k not in reported:
fused = self.reachable_nodes(fadj_, k)
for i in fused:
msg.append(fac * (depth + 1) * ' ' + '%s%s %s' % (sword, i, '(parallel)'))
reported.append(k)
print_wrapped('\n'.join(msg))
else:
print_g(fadj_, nadj_, k, depth + 1)
if nadj_[theroot] != []:
print_wrapped("Parallel region %s:" % region_id)
print_g(fadj_, nadj_, theroot, 0)
print("\n")
region_id = region_id + 1
return region_id
def print_fuse(ty, pf_id, adj, depth, region_id):
msg = []
print_wrapped("Parallel region %s:" % region_id)
msg.append(fac * depth * ' ' + '%s%s %s' % (sword, pf_id, '(parallel)'))
if adj[pf_id] != []:
fused = sorted(self.reachable_nodes(adj, pf_id))
for k in fused:
msg.append(fac * depth * ' ' + '%s%s %s' % (sword, k, '(parallel)'))
region_id = region_id + 1
print_wrapped('\n'.join(msg))
print("\n")
return region_id
# Walk the parfors by src line and print optimised structure
region_id = 0
reported = []
for line, info in sorted(lines.items()):
opt_ty, pf_id, adj = info
if opt_ty == 'fuse':
if pf_id not in reported:
region_id = print_fuse('f', pf_id, adj, 0, region_id)
elif opt_ty == 'nest':
region_id = print_nest(fadj, nadj, pf_id, reported, region_id)
else:
assert 0
def print_optimised(lines):
# This prints the optimised output based on the transforms that
# occurred during loop fusion and rewriting of loop nests
fac = len(sword)
summary = dict()
# region : {fused, serialized}
def print_nest(fadj_, nadj_, theroot, reported, region_id):
def print_g(fadj_, nadj_, nroot, depth):
for k in nadj_[nroot]:
msg = fac * depth * ' ' + '%s%s %s' % (sword, k, '(serial')
if nadj_[k] == []:
fused = []
if fadj_[k] != [] and k not in reported:
fused = sorted(self.reachable_nodes(fadj_, k))
msg += ", fused with loop(s): "
msg += ', '.join([str(x) for x in fused])
msg += ')'
reported.append(k)
print_wrapped(msg)
summary[region_id]['fused'] += len(fused)
else:
print_wrapped(msg + ')')
print_g(fadj_, nadj_, k, depth + 1)
summary[region_id]['serialized'] += 1
if nadj_[theroot] != []:
print_wrapped("Parallel region %s:" % region_id)
print_wrapped('%s%s %s' % (sword, theroot, '(parallel)'))
summary[region_id] = {'root': theroot, 'fused': 0, 'serialized': 0}
print_g(fadj_, nadj_, theroot, 1)
print("\n")
region_id = region_id + 1
return region_id
def print_fuse(ty, pf_id, adj, depth, region_id):
print_wrapped("Parallel region %s:" % region_id)
msg = fac * depth * ' ' + '%s%s %s' % (sword, pf_id, '(parallel')
fused = []
if adj[pf_id] != []:
fused = sorted(self.reachable_nodes(adj, pf_id))
msg += ", fused with loop(s): "
msg += ', '.join([str(x) for x in fused])
summary[region_id] = {'root': pf_id, 'fused': len(fused), 'serialized': 0}
msg += ')'
print_wrapped(msg)
print("\n")
region_id = region_id + 1
return region_id
# Walk the parfors by src line and print optimised structure
region_id = 0
reported = []
for line, info in sorted(lines.items()):
opt_ty, pf_id, adj = info
if opt_ty == 'fuse':
if pf_id not in reported:
region_id = print_fuse('f', pf_id, adj, 0, region_id)
elif opt_ty == 'nest':
region_id = print_nest(fadj, nadj, pf_id, reported, region_id)
else:
assert 0
# print the summary of the fuse/serialize rewrite
if summary:
for k, v in sorted(summary.items()):
msg = ('\n \nParallel region %s (loop #%s) had %s '
'loop(s) fused')
root = v['root']
fused = v['fused']
serialized = v['serialized']
if serialized != 0:
msg += (' and %s loop(s) '
'serialized as part of the larger '
'parallel loop (#%s).')
print_wrapped(msg % (k, root, fused, serialized, root))
else:
msg += '.'
print_wrapped(msg % (k, root, fused))
else:
print_wrapped("Parallel structure is already optimal.")
if print_pre_optimised:
print(' Before Optimisation '.center(_termwidth,'-'))
print_unoptimised(all_lines)
print(_termwidth * '-')
if print_post_optimised:
print(' After Optimisation '.center(_termwidth,'-'))
print_optimised(all_lines)
print(_termwidth * '-')
print_wrapped("")
print_wrapped(_termwidth * '-')
print_wrapped("\n ")
#----------- LICM section
if print_allocation_hoist or print_instruction_hoist:
print_wrapped('Loop invariant code motion'.center(80, '-'))
if print_allocation_hoist:
found = False
print('Allocation hoisting:')
for pf_id, data in self.hoist_info.items():
stmt = data.get('hoisted', [])
for inst in stmt:
if isinstance(inst.value, ir.Expr):
try:
attr = inst.value.attr
if attr == 'empty':
msg = ("The memory allocation derived from the "
"instruction at %s is hoisted out of the "
"parallel loop labelled #%s (it will be "
"performed before the loop is executed and "
"reused inside the loop):")
loc = inst.loc
print_wrapped(msg % (loc, pf_id))
try:
path = os.path.relpath(loc.filename)
except ValueError:
path = os.path.abspath(loc.filename)
lines = linecache.getlines(path)
if lines and loc.line:
print_wrapped(" Allocation:: " + lines[0 if loc.line < 2 else loc.line - 1].strip())
print_wrapped(" - numpy.empty() is used for the allocation.\n")
found = True
except (KeyError, AttributeError):
pass
if not found:
print_wrapped('No allocation hoisting found')
if print_instruction_hoist:
print("")
print('Instruction hoisting:')
hoist_info_printed = False
if self.hoist_info:
for pf_id, data in self.hoist_info.items():
hoisted = data.get('hoisted', None)
not_hoisted = data.get('not_hoisted', None)
if not hoisted and not not_hoisted:
print("loop #%s has nothing to hoist." % pf_id)
continue
print("loop #%s:" % pf_id)
if hoisted:
print(" Has the following hoisted:")
[print(" %s" % y) for y in hoisted]
hoist_info_printed = True
if not_hoisted:
print(" Failed to hoist the following:")
[print(" %s: %s" % (y, x)) for x, y in not_hoisted]
hoist_info_printed = True
if not hoist_info_printed:
print_wrapped('No instruction hoisting found')
print_wrapped(80 * '-')
else: # there are no parfors
print_wrapped('Function %s, %s, has no parallel for-loops.'.format(name, line))
def __str__(self):
r = "ParforDiagnostics:\n"
r += repr(self.replaced_fns)
return r
def __repr__(self):
r = "ParforDiagnostics"
return r
class PreParforPass(object):
"""Preprocessing for the Parfor pass. It mostly inlines parallel
implementations of numpy functions if available.
"""
def __init__(self, func_ir, typemap, calltypes, typingctx, options,
swapped={}):
self.func_ir = func_ir
self.typemap = typemap
self.calltypes = calltypes
self.typingctx = typingctx
self.options = options
# diagnostics
self.swapped = swapped
self.stats = {
'replaced_func': 0,
'replaced_dtype': 0,
}
def run(self):
"""Run pre-parfor processing pass.
"""
# e.g. convert A.sum() to np.sum(A) for easier match and optimization
canonicalize_array_math(self.func_ir, self.typemap,
self.calltypes, self.typingctx)
if self.options.numpy:
self._replace_parallel_functions(self.func_ir.blocks)
self.func_ir.blocks = simplify_CFG(self.func_ir.blocks)
def _replace_parallel_functions(self, blocks):
"""
Replace functions with their parallel implementation in
replace_functions_map if available.
The implementation code is inlined to enable more optimization.
"""
swapped = self.swapped
from numba.core.inline_closurecall import inline_closure_call
work_list = list(blocks.items())
while work_list:
label, block = work_list.pop()
for i, instr in enumerate(block.body):
if isinstance(instr, ir.Assign):
lhs = instr.target
lhs_typ = self.typemap[lhs.name]
expr = instr.value
if isinstance(expr, ir.Expr) and expr.op == 'call':
# Try and inline known calls with their parallel implementations
def replace_func():
func_def = get_definition(self.func_ir, expr.func)
callname = find_callname(self.func_ir, expr)
repl_func = replace_functions_map.get(callname, None)
# Handle method on array type
if (repl_func is None and
len(callname) == 2 and
isinstance(callname[1], ir.Var) and
isinstance(self.typemap[callname[1].name],
types.npytypes.Array)):
repl_func = replace_functions_ndarray.get(callname[0], None)
if repl_func is not None:
# Add the array that the method is on to the arg list.
expr.args.insert(0, callname[1])
require(repl_func is not None)
typs = tuple(self.typemap[x.name] for x in expr.args)
try:
new_func = repl_func(lhs_typ, *typs)
except:
new_func = None
require(new_func is not None)
g = copy.copy(self.func_ir.func_id.func.__globals__)
g['numba'] = numba
g['np'] = numpy
g['math'] = math
# if the function being inlined has a function
# checking the inputs, find it and add it to globals
check = replace_functions_checkers_map.get(callname,
None)
if check is not None:
g[check.name] = check.func
# inline the parallel implementation
new_blocks, _ = inline_closure_call(self.func_ir, g,
block, i, new_func, self.typingctx, typs,
self.typemap, self.calltypes, work_list)
call_table = get_call_table(new_blocks, topological_ordering=False)
# find the prange in the new blocks and record it for use in diagnostics
for call in call_table:
for k, v in call.items():
if v[0] == 'internal_prange':
swapped[k] = [callname, repl_func.__name__, func_def, block.body[i].loc]
break
return True
if guard(replace_func):
self.stats['replaced_func'] += 1
break
elif (isinstance(expr, ir.Expr) and expr.op == 'getattr' and
expr.attr == 'dtype'):
# Replace getattr call "A.dtype" with numpy.dtype(<actual type>).
# This helps remove superfluous dependencies from parfor.
typ = self.typemap[expr.value.name]
if isinstance(typ, types.npytypes.Array):
# Convert A.dtype to four statements.
# 1) Get numpy global.
# 2) Create var for known type of array as string
# constant. e.g. 'float64'
# 3) Get dtype function from numpy module.
# 4) Create var for numpy.dtype(var from #2).
# Create var for numpy module.
dtype = typ.dtype
scope = block.scope
loc = instr.loc
g_np_var = ir.Var(scope, mk_unique_var("$np_g_var"), loc)
self.typemap[g_np_var.name] = types.misc.Module(numpy)
g_np = ir.Global('np', numpy, loc)
g_np_assign = ir.Assign(g_np, g_np_var, loc)
# Create var for the inferred type of the array
# e.g., 'float64'
dtype_str = str(dtype)
if dtype_str == 'bool':
dtype_str = 'bool_'
typ_var = ir.Var(
scope, mk_unique_var("$np_typ_var"), loc)
self.typemap[typ_var.name] = types.StringLiteral(
dtype_str)
typ_var_assign = ir.Assign(
ir.Const(dtype_str, loc), typ_var, loc)
# Get the dtype function from the numpy module.
dtype_attr_var = ir.Var(scope, mk_unique_var("$dtype_attr_var"), loc)
temp = find_template(numpy.dtype)
tfunc = numba.core.types.Function(temp)
tfunc.get_call_type(self.typingctx, (self.typemap[typ_var.name],), {})
self.typemap[dtype_attr_var.name] = types.functions.Function(temp)
dtype_attr_getattr = ir.Expr.getattr(g_np_var, 'dtype', loc)
dtype_attr_assign = ir.Assign(dtype_attr_getattr, dtype_attr_var, loc)
# Call numpy.dtype on the statically coded type two steps above.
dtype_var = ir.Var(scope, mk_unique_var("$dtype_var"), loc)
self.typemap[dtype_var.name] = types.npytypes.DType(dtype)
dtype_getattr = ir.Expr.call(dtype_attr_var, [typ_var], (), loc)
dtype_assign = ir.Assign(dtype_getattr, dtype_var, loc)
self.calltypes[dtype_getattr] = signature(
self.typemap[dtype_var.name], self.typemap[typ_var.name])
# The original A.dtype rhs is replaced with result of this call.
instr.value = dtype_var
# Add statements to body of the code.
block.body.insert(0, dtype_assign)
block.body.insert(0, dtype_attr_assign)
block.body.insert(0, typ_var_assign)
block.body.insert(0, g_np_assign)
self.stats['replaced_dtype'] += 1
break
def find_template(op):
for ft in numba.core.typing.templates.builtin_registry.functions:
if ft.key == op:
return ft
class ParforPassStates:
"""This class encapsulates all internal states of the ParforPass.
"""
def __init__(self, func_ir, typemap, calltypes, return_type, typingctx,
options, flags, diagnostics=ParforDiagnostics()):
self.func_ir = func_ir
self.typemap = typemap
self.calltypes = calltypes
self.typingctx = typingctx
self.return_type = return_type
self.options = options
self.diagnostics = diagnostics
self.swapped_fns = diagnostics.replaced_fns
self.fusion_info = diagnostics.fusion_info
self.nested_fusion_info = diagnostics.nested_fusion_info
self.array_analysis = array_analysis.ArrayAnalysis(
self.typingctx, self.func_ir, self.typemap, self.calltypes,
)
ir_utils._max_label = max(func_ir.blocks.keys())
self.flags = flags
class ConvertInplaceBinop:
"""Parfor subpass to convert setitem on Arrays
"""
def __init__(self, pass_states):
"""
Parameters
----------
pass_states : ParforPassStates
"""
self.pass_states = pass_states
self.rewritten = []
def run(self, blocks):
pass_states = self.pass_states
# convert expressions like A += ... where A is an array.
topo_order = find_topo_order(blocks)
# variables available in the program so far (used for finding map
# functions in array_expr lowering)
for label in topo_order:
block = blocks[label]
new_body = []
equiv_set = pass_states.array_analysis.get_equiv_set(label)
for instr in block.body:
if isinstance(instr, ir.Assign):
lhs = instr.target
expr = instr.value
if isinstance(expr, ir.Expr) and expr.op == 'inplace_binop':
loc = expr.loc
target = expr.lhs
value = expr.rhs
target_typ = pass_states.typemap[target.name]
value_typ = pass_states.typemap[value.name]
# Handle A op= ...
if isinstance(target_typ, types.npytypes.Array):
# RHS is an array
if isinstance(value_typ, types.npytypes.Array):
new_instr = self._inplace_binop_to_parfor(equiv_set,
loc, expr.immutable_fn, target, value)
self.rewritten.append(
dict(old=instr, new=new_instr,
reason='inplace_binop'),
)
instr = [new_instr, ir.Assign(target, lhs, loc)]
if isinstance(instr, list):
new_body.extend(instr)
else:
new_body.append(instr)
block.body = new_body
def _inplace_binop_to_parfor(self, equiv_set, loc, op, target, value):
"""generate parfor from setitem node with a boolean or slice array indices.
The value can be either a scalar or an array variable, and if a boolean index
is used for the latter case, the same index must be used for the value too.
"""
pass_states = self.pass_states
scope = target.scope
arr_typ = pass_states.typemap[target.name]
el_typ = arr_typ.dtype
init_block = ir.Block(scope, loc)
value_typ = pass_states.typemap[value.name]
size_vars = equiv_set.get_shape(target)
# generate loopnests and size variables from target correlations
index_vars, loopnests = _mk_parfor_loops(pass_states.typemap, size_vars, scope, loc)
# generate body
body_label = next_label()
body_block = ir.Block(scope, loc)
index_var, index_var_typ = _make_index_var(
pass_states.typemap, scope, index_vars, body_block)
# Read value.
value_var = ir.Var(scope, mk_unique_var("$value_var"), loc)
pass_states.typemap[value_var.name] = value_typ.dtype
getitem_call = ir.Expr.getitem(value, index_var, loc)
pass_states.calltypes[getitem_call] = signature(
value_typ.dtype, value_typ, index_var_typ)
body_block.body.append(ir.Assign(getitem_call, value_var, loc))
# Read target
target_var = ir.Var(scope, mk_unique_var("$target_var"), loc)
pass_states.typemap[target_var.name] = el_typ
getitem_call = ir.Expr.getitem(target, index_var, loc)
pass_states.calltypes[getitem_call] = signature(
el_typ, arr_typ, index_var_typ)
body_block.body.append(ir.Assign(getitem_call, target_var, loc))
# Create temp to hold result.
expr_out_var = ir.Var(scope, mk_unique_var("$expr_out_var"), loc)
pass_states.typemap[expr_out_var.name] = el_typ
# Create binop and assign result to temporary.
binop_expr = ir.Expr.binop(op, target_var, value_var, loc)
body_block.body.append(ir.Assign(binop_expr, expr_out_var, loc))
unified_type = self.pass_states.typingctx.unify_pairs(el_typ, value_typ.dtype)
pass_states.calltypes[binop_expr] = signature(
unified_type, unified_type, unified_type)
# Write to target
setitem_node = ir.SetItem(target, index_var, expr_out_var, loc)
pass_states.calltypes[setitem_node] = signature(
types.none, arr_typ, index_var_typ, el_typ)
body_block.body.append(setitem_node)
parfor = Parfor(loopnests, init_block, {}, loc, index_var, equiv_set,
('inplace_binop', ''), pass_states.flags)
parfor.loop_body = {body_label: body_block}
if config.DEBUG_ARRAY_OPT >= 1:
print("parfor from inplace_binop")
parfor.dump()
return parfor
def _type_getitem(self, args):
fnty = operator.getitem
return self.pass_states.typingctx.resolve_function_type(fnty, tuple(args), {})
class ConvertSetItemPass:
"""Parfor subpass to convert setitem on Arrays
"""
def __init__(self, pass_states):
"""
Parameters
----------
pass_states : ParforPassStates
"""
self.pass_states = pass_states
self.rewritten = []
def run(self, blocks):
pass_states = self.pass_states
# convert setitem expressions like A[C] = c or A[C] = B[C] to parfor,
# where C is a boolean array.
topo_order = find_topo_order(blocks)
# variables available in the program so far (used for finding map
# functions in array_expr lowering)
for label in topo_order:
block = blocks[label]
new_body = []
equiv_set = pass_states.array_analysis.get_equiv_set(label)
for instr in block.body:
if isinstance(instr, ir.StaticSetItem) or isinstance(instr, ir.SetItem):
loc = instr.loc
target = instr.target
index = (instr.index if isinstance(instr, ir.SetItem)
else instr.index_var)
value = instr.value
target_typ = pass_states.typemap[target.name]
index_typ = pass_states.typemap[index.name]
value_typ = pass_states.typemap[value.name]
# Handle A[boolean_array] = <scalar or array>
if isinstance(target_typ, types.npytypes.Array):
if (isinstance(index_typ, types.npytypes.Array) and
isinstance(index_typ.dtype, types.Boolean) and
target_typ.ndim == index_typ.ndim):
# RHS is a scalar number
if isinstance(value_typ, types.Number):
new_instr = self._setitem_to_parfor(equiv_set,
loc, target, index, value)
self.rewritten.append(
dict(old=instr, new=new_instr,
reason='masked_assign_broadcast_scalar'),
)
instr = new_instr
# RHS is an array
elif isinstance(value_typ, types.npytypes.Array):
val_def = guard(get_definition, pass_states.func_ir,
value.name)
if (isinstance(val_def, ir.Expr) and
val_def.op == 'getitem' and
val_def.index.name == index.name):
new_instr = self._setitem_to_parfor(equiv_set,
loc, target, index, val_def.value)
self.rewritten.append(
dict(old=instr, new=new_instr,
reason='masked_assign_array'),
)
instr = new_instr
else:
# Handle A[:] = x
shape = equiv_set.get_shape(instr)
if shape is not None:
new_instr = self._setitem_to_parfor(equiv_set,
loc, target, index, value, shape=shape)
self.rewritten.append(
dict(old=instr, new=new_instr,
reason='slice'),
)
instr = new_instr
new_body.append(instr)
block.body = new_body
def _setitem_to_parfor(self, equiv_set, loc, target, index, value, shape=None):
"""generate parfor from setitem node with a boolean or slice array indices.
The value can be either a scalar or an array variable, and if a boolean index
is used for the latter case, the same index must be used for the value too.
"""
pass_states = self.pass_states
scope = target.scope
arr_typ = pass_states.typemap[target.name]
el_typ = arr_typ.dtype
index_typ = pass_states.typemap[index.name]
init_block = ir.Block(scope, loc)
if shape:
# Slice index is being used on the target array, we'll have to create
# a sub-array so that the target dimension matches the given shape.
assert(isinstance(index_typ, types.BaseTuple) or
isinstance(index_typ, types.SliceType))
# setitem has a custom target shape
size_vars = shape
# create a new target array via getitem
subarr_var = ir.Var(scope, mk_unique_var("$subarr"), loc)
getitem_call = ir.Expr.getitem(target, index, loc)
subarr_typ = typing.arraydecl.get_array_index_type( arr_typ, index_typ).result
pass_states.typemap[subarr_var.name] = subarr_typ
pass_states.calltypes[getitem_call] = self._type_getitem((arr_typ, index_typ))
init_block.append(ir.Assign(getitem_call, subarr_var, loc))
target = subarr_var
else:
# Otherwise it is a boolean array that is used as index.
assert(isinstance(index_typ, types.ArrayCompatible))
size_vars = equiv_set.get_shape(target)
bool_typ = index_typ.dtype
# generate loopnests and size variables from lhs correlations
loopnests = []
index_vars = []
for size_var in size_vars:
index_var = ir.Var(scope, mk_unique_var("parfor_index"), loc)
index_vars.append(index_var)
pass_states.typemap[index_var.name] = types.uintp
loopnests.append(LoopNest(index_var, 0, size_var, 1))
# generate body
body_label = next_label()
body_block = ir.Block(scope, loc)
index_var, index_var_typ = _make_index_var(
pass_states.typemap, scope, index_vars, body_block)
parfor = Parfor(loopnests, init_block, {}, loc, index_var, equiv_set,
('setitem', ''), pass_states.flags)
if shape:
# slice subarray
parfor.loop_body = {body_label: body_block}
true_block = body_block
end_label = None
else:
# boolean mask
true_label = next_label()
true_block = ir.Block(scope, loc)
end_label = next_label()
end_block = ir.Block(scope, loc)
parfor.loop_body = {body_label: body_block,
true_label: true_block,
end_label: end_block,
}
mask_var = ir.Var(scope, mk_unique_var("$mask_var"), loc)
pass_states.typemap[mask_var.name] = bool_typ
mask_val = ir.Expr.getitem(index, index_var, loc)
body_block.body.extend([
ir.Assign(mask_val, mask_var, loc),
ir.Branch(mask_var, true_label, end_label, loc)
])
value_typ = pass_states.typemap[value.name]
if isinstance(value_typ, types.npytypes.Array):
value_var = ir.Var(scope, mk_unique_var("$value_var"), loc)
pass_states.typemap[value_var.name] = value_typ.dtype
getitem_call = ir.Expr.getitem(value, index_var, loc)
pass_states.calltypes[getitem_call] = signature(
value_typ.dtype, value_typ, index_var_typ)
true_block.body.append(ir.Assign(getitem_call, value_var, loc))
else:
value_var = value
setitem_node = ir.SetItem(target, index_var, value_var, loc)
pass_states.calltypes[setitem_node] = signature(
types.none, pass_states.typemap[target.name], index_var_typ, el_typ)
true_block.body.append(setitem_node)
if end_label:
true_block.body.append(ir.Jump(end_label, loc))
if config.DEBUG_ARRAY_OPT >= 1:
print("parfor from setitem")
parfor.dump()
return parfor
def _type_getitem(self, args):
fnty = operator.getitem
return self.pass_states.typingctx.resolve_function_type(fnty, tuple(args), {})
def _make_index_var(typemap, scope, index_vars, body_block):
ndims = len(index_vars)
loc = body_block.loc
if ndims > 1:
tuple_var = ir.Var(scope, mk_unique_var(
"$parfor_index_tuple_var"), loc)
typemap[tuple_var.name] = types.containers.UniTuple(
types.uintp, ndims)
tuple_call = ir.Expr.build_tuple(list(index_vars), loc)
tuple_assign = ir.Assign(tuple_call, tuple_var, loc)
body_block.body.append(tuple_assign)
return tuple_var, types.containers.UniTuple(types.uintp, ndims)
elif ndims == 1:
return index_vars[0], types.uintp
else:
raise errors.UnsupportedRewriteError(
"Parfor does not handle arrays of dimension 0",
loc=loc,
)
def _mk_parfor_loops(typemap, size_vars, scope, loc):
"""
Create loop index variables and build LoopNest objects for a parfor.
"""
loopnests = []
index_vars = []
for size_var in size_vars:
index_var = ir.Var(scope, mk_unique_var("parfor_index"), loc)
index_vars.append(index_var)
typemap[index_var.name] = types.uintp
loopnests.append(LoopNest(index_var, 0, size_var, 1))
return index_vars, loopnests
class ConvertNumpyPass:
"""
Convert supported Numpy functions, as well as arrayexpr nodes, to
parfor nodes.
"""
def __init__(self, pass_states):
self.pass_states = pass_states
self.rewritten = []
def run(self, blocks):
pass_states = self.pass_states
topo_order = find_topo_order(blocks)
# variables available in the program so far (used for finding map
# functions in array_expr lowering)
avail_vars = []
for label in topo_order:
block = blocks[label]
new_body = []
equiv_set = pass_states.array_analysis.get_equiv_set(label)
for instr in block.body:
if isinstance(instr, ir.Assign):
expr = instr.value
lhs = instr.target
if self._is_C_order(lhs.name):
# only translate C order since we can't allocate F
if guard(self._is_supported_npycall, expr):
new_instr = self._numpy_to_parfor(equiv_set, lhs, expr)
if new_instr is not None:
self.rewritten.append(dict(
old=instr,
new=new_instr,
reason='numpy_allocator',
))
instr = new_instr
elif isinstance(expr, ir.Expr) and expr.op == 'arrayexpr':
new_instr = self._arrayexpr_to_parfor(
equiv_set, lhs, expr, avail_vars)
self.rewritten.append(dict(
old=instr,
new=new_instr,
reason='arrayexpr',
))
instr = new_instr
avail_vars.append(lhs.name)
new_body.append(instr)
block.body = new_body
def _is_C_order(self, arr_name):
typ = self.pass_states.typemap[arr_name]
return isinstance(typ, types.npytypes.Array) and typ.layout == 'C' and typ.ndim > 0
def _arrayexpr_to_parfor(self, equiv_set, lhs, arrayexpr, avail_vars):
"""generate parfor from arrayexpr node, which is essentially a
map with recursive tree.
"""
pass_states = self.pass_states
scope = lhs.scope
loc = lhs.loc
expr = arrayexpr.expr
arr_typ = pass_states.typemap[lhs.name]
el_typ = arr_typ.dtype
# generate loopnests and size variables from lhs correlations
size_vars = equiv_set.get_shape(lhs)
index_vars, loopnests = _mk_parfor_loops(pass_states.typemap, size_vars, scope, loc)
# generate init block and body
init_block = ir.Block(scope, loc)
init_block.body = mk_alloc(pass_states.typemap, pass_states.calltypes, lhs,
tuple(size_vars), el_typ, scope, loc)
body_label = next_label()
body_block = ir.Block(scope, loc)
expr_out_var = ir.Var(scope, mk_unique_var("$expr_out_var"), loc)
pass_states.typemap[expr_out_var.name] = el_typ
index_var, index_var_typ = _make_index_var(
pass_states.typemap, scope, index_vars, body_block)
body_block.body.extend(
_arrayexpr_tree_to_ir(
pass_states.func_ir,
pass_states.typingctx,
pass_states.typemap,
pass_states.calltypes,
equiv_set,
init_block,
expr_out_var,
expr,
index_var,
index_vars,
avail_vars))
pat = ('array expression {}'.format(repr_arrayexpr(arrayexpr.expr)),)
parfor = Parfor(loopnests, init_block, {}, loc, index_var, equiv_set, pat[0], pass_states.flags)
setitem_node = ir.SetItem(lhs, index_var, expr_out_var, loc)
pass_states.calltypes[setitem_node] = signature(
types.none, pass_states.typemap[lhs.name], index_var_typ, el_typ)
body_block.body.append(setitem_node)
parfor.loop_body = {body_label: body_block}
if config.DEBUG_ARRAY_OPT >= 1:
print("parfor from arrayexpr")
parfor.dump()
return parfor
def _is_supported_npycall(self, expr):
"""check if we support parfor translation for
this Numpy call.
"""
call_name, mod_name = find_callname(self.pass_states.func_ir, expr)
if not (isinstance(mod_name, str) and mod_name.startswith('numpy')):
return False
if call_name in ['zeros', 'ones']:
return True
if mod_name == 'numpy.random' and call_name in random_calls:
return True
# TODO: add more calls
return False
def _numpy_to_parfor(self, equiv_set, lhs, expr):
call_name, mod_name = find_callname(self.pass_states.func_ir, expr)
args = expr.args
kws = dict(expr.kws)
if call_name in ['zeros', 'ones'] or mod_name == 'numpy.random':
return self._numpy_map_to_parfor(equiv_set, call_name, lhs, args, kws, expr)
# return error if we couldn't handle it (avoid rewrite infinite loop)
raise errors.UnsupportedRewriteError(
f"parfor translation failed for {expr}", loc=expr.loc,
)
def _numpy_map_to_parfor(self, equiv_set, call_name, lhs, args, kws, expr):
"""generate parfor from Numpy calls that are maps.
"""
pass_states = self.pass_states
scope = lhs.scope
loc = lhs.loc
arr_typ = pass_states.typemap[lhs.name]
el_typ = arr_typ.dtype
# generate loopnests and size variables from lhs correlations
size_vars = equiv_set.get_shape(lhs)
if size_vars is None:
if config.DEBUG_ARRAY_OPT >= 1:
print("Could not convert numpy map to parfor, unknown size")
return None
index_vars, loopnests = _mk_parfor_loops(pass_states.typemap, size_vars, scope, loc)
# generate init block and body
init_block = ir.Block(scope, loc)
init_block.body = mk_alloc(pass_states.typemap, pass_states.calltypes, lhs,
tuple(size_vars), el_typ, scope, loc)
body_label = next_label()
body_block = ir.Block(scope, loc)
expr_out_var = ir.Var(scope, mk_unique_var("$expr_out_var"), loc)
pass_states.typemap[expr_out_var.name] = el_typ
index_var, index_var_typ = _make_index_var(
pass_states.typemap, scope, index_vars, body_block)
if call_name == 'zeros':
value = ir.Const(el_typ(0), loc)
elif call_name == 'ones':
value = ir.Const(el_typ(1), loc)
elif call_name in random_calls:
# remove size arg to reuse the call expr for single value
_remove_size_arg(call_name, expr)
# update expr type
new_arg_typs, new_kw_types = _get_call_arg_types(
expr, pass_states.typemap)
pass_states.calltypes.pop(expr)
pass_states.calltypes[expr] = pass_states.typemap[expr.func.name].get_call_type(
typing.Context(), new_arg_typs, new_kw_types)
value = expr
else:
raise NotImplementedError(
"Map of numpy.{} to parfor is not implemented".format(call_name))
value_assign = ir.Assign(value, expr_out_var, loc)
body_block.body.append(value_assign)
setitem_node = ir.SetItem(lhs, index_var, expr_out_var, loc)
pass_states.calltypes[setitem_node] = signature(
types.none, pass_states.typemap[lhs.name], index_var_typ, el_typ)
body_block.body.append(setitem_node)
parfor = Parfor(loopnests, init_block, {}, loc, index_var, equiv_set,
('{} function'.format(call_name,), 'NumPy mapping'),
pass_states.flags)
parfor.loop_body = {body_label: body_block}
if config.DEBUG_ARRAY_OPT >= 1:
print("generated parfor for numpy map:")
parfor.dump()
return parfor
class ConvertReducePass:
"""
Find reduce() calls and convert them to parfors.
"""
def __init__(self, pass_states):
self.pass_states = pass_states
self.rewritten = []
def run(self, blocks):
pass_states = self.pass_states
topo_order = find_topo_order(blocks)
for label in topo_order:
block = blocks[label]
new_body = []
equiv_set = pass_states.array_analysis.get_equiv_set(label)
for instr in block.body:
parfor = None
if isinstance(instr, ir.Assign):
loc = instr.loc
lhs = instr.target
expr = instr.value
callname = guard(find_callname, pass_states.func_ir, expr)
if (callname == ('reduce', 'builtins')
or callname == ('reduce', '_functools')):
# reduce function with generic function
parfor = guard(self._reduce_to_parfor, equiv_set, lhs,
expr.args, loc)
if parfor:
self.rewritten.append(dict(
new=parfor,
old=instr,
reason='reduce',
))
instr = parfor
new_body.append(instr)
block.body = new_body
return
def _reduce_to_parfor(self, equiv_set, lhs, args, loc):
"""
Convert a reduce call to a parfor.
The call arguments should be (call_name, array, init_value).
"""
pass_states = self.pass_states
scope = lhs.scope
call_name = args[0]
in_arr = args[1]
arr_def = get_definition(pass_states.func_ir, in_arr.name)
mask_var = None
mask_indices = None
# Search for array[boolean_mask]
mask_query_result = guard(_find_mask, pass_states.typemap, pass_states.func_ir, arr_def)
if mask_query_result:
in_arr, mask_var, mask_typ, mask_indices = mask_query_result
init_val = args[2]
size_vars = equiv_set.get_shape(in_arr if mask_indices is None else mask_var)
if size_vars is None:
return None
index_vars, loopnests = _mk_parfor_loops(pass_states.typemap, size_vars, scope, loc)
mask_index = index_vars
if mask_indices:
# the following is never tested
raise AssertionError("unreachable")
index_vars = tuple(x if x else index_vars[0] for x in mask_indices)
acc_var = lhs
# init block has to init the reduction variable
init_block = ir.Block(scope, loc)
init_block.body.append(ir.Assign(init_val, acc_var, loc))
# produce loop body
body_label = next_label()
index_var, loop_body = self._mk_reduction_body(call_name,
scope, loc, index_vars, in_arr, acc_var)
if mask_indices:
# the following is never tested
raise AssertionError("unreachable")
index_var = mask_index[0]
if mask_var is not None:
true_label = min(loop_body.keys())
false_label = max(loop_body.keys())
body_block = ir.Block(scope, loc)
loop_body[body_label] = body_block
mask = ir.Var(scope, mk_unique_var("$mask_val"), loc)
pass_states.typemap[mask.name] = mask_typ
mask_val = ir.Expr.getitem(mask_var, index_var, loc)
body_block.body.extend([
ir.Assign(mask_val, mask, loc),
ir.Branch(mask, true_label, false_label, loc)
])
parfor = Parfor(loopnests, init_block, loop_body, loc, index_var,
equiv_set, ('{} function'.format(call_name),
'reduction'), pass_states.flags)
if config.DEBUG_ARRAY_OPT >= 1:
print("parfor from reduction")
parfor.dump()
return parfor
def _mk_reduction_body(self, call_name, scope, loc,
index_vars, in_arr, acc_var):
"""
Produce the body blocks for a reduction function indicated by call_name.
"""
from numba.core.inline_closurecall import check_reduce_func
pass_states = self.pass_states
reduce_func = get_definition(pass_states.func_ir, call_name)
fcode = check_reduce_func(pass_states.func_ir, reduce_func)
arr_typ = pass_states.typemap[in_arr.name]
in_typ = arr_typ.dtype
body_block = ir.Block(scope, loc)
index_var, index_var_type = _make_index_var(
pass_states.typemap, scope, index_vars, body_block)
tmp_var = ir.Var(scope, mk_unique_var("$val"), loc)
pass_states.typemap[tmp_var.name] = in_typ
getitem_call = ir.Expr.getitem(in_arr, index_var, loc)
pass_states.calltypes[getitem_call] = signature(
in_typ, arr_typ, index_var_type)
body_block.append(ir.Assign(getitem_call, tmp_var, loc))
reduce_f_ir = compile_to_numba_ir(fcode,
pass_states.func_ir.func_id.func.__globals__,
pass_states.typingctx,
(in_typ, in_typ),
pass_states.typemap,
pass_states.calltypes)
loop_body = reduce_f_ir.blocks
end_label = next_label()
end_block = ir.Block(scope, loc)
loop_body[end_label] = end_block
first_reduce_label = min(reduce_f_ir.blocks.keys())
first_reduce_block = reduce_f_ir.blocks[first_reduce_label]
body_block.body.extend(first_reduce_block.body)
first_reduce_block.body = body_block.body
replace_arg_nodes(first_reduce_block, [acc_var, tmp_var])
replace_returns(loop_body, acc_var, end_label)
return index_var, loop_body
class ConvertLoopPass:
"""Build Parfor nodes from prange loops.
"""
def __init__(self, pass_states):
self.pass_states = pass_states
self.rewritten = []
def run(self, blocks):
pass_states = self.pass_states
call_table, _ = get_call_table(blocks)
cfg = compute_cfg_from_blocks(blocks)
usedefs = compute_use_defs(blocks)
live_map = compute_live_map(cfg, blocks, usedefs.usemap, usedefs.defmap)
loops = cfg.loops()
sized_loops = [(loops[k], len(loops[k].body)) for k in loops.keys()]
moved_blocks = []
# We go over all loops, smaller loops first (inner first)
for loop, s in sorted(sized_loops, key=lambda tup: tup[1]):
if len(loop.entries) != 1 or len(loop.exits) != 1:
continue
entry = list(loop.entries)[0]
for inst in blocks[entry].body:
# if prange or pndindex call
if (isinstance(inst, ir.Assign)
and isinstance(inst.value, ir.Expr)
and inst.value.op == 'call'
and self._is_parallel_loop(inst.value.func.name, call_table)):
body_labels = [ l for l in loop.body if
l in blocks and l != loop.header ]
args = inst.value.args
loop_kind, loop_replacing = self._get_loop_kind(inst.value.func.name,
call_table)
# Get the body of the header of the loops minus the branch terminator
# The general approach is to prepend the header block to the first
# body block and then let dead code removal handle removing unneeded
# statements. Not all statements in the header block are unnecessary.
header_body = blocks[loop.header].body[:-1]
# find loop index variable (pair_first in header block)
loop_index = None
for hbi, stmt in enumerate(header_body):
if (isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Expr)
and stmt.value.op == 'pair_first'):
loop_index = stmt.target.name
li_index = hbi
break
assert(loop_index is not None)
# Remove pair_first from header.
# We have to remove the pair_first by hand since it causes problems
# for some code below if we don't.
header_body = header_body[:li_index] + header_body[li_index+1:]
# loop_index may be assigned to other vars
# get header copies to find all of them
cps, _ = get_block_copies({0: blocks[loop.header]},
pass_states.typemap)
cps = cps[0]
loop_index_vars = set(t for t, v in cps