Tests passing, about to try something new on another branch

Author: connorjward

pyop3/array/harray.py

@@ -123,6 +123,10 @@ def __init__(
     ):
         super().__init__(name=name, prefix=prefix)
 
+        # debug
+        # if self.name == "t_0":
+        #     breakpoint()
+
         axes = as_axis_tree(axes)
 
         if isinstance(data, Buffer):

pyop3/axtree/layout.py

@@ -160,6 +160,7 @@ def collect_externally_indexed_axes(axes, axis=None, component=None, path=pmap()
                 }
             )
     else:
+        path_ = path | {axis.label: component.label}
         csize = component.count
         if isinstance(csize, HierarchicalArray):
             if csize.axes.is_empty:
@@ -168,14 +169,14 @@ def collect_externally_indexed_axes(axes, axis=None, component=None, path=pmap()
                 # is the path sufficient? i.e. do we have enough externally provided indices
                 # to correctly index the axis?
                 for caxis, ccpt in csize.axes.path_with_nodes(*csize.axes.leaf).items():
-                    if caxis.label in path:
-                        assert path[caxis.label] == ccpt, "Paths do not match"
+                    if caxis.label in path_:
+                        assert path_[caxis.label] == ccpt, "Paths do not match"
                     else:
+                        # also return an expr?
                         external_axes[caxis.label] = caxis
         else:
             assert isinstance(csize, numbers.Integral)
             if subaxis := axes.child(axis, component):
-                path_ = path | {axis.label: component.label}
                 for subcpt in subaxis.components:
                     external_axes.update(
                         {

pyop3/ir/lower.py

@@ -828,10 +828,23 @@ def parse_assignment_properly_this_time(
             (axis.id, component.label), pmap()
         )
 
+        # TODO move to register_extent
+        if isinstance(component.count, HierarchicalArray):
+            count_axes = component.count.axes
+            count_exprs = {}
+            for count_axis, count_cpt in count_axes.path_with_nodes(
+                *count_axes.leaf
+            ).items():
+                count_exprs.update(
+                    component.count.index_exprs.get((count_axis.id, count_cpt), {})
+                )
+        else:
+            count_exprs = {}
+
         extent_var = register_extent(
             component.count,
-            index_exprs[assignment.assignee] | loop_indices | domain_index_exprs,
-            iname_replace_map,
+            index_exprs[assignment.assignee] | count_exprs | domain_index_exprs,
+            iname_replace_map | loop_indices,
             codegen_context,
         )
         codegen_context.add_domain(iname, extent_var)

pyop3/lang.py

@@ -200,14 +200,19 @@ def kernel_arguments(self):
     @cached_property
     def _distarray_args(self):
         from pyop3.array import ContextSensitiveMultiArray, HierarchicalArray
+        from pyop3.buffer import DistributedBuffer
 
         arrays = {}
         for arg, intent in self.kernel_arguments:
             if isinstance(arg, ContextSensitiveMultiArray):
                 # take first
                 arg, *_ = arg.context_map.values()
 
-            if not isinstance(arg, HierarchicalArray) or not arg.buffer.is_distributed:
+            if (
+                not isinstance(arg, HierarchicalArray)
+                or not isinstance(arg.buffer, DistributedBuffer)
+                or arg.buffer.sf is None
+            ):
                 continue
 
             if arg.array not in arrays:

pyop3/transform.py

@@ -175,6 +175,7 @@ def _(self, assignment: Assignment):
                     data=NullBuffer(arg.dtype),  # does this need a size?
                     name=self._name_generator("t"),
                 )
+                breakpoint()
 
                 if intent == READ:
                     gathers.append(PetscMatLoad(arg, new_arg))