transformations: (csl-stencil) Add pass to handle async ops and enclosing cf #3192

n-io · 2024-09-18T16:20:35Z

Adds a pass to handle the async control flow of csl_stencil.apply and any enclosing loops by translating control flow into a csl.func call graph.

codecov · 2024-09-18T16:25:50Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 90.03%. Comparing base (25506b3) to head (3b011b2).
Report is 4 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3192      +/-   ##
==========================================
+ Coverage   90.01%   90.03%   +0.02%     
==========================================
  Files         431      432       +1     
  Lines       54282    54435     +153     
  Branches     8410     8443      +33     
==========================================
+ Hits        48863    49012     +149     
+ Misses       4061     4056       -5     
- Partials     1358     1367       +9

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

AntonLydike

Nice progress! I have some comments here and there

xdsl/transforms/csl_stencil_handle_async_flow.py

AntonLydike · 2024-09-19T08:18:45Z

xdsl/transforms/csl_stencil_handle_async_flow.py

+        )
+
+        no_params = FunctionType.from_lists([], [])
+        cond_task_id = self.counter + 1


Are you sure that this is a robust enough solution? I remember us thinking about having task-id allocation and referring to them by symbols in the meantime. What happened to that plan? (jsut curious, not necessarily criticism)

Yes and no, with our current comms library we have ids 1-3 free for local task ids. We could prob put an assert not to exceed that, but I think it'll simply not compile at a later stage. The question is how much knowledge this pass should have of things downstream? That said, it is the first thing in the pipeline to handle tasks and task ids, so it has a bit of freedom.

xdsl/transforms/csl_stencil_handle_async_flow.py

tests/filecheck/transforms/csl-stencil-handle-async-flow.mlir

AntonLydike · 2024-09-19T08:27:05Z

xdsl/transforms/csl_stencil_handle_async_flow.py

+        terminator = op.body.block.last_op
+        assert isinstance(terminator, scf.Yield)
+        assert all(
+            arg in op.body.block.args for arg in terminator.arguments
+        ), "Can only yield unmodified iter_args (in any order)"
+
+        # limitation: currently only loops built from arith.constant are supported
+        assert isinstance(op.lb, OpResult)
+        assert isinstance(op.ub, OpResult)
+        assert isinstance(op.step, OpResult)
+        assert isinstance(op.lb.op, arith.Constant)
+        assert isinstance(op.ub.op, arith.Constant)
+        assert isinstance(op.step.op, arith.Constant)
+        assert isa(op.lb.op.value, IntegerAttr[IndexType])
+        assert isa(op.ub.op.value, IntegerAttr[IndexType])
+        assert isa(op.step.op.value, IntegerAttr[IndexType])
+
+        # limitation: all iter_args must be memrefs (stencil buffers) and have the same data type
+        assert isa(op.iter_args[0].type, MemRefType[csl.ZerosOp.T])
+        element_type = op.iter_args[0].type.get_element_type()
+        assert all(
+            isa(a.type, MemRefType[csl.ZerosOp.T])
+            and element_type == a.type.get_element_type()
+            for a in op.iter_args
+        )


Nit: I'm generally not a fan of asserts in rewrite patterns, but I can see the appeal here. When the pass just aborts you may end up with silent failures, which are a pain to debug... I guess this shows that xDSL is lacking some sort of warning system for not-applied passes...

I think in MLIR there are separate match and apply functions, which separates out code that checks if a pattern should be applied. We also need quite a few of these asserts to make pyright happy, and the new make precommit forces us to split assert isinstance(op.lb, OpResult) and isinstance(op.lb.op, arith.Constant) into separate asserts.

xdsl/transforms/csl_stencil_handle_async_flow.py

dk949

Just a couple of things from me

dk949 · 2024-09-19T10:06:13Z

xdsl/dialects/csl/csl.py

@@ -400,7 +400,7 @@ def get_element_type(self):
        return self.res.type.get_element_type()

    @staticmethod
-    def from_type(child_type: TypeAttribute) -> VariableOp:


Should this not be a TypeAttribute?

Then pyright complains in this line: [csl.VariableOp.from_type(arg_t) for arg_t in op.iter_args.types]

This seems like the wrong place to fix this then..?

Although, @AntonLydike, is there a reason some types do not inherit TypeAttribute? In which case this should be an Attribute.

Well, good question. All types should inherit from TypeAttribute, as that will have them be printed as !dialect.name instead of #dialect.name (which is attribute specific and can't be used as a type according to MLIR).

xdsl/transforms/csl_stencil_handle_async_flow.py

n-io · 2024-09-19T15:31:07Z

I've added some functionality (and a test) to support non-loop, sequential stencil apply ops.

dk949

Looks good AFAICT

n-io added 6 commits September 5, 2024 15:04

update

445d73c

Merge branch 'main' into nicolai/csl-stencil-lower-enclosing-loops

10c195c

Merge branch 'main' into nicolai/csl-stencil-lower-enclosing-loops

92edadd

handle loops and single csl_stencil.apply

d7cb62b

Merge branch 'main' into nicolai/csl-stencil-lower-enclosing-loops

157cdef

filecheck

c686055

n-io added the transformations Changes or adds a transformatio label Sep 18, 2024

n-io requested review from AntonLydike and dk949 September 18, 2024 16:20

n-io self-assigned this Sep 18, 2024

AntonLydike reviewed Sep 19, 2024

View reviewed changes

Adding comments

7c8c99b

dk949 reviewed Sep 19, 2024

View reviewed changes

Support sequential applies

3b011b2

n-io requested review from AntonLydike and dk949 September 19, 2024 15:30

dk949 approved these changes Sep 20, 2024

View reviewed changes

dk949 mentioned this pull request Sep 20, 2024

transformations: (csl-stencil-to-csl-wrapper) Add unblock_cmd_stream call #3198

Merged

n-io merged commit 9dc299a into main Sep 20, 2024
14 checks passed

n-io deleted the nicolai/csl-stencil-lower-enclosing-loops branch September 20, 2024 10:48

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

transformations: (csl-stencil) Add pass to handle async ops and enclosing cf #3192

transformations: (csl-stencil) Add pass to handle async ops and enclosing cf #3192

n-io commented Sep 18, 2024

codecov bot commented Sep 18, 2024 •

edited

Loading

AntonLydike left a comment

AntonLydike Sep 19, 2024

n-io Sep 19, 2024

AntonLydike Sep 19, 2024

n-io Sep 19, 2024

dk949 left a comment

dk949 Sep 19, 2024

n-io Sep 19, 2024

dk949 Sep 19, 2024

AntonLydike Sep 19, 2024

n-io commented Sep 19, 2024

dk949 left a comment

transformations: (csl-stencil) Add pass to handle async ops and enclosing cf #3192

transformations: (csl-stencil) Add pass to handle async ops and enclosing cf #3192

Conversation

n-io commented Sep 18, 2024

codecov bot commented Sep 18, 2024 • edited Loading

Codecov Report

AntonLydike left a 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

dk949 left a 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

n-io commented Sep 19, 2024

dk949 left a comment

Choose a reason for hiding this comment

codecov bot commented Sep 18, 2024 •

edited

Loading