[C++] copy_gpu a transposed array

[zcbenz]

I'm doing something like this inside eval_gpu:

    array trans = swapaxes(in, axis, -1, s); 
    array in_copy(trans.shape(), trans.dtype(), nullptr, {});
    copy_gpu(trans, in_copy, CopyType::General, s); 

which does not work and it took me quite a while to realize that it is because trans has not been evaluated yet and it does not contain any data.

So made it work by adding an metal::eval:

    array trans = swapaxes(in, axis, -1, s);
    metal::eval(trans);
    array in_copy(trans.shape(), trans.dtype(), nullptr, {});
    copy_gpu(trans, in_copy, CopyType::General, s); 

which however makes me nervous.

Is there a better way to do this?

(For context: I'm doing this to pass the array to cub::DeviceSegmentedSort, which only sorts the innermost dimension of contiguous input.)

[awni]

which does not work and it took me quite a while to realize that it is because trans has not been evaluated yet and it does not contain any data

Yes you can't call MLX operations from inside the eval functions.

So made it work by adding an metal::eval

I'm glad this makes you nervous because that's a bad idea, we shouldn't evaluate anything inside an eval.

Is there a better way to do this?

There isn't a super clean way to do it. But you can see for example here where we transpose an array inside the eval_cpu.

A potentially cleaner option is to refactor the actual impentation of Transpose::eval so that you can call it from inside your eval_gpu. A lot of ops work like that where we have the safe to call version like copy_gpu or binary_op_gpu.

You have to be extra careful with temporary arrays like in_copy to make sure you hold on to them until the computation is finished. In the Metal and CPU backend we have a way to register temporaries with the device so it holds them until the command is finished.

[zcbenz]

Thanks!