transforms: add restrict flag to StencilShapeMinimize pass #3411

emmau678 · 2024-11-08T11:03:48Z

Add a restrict flag to the StencilShapeMinimize pass which restricts the grid size of a stencil computation.

superlopuh · 2024-11-08T11:04:20Z

Could you please complete the checklist in the description?

xdsl/transforms/stencil_shape_minimize.py

Co-authored-by: Nicolai Stawinoga <36768051+n-io@users.noreply.github.com>

codecov · 2024-11-08T14:01:04Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 90.17%. Comparing base (a731a28) to head (3bc2d9c).
Report is 19 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3411      +/-   ##
==========================================
+ Coverage   90.15%   90.17%   +0.01%     
==========================================
  Files         455      456       +1     
  Lines       57429    57562     +133     
  Branches     5530     5543      +13     
==========================================
+ Hits        51775    51905     +130     
- Misses       4195     4197       +2     
- Partials     1459     1460       +1

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

n-io

Great, while there are ways to combine this into one filecheck mlir file (if you grep for --check-prefix), you could also make a new one. Happy to help with either way.

xdsl/transforms/stencil_shape_minimize.py

Co-authored-by: Nicolai Stawinoga <36768051+n-io@users.noreply.github.com>

emmau678 · 2024-11-08T15:12:45Z

Great, while there are ways to combine this into one filecheck mlir file (if you grep for --check-prefix), you could also make a new one. Happy to help with either way.

any help on how to get started on this would be greatly appreciated!

n-io · 2024-11-08T16:17:29Z

Great, while there are ways to combine this into one filecheck mlir file (if you grep for --check-prefix), you could also make a new one. Happy to help with either way.

any help on how to get started on this would be greatly appreciated!

I just pushed a filecheck to your branch, but it will likely break. The last thing to do is to run shape inference (ShapeInferencePass().apply(ctx, op)) inside the if self.restrict conditional to tie everything together. Wanna do the honours? :)

superlopuh · 2024-11-08T18:11:36Z

xdsl/transforms/stencil_shape_minimize.py

@@ -91,6 +92,7 @@ def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
                    ]
                )
            ).rewrite_module(op)
+        ShapeInferencePass().apply(ctx, op)


I would recommend against this

@n-io, what's the motivation for doing this inside the pass and not as part of the pass pipeline?

The most immediate purpose is to undo the InvalidateTemps() rewrite pattern. The pass in general is super useful for restricting the stencil grid size to something small enough that for instance it can be run in a simulator. The idea is simple: invalidate temps, modify size on store ops, un-invalidate temps. The result will then be picked up by stencil-shape-minimize, which will handle the rest. The only alternative would be to separate this out from shape-minimize altogether, though running stencil-shape-minimize{restrict=30,30,30} seems to make a lot of sense in my head instead of having it as a fully independent pass, and I do think they have to be run together in order to produce a program that passes validation (which we want after every logical pass).

Ah OK, that makes sense. The main thing that triggered me was calling a pass inside of another pass, for the same reason, as the contract is that a pass leaves the IR in a valid state, if the IR is in a valid state to begin with. In that sense it would make sense to split out the body of the shape inference pass into a helper method, and to call into that from here.

I've seen this done in a few places, though was never sure why helper methods are preferred, including by post_walk_func?

My understanding is that it's preferable to expose the least restrictive API when possible. For example, there's no need to pass the context when performing shape inference. It's probably also to make things more flexible in general, like you may want to only perform shape inference on a subset of your module.

I'm not sure I follow - do I need to make some updates here?

Yes, please! Could pull out the body of the shape inference pass into a python function that takes any operation, and no MLContext, and call that helper function from here instead of the pass?

ok, thank you!

xdsl/transforms/stencil_shape_minimize.py

Co-authored-by: Nicolai Stawinoga <36768051+n-io@users.noreply.github.com>

xdsl/transforms/shape_inference.py

Co-authored-by: Sasha Lopoukhine <superlopuh@gmail.com>

superlopuh

It would probably be good to get @n-io's tick before merging, but LGTM!

n-io

LGTM, feel free to merge!

xdsl/transforms/stencil_shape_minimize.py

…ct#3411) Add a restrict flag to the StencilShapeMinimize pass which restricts the grid size of a stencil computation. --------- Co-authored-by: emmau678 <eu233@Emma-laptop> Co-authored-by: Nicolai Stawinoga <36768051+n-io@users.noreply.github.com> Co-authored-by: n-io <n-io@users.noreply.github.com> Co-authored-by: Sasha Lopoukhine <superlopuh@gmail.com>

transforms: add restrict flag to StencilShapeMinimize pass

