Now with tests in bazel

Author: wsmoses

.buildkite/pipeline.yml

@@ -62,7 +62,8 @@ steps:
       bazel --output_user_root=`pwd`/baztmp build :enzyme_ad
       cp bazel-bin/*.whl .
       python -m pip install *.whl
-      cd test && python -m pip install "jax[cpu]" && python test.py && python bench_vs_xla.py
+      python -m pip install "jax[cpu]"
+      bazel --output_user_root=`pwd`/baztmp test --test_output=errors ...
     artifact_paths:
       - "*.whl"
 

.github/workflows/build.yml

@@ -45,12 +45,7 @@ jobs:
       - name: test
         run: |
           python3 -m pip install --user --force-reinstall "jax[cpu]" *.whl
-          cd test
-          nm -C $(python3 -c "from enzyme_ad.jax import enzyme_call; print(enzyme_call.__file__)") | grep ExecutorCache::
-          python3 test.py
-          python3 bench_vs_xla.py
-          cd lit_tests
-          lit . --verbose
+          bazel test --test_output=errors ...
 
       - name: Upload Build
         uses: actions/upload-artifact@v3

test/BUILD

@@ -1,4 +1,31 @@
 load("@rules_python//python:py_test.bzl", "py_test")
+load("@llvm-project//llvm:lit_test.bzl", "package_path", "lit_test")
+
+
+exports_files(
+    ["lit_tests/lit.cfg.py"],
+    visibility = [":__subpackages__"],
+)
+
+[
+    lit_test(
+        name = "%s.test" % src,
+        srcs = [src],
+        data = [
+            "//test:lit_tests/lit.cfg.py",
+            "//src/enzyme_ad/jax:enzyme_jax_internal",
+            "@llvm-project//clang:builtin_headers_gen",
+            "@llvm-project//llvm:FileCheck",
+            "@llvm-project//llvm:count",
+            "@llvm-project//llvm:not",
+        ] + glob(["**/*.h"]),
+    )
+    for src in glob(
+        [
+            "**/*.pyt",
+        ],
+    )
+]
 
 py_test(
     name = "test",

test/lit_tests/ir.pyt

@@ -4,6 +4,7 @@ import jax
 import jax.numpy as jnp
 from enzyme_ad.jax import cpp_call
 
+argv = ("-I/usr/include/c++/11", "-I/usr/include/x86_64-linux-gnu/c++/11")
 
 def do_something(ones, twos):
     shape = jax.core.ShapedArray(tuple(3 * s for s in ones.shape), ones.dtype)
@@ -26,7 +27,7 @@ def do_something(ones, twos):
         }
         }
     }
-    """,
+    """, argv=argv,
         fn="myfn",
     )
     return a, b