Update all dependencies. #12

ingomueller-net · 2023-11-03T10:28:22Z

This PR depends on and therefore includes #10.

This PR updates the following dependencies in the WORKSPACE file to a
version from around November 2:

LLVM
XLA
JAX (even though there is no actual BUILD dependency)
Enzyme

The PR also adds direct dependencies to the following projects:

rules_cc (because LLVM requires a newer version than the one that XLA
imports)

For context: @ftynse pointed me to this repo to copy the Bazel recipe to another project that depends on JAX and I didn't make it work there yet. Since I got stuck on that project, I came back to this one to see if I can make it work here. Turns out I could. I hope that this will serve me as a good base to get the other project to work as well. I don't need this project currently, but since I grinded through making it work, you might as well benefit from it ;)

The current import was non-functional, and that only remained unnoticed because there is no BUILD dependency to JAX, so JAX isn't actually built. (The CI tests install JAX through pip, so the JAX dependency isn't needed there.) The fix consists in moving the configuration from JAX's WORKSPACE file (which is emptied when imported via `http_archive`) to the WORKSPACE file of this project (plus fixing how its dependency `ducc` is loaded). Now adding `@jax//jax` to some of the build targets actually (and successfully) builds JAX.

This PR updates the following dependencies in the WORKSPACE file to a version from around November 2: * LLVM * XLA * JAX (even though there is no actual BUILD dependency) * Enzyme The PR also adds direct dependencies to the following projects: * rules_cc (because LLVM requires a newer version than the one that XLA imports)

ingomueller-net

OK, I am done: CI passes :)

@ftynse: Could you or one of your collaborators have a look to check if I do anything stupid (in particular in how to use Bazel)?

ingomueller-net · 2023-11-03T16:14:31Z

WORKSPACE

@@ -14,10 +28,10 @@ http_archive(
 )

 load("@llvm-raw//utils/bazel:configure.bzl", "llvm_configure")
-llvm_configure(name = "llvm-project", targets = ["X86", "AArch64", "AMDGPU", "ARM", "NVPTX"])
+llvm_configure(name = "llvm-project", targets = LLVM_TARGETS)


Note that this ~intentionally breaks the regexes in the CI pipelines that replace the build targets depending on the CI target. This doesn't work anymore since XLA (I believe) always requires X86CodeGen to be present, so the X86 target is always required.

It might be possible to adapt the regexes accordingly but since (one of) the pipelines are (is) signed, I didn't do that change.

ingomueller-net · 2023-11-03T16:16:06Z

BUILD


 py_package(
    name = "enzyme_jax_data",
    # Only include these Python packages.
-    packages = ["@//enzyme_jax:enzyme_call.so", "@llvm-project//clang:builtin_headers_gen"],


Warning: I don't understand what packages is supposed to do. If I use the line as is, the .so file is not included in the wheel...

So it's supposed to force the so (And those headers) to be made available in the pip package.

It does look like CI thinks that it can no longer find them, so we may need something else here (it worked in the past so I'm not sure what would've changed, and I'm curious if putting it back fixes anything)

When it was there, CI broke. I could observe the same behavior locally: only the py files ended up in the wheel. I don't know what changed, though...

Ah yeah so the problem is that the package.bzl forces the dependencies to be included as well, so by getting rid of our version of package.bzl you've removed the ability to import the actual dependencies.

Revert back to our package.bzl and see if it works?

I have also tried that. I started looking into the wheel rule because it broke, even before updating rules_python. Also, the old package.bzl doesn't have filtering logic for all I can tell.

Is there anything in the current behavior that is undesirable?

ingomueller-net · 2023-11-03T16:17:14Z

BUILD

@@ -29,19 +27,15 @@ cc_library(
    ],
 )

-load("@rules_python//python:packaging.bzl", "py_wheel")
-
-load(":package.bzl", "py_package")


By reading the code, I had the feeling that the package.bzl from this repo had no effect other than calling the function of the same name in rules_python (plus allowing but not using the prefix parameter), so I removed it.

ingomueller-net · 2023-11-06T11:08:20Z

I went through the packing rule again. If I revert the last commit, i.e., all changes to that rule, then the resulting *.whl does not contain the .so file:

python -c "from zipfile import ZipFile; from pprint import pprint; pprint(ZipFile('bazel-bin/enzyme_jax-0.0.5-py3-none-manylinux2014_x86_64.whl').namelist())" | grep \\.so
# no output

Commenting out the prefix or packages arguments does not change that behavior. Only if I use the function from the @rules_python package instead of the one from the reverted package.bzl in addition to commenting out prefix and packages parameters does it work. I also tried changing the packages parameter to //enzyme_jax:enzyme_call instead of //enzyme_jax:enzyme_call.so (i.e., without file extension) -- doing that doesn't fix the behavior. Using the rule from @rules_python without packages parameter does work for the old version of rules_python, though. With the new version, no form of the packages argument leads in the resulting wheel containing the .so file.

In short, after the changes in the previous commits, only the fixes in the last commit lead to a valid wheel.

(I just force-pushed an updated last commit where I remove the comment for the also-removed packages argument.)

ingomueller-net · 2023-11-06T11:24:03Z

OK, I investigated more. Values in the packages arguments have all . replaced by / (here -- WTF??) and are used as prefix for matching. So with packages = ["enzyme_jax/enzyme_call"],, the wheel actually contains the *.so file. That filter, instead, filtered out the headers. But I guess that wasn't the intention: they where part of the packages filter originally. As I see it the original packages filter said to include all files that where already selected as deps, so even when they were working, they had no effect (and when they were not, they broke the wheel). TL;DR: I think we should drop the argument as done in the last commit.

ingomueller-net · 2023-11-08T10:03:09Z

OK, I fixed another bunch of problems. I am able to run test/test.py locally now 😄 Hopefully this will pass CI as well...

ingomueller-net · 2023-12-01T08:11:49Z

You made it work :) Kudos!

