test: Added an Extensive Set of Tests #21

philip-paul-mueller · 2024-09-24T13:27:05Z

This PR adds a large set of tests for various translators, especially the ones that where added by PR#20 but also other parts of the code are tested.
This is a very basic PR that does not add any functionality, except tests for the one that is already there.

This PR can only be merged after PR#20 has been merged.
Furthermore, for certain reasons this PR contains the whole development history of PR#20.

…o tests for them yet. `device_put` is actually a very powerfull operation, such as Memlets, where source and destination are on different devices. However, we do not support something like that yet, so we will keep it down yet.

It is not yet dynamic sclice, but soon.

…n value. This restirction was removed commit `05e4a885441c0cd`. However, I realized that it made some higher level translators impossible to write, such as `select_n`. Thus I removed this restriction again. Another solution would be to add another layer.

I also observe random failures for the `test_iota_broadcast()` test if I run all tests. However, if I only run it, then nothing happens, I have no idea why.

… the class description.

This should bring the development branch to the newest stage.

…Jax stuff about the primitives.

Jax adjust the start indexes if the window overruns, however, this is not done, instead an out of bound error happens.

During that work I also detected some [issue](spcl/dace#1579) in DaCe's simplification pipeline.

The cleaning was not correct. The function essentially created some deatached caches.

I just copied them and did not do a merge, which is not so nice. Furthermore, the tests are not yet there, in my view it makes sense to first have something that can be checked.

Before it was implementewd as a switch, for the case JAX would use the bool overload of XLA. The check for this was now essentially moved inside the function.

However, a similar issue (#1644) in DaCe is still open.

It is now better confiugured.

Now let's test if it works.

This is basically for testing them.

Before the `order` argument was a `Literal` but this caused more truble now it is a string.

I enabled the simplify pass in commit `411bd7bd` and it worked locally. However, this was because I was not running it inside nox and using my own version of DaCe. The bug in simplify was fixed in [PR#1603](spcl/dace#1603) which was merged _after_ 16.1 was released, thus the fix is not avaliable.

It seems JAX has updated the `make_jaxpr()` function and now that thing caches itself. This is now accounted for.

But the new function really tests if the loweing works, if teh strides are honored and infered.

Co-authored-by: Enrique González Paredes <enriqueg@cscs.ch>

This PR introduces a series of primitive translators, most of them are based on the prototype, with some improvements. I just copied them over from the development branch, which is not so nice, but was the simplest thing to to without also introducing the other stuff. It is important that the tests from the development branch were not added, to keep the PR small. Furthermore, we need something to test, so this PR must go first. For organizational reasons, the development history of this PR happened to be contained in [PR#21](#21). --------- Co-authored-by: Enrique González Paredes <enriqueg@cscs.ch>

codecov-commenter · 2024-10-01T07:13:37Z

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

All modified and coverable lines are covered by tests ✅

Please upload report for BASE (main@0a9f361). Learn more about missing BASE report.

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@           Coverage Diff           @@
##             main      #21   +/-   ##
=======================================
  Coverage        ?   88.66%           
=======================================
  Files           ?       31           
  Lines           ?     1235           
  Branches        ?      251           
=======================================
  Hits            ?     1095           
  Misses          ?       82           
  Partials        ?       58

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

egparedes

First round of review only looking at the test infrastructure.

egparedes · 2024-09-30T14:25:38Z

docs/main_differences.md

+
+- JaCe always traces with enabled `x64` mode.
+  This is a restriction that might be lifted in the future.
+- JAX returns scalars as zero-dimensional arrays, JaCe returns them as array with shape `(1, )`.


The main issue is that DaCe does not have a real concept of zero dimensional arrays, as far as I know.
Consider the following two functions

@dace.program def bar(a: dace.float64): return a + 1 @dace.program def baz(a: dace.float64[1]): return a + 1

if you pass a zero dimensional array to bar() then it will be casted to an scalar, if you pass it to baz() an error will happen.
Furthermore, the binary interface of the SDFG can not return scalars, return values have to be arrays there is no way around that without patching the code generator and making a lot of changes to handle special cases.
So I decided to follow PEP20 and decided that this case is not special enough to change the rule.
If you want this feature then please open an issue.

egparedes · 2024-09-30T14:28:09Z

docs/main_differences.md

+- JaCe always traces with enabled `x64` mode.
+  This is a restriction that might be lifted in the future.
+- JAX returns scalars as zero-dimensional arrays, JaCe returns them as array with shape `(1, )`.
+- In JAX parts of the computation runs on CPU parts on GPU, in JaCe everything runs (currently) either on CPU or GPU.


What do you mean here? Which parts run on CPU/GPU in JAX?

The JAX compiler, i.e. XLA can decide to do this.
The question is, if it really does it.

egparedes · 2024-09-30T14:28:53Z

docs/main_differences.md

+- JaCe does not return `jax.Array` instances, but NumPy/CuPy arrays.
+- The execution is not asynchronous.


These two points could be also fixed in the future, right?

Return Type: Return jax.Array #22

Asynchronous: Asynchronous Execution #23

egparedes · 2024-09-30T14:34:41Z

tests/common_fixture.py

Move fixtures from conftest.py here (in which case should be renamed to commom_fixtures.py) or remove this empty file.

egparedes · 2024-09-30T15:35:15Z

tests/conftest.py

+from jace import optimization, stages
+from jace.util import translation_cache as tcache
+
+


For all fixtures in this module:

remove the underscore of the name (all fixtures defined here are meant to be used outside of this module)
remove the autouse=true and follow the same approach we used in gt4py-dace (https://github.com/GridTools/gt4py/pull/1594/files#diff-4a5190762a7529af16fcf3d4f1726d2168726659f60f9d58495611d56eeb3008): an explicit pytestmark definition at a module level with the fixtures that should be automatically used for all tests in the module.

Additionally, I would define another fixture requesting all other fixtures expected in the standard case and use this group fixtures in the pytestmark at module level. Example:

# This file @pytest.fixture def standard_jace_test_settings(enable_x64_mode_in_jax, disable_jit, ...) -> ...: .... # Other test files pytestmark = pytest.mark.usefixtures("standard_jace_test_settings")

Finally, I'd create a simpler type alias for the return type of generator fixtures as suggested here :

T = TypeVar("T") YieldFixture = Generator[T, None, None] @pytest.fixture def foo() -> YieldFixture[str]: yield "foo"

egparedes · 2024-10-01T06:51:02Z

tests/util.py

+def make_array(
+    shape: Sequence[int] | int,
+    dtype: type = np.float64,
+    order: str = "C",


Suggested change

order: str = "C",

order: Literal["C", "F"] = "C",

This was the original, however, it has serious problems.
Whenever the order came from an argument (fixture to test if it works in both C and F order) then MyPy complained, and I had to add brainless # type: ignore[call-overload] marks (see commit 090c3a2).
I do not think that cluttering the code with these annotations just to get some tiny bit of type security is not worth it.

egparedes · 2024-10-01T06:59:51Z

tests/util.py

+    low: Any = None,
+    high: Any = None,


Why not?

Suggested change

low: Any = None,

high: Any = None,

low: int | float | np.number = 0,

high: int | float | np.number = 1,

Yours is a bit better.
But it must be still accept None, otherwise you will only generate 0 or 1 in the integer case.
Furthermore, technically it must also accept complex, but they are an edge case that we ignore.

egparedes · 2024-10-01T07:05:01Z

tests/util.py

+
+
+def make_array(
+    shape: Sequence[int] | int,


Suggested change

shape: Sequence[int] | int,

shape: Sequence[int],

This is optional, but I would stick to one single way to define the shape to increase the readability of the function and help the type checker to catch errors here. In my opinion, there is not too much value in providing automatic tuple conversion for ints.

I do think that this change would increase readability.
For my these ..., shape=(some_scalar_variable,), ... expression are much less readable than ..., shape=some_scalar_variable, ....
This typedef allows you to select the style that is most appropriate to the function.
Thus you do not have to replicate the "logic" that distinguish scalars from tuples everywhere you use this thing, but only at one central place.
Furthermore, it is the NumPy behaviour.

egparedes · 2024-10-01T07:06:01Z

tests/util.py

+__all__ = ["make_array"]
+
+
+def make_array(


Just a question: this utility function is dealing with the generation of random array values with a certain shape, dtype and range, but what about packing them in the the correct ndarray type (numpy, cupy, Jax)? Isn't it needed ?

At the end of the day, DaCe generates C code that operates on pointers to memory regions of a certain size.
Thus in the majority of cases you care about the data type and the size of the memory and not about the type of container that is used to store manage said memory it in the Python world.

However, you need to have tests to ensure that you can extract that memory from JAX, NumPy and CuPy arrays.
For JAX array we have:

tests/unit_tests/test_caching.py:test_caching_jax_numpy_array()

tests/unit_tests/test_jax_api.py:test_jax_array_as_input()

So in the end is composition at work.

philip-paul-mueller added 30 commits May 24, 2024 14:34

Made some small modifications.

979f4ec

Reloacated some code.

2d2c160

Added a transformation for slice.

8adae42

It is not yet dynamic sclice, but soon.

Implemented a test for teh select_n primitive.

a3ac868

I also observe random failures for the `test_iota_broadcast()` test if I run all tests. However, if I only run it, then nothing happens, I have no idea why.

First, chunk of work.

5f186f6

Small corrections.

df82038

Size corrections.

60b7acb

Moved the description of the constructor arguments from __init__ to…

2190faf

… the class description.

Moved the jax package.

5d6962a

Fixed the type annotation.

6c980a9

Added a todo.

6dd65f3

Merge branch 'initial_implementation' into initial_implementation_cont

a0996eb

This should bring the development branch to the newest stage.

Found something interesting in DaCe, but it does not fully match the …

f7fe13d

…Jax stuff about the primitives.

WIP: Dynamic slicing, it seems to work, but it is not nice.

7b7b7ef

Tidied up the dynamic slicing, however, it is not yet fully compatible.

28c77ca

Jax adjust the start indexes if the window overruns, however, this is not done, instead an out of bound error happens.

Dynamic slice now also supports the shifting of windows.

054bb22

During that work I also detected some [issue](spcl/dace#1579) in DaCe's simplification pipeline.

Some cleaning.

c7e8b9f

Removed the warning about the useless casting in the tests.

a28a9c8

Started to reorganize the tests.

5bb4b5e

WIP: Started to better organize the tests.

959849a

Applied a first version to the new code.

047bb46

Second round.

64135b9

Third rounds for updated configuration.

5756d1e

Fourth batch of adapting new configuration.

8b66af8

Fixed a bug in the cache.

a5826e9

The cleaning was not correct. The function essentially created some deatached caches.

Fixed a bug in the cache.

6f4c45b

The cleaning was not correct. The function essentially created some deatached caches.

Reorganized some stuff.

4735092

Improved coverage.

758adc8

philip-paul-mueller added 12 commits September 13, 2024 15:03

Added all translators from teh development branch.

d6265bc

I just copied them and did not do a merge, which is not so nice. Furthermore, the tests are not yet there, in my view it makes sense to first have something that can be checked.

Modified the condition primitive.

25616ae

Before it was implementewd as a switch, for the case JAX would use the bool overload of XLA. The check for this was now essentially moved inside the function.

Nobody needs that thest in this form.

24d97fb

Merge branch 'new_translators' into dev_branch

9f22bd2

Since the bug mention in DaCe issue |1595 was resolved we can enable it.

411bd7b

However, a similar issue (#1644) in DaCe is still open.

First batch of Enrique's suggestions, but still not done.

c8b7d86

Refactored the gather translator.

cb600d3

It is now better confiugured.

Some more corrections.

c29fc0d

Now let's test if it works.

Merge branch 'new_translators' into dev_branch

b0ac641

This is basically for testing them.

Made it run again.

3ee8dad

Fixed an error in the tests.

e787ee9

Fixed some errors.

846a345

philip-paul-mueller mentioned this pull request Sep 24, 2024

feat: First Batch of Translators #20

Merged

philip-paul-mueller requested a review from egparedes September 24, 2024 13:32

philip-paul-mueller added 7 commits September 24, 2024 15:36

Merge branch 'new_translators' into new_test_suite

a546932

Updated the make_array() function.

090c3a2

Before the `order` argument was a `Literal` but this caused more truble now it is a string.

Fixed an error in the tests.

6fe78b0

It seems JAX has updated the `make_jaxpr()` function and now that thing caches itself. This is now accounted for.

The test was not really testing anything, okay if it was working.

d88752a

But the new function really tests if the loweing works, if teh strides are honored and infered.

Applied Enriques primarly fixes.

2c7b3c8

Merge branch 'new_translators' into new_test_suite

5687f40

philip-paul-mueller force-pushed the new_test_suite branch from 71f9f86 to 5687f40 Compare September 26, 2024 13:41

philip-paul-mueller and others added 3 commits September 26, 2024 16:22

Splited a test into multiple one, for better pinpointing.

cc10a48

Apply suggestions from code review

18bdee9

Co-authored-by: Enrique González Paredes <enriqueg@cscs.ch>

Merge branch 'new_translators' into new_test_suite

3588a77

philip-paul-mueller added 2 commits September 27, 2024 11:12

The import error that panda is not found is now ignored.

2169470

Merge branch 'main' into new_test_suite

e17b4eb

egparedes requested changes Oct 1, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

test: Added an Extensive Set of Tests #21

test: Added an Extensive Set of Tests #21

philip-paul-mueller commented Sep 24, 2024 •

edited

Loading

codecov-commenter commented Oct 1, 2024 •

edited

Loading

egparedes left a comment

egparedes Sep 30, 2024

philip-paul-mueller Oct 2, 2024 •

edited

Loading

egparedes Sep 30, 2024

philip-paul-mueller Oct 2, 2024 •

edited

Loading

egparedes Sep 30, 2024 •

edited

Loading

philip-paul-mueller Oct 2, 2024

egparedes Sep 30, 2024

egparedes Sep 30, 2024 •

edited

Loading

egparedes Oct 1, 2024

philip-paul-mueller Oct 2, 2024 •

edited

Loading

egparedes Oct 1, 2024 •

edited

Loading

philip-paul-mueller Oct 2, 2024 •

edited

Loading

egparedes Oct 1, 2024

philip-paul-mueller Oct 2, 2024 •

edited

Loading

egparedes Oct 1, 2024

philip-paul-mueller Oct 2, 2024

		- JaCe does not return `jax.Array` instances, but NumPy/CuPy arrays.
		- The execution is not asynchronous.

		from jace import optimization, stages
		from jace.util import translation_cache as tcache

test: Added an Extensive Set of Tests #21

Are you sure you want to change the base?

test: Added an Extensive Set of Tests #21

Conversation

philip-paul-mueller commented Sep 24, 2024 • edited Loading

codecov-commenter commented Oct 1, 2024 • edited Loading

Codecov Report

egparedes left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

philip-paul-mueller Oct 2, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

philip-paul-mueller Oct 2, 2024 • edited Loading

Choose a reason for hiding this comment

egparedes Sep 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

egparedes Sep 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

philip-paul-mueller Oct 2, 2024 • edited Loading

Choose a reason for hiding this comment

egparedes Oct 1, 2024 • edited Loading

Choose a reason for hiding this comment

philip-paul-mueller Oct 2, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

philip-paul-mueller Oct 2, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

philip-paul-mueller commented Sep 24, 2024 •

edited

Loading

codecov-commenter commented Oct 1, 2024 •

edited

Loading

philip-paul-mueller Oct 2, 2024 •

edited

Loading

philip-paul-mueller Oct 2, 2024 •

edited

Loading

egparedes Sep 30, 2024 •

edited

Loading

egparedes Sep 30, 2024 •

edited

Loading

philip-paul-mueller Oct 2, 2024 •

edited

Loading

egparedes Oct 1, 2024 •

edited

Loading

philip-paul-mueller Oct 2, 2024 •

edited

Loading

philip-paul-mueller Oct 2, 2024 •

edited

Loading