Repository: flatironinstitute/jax-finufft (language: python)

Issue:
Crash during XLA compilation (and GPU runtime crash)
@Matematija has found a bug (or bugs) triggered by the following code:

```python
from functools import partial
import numpy as np
import jax
from jax import numpy as jnp
from jax_finufft import nufft1, nufft2

@jax.jit
@partial(jax.vmap, in_axes=(0, 0, None))
def cconv_test(f, xs, kernel):

    # f.shape = (n_grid, in_features)
    # x.shape = (n_grid, ndim)
    # kernel.shape = (*k_grid, in_features, out_features)

    ndim = xs.shape[-1]
    k_grid_shape = kernel.shape[:-2]

    f_ = f.transpose().astype(complex)
    coords = [xs[...,i] for i in range(ndim)]

    f_hat = nufft1(k_grid_shape, f_, *coords, iflag=-1)
    c_hat = jnp.einsum('a...,...ab->b...', f_hat, kernel)
    return nufft2(c_hat, *coords, iflag=1)

kernel = jnp.array(np.random.randn(32, 32, 32, 16, 16))
f = jnp.array(np.random.randn(8, 100, 16))
x = jnp.array(np.random.uniform(low=-np.pi, high=np.pi, size=(8, 100, 3)))

fconv = cconv_test(f, x, kernel) # crash!
```

It crashes both the CPU and GPU versions, but in very different ways, so I'll focus on CPU for now. The error is:

```console
$ JAX_PLATFORMS=cpu python scripts/jittest.py
[external/com_google_absl/absl/status/statusor.cc : 83] RAW: Attempting to fetch value instead of handling error INTERNAL: Could not convert JSON string to proto: Encountered non UTF-8 code points.
�7�5����
Aborted (core dumped)
```

The error is pretty sensitive to having `nufft2` using the output of `nufft1`, and using `jnp.einsum`, and `vmap`. There are lots of ways to tweak this function that make the crash go away, but it's not clear if doing so is just hiding an underlying problem.

The GDB traceback for this error suggests that the crash happens inside the XLA compilation, before execution even begins. Here's a partial traceback:
```
#0  0x00007ffff66baacf in raise () from /lib64/libc.so.6
#1  0x00007ffff668dea5 in abort () from /lib64/libc.so.6
#2  0x00007fffe2ed3d1c in absl::lts_20230802::raw_log_internal::(anonymous namespace)::RawLogVA(absl::lts_20230802::LogSeverity, char const*, int, char const*, __va_list_tag*) ()
   from /mnt/home/lgarrison/jax-finufft/venv-matija/lib/python3.10/site-packages/jaxlib/xla_extension.so
#3  0x00007fffe2ed3dae in absl::lts_20230802::raw_log_internal::RawLog(absl::lts_20230802::LogSeverity, char const*, int, char const*, ...) () from /mnt/home/lgarrison/jax-finufft/venv-matija/lib/python3.10/site-packages/jaxlib/xla_extension.so
#4  0x00007fffe2e95ffc in absl::lts_20230802::internal_statusor::Helper::Crash(absl::lts_20230802::Status const&) ()
   from /mnt/home/lgarrison/jax-finufft/venv-matija/lib/python3.10/site-packages/jaxlib/xla_extension.so
#5  0x00007fffdb210c7c in xla::cpu::EmitDotOperation(xla::HloInstruction const&, xla::llvm_ir::IrArray const&, xla::llvm_ir::IrArray const&, xla::llvm_ir::IrArray const&, xla::llvm_ir::IrArray const*, llvm::Value*, llvm::IRBuilder<llvm::ConstantFolder, llvm::IRBuilderDefaultInserter>*, mlir::MLIRContext*, xla::HloModuleConfig const&, xla::cpu::TargetMachineFeatures const&) ()
   from /mnt/home/lgarrison/jax-finufft/venv-matija/lib/python3.10/site-packages/jaxlib/xla_extension.so
#6  0x00007fffdb1ed68b in xla::cpu::IrEmitter::HandleDot(xla::HloInstruction*) ()
   from /mnt/home/lgarrison/jax-finufft/venv-matija/lib/python3.10/site-packages/jaxlib/xla_extension.so
#7  0x00007fffe2b98600 in absl::lts_20230802::Status xla::HloInstruction::Visit<xla::HloInstruction*>(xla::DfsHloVisitorBase<xla::HloInstruction*>*) () from /mnt/home/lgarrison/jax-finufft/venv-matija/lib/python3.10/site-packages/jaxlib/xla_extension.so
#8  0x00007fffdacaf158 in absl::lts_20230802::Status xla::HloComputation::AcceptOrdered<xla::HloInstruction*>(xla::DfsHloVisitorBase<xla::HloInstruction*>*, absl::lts_20230802::Span<xla::HloInstruction* const>) const ()
   from /mnt/home/lgarrison/jax-finufft/venv-matija/lib/python3.10/site-packages/jaxlib/xla_extension.so
#9  0x00007fffdb1f439c in xla::cpu::IrEmitter::EmitComputation(xla::HloComputation*, std::basic_string_view<char, std::char_traits<char> >, bool, absl::lts_20230802::Span<xla::HloInstruction* const>, bool) ()
   from /mnt/home/lgarrison/jax-finufft/venv-matija/lib/python3.10/site-packages/jaxlib/xla_extension.so
#10 0x00007fffdb1051ad in xla::cpu::CpuCompiler::CompileLegacyCpuExecutable(std::unique_ptr<xla::HloModule, std::default_delete<xla::HloModule> >) () from /mnt/home/lgarrison/jax-finufft/venv-matija/lib/python3.10/site-packages/jaxlib/xla_extension.so
```

