# [GT4Py/DaCe] adhoc: Fast path to Combined IR with DaCe <!-- Add the tag for the current cycle number in the top bar --> - Shaped by: Hannes - Appetite (FTEs, weeks): full cycle - Developers: Till, Enrique, Edoardo, Hannes ## Problem This is combining and parallelizing the work towards a combined IR ## Appetite <!-- Explain how much time we want to spend and how that constrains the solution --> ## Solution See the list of main tasks below. TODO link to the other documents ## Rabbit holes <!-- Details about the solution worth calling out to avoid problems --> ## No-gos <!-- Anything specifically excluded from the concept: functionality or use cases we intentionally aren’t covering to fit the ## appetite or make the problem tractable --> ## Progress <!-- Don't fill during shaping. This area is for collecting TODOs during building. As first task during building add a preliminary list of coarse-grained tasks for the project and refine them with finer-grained items when it makes sense as you work on them. --> - [ ] Adding new ITIR program structure and new builtin (Hannes) - [x] gtfn running on new structure - [x] pretty printer, parser - [ ] embedded execution (discuss short-cuts) - [ ] `as_fieldop` domain inference - [ ] New ITIR type system (Till) - [ ] Push fencil to program to top of `apply_common_transforms` (Till) - [ ] ... - [ ] Combined IR structure to DaCe (Edoardo) - [x] first test examples - [x] infrastructure - [x] as_fieldop translation - [x] select - [ ] shifts - [ ] reduce - [ ] concat_where - [ ] scan ## Scratchpad https://hackmd.io/_b9RLJX-QvmyGy21WVDngg ### Builtins #### Arithmetic No questions here really. ``` cast_ *ARITHMETIC_BUILTINS, *TYPEBUILTINS, ``` #### as_field_op ```python Stencil = Callable[[It, ...], Value] as_field_op(stencil: Stencil, domain: Domain)(*fields: Field | tuple[Field, ...]) ``` For now tuples of fields get translated into an iterator of tuples. We might want to change this. #### Shift ``` as_field_op(lambda it: deref(shift(V2E, 1)(it)), vertex_domain)(edge_field) ``` ``` premap(edge_field, V2E_1, vertex_domain) ``` No premap for now as we might change what arguments we pass, e.g. no tags, but connectivities. Dace backend pattern matches the stencil. ``` shift deref ``` TODO: Think about ICON case where in the interior we get a guarantee that we can access all neighbors without a `can_deref`. TODO: as_offset / dynamic offset ``` as_field_op(lambda it, offset_it: deref(shift(V2E, deref(offset_it)(it)), vertex_domain)(edge_field, offset_field) ``` #### Reduce-like ``` neighbors list_get map_ - we never combine list only fields so this should not be needed make_const_list reduce can_deref ``` ``` as_field_op(lambda it: neighbors(V2E, it), vertex_v2e_domain)(edge_field) ``` *Representations* ``` # representation 1 field_a: Field[[Vertex, V2E], dtype] = as_field_op(lambda it: neighbors(V2E, it), vertex_v2e_domain)(edge_field_a) # representation 2 - take this representation for now. it is incomplete # and does not match what we have in the frontend. this should be fixed by something along the lines of representation 3, but we want to design this properly considering all the cases. field_a: Field[[Vertex], List[dtype]] = as_field_op(lambda it: neighbors(V2E, it), vertex_domain)(edge_field_a) # representation 3 # conceptually neighbors return a LocalField[V2E, dtype] (itir.List extended with the dimension), which then gets collapsed after `as_field_op` into the mapped-over field field_a: Field[[Vertex, V2EDim], dtype]] = as_field_op(lambda it: neighbors(V2E, it), vertex_domain)(edge_field_a) ``` ``` premap(edge_field_a, V2E) -> Field[[Vertex], list[dtype]] ``` *make_const_list* ``` g = premap(1, V2E) -> Field[[Vertex], list[dtype]] g + 1 f: Field[Vertex, list[dtype]] = as_field_op(lambda: make_const_list(1), vertex_domain)() ``` <!--Argument to `reduce` as passed to dace with representation 1: ```python field_a = as_field_op(lambda it: neighbors(V2E, it), vertex_v2e_domain)(edge_field_a) field_b = as_field_op(lambda it: neighbors(V2E, it), vertex_v2e_domain)(edge_field_b) as_field_op(plus_stencil, vertex_v2e_field_a, vertex_v2e_field_b) ```--> TODO: continue discussion on reduce ``` field_a = as_field_op(lambda it: reduce(plus, 0)(deref(it)), vertex_domain)(vertex_list_field) neighbor_sum(vertex_v2e_field, axis=V2EDim) # frontend ``` reduction on `make_const_list` is not allowed. #### Conditionals / If ##### frontend where ``` itir.if_(cond: bool, true_branch: scalar | tuple, false_branch: scalar | tuple) # gtir where(cond_f, true_f false_f) # frontend as_field_op(lambda cond, true_, false_: if_(deref(cond), deref(true_), deref(false_)), domain)(cond_f, true_f, false_f) ``` ##### frontend ternary (or if-stmt translated to ternary) ``` cond(cond: bool, true_branch: Field | tuple, false_branch: Field | tuple) # gtir true_f if cond_ else false_f # frontend as_field_op(lambda cond, true_, false_: if(deref(cond), deref(true_), deref(false_)), domain)(cond_f, true_f, false_f) # we don't do this ``` TODO: Explain why this is needed. Explain transformation into IfStmt or local if above. Can not use local approach as above since tuples defined on different domains are possible. ``` as_field_op(...)(if_(is_lam, field_a, field_b)) ``` Introduce `gtir.cond(cond: Scalar, true_: Field, false_: Field)` #### Boundary conditions ```python! concat_where(domain: CartesianDomain | UnstructuredDomain, true_branch: Field, false_branch: Field) # frontend # for sure we will support this concat_where(IDim < 1, field_a, field_b) # IDim < 1 is a domain # maybe this concat_where(IDim < 1 & IDim > 10, field_a, field_b) # IDim < 1 & IDim > 10 is a domain # we will look at this to see what is ahead of us, but not implement it # the result must be rectangular (for now) concat_where(IDim < 10 & KDim > 20, field_a, field_b) ``` TODO: write down the conditions on domain ``` concat(field_0, ..., field_n) ``` #### Tuples ``` tuple_get(i, tuple_) make_tuple(el_0, el_n) ``` #### Scan ``` scan(scan_pass: Callable, init: tuple | scalar, direction: bool) ``` ## 23.04 minutes planning for next cycle Options: - Edoardo: dace lowering Till: PAST to GTIR, domain inference, GTIR passes (TODO: find out what dependencies are between the projects) Philip: Get infrastructure for JaCe then move to other things towards GT4Py 1 Project: Till, Edoardo, Sara