Parallel bits

Author: connorjward

pyop3/array/petsc.py

@@ -308,7 +308,6 @@ def _iter_nest_labels(
     @cached_property
     @PETSc.Log.EventDecorator()
     def maps(self):
-        print("HIT!")
         from pyop3.axtree.layout import my_product
 
         # TODO: Don't think these need to be lists here.
@@ -380,7 +379,8 @@ def maps(self):
                 # target_indices = {idx.index.id: idx.target_exprs for idx in idxs}
                 target_indices = merge_dicts([idx.replace_map for idx in idxs])
 
-                for p in self.raxes.iter(idxs):
+                # for p in self.raxes.iter(idxs):
+                for p in self.raxes.iter(idxs, include_ghost_points=True):  # seems to fix things
                     target_path = p.target_path
                     target_exprs = p.target_exprs
                     for key in dropped_rkeys:
@@ -390,6 +390,7 @@ def maps(self):
                     offset = orig_raxes.offset(
                         target_exprs, target_path, loop_exprs=target_indices
                     )
+
                     rmap.set_value(
                         p.source_exprs,
                         offset,
@@ -402,7 +403,8 @@ def maps(self):
                 # target_indices = {idx.index.id: idx.target_exprs for idx in idxs}
                 target_indices = merge_dicts([idx.replace_map for idx in idxs])
 
-                for p in self.caxes.iter(idxs):
+                # for p in self.caxes.iter(idxs):
+                for p in self.caxes.iter(idxs, include_ghost_points=True):  # seems to fix things
                     target_path = p.target_path
                     target_exprs = p.target_exprs
                     for key in dropped_ckeys:
@@ -419,6 +421,8 @@ def maps(self):
                         loop_exprs=target_indices,
                     )
 
+        # breakpoint()
+
         return (rmap, cmap)
 
     @property

pyop3/axtree/tree.py

@@ -209,18 +209,20 @@ def _collect_datamap(axis, *subdatamaps, axes):
 
 
 class AxisComponent(LabelledNodeComponent):
-    fields = LabelledNodeComponent.fields | {"size", "unit"}
+    fields = LabelledNodeComponent.fields | {"size", "unit", "rank_equal"}
 
     def __init__(
         self,
         size,
         label=None,
         *,
         unit=False,
+        rank_equal=False,
     ):
         from pyop3.array import HierarchicalArray
 
         if isinstance(size, collections.abc.Iterable):
+            assert not rank_equal  # nasty
             owned_count, count = map(int, size)
             distributed = True
             assert isinstance(owned_count, numbers.Integral) and isinstance(count, numbers.Integral)
@@ -247,8 +249,8 @@ def __init__(
         self.unit = unit
         self.distributed = distributed
 
-        # remove
-        self.rank_equal = not distributed
+        # cleanup
+        self.rank_equal = rank_equal
 
     # redone because otherwise getting a bizarre error (numpy types have confusing behaviour!)
     # def __eq__(self, other):
@@ -370,7 +372,7 @@ def from_serial(cls, serial: Axis, sf):
         # renumber the serial axis to store ghost entries at the end of the vector
         component_sizes, numbering = partition_ghost_points(serial, sf)
         components = [
-            c.copy(size=(size, c.count)) for c, size in checked_zip(serial.components, component_sizes)
+            c.copy(size=(size, c.count), rank_equal=False) for c, size in checked_zip(serial.components, component_sizes)
         ]
         return cls(components, serial.label, numbering=numbering, sf=sf)
 
@@ -811,14 +813,6 @@ def global_numbering(self):
 
         return numbering[self._buffer_indices_ghost]
 
-    @property
-    def comm(self):
-        paraxes = [axis for axis in self.nodes if axis.sf is not None]
-        if not paraxes:
-            return MPI.COMM_SELF
-        else:
-            return single_valued(ax.comm for ax in paraxes)
-
     @cached_property
     def leaf_target_paths(self):
         return tuple(
@@ -907,7 +901,7 @@ def _collect_owned_index_tree(self, axis=None):
                 slice_component = AffineSliceComponent(component.label)
             slice_components.append(slice_component)
 
-        slice_ = Slice(axis.label, slice_components)
+        slice_ = Slice(axis.label, slice_components, label=axis.label)
 
         index_tree = IndexTree(slice_)
         for component, slice_component in checked_zip(axis.components, slice_components):
@@ -1025,6 +1019,14 @@ def sf(self) -> StarForest:
             iremote = np.concatenate(iremotes)
             return StarForest.from_graph(self.size, nroots, ilocal, iremote, self.comm)
 
+    @property
+    def comm(self):
+        paraxes = [axis for axis in self.nodes if axis.sf is not None]
+        if not paraxes:
+            return MPI.COMM_SELF
+        else:
+            return single_valued(ax.comm for ax in paraxes)
+
     @cached_property
     def datamap(self):
         if self.is_empty:
@@ -1127,8 +1129,6 @@ def _buffer_indices_ghost(self):
         return slice(None)
 
 
-# are all of these necessary?
-# class IndexedAxisTree(Indexed, BaseAxisTree):
 class IndexedAxisTree(BaseAxisTree):
     def __init__(
         self,
@@ -1156,6 +1156,10 @@ def __init__(
     def unindexed(self):
         return self._unindexed
 
+    @property
+    def comm(self):
+        return self.unindexed.comm
+
     @property
     def target_paths(self):
         return self._target_paths

pyop3/itree/tree.py

@@ -1301,6 +1301,7 @@ def _(slice_: Slice, indices, *, target_path_acc, prev_axes, **kwargs):
                     size = target_cpt.count
                 else:
                     size = target_cpt.count[indices]
+                rank_equal = True
             else:
                 if subslice.stop is None:
                     stop = target_cpt.count
@@ -1313,22 +1314,31 @@ def _(slice_: Slice, indices, *, target_path_acc, prev_axes, **kwargs):
 
                     owned_count = min(target_cpt.owned_count, stop)
                     count = stop
-                    size = (owned_count, count)
+                    rank_equal = False
+
+                    if owned_count == count:
+                        size = count
+                    else:
+                        size = (owned_count, count)
                 else:
                     size = math.ceil((stop - subslice.start) / subslice.step)
+                    rank_equal = True
 
         else:
             assert isinstance(subslice, Subset)
             size = subslice.array.axes.leaf_component.count
 
+            # kind of misleading, the values may differ I think
+            rank_equal = True
+
             if target_cpt.distributed:
                 raise NotImplementedError
 
         if is_full_slice and subslice.label_was_none:
             mylabel = subslice.component
         else:
             mylabel = subslice.label
-        cpt = AxisComponent(size, label=mylabel, unit=target_cpt.unit)
+        cpt = AxisComponent(size, label=mylabel, unit=target_cpt.unit, rank_equal=rank_equal)
         components.append(cpt)
 
         target_path_per_subslice.append(pmap({slice_.axis: subslice.component}))