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.73% 90.72% -0.01%
==========================================
Files 467 467
Lines 58965 58973 +8
Branches 5640 5642 +2
==========================================
+ Hits 53500 53504 +4
- Misses 4020 4022 +2
- Partials 1445 1447 +2 ☔ 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.