[WIP] Stack direct build without Makefile

[i-chaochen]

This is a more aggressive change based on #243

The build is simple by one line without Makefile

python3 stack.py build --xla-dir=/your/xla --jax-dir=/your/jax

1. Direct Build Implementation: Refactored stack.py to handle the entire build and installation process (JAX, jaxlib, PJRT, and plugin) directly using Python subprocesses, removing the dependency on a Makefile.

2. Simplified Command Line:

   • Renamed --kernel-jax-dir to --jax-dir.
   • Enabled the use of the = syntax (e.g., --xla-dir=/path/to/xla).
   • Made --rebuild-makefile implicit for the build action.

3. New flag --debug: You can now add --debug to your build or develop commands.

• Debug configurations: When --debug is used, the build process will:
  • Enable --config=debug and --compilation_mode=dbg
  • Disable stripping of symbols (--strip=never).
  • Set maximum debug info level (-g3).
  • Disable optimizations (-O0) for both C and C++ code.
• Makefile support: Running python3 stack.py develop --debug will generate a Makefile configured to use the CFG_DEBUG options by default.
• Example usage:
  • To build everything in debug mode:

python3 stack.py build --xla-dir=/your/xla --jax-dir=/your/jax --debug

To just generate a Makefile with debug settings:

python3 stack.py develop --xla-dir=/your/xla --jax-dir=/your/jax --debug

If you really want Makefile you can still go with

python3 stack.py develop --xla-dir=/your/xla --jax-dir=/your/jax

4. Specific clang path and choose clang18 as default one

python3 stack.py build --xla-dir=/your/xla --jax-dir=/your/jax --clang-path=/usr/lib/llvm-18/bin/clang

[Ruturaj4]

LGTM @charleshofer can you also take a look?

[Arech8]

Direct Build Implementation: Refactored stack.py to handle the entire build and installation process (JAX, jaxlib, PJRT, and plugin) directly using Python subprocesses, removing the dependency on a Makefile.

I oppose this heavily as you didn't explain the rationale and I don't see a value of that, but only a burden of supporting yet another build sequence. No, full stop here, we should not add another layer to this onion.

[Arech8]

I don't think adding another build sequence to support is useful and valuable. Makefile route works flawlessly, why should we add another orthogonal way to build things? This creates nothing, but a maintenance hell...

IMO the PR should be discarded.

[i-chaochen]

I don't think adding another build sequence to support is useful and valuable. Makefile route works flawlessly, why should we add another orthogonal way to build things? This creates nothing, but a maintenance hell...

IIUC, previous, we need to have

1. python stack.py develop --rebuild-makefile
2. edit Makefile for xla and jax path
3. make clean dist install // which only has pjrt and plugin
4. make jaxlib_clean jaxlib jaxlib_install // to install jaxlib
5. go to jax repo and try pip install . // to install jax

Why we need such cumbersome procedure to build jax and xla? it's wasted everyone time every day to do these unnessary steps. You can still have Makefile if you want to, this is still there. Your previous procesdure is still there.

This is to simplfy the dev process.

[Arech8]

@i-chaochen who told you that you should use that procedure? It works, but that's not what anyone should do. Have you read @charleshofer descriptions of the process? The old process is as simple as doing:

1. python3 stack.py develop with correct flags once (flags to specify JAX and XLA existed there almost all the time, you must have seen them when modifying the file, didn't you?), + cd $YOUR_JAX_DIR && pip install . -e once.
2. make refresh refresh_jaxlib each time you need to install/update you jax instance.

That's it.

[i-chaochen]

Have you read @charleshofer descriptions of the process? The old process is as simple as doing:

which docs? not in https://github.com/ROCm/rocm-jax/blob/master/BUILDING.md and

https://github.com/ROCm/rocm-jax/blob/master/DEVSETUP.md#setting-up-a-docker-development-environment

is let me to do this thing

python3 stack.py develop --rebuild-makefile
cat jax_rocm_plugin/Makefile

so I edit my xla and jax path with what docs mentioned here

source .venv/bin/activate
make clean dist install

there are no use cases in the docs to what you said

1. python3 stack.py develop with correct flags once (flags to specify JAX and XLA existed there almost all the time, you must have seen them when modifying the file, didn't you?), + cd $YOUR_JAX_DIR && pip install . -e once.
2. make refresh refresh_jaxlib each time you need to install/update you jax instance.

[Arech8]

@i-chaochen

[i-chaochen]

python3 stack.py develop with correct flags once (flags to specify JAX and XLA existed there almost all the time, you must have seen them when modifying the file, didn't you?), + cd $YOUR_JAX_DIR && pip install . -e once.

so you mean we can just do like this?

python3 stack.py develop xla_dir=/my/own/xla/ test_jax_ref=/my/own/jax ?

BTW, why you got test_jax_ref and kernels_jax_dir ? cannot we just to have one only? is test_jax_ref is for jax repo and kernels_jax_dir is for jaxlib?

rocm-jax/stack.py

Lines 664 to 665 in 21840fc

	test_jax_ref=args.jax_ref,
	jax_dir=args.jax_dir,

2. make refresh refresh_jaxlib each time you need to install/update you jax instance.

I would like to add an use-case for make refresh refresh_jaxlib on https://github.com/ROCm/rocm-jax/blob/master/DEVSETUP.md to be more straightforward, is it ok?

so IIUC, the build procedure is these two?

1. python3 stack.py develop xla_dir=/my/own/xla/ test_jax_ref=/my/own/jax
2. make refresh refresh_jaxlib

[Arech8]

so you mean we can just do like this?
python3 stack.py develop xla_dir=/my/own/xla/ test_jax_ref=/my/own/jax ?
BTW, why you got test_jax_ref and kernels_jax_dir ? cannot we just to have one only? is test_jax_ref is for jax repo and kernels_jax_dir is for jaxlib?

@i-chaochen can you please finally have a look at stack.py CLI (that you proposed to modify, meaning you must already understand them) at https://github.com/ROCm/rocm-jax/blob/master/stack.py#L398 or just run stack.py -h and read the CLI arguments and their meaning? It's just a couple of dozen lines of simple readable code that contain all the answers. You can find correct arguments on your own, I believe in you!

I would like to add an use-case for ...

sure, please make a PR

[Arech8]

@i-chaochen please don't waste an effort on this PR, we will not merge a third parallel way to build the repo.

On top of that, as was said, @alekstheod is brewing a major refactoring of the build system to bazelize things as much as possible, so this will be obsolete in a very short time.

[i-chaochen]

@i-chaochen please don't waste an effort on this PR, we will not merge a third parallel way to build the repo.

On top of that, as was said, @alekstheod is brewing a major refactoring of the build system to bazelize things as much as possible, so this will be obsolete in a very short time.

my understanding is @alekstheod's work is about RBE in CI, I am not sure how it will be. For now, I'm using this for my team purpose to build.

[Arech8]

@alekstheod is working on the whole build system refactoring and wants to fully get rid of python wrappers as much as possible.

You are free to use whatever you want locally, but we will not support yet another build script. I highly suggest to learn using existing solution instead. Adding a burden of maintaining semantically duplicate code isn't a way to simplify working on a project.