Interestingly, error bisects to ec080e48b69558c7cbb5e798f0bc83f37d1c3b0a, where the XLA translation rules were replaced with MLIR lowering rules. So I think it's possible that a bug was introduced there.

But a counterpoint to that reasoning is that this script also crashes with the CUDA backend, and it still crashes even before that commit! However, the crash is very different than the CPU crash:
```console
$ JAX_PLATFORMS=cuda python scripts/jittest.py
2023-11-17 15:59:54.546455: W external/xla/xla/service/gpu/runtime/support.cc:58] Intercepted XLA runtime error:
UNKNOWN: nufft3d2f XLA extension have thrown an exception: exclusive_scan failed to synchronize: cudaErrorIllegalAddress: an illegal memory access was encountered
2023-11-17 15:59:54.546537: E external/xla/xla/pjrt/pjrt_stream_executor_client.cc:2716] Execution of replica 0 failed: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.custom_call' failed: nufft3d2f XLA extension have thrown an exception: exclusive_scan failed to synchronize: cudaErrorIllegalAddress: an illegal memory access was encountered; current tracing scope: custom-call.6; current profiling annotation: XlaModule:#hlo_module=jit_cconv_test,program_id=5#.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/mnt/home/lgarrison/jax-finufft/scripts/jittest.py", line 51, in <module>
    fconv = cconv_test(f, x, kernel) # crash!
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.custom_call' failed: nufft3d2f XLA extension have thrown an exception: exclusive_scan failed to synchronize: cudaErrorIllegalAddress: an illegal memory access was encountered; current tracing scope: custom-call.6; current profiling annotation: XlaModule:#hlo_module=jit_cconv_test,program_id=5#.
```

This appears to be a runtime, GPU-side crash, as opposed to the CPU XLA compilation crash. I can catch the error in `cuda-gdb`, so I'll continue to debug on that front.

So to summarize:
1) Maybe ec080e48b69558c7cbb5e798f0bc83f37d1c3b0a introduced a lowering rules bug that's triggering an XLA compilation crash, and there's also different GPU runtime bug.
2) Maybe the bugs have a common source, and ec080e48b69558c7cbb5e798f0bc83f37d1c3b0a simply moved things around enough to start triggering a CPU bug

