[schedule][mlir] consumer fusion

src/xtc/backends/jir/JIRScheduler.py

@@ -318,6 +318,11 @@ def fuse_producer_at(
         # TODO: not implemented for now
         pass
 
+    @override
+    def fuse_consumer_at(self, axis: str, root: str = DEFAULT_ROOT) -> None:
+        # TODO: not implemented for now
+        pass
+
     @override
     def define_memory_mesh(self, axes: dict[str, int]) -> None:
         # TODO: not implemented for now

src/xtc/backends/mlir/MlirCompilerPasses.py

@@ -33,7 +33,8 @@
 )
 from mlir.passmanager import PassManager
 from mlir.ir import Module
-import mlir.xtc_transform
+
+from mlir.xtc_transform import FuseConsumerOp
 
 # Import SDist if available
 try:
@@ -231,12 +232,17 @@ def _generate_scheduling(self) -> OpResult:
                     schedule=schedule,
                     root=list(schedule.permutation)[0],
                     handle=handle,
-                    fuse_axes=fused_producers.get(schedule.node_ident),
+                    producer_fuse_axes=fused_producers.get(schedule.node_ident),
                 )
                 if schedule.vectorization or self._always_vectorize:
                     self._post_vectorize(scheduling_state, schedule)
                 handle = scheduling_state.handle
 
+                if schedule.fused_consumers:
+                    self._fuse_consumers_into_loops(
+                        schedule, scheduling_state, unscheduled_handles
+                    )
+
         assert handle, "At least 1 operation should have been processed"
         return handle
 
@@ -308,7 +314,7 @@ def _generate_node_scheduling(
         schedule: MlirNodeSchedule,
         root: str,
         handle: OpResult,
-        fuse_axes: dict[str, list[str]] | None,
+        producer_fuse_axes: dict[str, list[str]] | None,
     ) -> SchedulingState:
         sched_state = SchedulingState({}, handle, None)
         split_state = SplitState(schedule.splits, root)
