Is the iteration of the jax version affecting the CUDA code?

[jing-alice]

I used my own cuda code to implement function similar to jax.at[].add(), it is worked for jax==0.4.8, jaxlib==0.4.7+cuda11.cudnn86.
This is my code:

import jax.numpy as jnp
import time
from neurailib import index_add
import jax
from jax import lax, jit
from functools import partial

size1 = 15
size2 = 10
arr = jnp.zeros((size1, size2))
index = jnp.array([0, 0, 1, 1, size1 - 1])
temp = jnp.ones(size2) * 1.0
temp1 = jnp.array([temp, temp, temp, temp, temp])
num_iterations = 700
temp_input = jnp.zeros((num_iterations, size2))
temp_input = temp_input.at[0].set(jnp.ones(size2))
start_at_time = time.time()


@partial(jit, static_argnums=1)
def update_arr(arr, num_iterations, temp_input):
  for i in range(num_iterations):
    arr = arr.at[i].add(temp_input[i])
    arr = index_add(arr, temp1, index)
  return arr


arr = update_arr(arr, num_iterations, temp_input)
print("new arr: ", arr)

and the result is:

However, when jax==0.4.13，jaxlib==0.4.13+cuda12.cudnn8, the result is incorrect and random, The following two figures show the result of running the code twice:

And this is my CUDA code:

#include "gpu_index_add.h"
#include <stdio.h>
#include <iostream>
#include <cuda_runtime.h>
using namespace std;

namespace neurai_lib
{

	namespace
	{
		template <typename T, typename I>
		__global__ void IndexAddKernel(
				T *result,
				T *input,
				const T *conn,
				const I *conn_index,
				const std::uint32_t num_rows,
				const std::uint32_t num_cols,
				const std::uint32_t num_conn)
		{

			const unsigned int ix = threadIdx.x + blockIdx.x * blockDim.x;
			const unsigned int iy = threadIdx.y + blockIdx.y * blockDim.y;
			const unsigned int idx = iy * num_cols + ix;

			if (iy < num_conn && ix < num_cols && conn_index[iy] < num_rows)
			{
				atomicAdd(&input[conn_index[iy] * num_cols + ix], conn[iy * num_cols + ix]);
			}

		}

		template <typename T, typename I>
		__global__ void IndexAddKernel1(
				T *result,
				T *input,
				const T *conn,
				const I *conn_index,
				const std::uint32_t num_rows,
				const std::uint32_t num_cols,
				const std::uint32_t num_conn)
		{
			const unsigned int ix = threadIdx.x + blockIdx.x * blockDim.x;
			const unsigned int iy = threadIdx.y + blockIdx.y * blockDim.y;
			const unsigned int idx = iy * num_cols + ix;
			if (ix < num_cols && iy < num_rows)
			{
				result[idx] = input[idx];
			}
		}


		template <typename T, typename I>
		inline void index_add(cudaStream_t stream, void **buffers, const char *opaque,
													std::size_t opaque_len)
		{
			const IndexAddDescriptor &d = *UnpackDescriptor<IndexAddDescriptor>(opaque, opaque_len);
			const std::uint32_t num_rows = d.num_rows;
			const std::uint32_t num_cols = d.num_cols;
			const std::uint32_t num_conn = d.num_conn;

			T *input = reinterpret_cast<T *>(buffers[0]);
			const T *conn = reinterpret_cast<const T *>(buffers[1]);
			const I *conn_index = reinterpret_cast<const I *>(buffers[2]);

			T *result = reinterpret_cast<T *>(buffers[3]);
			cudaMemset(result, 0, sizeof(T) * num_rows * num_cols);
			unsigned int block_size = 32;
			dim3 block(block_size, block_size);
			dim3 grid((num_cols + block.x - 1) / block.x, (num_rows + block.y - 1) / block.y);
			IndexAddKernel<T, I><<<grid, block>>>(
					result,
					input,
					conn,
					conn_index,
					num_rows,
					num_cols, num_conn);
			cudaDeviceSynchronize();
			IndexAddKernel1<T, I><<<grid, block>>>(
					result,
					input,
					conn,
					conn_index,
					num_rows,
					num_cols, num_conn);

			cudaDeviceSynchronize();
		}

	} // namespace

