pyop3 - #3318
Draft
connorjward wants to merge 1141 commits into
Draft
Conversation
Had to introduce nodal_axis_constraints and nodal_layout_axes so nodal axes work for mixed+Real.
* sample mlir generation and runnable
* demo to satisfy
* building plan and stepping through loopy.py to see inputs
* designing make_kernel function
* incorporating mlir -> loopy flag and building pymbolic to mlir pipeline
* MLIR generating stage - but MLIR typing is incorrect, need type inference
* refactoring pyop3/lower && integrating MLIR
- Goal: Create an interface to a code generation context (MLIR or Loopy)
- Status: Interface created, battling PETSc bug before cleaning more.
- Goal: Integrate MLIR for auto-generation
- Status: Was working but transitioning to pyop3->mlir pipeline as
opposed to pym->mlir. Refactoring process is ongoing.
* introducing dtypes for type-inference and testing case
* push before connorjward/pyop3 merge update attempt
requirements tracking pyop3 && typing hints
* Updating lowering paths for new context class
PyOP3 updates and path traversals reflected in the new lowering
structure
* folder for testing with mlir files
* bug fixed, delayed dtype integration
Still need to integrate dtype support for all op3 expressions.
Removed abstractmethod while testing.
* removing debugging python/text files
* cleaning up refactor, to add docstrings
* removing testing files from git
* removing offloading demo files from remote
* removing mlir class for merge
* adding docstrings
* removing dtype method for future pr
* removing debug flags
* removing unnecessary requirements-build update
* returning requirements.txt
* resolving simple PR comments
- rename compiler option `codegen` -> `backend`
- documenting `backend` compiler option
- functional variable renaming
* moving _compile to codegen
- moving dispatch _compile function from context.py -> codegen.py
- local import in _parse_loop to avoid circular import
- moving dispatch _collect_temporary_shapes from context.py ->
transform.py
* abstracting backend from pyop3 language
- codegen holds more pyop3 traversal.
- backend-specific functions introduced where necessary
notable flaw is bloated function call for add_petsc_mat.
* moving petsc_mat and exscan to codegen
- not implementing methods for MLIR, loopy-only.
* fix typo and move transform
as title
* restructuring lower_expr
- moving lower_expr to free function
- going to revisit this once I have a working MLIR lowering implementation
* Removing unnecessary imports
Co-authored-by: Connor Ward <c.ward20@imperial.ac.uk>
* cleaning and fixed bug with component.local size
---------
Co-authored-by: Connor Ward <c.ward20@imperial.ac.uk>
* adding dtype to pyop3.expr.Expression
done for later type inference with MLIR code generation
type resolution in this technique hinges on numpy.result_type
Key assumptions/points:
- Type resolution for binary and ternary operators use numpy.result_type
- This handles type promotion, and alike, as per rules here: https://numpy.org/devdocs/reference/arrays.promotion.html#arrays-promotion
- An assumption for numerals is made such that:
- integers are int32
- float are float64
- Another assumption for LoopIndexVar is that indexing variables
are int32.
* quick resolve using pyop3.dtypes and removing abstract
* adding generic np.number dispatch handler
* typo in handler
Memory leaks addressed but some caches still disabled
* relabel axis tree and indices on the way in and out * G-ADOPT test case passes with no memory leaking * Finally remove the grossness that was Map.name
* Cache exceptions as well as computed values * Stash the relabeled object
Key aim is to have efficient cache accesses for axis tree routines.
firedrake-check passes so back to connorjward/pyop3
Found a genuine bug in WeakKeyDictionary. Need to use a unique key for every cache entry. Need to check though that this works as expected and things are cleared up....
Turns out that stashing exceptions creates a tricky reference cycle. I will have to do something less magic next time I come to profile things.
Simple bugs from recent changes.
Got into a slight traversal mismatch. Fixed now.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Replace PyOP2 with pyop3.
Summary
Data structures for unstructured meshes are much more complicated than N-dimensional arrays. Firedrake therefore does an awful lot of very involved bookkeeping to make sure that things work.
The key contribution of pyop3 is a new abstraction for describing data layouts, called axis trees. Axis trees are able to describe mesh data layouts completely from a high-level specification. This has the dual set of benefits of: (a) being able to delete a lot of complex Firedrake code, and (b) we can now express a whole suite of algorithms that would otherwise be very difficult to write.
You can think of pyop3 as 'a JIT-compiled numpy for unstructured mesh data'.
For more information my thesis is useful reference.
See here for a list of the sorts of things that having pyop3 should enable Firedrake to now do.
In addition to enabling lots of new technologies, another very significant benefit to integrating pyop3 is the reduction in technical debt and code complexity. Many complex pieces of preexisting Firedrake code have been simplified and/or centralised into very few functions. For example, Real function spaces are now isomorphic to DG0 function spaces in all circumstances except in the specific 2 or 3 locations where the data layout is set up. This reduces special-casing throughout Firedrake, making everything more composable!
Major changes
PyOP2 is gone
Obviously. Anywhere that you explicitly use the PyOP2 API (i.e.
from pyop2 import op2) will have to be replaced with pyop3 (import pyop3 as op3). If you use PyOP2 indirectly (e.g.function.dat.data,mesh.cell_set,V.dof_dset) then some things will be expected to work (e.g.function.dat.data) but others have no equivalent pyop3 object and thus no longer exist.Extruded meshes are now DMPlexs
Extruded meshes are now fully fledged PETSc DMPlex objects, instead of hybrid objects where only Firedrake knew about the extrusion. This simplifies an awful lot of code and will enable things like submeshing extruded meshes.
Mixed functions now share a single underlying array, instead of one per field
Breaking API changes
TODO: prepare a pyop3 migration guide
mesh.interior_facetsandmesh.exterior_facetsare nowAxisTrees, not_FacetsSupport for variable layer extruded meshes is dropped
The functionality appeared not to be used by the community (e.g. thetisproject/thetis#405) and porting the functionality would have been a large amount of work.
Ensemble API for nonblocking communication returns a single request instead of a list of requests
This is because we no longer need to emit a request per subfunction.
API deprecations
.sub()expects a tuple instead of integer for tensor spaces.For example, to get the
(1, 1)component of a 2x2 tensor function space one should now passf.sub((1, 1))instead off.sub(3). The latter will continue to work for now but a warning is raised.