In terms of a "common source", I think it would have to be something in the reproducer script (like bad shapes that aren't caught properly) or a JAX bug (e.g. with `einsum`). Because the XLA crash happens so early, I think we could reproduce this outside of jax-finufft, with no compiled library code. I'll have to test this to be sure, though.

@dfm Do you know if a bug in the lowering rules could cause an XLA compilation crash like this?


Discussion / hints:
Oh wow, this is much deeper/weirder than I thought. Thank you for putting so much effort into this.

Fix patch:
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 7a15776..21e6615 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -2,6 +2,9 @@ cmake_minimum_required(VERSION 3.15...3.26)
 project(${SKBUILD_PROJECT_NAME} LANGUAGES C CXX)
 message(STATUS "Using CMake version: " ${CMAKE_VERSION})
 
+# for cuda-gdb and verbose PTXAS output
+# set(CMAKE_CUDA_FLAGS ${CMAKE_CUDA_FLAGS} "-g -G -Xptxas -v")
+
 # Workaround for LTO applied incorrectly to CUDA fatbin
 # https://github.com/pybind/pybind11/issues/4825
 set(CMAKE_INTERPROCEDURAL_OPTIMIZATION OFF)
diff --git a/lib/kernels.cc.cu b/lib/kernels.cc.cu
index ef4fcf5..38ee6fc 100644
--- a/lib/kernels.cc.cu
+++ b/lib/kernels.cc.cu
@@ -22,10 +22,11 @@ void run_nufft(int type, const NufftDescriptor<T> *descriptor, T *x, T *y, T *z,
   makeplan<T>(type, ndim, descriptor->n_k, descriptor->iflag, descriptor->n_transf,
               descriptor->eps, &plan, opts);
   for (int64_t index = 0; index < descriptor->n_tot; ++index) {
+    int64_t i = index * descriptor->n_j;
     int64_t j = index * descriptor->n_j * descriptor->n_transf;
     int64_t k = index * n_k * descriptor->n_transf;
 
-    setpts<T>(plan, descriptor->n_j, &(x[j]), y_index<ndim, T>(y, j), z_index<ndim, T>(z, j), 0,
+    setpts<T>(plan, descriptor->n_j, &(x[i]), y_index<ndim, T>(y, i), z_index<ndim, T>(z, i), 0,
               NULL, NULL, NULL);
 
     execute<T>(plan, &c[j], &F[k]);
diff --git a/pyproject.toml b/pyproject.toml
index 3da8250..1b818ea 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -23,6 +23,11 @@ wheel.install-dir = "jax_finufft"
 minimum-version = "0.5"
 build-dir = "build/{wheel_tag}"
 
+# For debugging:
+# cmake.build-type = "Debug"
+# cmake.verbose = true
+# install.strip = false
+
 [tool.setuptools_scm]
 version_file = "src/jax_finufft/jax_finufft_version.py"
 
diff --git a/vendor/finufft b/vendor/finufft
index d67dd92..968f90d 160000
--- a/vendor/finufft
+++ b/vendor/finufft
@@ -1,1 +1,1 @@
-Subproject commit d67dd92dea19912c631c069f3969184172bab52e
+Subproject commit 968f90d1214c9308a5ef2a647ad4ebb5ad8dffeb


Test patch:
diff --git a/tests/ops_test.py b/tests/ops_test.py
index 4a6d71a..fa88af0 100644
--- a/tests/ops_test.py
+++ b/tests/ops_test.py
@@ -312,10 +312,6 @@ def norm_nufft2(c, x):
 
 
 def test_issue37():
-    if jax.default_backend() != "cpu":
-        pytest.xfail("TODO: this test currently fails on the GPU")
-
-    @jax.jit
     @partial(jax.vmap, in_axes=(0, 0, None))
     def cconv_test(f, xs, kernel):
         # f.shape = (n_grid, in_features)
@@ -325,7 +321,9 @@ def cconv_test(f, xs, kernel):
         ndim = xs.shape[-1]
         k_grid_shape = kernel.shape[:-2]
 
-        f_ = f.transpose().astype(complex)
+        f_ = f.astype(
+            {np.float32: np.complex64, np.float64: np.complex128}[f.dtype.type]
+        ).transpose()
         coords = [xs[..., i] for i in range(ndim)]
 
         f_hat = nufft1(k_grid_shape, f_, *coords, iflag=-1)
@@ -336,4 +334,10 @@ def cconv_test(f, xs, kernel):
     f = jnp.array(np.random.randn(8, 100, 16))
     x = jnp.array(np.random.uniform(low=-np.pi, high=np.pi, size=(8, 100, 3)))
 
-    cconv_test(f, x, kernel)
+    a = cconv_test(f, x, kernel)
+    b = jax.jit(cconv_test)(f, x, kernel)
+
+    assert np.isfinite(a).all()
+    assert np.isfinite(b).all()
+
+    check_close(a, b)