@@ -361,9 +367,9 @@ def _generate_node_scheduling(
                 if loop_name in schedule.distribution:
                     self._distribute_loop(loop_name, schedule, sched_state)
             # Fuse the producers
-            if fuse_axes and loop_name in fuse_axes:
+            if producer_fuse_axes and loop_name in producer_fuse_axes:
                 self._fuse_producers_into_loop(
-                    loop_name, fuse_axes, schedule, sched_state
+                    loop_name, producer_fuse_axes, schedule, sched_state
                 )
 
         # For now on, the focus is on the outermost loop
@@ -376,6 +382,41 @@ def _generate_node_scheduling(
 
         return sched_state
 
+    def _fuse_consumers_into_loops(
+        self,
+        schedule: MlirNodeSchedule,
+        sched_state: SchedulingState,
+        unscheduled_handles: set[str | None],
+    ):
+        assert self._named_sequence is not None
+
+        fuse_root = parent_name(schedule.fused_consumers[0])
+        for fuse_axis in schedule.fused_consumers:
+            # derive handle of consumer
+            consumer_handles = find_consumer_handles(
+                self._mlir_program.mlir_module, schedule.node_ident
+            )
+            consumer_id = consumer_handles[0]
+            unscheduled_handles.add(consumer_id)
+            # fuse consumer into all loops until the fuse_axis
+            fuse_loops = []
+            for loop_dim in schedule.permutation[fuse_root]:
+                transform_result = sched_state.all_loops[loop_dim]
+                fuse_loops.append(transform_result)
+                if loop_dim == fuse_axis:
+                    break
+            consumer_handle = structured_match(
+                results_=transform.AnyOpType.get(),
+                target=self._named_sequence.bodyTarget,
+                op_attrs={consumer_id: UnitAttr.get()},
+            )
+            op = FuseConsumerOp(consumer_handle, fuse_loops)
+            # re-annotate the loops that were touched by the fusion
+            for i, loop_dim in enumerate(schedule.permutation[fuse_root]):
+                transform.AnnotateOp(op.new_loops[i], loop_dim)
+                if loop_dim == fuse_axis:
+                    break
+
     def _fuse_producers_into_loop(
         self,
         loop_name: str,
@@ -484,7 +525,10 @@ def _recursive_scheduling(
         self, schedule: MlirNodeSchedule, root: str, sched_state: SchedulingState
     ):
         inner_sched_state = self._generate_node_scheduling(
-            schedule=schedule, root=root, handle=sched_state.handle, fuse_axes=None
+            schedule=schedule,
+            root=root,
+            handle=sched_state.handle,
+            producer_fuse_axes=None,
         )
         sched_state.all_loops.update(inner_sched_state.all_loops)
         sched_state.handle = inner_sched_state.handle
@@ -645,35 +689,58 @@ def _pack_buffer(
                 )
 
     def _collect_fused_producers(self, unscheduled_handles: set[str | None]):
-        # maps each fused consumer op to the producer handles that must be
+        # maps each fused containing op to the producer handles that must be
         # fused through each loop dimension to reach their target fusion depth.
-        fused_producers = {}
+        fused_producer_handles = {}
 
         for schedule in self._nodes_schedules:
-            if schedule.fused:
+            if schedule.fused_producers:
                 prods = find_producer_handles(
                     self._mlir_program.mlir_module, schedule.node_ident
                 )
-                fuse_root = parent_name(schedule.fused[0][0])
+                fuse_root = parent_name(schedule.fused_producers[0][0])
                 unscheduled_handles.update(set(prods))
-                op_axes = {idx: ax for ax, idx in schedule.fused}
+                op_axes = {idx: ax for ax, idx in schedule.fused_producers}
 
                 fuse_destinations = {}
                 for idx, prod_handle in enumerate(prods):
                     if not prod_handle:
                         continue
                     if idx in op_axes:
                         fuse_destinations[prod_handle] = op_axes[idx]
-                # get outer dims to fuse, assumes fuse no splitting avove loop dim
+                # get outer dims to fuse, assumes fuse no splitting above loop dim
                 dim_fuse_handles: dict[str, list[str]] = {}
                 for fuse_handle, fuse_dest in fuse_destinations.items():
                     for dim in schedule.permutation[fuse_root]:
                         dim_fuse_handles.setdefault(dim, []).append(fuse_handle)
                         if dim == fuse_dest:
                             break
-                fused_producers[schedule.node_ident] = dim_fuse_handles
+                fused_producer_handles[schedule.node_ident] = dim_fuse_handles
+
+        return fused_producer_handles
+
+
+def find_consumer_handles(module: Module, root_handle: str) -> list[str | None]:
+    # returns the handles for each consumer op of the operation specified by root_handle
+    consumer_handles: list[str | None] = []
+    root_op = None
+    for func_op in module.body.operations:
+        for op in func_op.regions[0].blocks[0].operations:
+            if root_handle in op.attributes:
+                root_op = op
+                break
+        if root_op:
+            break
+
+    if not root_op:
+        return consumer_handles
 
-        return fused_producers
+    for use in root_op.results[0].uses:
+        consumer_op = use.owner
+        for attr in consumer_op.attributes:
+            if attr.startswith("__xtc_id_"):
+                consumer_handles.append(attr)
+    return consumer_handles
 
 
 def find_producer_handles(module: Module, root_handle: str) -> list[str | None]:
@@ -751,7 +818,6 @@ def run(self, pass_names: list[str]) -> None:
 
 
 def apply_bufferization_passes(mlir_program: RawMlirProgram, mlir_install_dir: str):
-    assert mlir.xtc_transform
     bufferize_options = [
         "bufferize-function-boundaries",
         "function-boundary-type-conversion=identity-layout-map",

src/xtc/backends/mlir/MlirNodeScheduler.py

@@ -31,7 +31,8 @@ class MlirNodeSchedule:
     processor_mesh: dict[str, int]
     distribution: dict[str, str]
     distributed_buffers: dict[str, dict]
-    fused: list[tuple[str, int]]
+    fused_producers: list[tuple[str, int]]
+    fused_consumers: list[str]
 
     def index_of_dim(self, dim: str) -> int:
         return list(self.dims).index(dim)
@@ -92,13 +93,14 @@ def __init__(
         self.processor_mesh: dict[str, int] = {}
         self.distribution: dict[str, str] = {}
         self.distributed_buffers: dict[str, dict] = {}
-        self.fused: list[tuple[str, int]] = []
+        self.fused_producers: list[tuple[str, int]] = []
+        self.fused_consumers: list[str] = []
 
     def mlir_node_schedule(self) -> MlirNodeSchedule:
         if not self.permutation:
             self.permutation[DEFAULT_ROOT] = self.get_default_interchange(DEFAULT_ROOT)
 
-        for fuse_axis in self.fused:
+        for fuse_axis in self.fused_producers:
             assert fuse_axis[0] in self.permutation[next(iter(self.permutation))], (
                 "Fusion must be to an axis in the base root not the result of a split."
             )
@@ -119,7 +121,8 @@ def mlir_node_schedule(self) -> MlirNodeSchedule:
             processor_mesh=self.processor_mesh,
             distribution=self.distribution,
             distributed_buffers=self.distributed_buffers,
-            fused=self.fused,
+            fused_producers=self.fused_producers,
+            fused_consumers=self.fused_consumers,
         )
 
     @override
@@ -229,4 +232,7 @@ def distributed_buffer_at(
     def fuse_producer_at(
         self, axis: str, input_idx: int, root: str = DEFAULT_ROOT
     ) -> None:
-        self.fused.append((make_loop_name(root, axis), input_idx))
+        self.fused_producers.append((make_loop_name(root, axis), input_idx))
+
+    def fuse_consumer_at(self, axis: str, root: str = DEFAULT_ROOT) -> None:
+        self.fused_consumers.append(make_loop_name(root, axis))

src/xtc/backends/mlir/MlirScheduler.py

@@ -174,6 +174,10 @@ def fuse_producer_at(
     ) -> None:
         self._current_scheduler.fuse_producer_at(axis, input_idx, root=root)
 
+    @override
+    def fuse_consumer_at(self, axis: str, root: str = DEFAULT_ROOT) -> None:
+        self._current_scheduler.fuse_consumer_at(axis, root=root)
+
     @override
     def define_memory_mesh(self, axes: dict[str, int]) -> None:
         self._require_extension("sdist")

src/xtc/backends/tvm/TVMScheduler.py

@@ -467,6 +467,11 @@ def fuse_producer_at(
         assert input_idx >= 0 and input_idx < len(self._op.np_inputs_spec())
         self.fused.append((axis, input_idx))
 
+    @override
+    def fuse_consumer_at(self, axis: str, root: str = DEFAULT_ROOT) -> None:
+        # TODO: not implemented for now
+        pass
+
     @override
     def define_memory_mesh(self, axes: dict[str, int]) -> None:
         # TODO: not implemented for now

src/xtc/itf/schd/scheduler.py

@@ -232,6 +232,19 @@ def fuse_producer_at(
         """
         ...
 
+    @abstractmethod
+    def fuse_consumer_at(self, axis: str, root: str = DEFAULT_ROOT) -> None:
+        """Fuse the consumer computation at the given producer location.
+
+        The consumer of output zero is fused at the given scheduled producer
+        axis. Other outputs are not currently supported.
+
+        Args:
+            axis: localisation of the fusion in the producer
+            root: the parent split (or the operator's absolute root)
+        """
+        ...
+
     @abstractmethod
     def define_memory_mesh(self, axes: dict[str, int]) -> None:
         """Define a memory mesh.