def ompfy(stmt, **options):
"""
Converts some statements to OpenMP statments.
stmt: stmt, list
statement or a list of statements
"""
if isinstance(stmt, (list, tuple, Tuple)):
return [ompfy(i, **options) for i in stmt]
if isinstance(stmt, Tensor):
# TODO to implement
return stmt
if isinstance(stmt, ForIterator):
iterable = ompfy(stmt.iterable, **options)
target = stmt.target
body = ompfy(stmt.body, **options)
info, clauses = get_for_clauses(iterable)
if (clauses is None):
return ForIterator(target, iterable, body, strict=False)
else:
loop = ForIterator(target, iterable, body, strict=False)
nowait = info['nowait']
return OMP_For(loop, clauses, nowait)
if isinstance(stmt, For):
iterable = ompfy(stmt.iterable, **options)
target = stmt.target
body = ompfy(stmt.body, **options)
return For(target, iterable, body, strict=False)
if isinstance(stmt, list):
return [ompfy(a, **options) for a in stmt]
if isinstance(stmt, While):
test = ompfy(stmt.test, **options)
body = ompfy(stmt.body, **options)
return While(test, body)
if isinstance(stmt, With):
test = ompfy(stmt.test, **options)
body = ompfy(stmt.body, **options)
settings = ompfy(stmt.settings, **options)
clauses = get_with_clauses(test)
if (clauses is None):
return With(test, body, settings)
else:
# TODO to be defined
variables = []
return OMP_Parallel(clauses, variables, body)
if isinstance(stmt, If):
args = []
for block in stmt.args:
test = block[0]
stmts = block[1]
t = ompfy(test, **options)
s = ompfy(stmts, **options)
args.append((t,s))
return If(*args)
if isinstance(stmt, FunctionDef):
name = ompfy(stmt.name, **options)
arguments = ompfy(stmt.arguments, **options)
results = ompfy(stmt.results, **options)
body = ompfy(stmt.body, **options)
local_vars = ompfy(stmt.local_vars, **options)
global_vars = ompfy(stmt.global_vars, **options)
return FunctionDef(name, arguments, results,
body, local_vars, global_vars)
if isinstance(stmt, ClassDef):
name = ompfy(stmt.name, **options)
attributs = ompfy(stmt.attributs, **options)
methods = ompfy(stmt.methods, **options)
options = ompfy(stmt.options, **options)
return ClassDef(name, attributs, methods, options)
if isinstance(stmt, Module):
name = ompfy(stmt.name, **options)
variables = ompfy(stmt.variables, **options)
funcs = ompfy(stmt.funcs , **options)
classes = ompfy(stmt.classes , **options)
imports = ompfy(stmt.imports , **options)
imports += [Import('omp_lib')]
return Module(name, variables, funcs, classes,
imports=imports)
if isinstance(stmt, Program):
name = ompfy(stmt.name, **options)
variables = ompfy(stmt.variables, **options)
funcs = ompfy(stmt.funcs , **options)
classes = ompfy(stmt.classes , **options)
imports = ompfy(stmt.imports , **options)
body = ompfy(stmt.body , **options)
modules = ompfy(stmt.modules , **options)
imports += [Import('omp_lib')]
return Program(name, variables, funcs, classes, body,
imports=imports, modules=modules)
if isinstance(stmt, ParallelBlock):
variables = stmt.variables
body = stmt.body
clauses = stmt.clauses
return OMP_Parallel(clauses, variables, body)
return stmt
def _visit_While(self, stmt):
test = self._visit(stmt.test)
body = self._visit(stmt.body)
return While(test, body)
def mpify(stmt, **options):
"""
Converts some statements to MPI statments.
stmt: stmt, list
statement or a list of statements
"""
if isinstance(stmt, (list, tuple, Tuple)):
return [mpify(i, **options) for i in stmt]
if isinstance(stmt, MPI):
return stmt
if isinstance(stmt, Tensor):
options['label'] = stmt.name
return stmt
if isinstance(stmt, ForIterator):
iterable = mpify(stmt.iterable, **options)
target = stmt.target
body = mpify(stmt.body, **options)
return ForIterator(target, iterable, body, strict=False)
if isinstance(stmt, For):
iterable = mpify(stmt.iterable, **options)
target = stmt.target
body = mpify(stmt.body, **options)
return For(target, iterable, body, strict=False)
if isinstance(stmt, list):
return [mpify(a, **options) for a in stmt]
if isinstance(stmt, While):
test = mpify(stmt.test, **options)
body = mpify(stmt.body, **options)
return While(test, body)
if isinstance(stmt, If):
args = []
for block in stmt.args:
test = block[0]
stmts = block[1]
t = mpify(test, **options)
s = mpify(stmts, **options)
args.append((t, s))
return If(*args)
if isinstance(stmt, FunctionDef):
return stmt
# TODO uncomment this
# name = mpify(stmt.name, **options)
# arguments = mpify(stmt.arguments, **options)
# results = mpify(stmt.results, **options)
# body = mpify(stmt.body, **options)
# local_vars = mpify(stmt.local_vars, **options)
# global_vars = mpify(stmt.global_vars, **options)
#
# return FunctionDef(name, arguments, results, \
# body, local_vars, global_vars)
if isinstance(stmt, ClassDef):
name = mpify(stmt.name, **options)
attributs = mpify(stmt.attributs, **options)
methods = mpify(stmt.methods, **options)
options = mpify(stmt.options, **options)
return ClassDef(name, attributs, methods, options)
if isinstance(stmt, Assign):
if isinstance(stmt.rhs, Tensor):
lhs = stmt.lhs
options['label'] = lhs.name
rhs = mpify(stmt.rhs, **options)
return Assign(lhs, rhs, \
strict=stmt.strict, \
status=stmt.status, \
like=stmt.like)
if isinstance(stmt, Del):
variables = [mpify(a, **options) for a in stmt.variables]
return Del(variables)
if isinstance(stmt, Ones):
if stmt.grid:
lhs = stmt.lhs
shape = stmt.shape
grid = mpify(stmt.grid, **options)
return Ones(lhs, grid=grid)
if isinstance(stmt, Zeros):
if stmt.grid:
lhs = stmt.lhs
shape = stmt.shape
grid = mpify(stmt.grid, **options)
return Zeros(lhs, grid=grid)
if isinstance(stmt, Module):
name = mpify(stmt.name, **options)
variables = mpify(stmt.variables, **options)
funcs = mpify(stmt.funcs, **options)
classes = mpify(stmt.classes, **options)
imports = mpify(stmt.imports, **options)
imports += [Import('mpi')]
# TODO add stdlib_parallel_mpi module
return Module(name, variables, funcs, classes, imports=imports)
if isinstance(stmt, Program):
name = mpify(stmt.name, **options)
variables = mpify(stmt.variables, **options)
funcs = mpify(stmt.funcs, **options)
classes = mpify(stmt.classes, **options)
imports = mpify(stmt.imports, **options)
body = mpify(stmt.body, **options)
modules = mpify(stmt.modules, **options)
imports += [Import('mpi')]
# TODO improve this import, without writing 'mod_...'
# maybe we should create a new class for this import
imports += [Import('mod_pyccel_stdlib_parallel_mpi')]
return Program(name,
variables,
funcs,
classes,
body,
imports=imports,
modules=modules)
return stmt