88cf605

emmau678 self-assigned this Nov 8, 2024

emmau678 requested a review from n-io November 8, 2024 11:13

emmau678 added the transformations Changes or adds a transformatio label Nov 8, 2024

n-io reviewed Nov 8, 2024

View reviewed changes

emmau678 and others added 3 commits November 8, 2024 13:30

Update xdsl/transforms/stencil_shape_minimize.py

2126efc

Co-authored-by: Nicolai Stawinoga <36768051+n-io@users.noreply.github.com>

Update xdsl/transforms/stencil_shape_minimize.py

637ce69

Co-authored-by: Nicolai Stawinoga <36768051+n-io@users.noreply.github.com>

add changes from review

12a062b

Merge branch 'main' into add_restrict_flag

d633879

emmau678 marked this pull request as ready for review November 8, 2024 14:15

emmau678 requested a review from n-io November 8, 2024 14:16

n-io reviewed Nov 8, 2024

View reviewed changes

xdsl/transforms/stencil_shape_minimize.py Outdated Show resolved Hide resolved

xdsl/transforms/stencil_shape_minimize.py Outdated Show resolved Hide resolved

emmau678 and others added 2 commits November 8, 2024 14:48

Update xdsl/transforms/stencil_shape_minimize.py

c99f37c

Co-authored-by: Nicolai Stawinoga <36768051+n-io@users.noreply.github.com>

improve handling of rewriting bounds

b7ded38

update filecheck

d8106af

run shape inference

bbcccc3

superlopuh reviewed Nov 8, 2024

View reviewed changes

n-io reviewed Nov 8, 2024

View reviewed changes

xdsl/transforms/stencil_shape_minimize.py Outdated Show resolved Hide resolved

emmau678 and others added 2 commits November 11, 2024 09:07

Update xdsl/transforms/stencil_shape_minimize.py

37d2853

Co-authored-by: Nicolai Stawinoga <36768051+n-io@users.noreply.github.com>

add helper function for ShapeInferencePass

1a54263

superlopuh reviewed Nov 11, 2024

View reviewed changes

xdsl/transforms/shape_inference.py Outdated Show resolved Hide resolved

emmau678 and others added 3 commits November 11, 2024 22:09

Update xdsl/transforms/shape_inference.py

b604103

Co-authored-by: Sasha Lopoukhine <superlopuh@gmail.com>

minor fixes to helper function handling

b92632c

change name of helper function

3bc2d9c

superlopuh approved these changes Nov 11, 2024

View reviewed changes

n-io approved these changes Nov 12, 2024

View reviewed changes

xdsl/transforms/stencil_shape_minimize.py Show resolved Hide resolved

emmau678 merged commit e971815 into main Nov 12, 2024
15 checks passed

emmau678 deleted the add_restrict_flag branch November 12, 2024 14:07

n-io mentioned this pull request Nov 13, 2024

Steps needed to compile an end-to-end stencil CSL example #3200

Open

50 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

transforms: add restrict flag to StencilShapeMinimize pass #3411

transforms: add restrict flag to StencilShapeMinimize pass #3411

emmau678 commented Nov 8, 2024 •

edited

Loading

superlopuh commented Nov 8, 2024

codecov bot commented Nov 8, 2024 •

edited

Loading

n-io left a comment

emmau678 commented Nov 8, 2024

n-io commented Nov 8, 2024 •

edited

Loading

superlopuh Nov 8, 2024

superlopuh Nov 8, 2024

n-io Nov 8, 2024 •

edited

Loading

superlopuh Nov 8, 2024

n-io Nov 8, 2024

superlopuh Nov 8, 2024

emmau678 Nov 11, 2024

superlopuh Nov 11, 2024

emmau678 Nov 11, 2024

superlopuh left a comment

n-io left a comment

transforms: add restrict flag to StencilShapeMinimize pass #3411

transforms: add restrict flag to StencilShapeMinimize pass #3411

Conversation

emmau678 commented Nov 8, 2024 • edited Loading

superlopuh commented Nov 8, 2024

codecov bot commented Nov 8, 2024 • edited Loading

Codecov Report

n-io left a comment

Choose a reason for hiding this comment

emmau678 commented Nov 8, 2024

n-io commented Nov 8, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

n-io Nov 8, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

superlopuh left a comment

Choose a reason for hiding this comment

n-io left a comment

Choose a reason for hiding this comment

emmau678 commented Nov 8, 2024 •

edited

Loading

codecov bot commented Nov 8, 2024 •

edited

Loading

n-io commented Nov 8, 2024 •

edited

Loading

n-io Nov 8, 2024 •

edited

Loading