Fix segfault (#35)

* Fix segfault

* add test build file

* Now with tests in bazel

* fix lit

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

WORKSPACE

@@ -60,8 +60,8 @@ load("@rules_python//python/pip_install:repositories.bzl", "pip_install_dependen
 
 pip_install_dependencies()
 
-ENZYME_COMMIT = "41e16dee0efd7f4c81e552e5899fe0068576573f"
-ENZYME_SHA256 = "fbfa6707db19b96ac8bd7f5edb58c4d028e11301873305fd7df7d7cd52abe66d"
+ENZYME_COMMIT = "e8ca2b1de3b770c767145d027b357bed97178bb0"
+ENZYME_SHA256 = "dd3789b8e749ed989d7a1a4956880826d9fd2501283dc60be468adec1186e63f"
 
 http_archive(
     name = "enzyme",

src/enzyme_ad/jax/enzyme_call.cc

@@ -972,7 +972,6 @@ class CpuKernel {
     void **outs = num_out > 1 ? reinterpret_cast<void **>(out) : &out;
     for (int i = 0; i < num_out; i++) {
       void *data = outs[i];
-      *(void **)(data) = 0;
     }
     auto fn = (void (*)(void **outs, void **ins))addr;
     fn(outs, ins);

test/BUILD

@@ -0,0 +1,65 @@
+load("@rules_python//python:py_test.bzl", "py_test")
+load("@llvm-project//llvm:lit_test.bzl", "package_path", "lit_test")
+load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
+
+expand_template(
+    name = "lit_site_cfg_py",
+    testonly = True,
+    out = "lit.site.cfg.py",
+    substitutions = {
+        "@LIT_SITE_CFG_IN_HEADER@": "# Autogenerated, do not edit.",
+        "@LLVM_TOOLS_BINARY_DIR@": package_path("@llvm-project//llvm:BUILD"),
+        "@LLVM_LIBS_DIR@": package_path("@llvm-project//llvm:BUILD"),
+        "@ENZYME_SOURCE_DIR@": "",
+        "@ENZYME_BINARY_DIR@": "",
+    },
+    template = "lit.site.cfg.py.in",
+    visibility = [":__subpackages__"],
+)
+
+exports_files(
+    ["lit.cfg.py"],
+    visibility = [":__subpackages__"],
+)
+
+[
+    lit_test(
+        name = "%s.test" % src,
+        srcs = [src],
+        data = [
+            ":lit.cfg.py",
+            ":lit_site_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",
+    srcs = [
+        "test.py",
+    ],
+    deps = [
+        "//src/enzyme_ad/jax:enzyme_jax_internal",
+    ],
+)
+
+py_test(
+    name = "bench_vs_xla",
+    srcs = [
+        "bench_vs_xla.py",
+    ],
+    deps = [
+        "//src/enzyme_ad/jax:enzyme_jax_internal",
+    ],
+)
+

test/lit_tests/lit.cfg.py renamed to test/lit.cfg.py

@@ -34,15 +34,7 @@
 
 # Tweak the PATH to include the tools dir and the scripts dir.
 base_paths = [
-    os.path.join(
-        os.path.dirname(__file__),
-        "..",
-        "..",
-        "bazel-bin",
-        "external",
-        "llvm-project",
-        "llvm",
-    ),
+    config.llvm_tools_dir,
     config.environment["PATH"],
 ]
 path = os.path.pathsep.join(base_paths)  # + config.extra_paths)

test/lit.site.cfg.py.in

@@ -0,0 +1,11 @@
+@LIT_SITE_CFG_IN_HEADER@
+
+import os
+config.llvm_tools_dir = "@LLVM_TOOLS_BINARY_DIR@"
+
+
+if len("@ENZYME_BINARY_DIR@") == 0:
+  config.llvm_tools_dir = os.getcwd() + "/" + config.llvm_tools_dir
+
+cfgfile = os.path.dirname(os.path.abspath(__file__)) + "/lit.cfg.py"
+lit_config.load_config(config, cfgfile)

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