Skip to content
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.

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

add optional output stationary check #354

Merged
merged 4 commits into from
Jan 29, 2025
Merged

add optional output stationary check #354

merged 4 commits into from
Jan 29, 2025

Conversation

jorendumoulin
Copy link
Contributor

stacked on #350

This one adds the capability to the scheduler to add extra checks/ constraints to filter out unwanted schedules, and adds a check to only accept pure output stationary (= never send partial results back to memory) schedules.

This is a useful check for for example accelerators which do not have the capability to preload partial results. In our case the codegen further down the pipeline is not yet capable of non-output stationary schedules.

Currently, this constraint is embedded in accelerator templates in a hacky way. This way of handling things makes it more explicit and controllable.

Copy link
Contributor

@JosseVanDelm JosseVanDelm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is beautiful

"""
Checks whether a schedule, outside of the template, is fully output
stationary. This is determined by making sure all parallel dimensions
preced the reduction dimensions in the output operand (last operand).
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
preced the reduction dimensions in the output operand (last operand).
precede the reduction dimensions in the output operand (last operand).

@@ -7,6 +9,7 @@ def scheduler_backtrack(
template: Template,
schedule: Schedule,
inner_dims: int = 1,
extra_checks: Sequence[Callable[[Template, Schedule], bool]] = [],
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🥹

@jorendumoulin jorendumoulin force-pushed the joren/backtracking-scheduler branch from e5362ce to 8eb7eab Compare January 29, 2025 07:30
Base automatically changed from joren/backtracking-scheduler to main January 29, 2025 07:31
@jorendumoulin jorendumoulin merged commit 5c037a4 into main Jan 29, 2025
14 checks passed
@jorendumoulin jorendumoulin deleted the joren/os-check branch January 29, 2025 07:42
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants