Implement Transpose primitives #17

ASKabalan · 2024-05-02T16:12:06Z

First pull request for Tranpose ops

They work for cubes for now, (Pencils and slabs)

ASKabalan · 2024-05-02T16:13:33Z

scripts/test_fft3d.py

@@ -9,6 +9,7 @@
 import jaxdecomp

 # Initialize jax distributed to instruct jax local process which GPU to use
+jaxdecomp.init()


This is needed on JZ for now

Ok, that's fine :-)

ASKabalan · 2024-05-02T16:16:34Z

tests/test_transpose.py

+  print(f"jd_tranposed_zy shape {jd_tranposed_zy.shape}")
+  print(f"jd_tranposed_yx shape {jd_tranposed_yx.shape}")
+
+  if pdims[1] == 1:


Make sure that the grid is transposed and traposed back every 2 steps

ASKabalan · 2024-05-02T16:19:58Z

jaxdecomp/_src/transpose.py

+      case 'x_y' | 'y_z':
+        transpose_shape = (2, 0, 1)
+      case 'y_x' | 'z_y':
+        transpose_shape = (1, 2, 0)


Perhaps the global shape has to follow the col to row tranposition that cudecomp does?
This is fine for cubes, for now

what happens for non cubes?

I solved this in the last commit

ASKabalan · 2024-05-02T16:55:25Z

src/jaxdecomp.cc

@@ -118,129 +119,23 @@ decompGridDescConfig_t getAutotunedGridConfig(decompGridDescConfig_t grid_config
  return output_config;
 };

-/// XLA interface ops


Instead of three custom callable funtion, I added only one and it uses an enum to know which transposition to do

ASKabalan · 2024-05-02T16:57:13Z

jaxdecomp/_src/transpose.py


+# Y to X


Derivatives are actually just residuals because tranposition is linear.
I defined a VJP anyways using the back tranposes as bwd passes for the forrward ones and vice versa

ASKabalan · 2024-05-02T16:59:01Z

jaxdecomp/__init__.py

@@ -4,13 +4,15 @@
 import jaxdecomp.fft as fft
 from jaxdecomp.fft import pfft3d, pifft3d

