Skip to content

[cute] Omit redundant persistent shared loop #3293

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
yushangdi wants to merge 1 commit into from
+164 −123
Closed

[cute] Omit redundant persistent shared loop #3293

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 43 additions & 21 deletions helion/_compiler/program_id.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,8 @@ def _stmt_name_uses(stmt: ast.AST) -> tuple[set[str], set[str]]:
writes.add(node.id)
else:
reads.add(node.id)
if isinstance(node, ast.AugAssign) and isinstance(node.target, ast.Name):
reads.add(node.target.id)
return reads, writes


Expand Down Expand Up @@ -2118,7 +2120,8 @@ def _setup_tcgen05_persistent_kernel(
full_role_local_body
and not is_multi_root
and (
layout.cluster_m > 1
use_validated_cluster_m1_role_local_body
or layout.cluster_m > 1
or self._tcgen05_has_scheduler_warp()
or self._tcgen05_uses_grouped_static_persistent()
)
Expand Down Expand Up @@ -2150,9 +2153,10 @@ def _setup_tcgen05_persistent_kernel(
)

setup: list[ast.stmt] = []
# Fully role-local CtaGroup.TWO does not consume the shared work-tile
# SMEM handoff. Validated CtaGroup.TWO skips the shared scheduler;
# each role owns a scheduler loop over the capped persistent grid.
# Fully role-local codegen does not consume the shared work-tile SMEM
# handoff. Each role owns a scheduler loop over the capped persistent
# grid, so validated cluster_m=1 and CtaGroup.TWO skip the shared
# scheduler and residual loop.
if not omit_shared_loop:
setup.extend(self._build_tcgen05_persistent_prelude(layout))
elif self._tcgen05_has_scheduler_warp():
Expand All @@ -2173,6 +2177,7 @@ def _setup_tcgen05_persistent_kernel(
partition,
build_shared_tile_body=False,
epi_role_prelude_stmts=epi_role_prelude_stmts,
post_loop_stmts=post_loop_stmts,
)
)
else:
Expand All @@ -2186,11 +2191,10 @@ def _setup_tcgen05_persistent_kernel(
)
setup.extend(role_local_whiles)
if not omit_shared_loop:
# Validated cluster_m=1 and guarded partial/multi-root
# role-local shapes still rejoin the shared loop so existing
# CTA-wide barriers remain valid. Fully role-local CtaGroup.TWO
# codegen skips this residual loop; its work is already owned
# by role-local schedulers and cross-role pipelines.
# Partial and multi-root role-local shapes still rejoin the
# shared loop. Validated fully role-local codegen skips this
# residual loop; its work is already owned by role-local
# schedulers and cross-role pipelines.
setup.append(
create(
ast.While,
Expand Down Expand Up @@ -2751,9 +2755,9 @@ def _build_role_local_while(
acc pipelines (the existing pipeline barriers carry the data
dependency); no ``cute.arch.sync_threads()`` is emitted inside
the role-local loop. The caller decides whether to append a residual
shared loop after these role-local loops; validated cluster_m=1 keeps
it for existing CTA-wide barriers, while guarded fully role-local
CtaGroup.TWO omits it.
shared loop after these role-local loops. It is omitted when the
residual shared body contains only cloned dependency setup and legacy
barriers that no longer protect shared work.

The returned statement is the role-local ``while`` itself,
wrapped in ``if {role_predicate}:`` so only the matching warps
Expand Down Expand Up @@ -5499,8 +5503,8 @@ def _tcgen05_grouped_stmt_safe_to_omit(
def _tcgen05_shared_stmt_safe_to_omit(cls, stmt: ast.stmt) -> bool:
"""Return whether a removed shared stmt is dependency-only setup.

Fully role-local CtaGroup.TWO codegen intentionally omits the residual
shared ``while``. The remaining shared view may still contain scalar
Fully role-local codegen intentionally omits the residual shared
``while``. The remaining shared view may still contain scalar
PID/offset/view setup that role-local loops clone through dependency
extraction, plus legacy bare ``sync_threads`` barriers that no longer
bracket shared work after every role has moved out. Other observable
Expand Down Expand Up @@ -5543,7 +5547,9 @@ def _tcgen05_shared_stmt_safe_to_omit(cls, stmt: ast.stmt) -> bool:
return isinstance(stmt, ast.Pass)

def _assert_tcgen05_omit_shared_loop_safe(
self, partition: Tcgen05PersistentProgramIDs._PartitionedRoleBody
self,
partition: Tcgen05PersistentProgramIDs._PartitionedRoleBody,
post_loop_stmts: list[ast.stmt] | None = None,
) -> None:
unsafe = [
ast.unparse(stmt)
Expand All @@ -5554,6 +5560,19 @@ def _assert_tcgen05_omit_shared_loop_safe(
"tcgen05 fully role-local codegen would discard observable shared "
"statement(s) while omitting the residual shared loop: " + "; ".join(unsafe)
)
shared_writes: set[str] = set()
for stmt in partition.shared_body_extracted:
_, writes = _stmt_name_uses(stmt)
shared_writes.update(writes)
post_loop_reads: set[str] = set()
for stmt in post_loop_stmts or ():
reads, _ = _stmt_name_uses(stmt)
post_loop_reads.update(reads)
dependencies = shared_writes & post_loop_reads
assert not dependencies, (
"tcgen05 fully role-local codegen would discard shared definition(s) "
"used by post-loop cleanup: " + ", ".join(sorted(dependencies))
)

def _assert_tcgen05_grouped_omit_shared_loop_safe(
self, partition: Tcgen05PersistentProgramIDs._PartitionedRoleBody
Expand Down Expand Up @@ -5609,6 +5628,7 @@ def _build_tcgen05_persistent_tile_body_role_local(
*,
build_shared_tile_body: bool = True,
epi_role_prelude_stmts: list[ast.stmt] | None = None,
post_loop_stmts: list[ast.stmt] | None = None,
) -> tuple[list[ast.stmt], list[ast.stmt]]:
"""Build the per-tile body in role-local-while form.

Expand All @@ -5628,11 +5648,9 @@ def _build_tcgen05_persistent_tile_body_role_local(
- ``shared_tile_body`` is the optional per-tile body for the shared
``while`` (the work-tile body without the extracted role blocks).
Built via :meth:`_build_tcgen05_persistent_tile_body` with existing
``cute.arch.sync_threads()`` calls preserved. Validated cluster_m=1
role-local kernels still append this loop after role-local work so
those CTA-wide barriers remain valid for epilogue synchronization
and work-tile metadata publication. Guarded fully role-local
CtaGroup.TWO codegen omits the residual shared loop in the caller.
``cute.arch.sync_threads()`` calls preserved. The caller omits this
loop only when the residual statements are dependency-only setup or
legacy barriers that no longer protect shared work.

Caller wires both into the persistent kernel as siblings of
each other inside the same setup list when the residual shared loop
Expand Down Expand Up @@ -5663,10 +5681,14 @@ def _build_tcgen05_persistent_tile_body_role_local(
layout, shared_role_blocks
)
else:
assert post_loop_stmts is not None, (
"omitting the tcgen05 shared loop requires explicit post-loop "
"dependency validation"
)
if self._tcgen05_uses_grouped_static_persistent():
self._assert_tcgen05_grouped_omit_shared_loop_safe(partition)
else:
self._assert_tcgen05_omit_shared_loop_safe(partition)
self._assert_tcgen05_omit_shared_loop_safe(partition, post_loop_stmts)
shared_tile_body = []
# Merge extracted blocks by ``role_predicate`` so each predicate
# gets one role-local loop carrying all of its per-tile
Expand Down
Loading
Loading