def _infer_h5_typ(self, rhs):
# infer the type if it is of the from f['A']['B'][:]
# with constant filename
# TODO: static_getitem has index_var for sure?
# make sure it's slice, TODO: support non-slice like integer
require(rhs.op in ('getitem', 'static_getitem'))
index_var = rhs.index if rhs.op == 'getitem' else rhs.index_var
index_def = get_definition(self.func_ir, index_var)
require(isinstance(index_def, ir.Expr) and index_def.op == 'call')
require(
find_callname(self.func_ir, index_def) == ('slice', 'builtins'))
# collect object names until the call
val_def = rhs
obj_name_list = []
while True:
val_def = get_definition(self.func_ir, val_def.value)
require(isinstance(val_def, ir.Expr))
if val_def.op == 'call':
return self._get_h5_type_file(val_def, obj_name_list)
# object_name should be constant str
require(val_def.op in ('getitem', 'static_getitem'))
val_index_var = val_def.index if val_def.op == 'getitem' else val_def.index_var
obj_name = find_const(self.func_ir, val_index_var)
require(isinstance(obj_name, str))
obj_name_list.append(obj_name)
def is_const_slice(typemap, func_ir, var, accept_stride=False):
""" return True if var can be determined to be a constant size slice """
require(typemap[var.name] == types.slice2_type
or (accept_stride and typemap[var.name] == types.slice3_type))
call_expr = get_definition(func_ir, var)
require(isinstance(call_expr, ir.Expr) and call_expr.op == 'call')
assert (len(call_expr.args) == 2
or (accept_stride and len(call_expr.args) == 3))
assert find_callname(func_ir, call_expr) == ('slice', 'builtins')
arg0_def = get_definition(func_ir, call_expr.args[0])
require(isinstance(arg0_def, ir.Const) and arg0_def.value is None)
size_const = find_const(func_ir, call_expr.args[1])
require(isinstance(size_const, int))
return True
def _get_h5_type_file(self, val_def, obj_name_list):
require(len(obj_name_list) > 0)
require(find_callname(self.func_ir, val_def) == ('File', 'h5py'))
require(len(val_def.args) > 0)
f_name = find_const(self.func_ir, val_def.args[0])
require(isinstance(f_name, str))
obj_name_list.reverse()
import h5py
f = h5py.File(f_name, 'r')
obj = f
for obj_name in obj_name_list:
obj = obj[obj_name]
ndims = len(obj.shape)
numba_dtype = numba.numpy_support.from_dtype(obj.dtype)
f.close()
return types.Array(numba_dtype, ndims, 'C')
def _analyze_call_np(self, lhs, func_name, args, array_dists):
"""analyze distributions of numpy functions (np.func_name)
"""
if func_name == 'ascontiguousarray':
self._meet_array_dists(lhs, args[0].name, array_dists)
return
if func_name == 'ravel':
self._meet_array_dists(lhs, args[0].name, array_dists)
return
if func_name == 'concatenate':
self._analyze_call_np_concatenate(lhs, args, array_dists)
return
if func_name == 'array' and is_array(self.typemap, args[0].name):
self._meet_array_dists(lhs, args[0].name, array_dists)
return
# sum over the first axis is distributed, A.sum(0)
if func_name == 'sum' and len(args) == 2:
axis_def = guard(get_definition, self.func_ir, args[1])
if isinstance(axis_def, ir.Const) and axis_def.value == 0:
array_dists[lhs] = Distribution.REP
return
if func_name == 'dot':
self._analyze_call_np_dot(lhs, args, array_dists)
return
# used in df.values
if func_name == 'stack':
seq_info = guard(find_build_sequence, self.func_ir, args[0])
if seq_info is None:
self._analyze_call_set_REP(lhs, args, array_dists,
'np.' + func_name)
return
in_arrs, _ = seq_info
axis = 0
# TODO: support kws
# if 'axis' in kws:
# axis = find_const(self.func_ir, kws['axis'])
if len(args) > 1:
axis = find_const(self.func_ir, args[1])
# parallel if args are 1D and output is 2D and axis == 1
if axis is not None and axis == 1 and self.typemap[lhs].ndim == 2:
for v in in_arrs:
self._meet_array_dists(lhs, v.name, array_dists)
return
if (func_name in [
'cumsum', 'cumprod', 'empty_like', 'zeros_like', 'ones_like',
'full_like', 'copy'
]):
in_arr = args[0].name
self._meet_array_dists(lhs, in_arr, array_dists)
return
# set REP if not found
self._analyze_call_set_REP(lhs, args, array_dists, 'np.' + func_name)
