feat[next]: Inline dynamic shifts #1738
Conversation
…p' into gtir_fuse_as_fieldop
|
I run icon4py stencils and I found that one small change is needed in dace backend. It makes sense to include it in this PR: |
|
Feel free to push your Dace changes (as long as they are isolated to dace parts otherwise it'll be to confusing). |
…ets' into inline_dynamic_offsets
| @@ -67,6 +67,107 @@ def _is_tuple_expr_of_literals(expr: itir.Expr): | |||
| return isinstance(expr, itir.Literal) | |||
There was a problem hiding this comment.
All changes in this file are merely a small refactoring that extracts the field operator fusion for a single as_fieldop from the FuseAsFieldOp pass into a dedicated function fuse_as_fieldop used by the InlineDynamicShifts pass.
| @@ -10,7 +10,7 @@ | |||
|
|
|||
There was a problem hiding this comment.
@SF-N could you review the changes in this file?
| return expr, accessed_domains | ||
|
|
||
|
|
||
| def _infer_stmt( | ||
| stmt: itir.Stmt, | ||
| offset_provider: common.OffsetProvider, | ||
| symbolic_domain_sizes: Optional[dict[str, str]], | ||
| allowed_unknown_domains: list[str], |
There was a problem hiding this comment.
In general, I am confused why this is introduced...
Are all program parameters allowed to be unknown, but the rest isn't?
There was a problem hiding this comment.
If I understand the logic, then I would remove it from here and rather write a separate path for checking this and avoid uglifying the this interface.
There was a problem hiding this comment.
I minimized cluttering the interface by introducing a typed dict InferenceOptions that handles forwarding of the arguments. I tried a separate path for checking, but this felt more error prone and ugly (let's discuss if you want more details why). The typed dict feels like a good compromise to me (also allows setting a breakpoint at the place where the error occurs), but I don't insist on it :-P.
| @@ -157,5 +161,8 @@ def apply_fieldview_transforms( | |||
| ir = inline_fundefs.prune_unreferenced_fundefs(ir) | |||
| ir = InlineLambdas.apply(ir, opcount_preserving=True, force_inline_lambda_args=True) | |||
| ir = CollapseTuple.apply(ir, offset_provider=offset_provider) # type: ignore[assignment] # type is still `itir.Program` | |||
| ir = inline_dynamic_shifts.InlineDynamicShifts.apply( | |||
| ir | |||
| ) # domain inference does not support dynamic offsets yet | |||
| ir = infer_domain.infer_program(ir, offset_provider=offset_provider) | |||
There was a problem hiding this comment.
Would it make sense to support unknown domain in the type inference for as_offset?
There was a problem hiding this comment.
Let's discuss quickly, I'm not sure I understand the question in its entirety.
Dynamic shifts are not supported in the domain inference. In order to make them work nonetheless this PR aggressively inlines all arguments to
as_fieldopuntil they contain only references toitir.Programparams. Additionally the domain inference is extended to tolerate suchas_fieldopby introducing a special domain marker that signifies a domain is unknown.