Add fast operator evaluation for tensor-product discretizations#362
Draft
a-alveyblanc wants to merge 65 commits intoinducer:mainfrom
Draft
Add fast operator evaluation for tensor-product discretizations#362a-alveyblanc wants to merge 65 commits intoinducer:mainfrom
a-alveyblanc wants to merge 65 commits intoinducer:mainfrom
Conversation
…product-operators
Owner
|
Should this be marked draft given that there are pending TODOs? |
…product-operators
inducer
reviewed
Nov 25, 2024
| # {{{ BilinearForm base class | ||
|
|
||
| @dataclass | ||
| class _BilinearForm: |
| new_args = [] | ||
| new_access_descrs = [] | ||
| for iarg, arg in enumerate(expr.args): | ||
| # we assume that the 0th argument will be the one with this tag |
Owner
There was a problem hiding this comment.
- So check that?
- Generally, beef up the matching. (maybe come up with some tools,
matchstatement)
| included in the original einsum), and properly reshapes to and from | ||
| tensor-product form to apply the 1D mass operator. | ||
| """ | ||
| def map_einsum(self, expr, *args, **kwargs): |
| return expr.copy(args=tuple(new_args)) | ||
|
|
||
|
|
||
| class InverseMassDistributor(CopyMapperWithExtraArgs): |
Owner
There was a problem hiding this comment.
- Split into two mappers (with the inner only handling what's allowed as part of distribution)
- Bonus points for generalizing to simplex
| return expr.copy(args=tuple(new_args)) | ||
|
|
||
|
|
||
| class RedundantMassTimesMassInverseRemover(CopyMapperWithExtraArgs): |
| return expr.copy(args=tuple(new_args)) | ||
|
|
||
|
|
||
| class RedundantMassTimesMassInverseRemover(CopyMapperWithExtraArgs): |
inducer
reviewed
Dec 16, 2024
| in_element_group: InterpolatoryElementGroupBase): | ||
| in_element_group: InterpolatoryElementGroupBase) -> ArrayOrContainer: | ||
|
|
||
| @keyed_memoize_in( |
Owner
There was a problem hiding this comment.
Let's have @keyed_memoize_in_first_arg.
| else: | ||
| quadrature_rule = in_grp.quadrature_rule() | ||
|
|
||
| if isinstance(in_grp, SimplexElementGroupBase): |
Owner
There was a problem hiding this comment.
Maybe the TP and full versions want to be separate functions?
|
|
||
| def _strong_scalar_grad(dcoll, dd_in, vec): | ||
| assert isinstance(dd_in.domain_tag, VolumeDomainTag) | ||
| def _reference_stiffness_transpose_matrices( |
Owner
There was a problem hiding this comment.
Unify with above via basis_getter?
|
|
||
| discr = dcoll.discr_from_dd(dd_in) | ||
| actx = vec.array_context | ||
| if in_grp == out_grp: |
Owner
There was a problem hiding this comment.
Suggested change
| if in_grp == out_grp: | |
| if isinstance(in_grp, Quadrature) |
?
Suggested change
| if in_grp == out_grp: | |
| if not isinstance(in_grp, Interpolatory) |
?
| out_grp.shape | ||
| ) | ||
| else: | ||
| quadrature_rule = in_grp.quadrature_rule() |
Owner
There was a problem hiding this comment.
Check that quadrature_rule provides sufficient exactness.
| return get_reference_stiffness_transpose_matrices(out_grp, in_grp) | ||
|
|
||
|
|
||
| def reference_mass_matrix(actx: ArrayContext, |
| """ | ||
|
|
||
| per_group_grads = [] | ||
| for out_grp, in_grp, vec_i, ijm_i in zip( |
inducer
reviewed
Jan 8, 2025
| use_tensor_product_fast_eval=use_tensor_product_fast_eval | ||
| ) | ||
|
|
||
| if input_group == output_group: |
Owner
There was a problem hiding this comment.
if isinstance(input_group, Interpolatory):
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
2 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.
Supersedes #313, #354
Adds:
Non-essential transformations will be added in a later PR. This is to keep this (already large) PR manageable.The plan is to break this PR up into a sequence of smaller PRs to ease the review process. Keeping this up until that starts happening.TODOs:
op.pyto use generic bilinear form evaluationcc @MTCam