def _alloc_array_on_high_bw_mem(self, site, obj, storage):
_storage = Storage()
super()._alloc_array_on_high_bw_mem(site, obj, _storage)
allocs = _storage[site].allocs + [self._Parallelizer._map_alloc(obj)]
frees = [self._Parallelizer._map_delete(obj)] + _storage[site].frees
storage.update(obj, site, allocs=allocs, frees=frees)
def _(iet):
# Collect written and read-only symbols
writes = set()
reads = set()
for i, v in MapExprStmts().visit(iet).items():
if not i.is_Expression:
# No-op
continue
if not any(isinstance(j, self._Parallelizer._Iteration) for j in v):
# Not an offloaded Iteration tree
continue
if i.write.is_DiscreteFunction:
writes.add(i.write)
reads = (reads | {r for r in i.reads if r.is_DiscreteFunction}) - writes
# Populate `storage`
storage = Storage()
for i in filter_sorted(writes):
self._map_function_on_high_bw_mem(iet, i, storage)
for i in filter_sorted(reads):
self._map_function_on_high_bw_mem(iet, i, storage, read_only=True)
iet = self._dump_storage(iet, storage)
return iet
def place_ondevice(self, iet, **kwargs):
efuncs = kwargs['efuncs']
storage = Storage()
if iet.is_ElementalFunction:
return iet, {}
# Collect written and read-only symbols
writes = set()
reads = set()
for efunc in efuncs:
for i, v in MapExprStmts().visit(efunc).items():
if not i.is_Expression:
# No-op
continue
if not any(isinstance(j, DeviceParallelIteration) for j in v):
# Not an offloaded Iteration tree
continue
if i.write.is_DiscreteFunction:
writes.add(i.write)
reads = (reads | {r
for r in i.reads
if r.is_DiscreteFunction}) - writes
# Update `storage`
for i in filter_sorted(writes):
self._map_function_on_high_bw_mem(i, storage)
for i in filter_sorted(reads):
self._map_function_on_high_bw_mem(i, storage, read_only=True)
iet = self._dump_storage(iet, storage)
return iet, {}