refactor[next]: embedded with itir.Program #1530
Conversation
Co-authored-by: Till Ehrengruber <t.ehrengruber@me.com>
Co-authored-by: Till Ehrengruber <t.ehrengruber@me.com>
src/gt4py/next/iterator/ir.py
Outdated
There was a problem hiding this comment.
I have a question, not related to this PR but in general to the new IR and specifically to this type:
class SetAt(Stmt): # from JAX array.at[...].set()
expr: Expr # only `as_fieldop(stencil)(inp0, ...)` in first refactoring
domain: Expr
target: Expr # `make_tuple` or SymRef
Do we really need to support make_tuple as target expression?
There was a problem hiding this comment.
it's the current representation for something like
@fundef
def prog(a,b,c)
setat(as_field_op(lambda x: make_tuple(deref(x)+1, deref(x)+2))(c), domain, make_tuple(a,b))
There was a problem hiding this comment.
Should the fieldview backend support that representation as is, I mean with inlined make_tuple? or just this one?
@fundef
def prog(a,b,c)
setat(as_field_op(lambda x, y: make_tuple(x, y))(as_field_op(lambda x: x+1)(c), as_field_op(lambda x: x+2)(c), ), domain, make_tuple(a,b))
There was a problem hiding this comment.
I am asking because from a lowering perspective, make_tuple and tuple_get do not implement any kind of computation on fields. Therefore, it is difficult to represent them in my map-tasklet graph. Would it be too strange to treat these builtins on the same level as as_field_op?
There was a problem hiding this comment.
@fundef
def prog(a,b,c)
setat(make_tuple(as_field_op(lambda x: x+1)(c), as_field_op(lambda x: x+2)(c)), domain, make_tuple(a,b))
| ) -> Callable[[Callable[_P, _R]], Callable[..., _R | tuple[_R | tuple, ...]]]: ... | ||
|
|
||
|
|
||
| def tree_map( |
There was a problem hiding this comment.
I think it is time to remove apply_to_primitive_constituents and absorb it in this function. I'll check if @SF-N has capacity to work on this.
There was a problem hiding this comment.
Sara will extract this into a new PR and also remove the apply_to_primitive_constituents cases there. I'll keep you in the loop.
tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py
Outdated
Show resolved
Hide resolved
|
cscs-ci run |
|
|
||
| return impl | ||
|
|
||
|
|
||
| @runtime.closure.register(EMBEDDED) | ||
| def closure( |
There was a problem hiding this comment.
The type annotation of ins appears to be wrong as there is the promotion of scalars below.
src/gt4py/next/iterator/embedded.py
Outdated
| for inp in promoted_ins | ||
| ) | ||
| res = sten(*ins_iters) | ||
| res = _compute_point(sten, ins, pos, column_range) |
There was a problem hiding this comment.
| res = _compute_point(sten, ins, pos, column_range) | |
| res = _compute_at_point(sten, ins, pos, column_range) |
Just a suggestion.
There was a problem hiding this comment.
decided to go with _compute_at_position
|
cscs-ci run |
Updates itir.embedded to work with
itir.Programs, i.e.set_atandas_fieldop.For programs to be able to run in embedded, the domain needs to be provided as second argument to
as_fieldop.Introduces a
DimensionKindtoitir.AxisLiteralto be able to reconstruct the kind from the IR. This is needed now as theset_atassigns from field to field, which requires matching dimensions. However, previously the python program generated from IR would always construct horizontal dimensions (but the information would not be used).