transformations: (convert-stencil-to-csl-stencil) Make pass runnable before tensorization #3220
Conversation
…before tensorization
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #3220 +/- ##
==========================================
- Coverage 90.03% 89.98% -0.06%
==========================================
Files 433 433
Lines 54568 54669 +101
Branches 8465 8499 +34
==========================================
+ Hits 49133 49196 +63
- Misses 4061 4099 +38
Partials 1374 1374 ☔ View full report in Codecov by Sentry. |
| if isa(self.swaps, ArrayAttr[dmp.ExchangeDeclarationAttr]) | ||
| else type(self.accumulator.type)( | ||
| self.accumulator.type.get_element_type(), | ||
| ( | ||
| len(self.swaps), | ||
| self.accumulator.type.get_shape()[0] | ||
| // self.num_chunks.value.data, | ||
| ), | ||
| ), |
There was a problem hiding this comment.
Maybe this could be constructed outside of the list? It's quite hard to follow where one element of the list begins and another ends.
| for use in uses: | ||
| if not isinstance(use.operation, stencil.ApplyOp): | ||
| continue | ||
| apply_op = use.operation |
There was a problem hiding this comment.
Will there always be only one apply in this situation? If so we could assert apply_op is None (or something similar).
Or alternatively:
for use in uses:
if isinstance(use.operation, stencil.ApplyOp):
apply_op = use.operation
break
else:
raise RuntimeError("Expected ApplyOp")| # output_size = temp.bounds | ||
| # if isinstance(output_size, StencilBoundsAttr): | ||
| # return | ||
| # | ||
| # update_result_size(op.res, temp.bounds | op.field.type.bounds, rewriter) |
| for use in uses: | ||
| if not isinstance(use.operation, stencil.ApplyOp): | ||
| continue | ||
| apply_op = use.operation |
There was a problem hiding this comment.
Is this not guaranteed to operate on the apply_op declared above?
| assert all( | ||
| swap.size[2] == uniform_size for swap in op.swaps | ||
| ), "all swaps need to be of uniform size" | ||
| # uses have to be retrieved *before* the loop because of the rewriting happening inside the loop |
There was a problem hiding this comment.
This comment is somewhat misleading, since there is no rewriting happening in this loop
| @@ -41,6 +42,13 @@ def infer_core_size(op: LoadOp) -> tuple[IndexAttr, IndexAttr]: | |||
| from an LoadOp by walking the def-use chain down to the `apply` | |||
| """ | |||
| applies: list[ApplyOp] = list(all_matching_uses([op.res], ApplyOp)) | |||
| if len(applies) == 0: | |||
| dmp_swaps: list[dmp.SwapOp] = list(all_matching_uses([op.res], dmp.SwapOp)) | |||
There was a problem hiding this comment.
nit: do we need to materialise this dmp_swaps?
Making
convert-stencil-to-csl-stencilrunnable before tensorisation should address a variety of issues that we have with the tensorisation pass, and provide us with the opportunity to store important bits and blobs in our csl_stencil dialect.Currently, the pass is runnable both before and after tensorisation, which results in some degree of code duplication.