From 44e1ff26892ced708dc19af16cbb134d06e1582f Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 22 Apr 2024 16:57:23 +0100 Subject: [PATCH 1/2] Tabulated Offsets --- pyop3/axtree/tree.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/pyop3/axtree/tree.py b/pyop3/axtree/tree.py index 8cc0d6d..e4fa082 100644 --- a/pyop3/axtree/tree.py +++ b/pyop3/axtree/tree.py @@ -1286,6 +1286,30 @@ def _collect_buffer_indices(self, *, include_ghost_points: bool): else: return indices + @cached_property + def tabulated_offsets(self): + from pyop3.array import HierarchicalArray + + loop_index = just_one(self.outer_loops) + iterset = AxisTree(loop_index.iterset.node_map) + rmap_axes = iterset.add_subtree(self, *iterset.leaf) + rmap = HierarchicalArray(rmap_axes, dtype=IntType) + rmap = rmap[loop_index.local_index] + for idx in loop_index.iter(): + target_indices = idx.replace_map + # for p in self.iter(idxs): + for p in self.iter([idx], include_ghost_points=True): # seems to fix thing + offset = self.axes.unindexed.offset( + p.target_exprs, p.target_path, loop_exprs=target_indices + ) + rmap.set_value( + p.source_exprs, + offset, + p.source_path, + loop_exprs=target_indices, + ) + return rmap + class ContextSensitiveAxisTree(ContextSensitiveLoopIterable): def __getitem__(self, indices) -> ContextSensitiveAxisTree: From a910cb2dcce262abddcf91d35b4917a950b09051 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 25 Apr 2024 13:28:08 +0100 Subject: [PATCH 2/2] More useful message when caling a Map --- pyop3/ir/lower.py | 4 ++-- pyop3/itree/tree.py | 13 ++++++++----- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/pyop3/ir/lower.py b/pyop3/ir/lower.py index f042062..0b915ce 100644 --- a/pyop3/ir/lower.py +++ b/pyop3/ir/lower.py @@ -450,8 +450,8 @@ def compile(expr: Instruction, name="mykernel"): textwrap.dedent( """ #include - - + + int32_t cmpfunc(const void * a, const void * b) { return ( *(int32_t*)a - *(int32_t*)b ); } diff --git a/pyop3/itree/tree.py b/pyop3/itree/tree.py index 6340a24..4c415aa 100644 --- a/pyop3/itree/tree.py +++ b/pyop3/itree/tree.py @@ -567,11 +567,14 @@ def __init__(self, connectivity, name=None, *, numbering=None) -> None: def __call__(self, index): if isinstance(index, (ContextFreeIndex, ContextFreeCalledMap)): - leaf_target_paths = tuple( - freeze({mcpt.target_axis: mcpt.target_component}) - for path in index.leaf_target_paths - for mcpt in self.connectivity[path] - ) + try: + leaf_target_paths = tuple( + freeze({mcpt.target_axis: mcpt.target_component}) + for path in index.leaf_target_paths + for mcpt in self.connectivity[path] + ) + except KeyError: + raise KeyError(f"Map does not have a suitable 'from index' that matches the provided index: {index.leaf_target_paths}") return ContextFreeCalledMap(self, index, leaf_target_paths) else: return CalledMap(self, index)