def get_predefined_types(precomp):
"""Return a dictionary that can be used by a CythonGenerator for
the precomputed symbols.
"""
result = {'dt': 0.0,
't': 0.0,
'dst': KnownType('object'),
'NBRS': KnownType('unsigned int*'),
'N_NBRS': KnownType('int'),
'src': KnownType('ParticleArrayWrapper')}
for sym, value in precomp.items():
result[sym] = value.context[sym]
return result
def get_equation_wrappers(self, known_types={}):
classes = defaultdict(lambda: 0)
eqs = {}
for equation in self.equations:
cls = equation.__class__.__name__
n = classes[cls]
equation.var_name = '%s%d' % (camel_to_underscore(
equation.name), n)
classes[cls] += 1
eqs[cls] = equation
wrappers = []
predefined = dict(get_predefined_types(self.pre_comp))
predefined.update(known_types)
predefined['NBRS'] = KnownType('GLOBAL_MEM unsigned int*')
use_local_memory = get_config().use_local_memory
modified_classes = []
if use_local_memory:
modified_classes = self._update_for_local_memory(predefined, eqs)
code_gen = self._Converter_Class(known_types=predefined)
ignore = ['reduce', 'converged']
for cls in sorted(classes.keys()):
src = code_gen.parse_instance(eqs[cls], ignore_methods=ignore)
wrappers.append(src)
if use_local_memory:
# Remove the added annotations
for cls in modified_classes:
self._set_loop_annotation(cls.loop, {})
return '\n'.join(wrappers)
def test_array_declarations_with_known_types(self):
# Given
g = self.group
known_types = {'d_x': KnownType('float*')}
# When
result = g.get_array_declarations(['d_x'], known_types)
# Then.
expect = 'cdef float* d_x'
self.assertEqual(result, expect)
def _update_for_local_memory(self, predefined, eqs):
modified_classes = []
loop_ann = predefined.copy()
for k in loop_ann.keys():
if 's_' in k:
# TODO: Make each argument have their own KnownType
# right from the start
loop_ann[k] = KnownType(loop_ann[k].type.replace(
'__global', '__local'))
for eq in eqs.values():
cls = eq.__class__
loop = getattr(cls, 'loop', None)
if loop is not None:
self._set_loop_annotation(loop, loop_ann)
modified_classes.append(cls)
return modified_classes
def test_that_all_types_are_detected_correctly(self):
x = np.linspace(0, 1, 10)
pa = ParticleArray(name='f', x=x)
pa.remove_property('pid')
info = get_all_array_names([pa])
result = get_known_types_for_arrays(info)
expect = {
'd_gid': KnownType("unsigned int*"),
'd_tag': KnownType("int*"),
'd_x': KnownType("double*"),
's_gid': KnownType("unsigned int*"),
's_tag': KnownType("int*"),
's_x': KnownType("double*")
}
for key in expect:
self.assertEqual(repr(result[key]), repr(expect[key]))