inducer / inducer/pytato

[RFC, generate_loopy] Propagating tags to `lp.Instruction`s

Open
#361 1 comment 0 reactions 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.