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, {}