def get_dims_and_symbol(expr):
if isinstance(expr, Indexed):
dims = tuple([(S.Zero, dim - 1) for dim in expr.shape])
symbol = expr.base.label
elif isinstance(expr, Symbol):
dims = []
symbol = expr
elif isinstance(expr, MatrixSymbol):
dims = tuple([(S.Zero, dim - 1) for dim in expr.shape if dim != 1])
symbol = expr
elif isinstance(expr, MatrixBase):
# if we have a Matrix, we set line by line the code
# useful when you have indexes like
# A[i, j, 0] = B[i, j, 4]
# todo: regarder que les matrices pour le in et out ont la meme taille
# todo: regarder si le symbole est toujours le meme dans le terme de gauche (expr)
for i in range(expr.shape[0]):
symbol = expr[i].base.label
dims = tuple([(S.Zero, dim - 1) for dim in expr[i].base.shape
if dim != 1])
elif isinstance(expr, MatrixSlice):
symbol = expr.parent
dims = tuple([(S.Zero, dim - 1) for dim in symbol.shape if dim != 1])
else:
raise CodeGenError("Only Indexed, Symbol, or MatrixSymbol "
"can define output arguments.")
return dims, symbol
def _get_routine_opening(self, routine):
"""Returns the opening statements of the routine."""
code_list = []
export = True
# export = self.settings.pop('export', True)
if export:
code_list.append("def ")
else:
code_list.append("cdef void ")
# Inputs
args = []
for i, arg in enumerate(routine.arguments):
if isinstance(arg, OutputArgument):
raise CodeGenError("Cython: invalid argument of type %s" %
str(type(arg)))
if isinstance(arg, (InputArgument, InOutArgument)):
name = self._get_symbol(arg.name)
if not arg.dimensions:
# If it is a scalar
if isinstance(arg, ResultBase):
# if it is an output
args.append((arg.get_datatype('C'), "*%s" % name))
else:
# if it is an output
args.append((arg.get_datatype('C'), name))
else:
if not export and len(arg.dimensions) == 1:
# if the dimension is 1
args.append((arg.get_datatype('C'), "*%s" % name))
else:
args.append(
(arg.get_datatype('C') + '[' +
', '.join([':'] * len(arg.dimensions)) + ':1]',
"%s" % name))
args = ", ".join(["%s %s" % t for t in args])
code_list.append("%s(%s)%s\n" %
(routine.name, args, ":" if export else " nogil:"))
code_list = ["".join(code_list)]
return code_list
def _get_routine_opening(self, routine):
"""Returns the opening statements of the routine."""
code_list = []
code_list.append("def ")
# Inputs
args = []
for i, arg in enumerate(routine.arguments):
if isinstance(arg, OutputArgument):
raise CodeGenError("Numpy: invalid argument of type %s" %
str(type(arg)))
if isinstance(arg, (InputArgument, InOutArgument)):
name = self._get_symbol(arg.name)
args.append(name)
args = ", ".join(args)
code_list.append("%s(%s):\n" % (routine.name, args))
code_list = ["".join(code_list)]
return code_list
def routine(self, name, expr, argument_sequence, global_vars):
"""Creates an Routine object that is appropriate for this language.
This implementation is appropriate for at least C/Fortran. Subclasses
can override this if necessary.
Here, we assume at most one return value (the l-value) which must be
scalar. Additional outputs are OutputArguments (e.g., pointers on
right-hand-side or pass-by-reference). Matrices are always returned
via OutputArguments. If ``argument_sequence`` is None, arguments will
be ordered alphabetically, but with all InputArguments first, and then
OutputArgument and InOutArguments.
This implementation is almost the same as the CodeGen class, but
expensive calls to Basic.atoms() have been replaced with
cheaper equivalents.
"""
if is_sequence(expr) and not isinstance(expr, (MatrixBase, MatrixExpr)):
if not expr:
raise ValueError("No expression given")
expressions = Tuple(*expr)
else:
expressions = Tuple(expr)
expr_free_symbols = expressions.free_symbols
# local variables
local_vars = {i.label for i in expr_free_symbols if isinstance(i, Idx)}
# global variables
global_vars = set() if global_vars is None else set(global_vars)
# symbols that should be arguments
symbols = expr_free_symbols - local_vars - global_vars
new_symbols = set([])
new_symbols.update(symbols)
for symbol in symbols:
if isinstance(symbol, Idx):
new_symbols.remove(symbol)
new_symbols.update(symbol.args[1].free_symbols)
if isinstance(symbol, Indexed):
new_symbols.remove(symbol)
symbols = new_symbols
# Decide whether to use output argument or return value
return_val = []
output_args = []
for expr in expressions:
if isinstance(expr, Equality):
out_arg = expr.lhs
expr = expr.rhs
if isinstance(out_arg, Indexed):
dims = tuple([ (S.Zero, dim - 1) for dim in out_arg.shape])
symbol = out_arg.base.label
elif isinstance(out_arg, Symbol):
dims = []
symbol = out_arg
elif isinstance(out_arg, MatrixSymbol):
dims = tuple([ (S.Zero, dim - 1) for dim in out_arg.shape])
symbol = out_arg
else:
raise CodeGenError("Only Indexed, Symbol, or MatrixSymbol "
"can define output arguments.")
if expr.has(symbol):
output_args.append(
InOutArgument(symbol, out_arg, expr, dimensions=dims))
else:
output_args.append(
OutputArgument(symbol, out_arg, expr, dimensions=dims))
# avoid duplicate arguments
symbols.remove(symbol)
elif isinstance(expr, (ImmutableMatrix, MatrixSlice)):
# Create a "dummy" MatrixSymbol to use as the Output arg
out_arg = MatrixSymbol('out_%s' % abs(hash(expr)), *expr.shape)
dims = tuple([(S.Zero, dim - 1) for dim in out_arg.shape])
output_args.append(
OutputArgument(out_arg, out_arg, expr, dimensions=dims))
else:
return_val.append(Result(expr))
arg_list = []
# setup input argument list
array_symbols = {}
for array in [i for i in expr_free_symbols if isinstance(i, Indexed)]:
array_symbols[array.base.label] = array
for array in [i for i in expr_free_symbols if isinstance(i, MatrixSymbol)]:
array_symbols[array] = array
for symbol in sorted(symbols, key=str):
if symbol in array_symbols:
dims = []
array = array_symbols[symbol]
for dim in array.shape:
dims.append((S.Zero, dim - 1))
metadata = {'dimensions': dims}
else:
metadata = {}
arg_list.append(InputArgument(symbol, **metadata))
output_args.sort(key=lambda x: str(x.name))
arg_list.extend(output_args)
if argument_sequence is not None:
# if the user has supplied IndexedBase instances, we'll accept that
new_sequence = []
for arg in argument_sequence:
if isinstance(arg, IndexedBase):
new_sequence.append(arg.label)
else:
new_sequence.append(arg)
argument_sequence = new_sequence
missing = [x for x in arg_list if x.name not in argument_sequence]
if missing:
msg = "Argument list didn't specify: {0} "
msg = msg.format(", ".join([str(m.name) for m in missing]))
raise CodeGenArgumentListError(msg, missing)
# create redundant arguments to produce the requested sequence
name_arg_dict = {x.name: x for x in arg_list}
new_args = []
for symbol in argument_sequence:
try:
new_args.append(name_arg_dict[symbol])
except KeyError:
new_args.append(InputArgument(symbol))
arg_list = new_args
return Routine(name, arg_list, return_val, local_vars, global_vars)
def _get_routine_ending(self, routine):
code_list = []
code_list.append('""",')
# Inputs
args = []
dtypes = []
for i, arg in enumerate(routine.arguments):
if isinstance(arg, OutputArgument):
raise CodeGenError("Loopy: invalid argument of type %s" %
str(type(arg)))
if isinstance(arg, (InputArgument, InOutArgument)):
name = self._get_symbol(arg.name)
if arg.dimensions:
dims = [
"{}".format(d[1] - d[0] + 1) for d in arg.dimensions
]
dtype = arg.get_datatype('PYTHON')
if dtype == 'int':
dtype = 'np.int32'
args.append(
'lp.GlobalArg("{name}", dtype={dtype}, shape="{shape}")'
.format(name=name, dtype=dtype, shape=", ".join(dims)))
else:
args.append('lp.ValueArg("{name}", dtype={dtype})'.format(
name=name, dtype=arg.get_datatype('PYTHON')))
for i, arg in enumerate(routine.local_vars):
if isinstance(arg, Symbol):
args.append(
'lp.TemporaryVariable("{name}", dtype=float)'.format(
name=self._get_symbol(arg)))
else:
dims = [d for d in arg.shape if d != 1]
args.append(
'lp.TemporaryVariable("{name}", dtype=float, shape="{shape}")'
.format(name=self._get_symbol(arg),
shape=','.join("%s" % s for s in dims)))
code_list.append('[')
args = ",\n".join(args)
code_list.append(args)
code_list.append('])#endArg\n')
# add type
dim = len(routine.idx_vars)
if dim == 1:
block_size = [256]
if dim == 2:
block_size = [16, 16]
if dim == 3:
block_size = [4, 4, 4]
for i, idx in enumerate(routine.idx_vars[-1::-1]):
code_list.append(
'{name} = lp.split_iname({name}, "{label}", {block}, outer_tag="g.{ilabel}", inner_tag="l.{ilabel}")'
.format(name=routine.name,
label="%s_" % idx.label,
ilabel=i,
block=block_size[i]))
code_list.append(
'{name} = lp.expand_subst({name})\n'.format(name=routine.name))
code_list.append(
'{name} = lp.set_options({name}, no_numpy = True)\n'.format(
name=routine.name))
prefetch = routine.settings.get("prefetch", None)
if prefetch:
for var in prefetch:
indices = []
for idx in routine.idx_vars:
indices.append("%s__inner" % idx.label)
# for i in range(var.rank):
# if isinstance(var.indices[i], Idx):
# indices.append("%s__inner"%var.indices[i].label)
code_list.append(
'{name} = lp.add_prefetch({name}, "{var}", "{label}", fetch_bounding_box=True)\n'
.format(name=routine.name,
var=var.base.label,
label=",".join(indices)))
# print("PREFETCH")
# label = []
# for var in self._settings["prefetch"]:
# indices = []
# for i in var.indices:
# if isinstance(i, Idx):
# indices.append("%s__inner"%i.label)
# code_list.append('{name} = lp.add_prefetch({name}, "{var}", "{label}", fetch_bounding_box=True)'.format(name=routine.name, var=var.label, label=",".join(indices)))
#print(LoopyCodePrinter()._sort_optimized(routine.local_vars, routine.instructions))
# one_time_step = lp.split_iname(one_time_step, "ii", 16, outer_tag="g.1", inner_tag="l.1")
# one_time_step = lp.split_iname(one_time_step, "jj", 16, outer_tag="g.0", inner_tag="l.0")
# one_time_step = lp.expand_subst(one_time_step)
# one_time_step = lp.add_prefetch(one_time_step, "f", "ii_inner,jj_inner", fetch_bounding_box=True)
code_list = ["\n".join(code_list)]
return code_list