Skip to content

Flatten entire array containers#91

Merged
thomasgibson merged 23 commits into
inducer:mainfrom
alexfikl:flatten-to-numpy
Oct 19, 2021
Merged

Flatten entire array containers#91
thomasgibson merged 23 commits into
inducer:mainfrom
alexfikl:flatten-to-numpy

Conversation

@alexfikl

@alexfikl alexfikl commented Sep 26, 2021

Copy link
Copy Markdown
Collaborator

Adds two functions, flatten and unflatten, that go between an array container and a 1D array (as supported by the array context).

The couple of gotchas here:

  • serialize_container and deserialize_container should not be doing anything crazy, i.e. two calls will return the components in the same order, which is the case for all the containers we have.
  • As it doesn't impose its own order, it's not safe for long time storage (i.e. in case something changes order in serialiaze_container), but that shouldn't be too much of an issue.
  • Naming? The same functions in meshmode basically just flatten DOFArrays and leave the rest of the array container the same.

Before merging:

@thomasgibson Would something like this do what you need?

@alexfikl alexfikl changed the title Flatten entire array container Flatten entire array containers Sep 26, 2021
@alexfikl alexfikl force-pushed the flatten-to-numpy branch 2 times, most recently from e9450ed to e133b9a Compare September 26, 2021 00:48

@thomasgibson thomasgibson left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is awesome. Thanks for doing this! I tried them out in inducer/grudge#154 and they appear to work perfectly. I left some comments below, but I'm really keen to have this functionality!

Comment thread arraycontext/container/traversal.py Outdated
Comment on lines +30 to +31
.. autofunction:: flatten_to_numpy
.. autofunction:: unflatten_from_numpy

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we should try to distinguish {flatten, unflatten}_{to,from}_numpy from their named equivalents in meshmode.dof_array. Though in my opinion, I think that {flatten, unflatten}_{to,from}_numpy makes way more sense here because it's actually flattening the entire container (rather than only flattening components as is the case of meshmode.dof_array.{flatten, unflatten}_{to,from}_numpy.

However, I am fine with just renaming these functions too. How about rec_flatten_to_numpy and similar? rec_ because it recursively flattens?

@thomasgibson thomasgibson Sep 26, 2021

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

flatten_container_to_numpy and unflatten_numpy_to_container actually could work too.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Those sound good to me! I'll wait for some more feedback, but otherwise we can rename them like that.

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agree with @thomasgibson's point: We should clarify the relationship between these and the DOFArray flatteners. Given that these do just fine without resorting to the latter, why don't we just deprecate those and define flattening generically, here? Then we wouldn't be obliged to change the name either, since it's clear that these are simply a more generic version of the former. Passing a template seems nicer anyway than passing a discretization.

@alexfikl alexfikl Sep 27, 2021

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed. One issue is that they do slightly different things:

  • these flatten the whole container into one array
  • the ones in meshmode.dof_array traverse the container and only flatten the DOFArrays.

Not sure if anyone depends on that functionality though? So it could be safe to deprecate and point to these functions instead.

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for spotting that! I don't think anyone is relying on the "just the DOFArrays" aspect of those. (I'm not even sure in what scenario that would be useful.) Plus, if desired, it can be replicated with relative ease atop these new functions.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, that sounds good! I'll add some deprecation warnings in there after this is in.

Comment thread arraycontext/container/traversal.py Outdated

@inducer inducer left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for working on this! This actually contains some interesting questions about the abstraction as a whole.

Comment thread arraycontext/container/traversal.py Outdated
Comment on lines +30 to +31
.. autofunction:: flatten_to_numpy
.. autofunction:: unflatten_from_numpy

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agree with @thomasgibson's point: We should clarify the relationship between these and the DOFArray flatteners. Given that these do just fine without resorting to the latter, why don't we just deprecate those and define flattening generically, here? Then we wouldn't be obliged to change the name either, since it's clear that these are simply a more generic version of the former. Passing a template seems nicer anyway than passing a discretization.

Comment thread arraycontext/container/traversal.py Outdated
Comment thread arraycontext/container/traversal.py Outdated
Comment thread arraycontext/container/traversal.py
Comment thread arraycontext/container/traversal.py Outdated
Comment thread arraycontext/container/traversal.py Outdated
Comment thread arraycontext/container/traversal.py Outdated

@inducer inducer left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! I think this on a good path. A few more comments below.