	// Index Add.
	pybind11::bytes build_index_add_descriptor(std::uint32_t num_rows, std::uint32_t num_cols, std::uint32_t num_conn)
	{
		return PackDescriptor(IndexAddDescriptor{num_rows, num_cols, num_conn});
	}

	void gpu_index_add(cudaStream_t stream, void **buffers, const char *opaque, std::size_t opaque_len)
	{
		index_add<float, std::uint32_t>(stream, buffers, opaque, opaque_len);
	}
}

This is the code how to use CUDA in JAX:

 from functools import partial
from jax import core
from jax.interpreters import mlir, xla
from jaxlib.hlo_helpers import custom_call
import jaxlib.mlir.ir as ir
from .register_custom_calls import default_layouts, gpu_ops
from jax._src import abstract_arrays
import numpy as np

__all__ = ["index_add"]

# Create index_add_p for index_add operation.
_index_add_p = core.Primitive('index_add')
_index_add_p.multiple_results = False
_index_add_p.def_impl(partial(xla.apply_primitive, _index_add_p))


def index_add(input, conn, conn_index):
  output = _index_add_p.bind(input, conn, conn_index)
  return output


def _index_add_cuda_lowing(ctx, input, conn, conn_index):
  input_type = ir.RankedTensorType(input.type)
  input_shape = input_type.shape
  conn_type = ir.RankedTensorType(conn.type)
  conn_shape = conn_type.shape
  conn_index_type = ir.RankedTensorType(conn_index.type)
  conn_index_shape = conn_index_type.shape

  num_rows = np.array(input_shape[0], dtype=np.uint32)
  num_cols = np.array(input_shape[1], dtype=np.uint32)
  num_conn = np.array(conn_index_shape[0], dtype=np.uint32)

  if gpu_ops is None:
    raise ValueError('Cannot find compiled gpu wheels. Please install neurailib setup_cuda.')

  opaque = gpu_ops.build_index_add_descriptor(num_rows, num_cols, num_conn)
  out = custom_call(
    b"gpu_index_add",
    out_types=[ir.RankedTensorType.get(input_shape, input_type.element_type)],
    operands=[input, conn, conn_index],
    backend_config=opaque,
    operand_layouts=default_layouts(input_shape, conn_shape, conn_index_shape),
    result_layouts=default_layouts(input_shape))
  return out


mlir.register_lowering(_index_add_p, _index_add_cuda_lowing, platform='gpu')

#######################
# Abstract evaluation #
#######################


def _index_add_abstract(input, conn, conn_index):
  return abstract_arrays.ShapedArray(input.shape, input.dtype)


_index_add_p.def_abstract_eval(_index_add_abstract)

[hawkinsp]

I think the issue is that you are updating the input array, but you should be updating only output array instead. You must not alter inputs to the custom call, with one exception below. If you do update the input in any way the behavior is undefined.

If you want the input and output to be the same array so you can do updates in-place, use the input_output_aliases argument to custom_call.

(I should note that jaxlib.hlo_helpers is technically an internal function, but there aren't great alternatives available yet.)

[jing-alice]

Thank you for answering my question. However, I still have doubts. When jax jax==0.4.8, jaxlib==0.4.7+cuda11.cudnn86, the code runs correctly.
When jax==0.4.13, jaxlib==0.4.13+cuda12.cudnn8, I tried something else. To avoid the affect of jit, I moved arr into the function body. This is the code.

import jax.numpy as jnp
import time
from neurailib import index_add
import jax
from jax import lax, jit
from functools import partial

size1 = 15
size2 = 10
index = jnp.array([0, 0, 1, 1, size1 - 1])
temp = jnp.ones(size2) * 1.0
temp1 = jnp.array([temp, temp, temp, temp, temp])
num_iterations = 700
temp_input = jnp.zeros((num_iterations, size2))
temp_input = temp_input.at[0].set(jnp.ones(size2))
start_at_time = time.time()


@partial(jit, static_argnums=0)
def update_arr(num_iterations, temp_input):
  arr = jnp.zeros((size1, size2))
  for i in range(num_iterations):
    arr = arr.at[i].add(temp_input[i])
    jax.debug.print("arr: {}", arr)
    arr = index_add(arr, temp1, index)
  return arr


arr = update_arr(num_iterations, temp_input)
print("new arr: ", arr)

But then something magical happened, why I added jax.debug.print("arr: {}", arr), in figure 1 the result is correct. When commenting out, in figure 2, the result is random and incorrect again.

figure1:

figure 2：