Skip to content

Commit

Permalink
Tracer: don't fuse reduce or scan workunits with fors as it is causin…
Browse files Browse the repository at this point in the history
…g slowdowns in ExaMiniMD (#286)
  • Loading branch information
NaderAlAwar authored Aug 18, 2024
1 parent e187e16 commit 5d645be
Showing 1 changed file with 14 additions and 0 deletions.
14 changes: 14 additions & 0 deletions pykokkos/core/fusion/trace.py
Original file line number Diff line number Diff line change
Expand Up @@ -335,6 +335,20 @@ def fuse_naive(self, operations: List[TracerOperation]) -> List[TracerOperation]
fused_ops.append(op)
continue

if op.operation in {"reduce", "scan"}:
if len(ops_to_fuse) > 0:
ops_to_fuse.reverse()
fused_ops.append(self.fuse_operations(ops_to_fuse, fused_safety_info))
ops_to_fuse.clear()
ops_to_fuse_views.clear()
fused_safety_info = {}
fused_range = None

# Don't fuse reduce or scan now because that might
# cause slowdowns.
fused_ops.append(op)
continue

current_range: Tuple[int, int] = (op.policy.begin, op.policy.end)
if fused_range is None:
fused_range = current_range
Expand Down

0 comments on commit 5d645be

Please sign in to comment.