Comment thread arraycontext/container/__init__.py Outdated
Comment thread arraycontext/container/traversal.py Outdated
Comment thread arraycontext/container/traversal.py Outdated
Comment thread arraycontext/container/traversal.py
Comment thread arraycontext/container/traversal.py Outdated
Comment thread arraycontext/container/traversal.py Outdated

@inducer inducer left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! This looks great now. Just two final wrinkles.

Comment thread arraycontext/container/traversal.py Outdated
Comment thread arraycontext/container/traversal.py Outdated

@inducer inducer left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! Spotted just a few more things. I think those are it though. :)

Comment thread arraycontext/container/traversal.py Outdated
iterable = serialize_container(template_subary)
except TypeError:
# NOTE: the max is needed to handle device scalars with size == 0
offset += max(1, template_subary.size)

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Upon thinking about it more, I don't think this is right. Device scalars report a size of 1:

>>> import numpy as np
>>> np.zeros(()).size
1
>>> a = np.zeros(())
>>> import pyopencl as cl
>>> ctx = cl.create_some_context(interactive=False)
>>> queue = cl.CommandQueue(ctx)
>>> import pyopencl.array
>>> a_dev = cl.array.to_device(queue, a)
>>> a_dev.size
1

The tests should cover device scalars and arrays that genuinely have zero size (e.g. shape=(3, 0, 4))

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, you're completely right! I removed that and added some tests in 27b894a.

A couple of wiggles there:

  • for 0-sized arrays, e.g. (18, 0) in the test, when reshaped the strides were (0, 8), but the template had (8, 8). Not sure who's fault that is?
  • pytato didn't like reshaping scalars or 0-sized arrays. I'll make a small example and complain there.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For the first part, it seems like it's something in pyopencl.

Hopefully inducer/compyte#36

Comment thread arraycontext/container/traversal.py Outdated
Comment thread arraycontext/container/traversal.py Outdated

@thomasgibson thomasgibson left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I looked through this PR after the latest changes and I think this looks even better. Can also confirm it works like a charm for inducer/grudge#154. I'll let @inducer have a chance to read over it but I wanted to leave my approval now. 😄

@inducer inducer left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, aside from these final wrinkles. If you're OK with accepting all these suggestions, then this is ready to go from my end.

Comment thread arraycontext/container/__init__.py Outdated
Comment thread arraycontext/container/traversal.py Outdated
Comment thread arraycontext/container/traversal.py Outdated
Comment thread arraycontext/container/traversal.py Outdated
@alexfikl

alexfikl commented Oct 7, 2021

Copy link
Copy Markdown
Collaborator Author

LGTM, aside from these final wrinkles. If you're OK with accepting all these suggestions, then this is ready to go from my end.

They look good to me! This still depends on the fixes from inducer/compyte#36 and inducer/pyopencl#514 to get the tests to pass without the added xfail. We can remove that later though.

@inducer

inducer commented Oct 7, 2021

Copy link
Copy Markdown
Owner

Those are close, too. So let's just wait on them.

@thomasgibson

Copy link
Copy Markdown
Collaborator

What's the status of this PR? Are there still issues remaining?

@alexfikl

Copy link
Copy Markdown
Collaborator Author

What's the status of this PR? Are there still issues remaining?

The general case should work as advertised. We hit an edge case with arrays that have a 0-size in one of the axes and don't work quite yet on pytato due to inducer/loopy#497.

@inducer

inducer commented Oct 12, 2021

Copy link
Copy Markdown
Owner

And I've been just slow to review. Sorry!

@thomasgibson

Copy link
Copy Markdown
Collaborator

No worries! I just wanted to make sure there wasn't a "gotcha" with this change

@inducer

inducer commented Oct 18, 2021

Copy link
Copy Markdown
Owner

I wonder whether we should just xfail the loopy-dependent zero-size bits for now and bump them to an issue, so @thomasgibson can start using this... what do you all think?

@alexfikl

alexfikl commented Oct 18, 2021

Copy link
Copy Markdown
Collaborator Author

I wonder whether we should just xfail the loopy-dependent zero-size bits for now and bump them to an issue, so @thomasgibson can start using this... what do you all think?

Yeah, definitely agree with that. The loopy fix might take a while and it's not breaking any actual code at the moment, just an edge case.

EDIT: Updated the url in c50ee3e to point to the loopy PR.

Comment thread arraycontext/impl/pyopencl/fake_numpy.py Outdated
@thomasgibson

Copy link
Copy Markdown
Collaborator

I wonder whether we should just xfail the loopy-dependent zero-size bits for now and bump them to an issue

Yes, sounds good to me.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants