def _parse_variables(self):
for ix, (sv_id, sv) in enumerate(self.scope_variables.items()):
full_tag = d_u(sv.id)
self.values_order[full_tag] = ix
self._parse_variable(full_tag, sv, sv_id)
for ix, (sv_id, sv) in enumerate(self.temporary_variables.items()):
full_tag = d_u(sv.id)
self._parse_variable(full_tag, sv, sv_id)
def create_assignments(self, variables):
from tqdm import tqdm
temp_variables = {}
for ii, n in tqdm(
enumerate(
self.get_where_node_attr('node_type',
NodeTypes.EQUATION))):
for i, e in self.get_edges_for_node(start_node=n):
va = e[1].copy()
if va in self.vars_assignments and len(
self.vars_assignments[va]) > 1:
# Make new temp var
sv = self.get(e[1], 'scope_var')
tmp_key = sv.tag + str(self.key_map[va]) + '_tmp'
tmp_label = sv.tag + variables[str(
self.key_map[va])].path.primary_path + '_tmp'
# Create fake scope variables for tmp setvar
fake_sv = {}
svf = None
if isinstance(sv, SetOfVariables):
tmp_var_counter = 0
tsv = TemporarySetVar(tmp_key, sv)
for svi in sv.variables.values():
tmp_var_counter += 1
svf = TemporaryVar(
'tmp_var_' + str(tmp_var_counter), svi,
svi.tag + '_' + str(sv.id), svi.set_var,
svi.set_var_ix)
tsv.tmp_vars.append(svf)
fake_sv[tsv.id] = tsv
else:
svf = TemporaryVar(d_u(tmp_key), sv, tmp_key, None,
None)
fake_sv[d_u(svf.get_path_dot())] = svf
temp_variables.update(fake_sv)
tmp = self.add_node(Node(key=tmp_key,
node_type=NodeTypes.TMP,
name=tmp_key,
file='sum',
label=tmp_label,
ln=0,
scope_var=svf),
ignore_existing=False)
# Add temp var to Equation target
self.add_edge(
Edge(n,
tmp,
e_type=EdgeType.TARGET,
arg_local=self.edges_c[i[0]].arg_local))
# Add temp var in var assignments
self.vars_assignments_mappings[va][(
nix := self.vars_assignments[va].index(n))] = ':'
self.vars_assignments[va][nix] = tmp
return temp_variables
def __init__(self,
filename,
equation_graph,
scope_variables,
equations,
scoped_equations,
temporary_variables,
system_tag="",
use_llvm=True,
imports=None):
self.filename = filename
self.imports = imports
self.system_tag = system_tag
self.scope_variables = scope_variables
self.set_variables = {}
self.states = []
self.deriv = []
for ix, (sv_id, sv) in enumerate(self.scope_variables.items()):
full_tag = d_u(sv.id)
if sv.type == VariableType.STATE:
self.states.append(full_tag)
elif sv.type == VariableType.DERIVATIVE:
self.deriv.append(full_tag)
for k, var in temporary_variables.items():
if var.type == VariableType.TMP_PARAMETER_SET:
self.set_variables.update({k: var})
new_sv = {}
tail = {}
for k, v in self.scope_variables.items():
if k in var.set_var.variables:
tail.update({k: v})
new_sv.update({
var.tmp_vars[v.set_var_ix].id:
var.tmp_vars[v.set_var_ix]
})
else:
new_sv.update({k: v})
self.scope_variables = dict(new_sv, **tail)
if var.type == VariableType.TMP_PARAMETER:
new_sv = {}
tail = {}
for k, v in self.scope_variables.items():
if k == var.scope_var_id:
tail.update({k: v})
new_sv.update({var.id: var})
else:
new_sv.update({k: v})
self.scope_variables = dict(new_sv, **tail)
self.scoped_equations = scoped_equations
self.temporary_variables = temporary_variables
self.values_order = {}
self.scope_var_node = {}
self.scalar_variables = {}
self._parse_variables()
# Sort the graph topologically to start generating code
self.topo_sorted_nodes = equation_graph.topological_nodes()
self.equation_graph = equation_graph.clean()
self.number_of_states = len(self.states)
self.number_of_derivatives = len(self.deriv)
# Initialize llvm builder - will be a list of intermediate llvm instructions to be lowered in generate
self.llvm = use_llvm
if self.llvm:
self.generated_program = LLVMBuilder(np.ascontiguousarray(
[x.value for x in self.scope_variables.values()],
dtype=np.float64),
self.values_order,
self.states,
self.deriv,
system_tag=self.system_tag)
else:
self.generated_program = ASTBuilder(np.ascontiguousarray(
[x.value for x in self.scope_variables.values()],
dtype=np.float64),
self.values_order,
self.states,
self.deriv,
system_tag=self.system_tag)
self.mod_body = []
# Create a kernel of assignments and calls
self.eq_vardefs = {}
# Loop over equation functions and generate code for each equation.
self._parse_equations(equations)
self.all_targeted = []
self.all_read = []
self.all_targeted_set_vars = []
self.all_read_set_vars = []
class EquationGenerator:
def __init__(self,
filename,
equation_graph,
scope_variables,
equations,
scoped_equations,
temporary_variables,
system_tag="",
use_llvm=True,
imports=None):
self.filename = filename
self.imports = imports
self.system_tag = system_tag
self.scope_variables = scope_variables
self.set_variables = {}
self.states = []
self.deriv = []
for ix, (sv_id, sv) in enumerate(self.scope_variables.items()):
full_tag = d_u(sv.id)
if sv.type == VariableType.STATE:
self.states.append(full_tag)
elif sv.type == VariableType.DERIVATIVE:
self.deriv.append(full_tag)
for k, var in temporary_variables.items():
if var.type == VariableType.TMP_PARAMETER_SET:
self.set_variables.update({k: var})
new_sv = {}
tail = {}
for k, v in self.scope_variables.items():
if k in var.set_var.variables:
tail.update({k: v})
new_sv.update({
var.tmp_vars[v.set_var_ix].id:
var.tmp_vars[v.set_var_ix]
})
else:
new_sv.update({k: v})
self.scope_variables = dict(new_sv, **tail)
if var.type == VariableType.TMP_PARAMETER:
new_sv = {}
tail = {}
for k, v in self.scope_variables.items():
if k == var.scope_var_id:
tail.update({k: v})
new_sv.update({var.id: var})
else:
new_sv.update({k: v})
self.scope_variables = dict(new_sv, **tail)
self.scoped_equations = scoped_equations
self.temporary_variables = temporary_variables
self.values_order = {}
self.scope_var_node = {}
self.scalar_variables = {}
self._parse_variables()
# Sort the graph topologically to start generating code
self.topo_sorted_nodes = equation_graph.topological_nodes()
self.equation_graph = equation_graph.clean()
self.number_of_states = len(self.states)
self.number_of_derivatives = len(self.deriv)
# Initialize llvm builder - will be a list of intermediate llvm instructions to be lowered in generate
self.llvm = use_llvm
if self.llvm:
self.generated_program = LLVMBuilder(np.ascontiguousarray(
[x.value for x in self.scope_variables.values()],
dtype=np.float64),
self.values_order,
self.states,
self.deriv,
system_tag=self.system_tag)
else:
self.generated_program = ASTBuilder(np.ascontiguousarray(
[x.value for x in self.scope_variables.values()],
dtype=np.float64),
self.values_order,
self.states,
self.deriv,
system_tag=self.system_tag)
self.mod_body = []
# Create a kernel of assignments and calls
self.eq_vardefs = {}
# Loop over equation functions and generate code for each equation.
self._parse_equations(equations)
self.all_targeted = []
self.all_read = []
self.all_targeted_set_vars = []
self.all_read_set_vars = []
def _parse_variable(self, full_tag, sv, sv_id):
if full_tag not in self.scope_var_node:
self.scope_var_node[full_tag] = sv
# If a scope_variable is part of a set it should be referenced alone
if sv.set_var:
if not sv.set_var.id in self.set_variables:
self.set_variables[sv.set_var.id] = sv.set_var
else:
self.scalar_variables[full_tag] = sv
def _parse_variables(self):
for ix, (sv_id, sv) in enumerate(self.scope_variables.items()):
full_tag = d_u(sv.id)
self.values_order[full_tag] = ix
self._parse_variable(full_tag, sv, sv_id)
for ix, (sv_id, sv) in enumerate(self.temporary_variables.items()):
full_tag = d_u(sv.id)
self._parse_variable(full_tag, sv, sv_id)
def get_external_function_name(self, ext_func):
if self.llvm:
return self._llvm_func_name(ext_func)
return ext_func
def _llvm_func_name(self, ext_func):
return ext_func + '_llvm1.<locals>.' + ext_func + '_llvm'
def _parse_equations(self, equations):
logging.info('make equations for compilation')
for eq_key, eq in equations.items():
vardef = Vardef(llvm=self.llvm)
eq[2].lower_graph = None
if self.llvm:
func_llvm, signature, args, target_ids = compiled_function_from_graph_generic_llvm(
eq[2],
imports=self.imports,
var_def_=Vardef(llvm=self.llvm),
compiled_function=True)
self.generated_program.add_external_function(
func_llvm, signature, len(args), target_ids)
else:
func, args, target_ids = function_from_graph_generic(
eq[2], var_def_=vardef, arg_metadata=eq[2].arg_metadata)
self.generated_program.add_external_function(
func, None, len(args), target_ids)
vardef.llvm_target_ids = target_ids
vardef.args_order = args
self.eq_vardefs[eq_key] = vardef
def search_in_item_scope(self, var_id, item_id):
for var in self.scope_variables.values():
##TODO add namespacecheck
if var.item.id == item_id and var.tag == self.scope_variables[
var_id].tag:
return var.id
raise ValueError("No variable found for id {}", var_id)
def _process_equation_node(self, n):
eq_key = self.scoped_equations[self.equation_graph.key_map[n]]
# Define the function to call for this eq
ext_func = recurse_Attribute(self.equation_graph.get(n, 'func'))
item_id = self.equation_graph.get(n, 'item_id')
vardef = self.eq_vardefs[eq_key]
# Find the arguments by looking for edges of arg type
a_indcs, a_edges = list(
self.equation_graph.get_edges_for_node_filter(
end_node=n, attr='e_type', val=EdgeType.ARGUMENT))
# Determine the local arguments names
args_local = [
self.equation_graph.key_map[ae[0]]
for i, ae in zip(a_indcs, a_edges)
if not self.equation_graph.edges_c[i].arg_local == 'local'
]
# Determine the local arguments names
args_scope_var = [
self.equation_graph.edges_c[i].arg_local
for i, ae in zip(a_indcs, a_edges)
if not self.equation_graph.edges_c[i].arg_local == 'local'
]
# Find the targets by looking for target edges
t_indcs, t_edges = list(
self.equation_graph.get_edges_for_node_filter(start_node=n,
attr='e_type',
val=EdgeType.TARGET))
targets_local = [
self.equation_graph.key_map[te[1]]
for i, te in zip(t_indcs, t_edges)
if not self.equation_graph.edges_c[i].arg_local == 'local'
]
targets_scope_var = [
self.equation_graph.edges_c[i].arg_local
for i, ae in zip(t_indcs, t_edges)
if not self.equation_graph.edges_c[i].arg_local == 'local'
]
set_size = 0
# Record targeted and read variables
if self.equation_graph.get(n, 'vectorized'):
# Map of scope.?? vars and set variable names
scope_vars = {
'scope.' + self.set_variables[k].tag: v
for k, v in zip(args_scope_var +
targets_scope_var, args_local + targets_local)
}
# Put the information of args and targets in the scope_var attr of the graph node for those equation
self.equation_graph.nodes[n].scope_var = {
'args': [scope_vars[a] for a in vardef.args],
'targets': [scope_vars[a] for a in vardef.targets]
}
# Record all targeted variables
for t in vardef.targets:
self.all_targeted_set_vars.append(scope_vars[t])
##TODO check that they all the same size
set_size = self.set_variables[scope_vars[t]].get_size()
# Record all read variables
for a in vardef.args:
self.all_read_set_vars.append(scope_vars[a])
else:
# Map of scope.?? vars and global-scope variable names
scope_vars = {
'scope.' + self.scope_variables[k].tag: v
for k, v in zip(args_scope_var +
targets_scope_var, args_local + targets_local)
}
# Put the information of args and targets in the scope_var attr of the graph node for those equation
self.equation_graph.nodes[n].scope_var = {
'args': [scope_vars[a] for a in vardef.args],
'targets': [scope_vars[a] for a in vardef.targets]
}
for a in vardef.args:
if (sva := scope_vars[a]) in self.set_variables:
self.all_read_set_vars.append(sva)
else:
self.all_read.append(sva)
self.all_targeted += [scope_vars[t] for t in vardef.targets]
# Generate ast for this equation callcompiled_function_from_graph_generic_llvm
if self.equation_graph.get(n, 'vectorized'):
llvm_args = []
for t in vardef.args_order:
llvm_args_ = []
set_var = self.set_variables[scope_vars[t]]
for i in range(set_var.get_size()):
llvm_args_.append(set_var.get_var_by_idx(i).id)
llvm_args.append(llvm_args_)
##reshape to correct format
llvm_args = [list(x) for x in zip(*llvm_args)]
self.generated_program.add_set_call(
self.get_external_function_name(ext_func), llvm_args,
vardef.llvm_target_ids)
else:
# Generate llvm arguments
args = []
for a in vardef.args_order:
if a in scope_vars:
args.append(d_u(scope_vars[a]))
else:
args.append(self.search_in_item_scope(a, item_id))
# Add this eq to the llvm_program
self.generated_program.add_call(
self.get_external_function_name(ext_func), args,
vardef.llvm_target_ids)
self.equation_graph.key_map[t]])):
target_indcs_map[mi].append((v[0], mi))
else:
target_indcs_map[0].append((v[0], None))
else:
for mi in maps:
target_indcs_map[mi[1] if mi[1] else 0].append(
(v[0], mi[0]))
target_var = self.equation_graph.key_map[t]
# Generate llvm/ast
mapping_dict = {}
for values in target_indcs_map:
for v in values:
var_name = d_u(self.equation_graph.key_map[v[0]])
if var_name in self.scope_variables:
if target_var in mapping_dict:
mapping_dict[target_var].append(var_name)
else:
mapping_dict[target_var] = [var_name]
else:
if var_name in self.set_variables:
if var_name in mapping_dict:
mapping_dict[target_var].append(
self.set_variables[var_name].
get_var_by_idx(v[1]).id)
else:
mapping_dict[target_var] = [
self.set_variables[var_name].
get_var_by_idx(v[1]).id