feat: Added Auto Opt, GPU and `jax.Array` #26

philip-paul-mueller · 2024-10-02T12:47:04Z

This commit addresses the following issues:

It adds an implementation for auto optimizer (Add auto_optimize to the options #24).
The current implementation is not likely to stay, since it essentially uses DaCe's version, which is known to have problems with JaCe's SDFGs.
It allows to run stuff on GPU (Add GPU Capabilities. #25).
While it is possible it is still needed that the user explicitly specify it, JAX does an auto detection.
Instead of returning NumPy arrays JaCe now returns jax.Array objects (Return jax.Array #22).
This goes in tandem with a reworking of the type annotation, which was wrong before (and can not be correctly made).

The auto optimizer is disabled and must be enabled explicitly. This will be the behaviour will be maintained at least until the map fusion PR is merged. Essentially this commit addresses issues: - [issue#24](GridTools#24) - [issue#25](GridTools#25)

As explained elsewhere the annotation was wrong. Simple example: ```python @jace.jit def foo(a: np.ndarray) -> np.float64: return np.float64(a.sum()) ``` The actual return value would be `np.ndarray`, although one dimension. This limitation is also present in JAX.

This addresses [issue#22](GridTools#22).

codecov-commenter · 2024-10-02T12:50:06Z

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

Codecov Report

Attention: Patch coverage is 68.91892% with 23 lines in your changes missing coverage. Please review.

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

Files with missing lines	Patch %	Lines
src/jace/util/jax_helper.py	33.33%	9 Missing and 3 partials ⚠️
src/jace/api.py	63.63%	2 Missing and 2 partials ⚠️
src/jace/optimization.py	76.47%	2 Missing and 2 partials ⚠️
src/jace/__init__.py	33.33%	1 Missing and 1 partial ⚠️
src/jace/stages.py	94.11%	1 Missing ⚠️

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

Additional details and impacted files

@@           Coverage Diff           @@
##             main      #26   +/-   ##
=======================================
  Coverage        ?   68.36%           
=======================================
  Files           ?       31           
  Lines           ?     1277           
  Branches        ?      261           
=======================================
  Hits            ?      873           
  Misses          ?      329           
  Partials        ?       75

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

egparedes

Looks good, I have some minor comments. (BTW, Iit is not by any means "a super small PR")

src/jace/api.py

egparedes · 2024-10-02T13:17:16Z

src/jace/api.py

    """

+    backend: str


Suggested change

backend: str

backend: Literal['cpu', 'gpu']

I do not think that enforcing this through the annotations and MyPy is the right way.

Why not? What's the use case for the Literal type annotation then? Using Literal here also helps in documenting the supported values in a more prominent way.

First a practical reason, if you read the backend string say from a file then you will get a str, so you will have to cast it.
It also does not really helps documenting since in general it does not explain what a particular value mean or does, you still have to read the doc anyway.

egparedes · 2024-10-02T13:18:00Z

src/jace/api.py

@@ -72,18 +75,19 @@ def jit(
        fun: Function to wrap.
        primitive_translators: Use these primitive translators for the lowering to SDFG.
            If not specified the translators in the global registry are used.
-        kwargs: Jit arguments.
+        kwargs: Jit arguments, see `JITOptions` for more.


Suggested change

kwargs: Jit arguments, see `JITOptions` for more.

kwargs: jit arguments, see `JITOptions` for more.

We should be consistent and write JIT capital or Jit as it is the abbreviation of "Just in time" and we are at the beginning of a sentence.

I'd rather use JIT then but it's up to you...

src/jace/api.py

egparedes · 2024-10-02T13:50:18Z

src/jace/util/jax_helper.py

+    if isinstance(arr, jax.Array):
+        return arr
+    # In newer version it is no longer needed to pass a capsule.
+    return jax_dlpack.from_dlpack(arr, copy=copy)  # type: ignore[attr-defined]  # `from_dlpack` is not found.


Why is this type ignore needed if jax is included in our lenient mypy settings (and also the function should actually exist)?

I have no idea.

egparedes · 2024-10-02T13:51:49Z

src/jace/util/jax_helper.py

+try:
+    import cupy as cp  # type: ignore[import-not-found]
+except ImportError:
+    cp = None
+
+


Why is this needed?

In case you want to use JaCe one a CPU only machine as the CI.

I meant, why is cupy needed here? I can't see where it is used in this module.

Before the update of the newest JAX we had to test for the array type, since for cupy we have to call .toDLPack() but for numpy it was necessary to call .__dlpack__().
I simply forgot that.

egparedes · 2024-10-02T13:52:53Z

src/jace/tracing.py

+    # NOTE: `jax.make_jaxpr()` to current tracing backend we use, never supported
+    #   all arguments `jax.jit()` supported, which is strange. But in JaCe this
+    #   _will_ not be the case. To make things work, until we have a better working
+    #   backend we have to filter, i.e. clearing the `trace_options`.


Check and rewrite, I don't understand it.

You are right it was written very strangely.

egparedes · 2024-10-02T13:53:27Z

src/jace/tracing.py

+            # For some reasons MyPy seems to think that `jax.make_jaxpr()` is the same
+            #  as `jace.make_jaxpr()` so we have to ignore the error.


Sounds suspicious...

I know, but I have no idea why.

src/jace/translator/post_translation.py

src/jace/api.py

egparedes · 2024-10-03T11:00:54Z

src/jace/api.py

    """

+    backend: str


Why not? What's the use case for the Literal type annotation then? Using Literal here also helps in documenting the supported values in a more prominent way.

egparedes

I pressed the send button by mistake with only half of my comments. Here is the second half.

src/jace/tracing.py

egparedes · 2024-10-03T11:13:19Z

src/jace/util/jax_helper.py

+try:
+    import cupy as cp  # type: ignore[import-not-found]
+except ImportError:
+    cp = None
+
+


I meant, why is cupy needed here? I can't see where it is used in this module.

src/jace/util/jax_helper.py

egparedes

LGTM. I have just a comment about an error message but I don't need to review it again so I'm already approving the PR now.

egparedes · 2024-10-03T12:04:38Z

src/jace/util/jax_helper.py

+        case "tpu" | "TPU":
+            raise NotImplementedError("TPU are not supported.")
+        case _:
+            raise ValueError(f"Could not parse the backend '{backend}'.")


Rewrite the error message to reflect the new name of the function.

You are absolutely right.

philip-paul-mueller added 4 commits October 2, 2024 14:41

Forcing a modern version of JAX.

e2aa76f

Now JaCe uses the jax.Array as return type.

255267d

This addresses [issue#22](GridTools#22).

philip-paul-mueller requested a review from egparedes October 2, 2024 12:47

This was referenced Oct 2, 2024

Return jax.Array #22

Closed

Add auto_optimize to the options #24

Closed

Add GPU Capabilities. #25

Closed

egparedes reviewed Oct 2, 2024

View reviewed changes

Addressed Enriques comments.

29d8452

philip-paul-mueller requested a review from egparedes October 3, 2024 09:10

egparedes reviewed Oct 3, 2024

View reviewed changes

Added Enrique's fixes.

95df5e2

philip-paul-mueller requested a review from egparedes October 3, 2024 11:56

egparedes approved these changes Oct 3, 2024

View reviewed changes

philip-paul-mueller added 2 commits October 4, 2024 07:52

Fixed an error message.

dff8039

Made a comment about the memory allocation of the output arrays.

1f48a1e

philip-paul-mueller merged commit 3be9f36 into GridTools:main Oct 4, 2024
4 checks passed

philip-paul-mueller deleted the auto_opt_and_gpu branch October 7, 2024 13:42

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Added Auto Opt, GPU and `jax.Array` #26

feat: Added Auto Opt, GPU and `jax.Array` #26

philip-paul-mueller commented Oct 2, 2024

codecov-commenter commented Oct 2, 2024 •

edited

Loading

egparedes left a comment

egparedes Oct 2, 2024

philip-paul-mueller Oct 3, 2024 •

edited

Loading

egparedes Oct 3, 2024

philip-paul-mueller Oct 3, 2024 •

edited

Loading

egparedes Oct 2, 2024

philip-paul-mueller Oct 3, 2024 •

edited

Loading

egparedes Oct 3, 2024

egparedes Oct 2, 2024

philip-paul-mueller Oct 3, 2024

egparedes Oct 2, 2024

philip-paul-mueller Oct 3, 2024

egparedes Oct 3, 2024

philip-paul-mueller Oct 3, 2024

egparedes Oct 2, 2024

philip-paul-mueller Oct 3, 2024

egparedes Oct 2, 2024

philip-paul-mueller Oct 3, 2024

egparedes Oct 3, 2024

egparedes left a comment

egparedes Oct 3, 2024

egparedes left a comment

egparedes Oct 3, 2024

philip-paul-mueller Oct 4, 2024

	kwargs: Jit arguments, see `JITOptions` for more.
	kwargs: jit arguments, see `JITOptions` for more.

		# For some reasons MyPy seems to think that `jax.make_jaxpr()` is the same
		# as `jace.make_jaxpr()` so we have to ignore the error.

feat: Added Auto Opt, GPU and jax.Array #26

feat: Added Auto Opt, GPU and jax.Array #26

Conversation

philip-paul-mueller commented Oct 2, 2024

codecov-commenter commented Oct 2, 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 3, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

egparedes left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

egparedes left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

feat: Added Auto Opt, GPU and `jax.Array` #26

feat: Added Auto Opt, GPU and `jax.Array` #26

codecov-commenter commented Oct 2, 2024 •

edited

Loading

philip-paul-mueller Oct 3, 2024 •

edited

Loading

philip-paul-mueller Oct 3, 2024 •

edited

Loading

philip-paul-mueller Oct 3, 2024 •

edited

Loading