[RFC, generate_loopy] Propagating tags to `lp.Instruction`s
- Dominant language
- Python
- Stars
- 15
- Forks
- 16
- Avg merge
- 17m
- Merged PRs (30d)
- 1
Description
In the current design, the tags attached to a materialized array are propagated to the corresponding variable's tags in the loopy kernel. However, no API is present to propagate any tags to the loopy instructions.
Here's what @matthiasdiener and I came up with: Make `generate_loopy` take in another argument that tells us which `Array.tags` should make into `lp.Instruction.tags`. The array expressions that are evaluated/read by a statement contribute toward this propagation. This way we could even capture any tags that were attached to unmaterialized arrays.
```python
class OperationDefiningTag(Tag): pass
class FluxComputingStage(OperationDefiningTag): pass
class QuadratureStage(OperationDefiningTag): pass
class ScalarOpOnBoundary(OperationDefiningTag): pass
def flux(x, y):
a1 = x[y].tagged(FluxComputingStage())
a2 = pt.sin(a1).tagged(ScalarOpOnBoundary())
return 2*a2
def quadrature(D, u):
a1 = pt.einsum("ij,ej->ei", D, u)
a2 = 2 * a1.tagged(QuadratureStage())
return pt.sin(a2)
def face_mass(..) : ...
def main():
u = ... (DOF values)
map_ = ... (vol. dof to face dof mapping)
D = ... (derivative matrix)
flux_values = flux(u, map_).tagged(ImplStored())
integral = quadrature(D, u).tagged(ImplStored())
result = face_mass(flux_value) + integral
t_unit = pt.generate_loopy(result, insn_tag_t_to_propagate=OperationDefiningTag)
```
For the above expression, the generated loopy kernel will have 3 statements: one for `flux_values`, one for `integral` and one for `result`. The statement writing `flux_values` would be tagged with `FluxComputingStage` and `ScalarOpOnBoundary`, the statement writing `integral` would be tagged with `QuadratureStage`, and the instruction writing `result` will have zero tags.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.