How to implement in-place logic in a CUDA custom call ?

[Dong-Jiahuan]

I want to register a CUDA function into JAX and this CUDA function will do something to an operand which passed from python.
For an example , in python, I use custom_call('my_fun', operands=[BUFFER], ...). So, in C++, my_fun receives
( cudaStream_t stream,
void **buffers,
const char *opaque,
std::size_t opaque_len)

I get BUFFER from buffers[0] and buffers[1] is the output.

My CUDA function wants do something to BUFFER and set BUFFER as the last element of buffers. (I think this is where the problem lies in.)

This example is not my real case and I didn't test this example. In real case, BUFFER is one of many operands and I still want to implement in-place logic to BUFFER.

After I register my real CUDA function to JAX and use it, I got an error :

jaxlib.mlir._mlir_libs.MLIRError: Verification failed:
error: "XXXXXXXX"("XXXXXXXX.py":46:0): 'stablehlo.custom_call' op incorrect layout dense<0> : tensor<1xindex> for type 'tensor', layout must be a permutation of [0, 0)