-from ._src import (  # transposeXtoY, transposeYtoX, transposeYtoZ, transposeZtoY
+from ._src import (  # transposeXtoY, transposeYtoX, transposeYtoZ, transposeZtoY,


without this commented line, isort and yapf are stuck in an endless battle 🥲 during a precommit run

aboucaud

Looks like a nice job !

From my perspective, what is missing from such PR is a bit more documentation.
There are comments here and there but no docstrings or parts that express the reasoning behind what is done.
I also like the fact you tried to add some typing in Python, but it's rather sparse so far, hence I don't really see the point.

To be clear, making this work and behaving as expected is the priority. But then, I think all this very clever piece of code could benefit from statements / explanations on the design and strategy. This could be in the code itself or somewhere next to it. Feel free to disagree.

ASKabalan · 2024-05-03T21:52:42Z

Thank you @aboucaud .
The tranpose primitive uses a much better structure than the others, I am gradually trying to improve readability.
But it still missing some stuff that I am very close to solving (non-cubes and understand fft forward axis order).

That said, any feedback is much appreciated.
Once this PR is ready I would highly appreciate feed back on code readability and maintability.

ASKabalan · 2024-05-04T08:25:53Z

@aboucaud @EiffL .
This is good to go.
Now works for cubes non cubes , slabs XY, slabs YZ and pencils.

ASKabalan · 2024-05-04T08:28:16Z

jaxdecomp/_src/transpose.py

@@ -44,16 +44,21 @@ def abstract(x, kind, pdims, global_shape):
    assert kind in ['x_y', 'y_z', 'z_y', 'y_x']
    match kind:
    # From X to Y the axis are rolled by 1 and pdims are swapped wrt to the input pdims


The process dims are being transposed at every step.
But I need the original pdims for my compiled primitive (the grid descriptor and the tranpose operator are indexed using gdim and pdim ... so even after a tranpose I need them back to find my primitive in m_TransposeDescriptors32 or m_TransposeDescriptors64)

ASKabalan · 2024-05-04T08:28:55Z

jaxdecomp/_src/transpose.py

@@ -141,7 +154,10 @@ def infer_sharding_from_operands(kind: str, mesh: Mesh,
                                   arg_infos: Tuple[ShapeDtypeStruct],
                                   result_infos: Tuple[ShapedArray]):
    input_sharding = arg_infos[0].sharding


Every step the processor grid is being switched from col major to row major
A test is added for this

ASKabalan · 2024-05-04T08:29:13Z

jaxdecomp/_src/transpose.py

        transpose_type = _jaxdecomp.TRANSPOSE_ZY
      case 'y_x':
+        transpose_shape = (1, 2, 0)
        transpose_type = _jaxdecomp.TRANSPOSE_YX
      case _:
        raise ValueError("Invalid kind")



Get the original shape for my indexed operator (same as above)

ASKabalan · 2024-05-04T08:29:46Z

tests/test_transpose.py

@@ -98,6 +98,16 @@ def test_tranpose(pdims):

  print(f"JD tranposed yz sharding {jd_tranposed_yz.sharding.spec}")
  print(f"JD tranposed xy sharding {jd_tranposed_xy.sharding.spec}")
+  print(f"JD tranposed zy sharding {jd_tranposed_zy.sharding.spec}")
+  print(f"JD tranposed yx sharding {jd_tranposed_yx.sharding.spec}")
+


Test that the sharding is being transposed

ASKabalan · 2024-05-04T08:37:30Z

jaxdecomp/_src/transpose.py

@@ -150,7 +166,14 @@ def partition(kind: str, mesh: Mesh, arg_infos: Tuple[ShapeDtypeStruct],
    input_sharding = NamedSharding(mesh, P(*arg_infos[0].sharding.spec))
    output_sharding = NamedSharding(mesh, P(*result_infos.sharding.spec))
    global_shape = arg_infos[0].shape
-    pdims = (get_axis_size(input_sharding, 0), get_axis_size(input_sharding, 1))


Make sure that pdims provided to the c++ primitive is the original one (same reason as above)
But in the JAX abstract eval we need the reversed pdims because it changes the pencil shape
So c++ pdims are the original and cudecomp knows how to handle them
We mirror cudecomp reversal in our abstract eval

ASKabalan · 2024-05-04T08:38:29Z

tests/test_transpose.py

+

Cartesian product testing of cubes non cubes, prime factors with pencils and slabs

ASKabalan · 2024-05-07T18:44:42Z

jaxdecomp/_src/transpose.py

+        transpose_pdims = pdims[::-1]
+      case 'y_x':
+        transpose_shape = (1, 2, 0)
+        transpose_pdims = pdims



A forward transpose cycle shifts the axis right for us (2 0 1)
A backward transpose cycle shifts the axis left (1 2 0)

Every transpose the pdims are switched (implicitly by cudecomp)
I need the original pdims for the grid descriptor
So I reverse the pdims every other transpose

EiffL

LGTM

Add back init to init the MPI comm (needed for JZ for now)

50232fb

ASKabalan linked an issue May 2, 2024 that may be closed by this pull request

Things to update for version 0.0.1 #10

Closed

7 tasks

ASKabalan force-pushed the u/ASKabalan/transpose_ops branch 2 times, most recently from 68b1b73 to 822f657 Compare May 2, 2024 16:18

ASKabalan commented May 2, 2024

View reviewed changes

ASKabalan requested a review from EiffL May 2, 2024 17:02

aboucaud reviewed May 3, 2024

View reviewed changes

ASKabalan force-pushed the u/ASKabalan/transpose_ops branch 4 times, most recently from 74f1656 to b5fe313 Compare May 4, 2024 08:23

ASKabalan commented May 4, 2024

View reviewed changes

ASKabalan requested a review from aboucaud May 4, 2024 08:32

ASKabalan force-pushed the u/ASKabalan/transpose_ops branch from b5fe313 to d69c2a0 Compare May 4, 2024 08:33

ASKabalan commented May 4, 2024

View reviewed changes

ASKabalan force-pushed the u/ASKabalan/transpose_ops branch from d69c2a0 to 48d9b49 Compare May 4, 2024 09:11

This was referenced May 4, 2024

Things to update for version 0.0.1 #10

Closed

Fix slice unpad #18

Merged

ASKabalan changed the title ~~U/as kabalan/transpose ops~~ Implement Transpose primitives May 4, 2024

ASKabalan added 3 commits May 4, 2024 16:52

Adjusted abstract evals and sharding rules for PFFT

9d0c539

Added full test for PFFTs

1b3e458

Adjust grad test for new yz slab flag

bd467cc

ASKabalan commented May 7, 2024

View reviewed changes

ASKabalan force-pushed the u/ASKabalan/transpose_ops branch from 48d9b49 to 7126c3d Compare May 30, 2024 18:02

ASKabalan added 4 commits May 31, 2024 14:29

Roll back yz flag and add an extra JAX transposition

ee8e102

Add extra transpose in cuda code for YZ slab

83554e0

Implement cu/c++ side

671540c

Implement JAX Side for tranpositions

aa3529b

ASKabalan added 4 commits May 31, 2024 17:36

Add test_transpose all OK for cubes only

ba0b76c

Add Gradient test for transpose_ops

39d62e6

Handle non-cube shapes and primes in the tranpose primitive

af3953d

Add tests for non-cubes and prime*size dimensions

1d1a779

ASKabalan force-pushed the u/ASKabalan/transpose_ops branch from 7126c3d to 1d1a779 Compare May 31, 2024 15:36

EiffL approved these changes Jun 11, 2024

View reviewed changes

EiffL merged commit 80172aa into main Jun 11, 2024
1 check passed

ASKabalan deleted the u/ASKabalan/transpose_ops branch June 19, 2024 09:07

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement Transpose primitives #17

Implement Transpose primitives #17

ASKabalan commented May 2, 2024

ASKabalan May 2, 2024

EiffL May 2, 2024

ASKabalan May 2, 2024

ASKabalan May 2, 2024

EiffL May 2, 2024

ASKabalan May 4, 2024

ASKabalan May 2, 2024

ASKabalan May 2, 2024

ASKabalan May 2, 2024

aboucaud left a comment

ASKabalan commented May 3, 2024

ASKabalan commented May 4, 2024

ASKabalan May 4, 2024

ASKabalan May 4, 2024

ASKabalan May 4, 2024

ASKabalan May 4, 2024

ASKabalan May 4, 2024

ASKabalan May 4, 2024

ASKabalan May 7, 2024 •

edited

Loading

EiffL left a comment

Implement Transpose primitives #17

Implement Transpose primitives #17

Conversation

ASKabalan commented May 2, 2024

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

aboucaud left a comment

Choose a reason for hiding this comment

ASKabalan commented May 3, 2024

ASKabalan commented May 4, 2024

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

ASKabalan May 7, 2024 • edited Loading

Choose a reason for hiding this comment

EiffL left a comment

Choose a reason for hiding this comment

ASKabalan May 7, 2024 •

edited

Loading