bug: (csl-lowering) Make multi-apply lowering work #3614
Conversation
n-io
added
bug
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #3614 +/- ##
=======================================
Coverage 90.47% 90.47%
=======================================
Files 472 472
Lines 59284 59289 +5
Branches 5638 5639 +1
=======================================
+ Hits 53639 53644 +5
Misses 4206 4206
Partials 1439 1439 ☔ View full report in Codecov by Sentry. |
| @@ -385,6 +388,9 @@ def match_and_rewrite(self, op: arith.ConstantOp, rewriter: PatternRewriter, /): | |||
| class InjectApplyOutsIntoLinalgOuts(RewritePattern): | |||
| @op_type_rewrite_pattern | |||
| def match_and_rewrite(self, op: csl_stencil.ApplyOp, rewriter: PatternRewriter, /): | |||
| if not op.dest: | |||
| return | |||
|
|
|||
There was a problem hiding this comment.
Adding support for zero output apply ops.
| @@ -94,6 +94,9 @@ def match_and_rewrite(self, op: csl_stencil.ApplyOp, rewriter: PatternRewriter, | |||
| # convert args | |||
| buf_args: list[SSAValue] = [] | |||
| to_memrefs: list[Operation] = [buf_iter_arg := to_memref_op(op.accumulator)] | |||
| op.accumulator.replace_by_if( | |||
| buf_iter_arg.memref, lambda use: use.operation != buf_iter_arg | |||
| ) | |||
There was a problem hiding this comment.
Allowing the bufferization.to_memref (buf_iter_arg) to be used by further apply ops, in case there is more than one.
| ) | ||
| assert isa(ac_op.result.type, AnyTensorType) | ||
| dest = tensor.InsertSliceOp.get( | ||
| source=ac_op.result, | ||
| dest=dest, | ||
| static_sizes=ac_op.result.type.get_shape(), | ||
| static_sizes=[1, *ac_op.result.type.get_shape()], |
There was a problem hiding this comment.
This needs 1 in the 1st dimension to insert 1 slice of z-values (ac_op.result.type.get_shape() are the shape of the z-values)
| dest = acc | ||
| for i, acc_offset in enumerate(offsets): | ||
| ac_op = csl_stencil.AccessOp( | ||
| dest, stencil.IndexAttr.get(*acc_offset), chunk_t | ||
| buf, stencil.IndexAttr.get(*acc_offset), chunk_t |
There was a problem hiding this comment.
Read from the buf, not from the dest (which is the accumulator).
| return TensorType( | ||
| use.operation.result.type.get_element_type(), | ||
| static_sizes, | ||
| static_sizes[dimdiff:], |
There was a problem hiding this comment.
Small fix for inserting a 1d slice into a 2d tensor.
That does actually make sense in most cases, thanks. |
This PR includes a few small fixes, described below.