ingomueller-net added 2 commits November 2, 2023 12:41

Fix prefix in JAX BUILD rule. Add SH256.

26ec69b

ingomueller-net force-pushed the update-dependencies branch 3 times, most recently from acaa8c7 to 3cba4bf Compare November 3, 2023 12:58

ingomueller-net force-pushed the update-dependencies branch from 3cba4bf to 68deb3f Compare November 3, 2023 16:10

ingomueller-net commented Nov 3, 2023

View reviewed changes

ingomueller-net force-pushed the update-dependencies branch from 68deb3f to 10f8438 Compare November 6, 2023 11:07

ingomueller-net added 2 commits November 7, 2023 10:27

Fix problem of circular dependency.

26502cb

Update rules_python. Simplify and fix packaging of wheel.

bdcd27b

ingomueller-net force-pushed the update-dependencies branch from 10f8438 to bdcd27b Compare November 7, 2023 10:28

Add missing dependencies. Fix pybind related rules.

2b1b359

ingomueller-net mentioned this pull request Nov 8, 2023

What's the right way to consume XLA with bazel? openxla/xla#6845

Open

ingomueller-net force-pushed the update-dependencies branch 4 times, most recently from ebe784a to 12c4999 Compare November 20, 2023 10:33

xxx-ci-debug

0a6155d

ingomueller-net force-pushed the update-dependencies branch from 12c4999 to 0a6155d Compare November 20, 2023 13:03

wsmoses added 3 commits November 29, 2023 12:30

fixup

cbb1170

bump xla

dc2f014

fix test

8ac8bd8

wsmoses merged commit 5f7f3de into EnzymeAD:main Nov 30, 2023
1 of 2 checks passed

ingomueller-net deleted the update-dependencies branch December 1, 2023 08:11

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Update all dependencies. #12

Update all dependencies. #12

ingomueller-net commented Nov 3, 2023

ingomueller-net left a comment

ingomueller-net Nov 3, 2023

ingomueller-net Nov 3, 2023

wsmoses Nov 3, 2023

wsmoses Nov 3, 2023

ingomueller-net Nov 3, 2023

wsmoses Nov 3, 2023

ingomueller-net Nov 4, 2023

ingomueller-net Nov 3, 2023

ingomueller-net commented Nov 6, 2023

ingomueller-net commented Nov 6, 2023

ingomueller-net commented Nov 8, 2023

ingomueller-net commented Dec 1, 2023

Update all dependencies. #12

Update all dependencies. #12

Conversation

ingomueller-net commented Nov 3, 2023

ingomueller-net left a 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

ingomueller-net commented Nov 6, 2023

ingomueller-net commented Nov 6, 2023

ingomueller-net commented Nov 8, 2023

ingomueller-net commented Dec 1, 2023