mapreducedim! is not implemented for AnyROCArray Types

[matinraayai]

Hello,
GPUArray.mapreducedim! should be implemented for AnyROCArray instead of the more specific type of ROCArray. Is there a way for AnyROCArray to be adapted to ROCArray, so that this function is implemented for AnyROCArray?
Thanks

[jpsamaroo]

As a stop-gap we can just convert to a dense ROCArray and then switch back the type after. @maleadt any recommendations?

[maleadt]

Why not implement mapreducedim! for AnyROCArray? Works for CUDA: https://github.com/JuliaGPU/CUDA.jl/blob/0060ea6f29b7dcfc5563dae868c11feb5d26281e/src/mapreduce.jl#L169-L171=, as these types are 'guaranteed' to work in a kernel.

[matinraayai]

@maleadt should changing the function prototype to accept AnyROCArray work directly in theory?

[maleadt]

I think so, but I'm not very familiar with AMDGPU.jl or it's mapreduce implementation. Try it out and see what breaks?

[matinraayai]

@maleadt @jpsamaroo I've tried it before, but the problem is that the AMDGPU implementation queries the ROCm device associated with the ROCArray using its buffer field here:

AMDGPU.jl/src/mapreduce.jl

Line 133 in 99967b7

agent = R.buf.agent

I have to look into it, but I don't think it will be as simple as changing the function prototype.

[maleadt]

Then JuliaGPU/Adapt.jl#52 would probably be useful.

[matinraayai]

@maleadt it looks good, but do you have an ETA on when it will merge?
Also it seems like the CUDA version also queries the current thread's device, which might not correspond to the actual device the array resides on. I think when merged, we can use this feature to fix bugs in both CUDA.jl and AMDGPU.jl.
https://github.com/JuliaGPU/CUDA.jl/blob/c1ef162c06d78c4b5cf31914fa1acf7db34eb7f5/src/mapreduce.jl#L174

[maleadt]

No ETA on that change. You should probably use an agent function instead and define it on a couple of common wrappers, like oneAPI.jl does for its device function: https://github.com/JuliaGPU/oneAPI.jl/blob/9a43fb6fbf0ef20ffba34daef4b7ff44e7ed01c8/src/array.jl#L370-L371=

Also it seems like the CUDA version also queries the current thread's device, which might not correspond to the actual device the array resides on.

It's currently up to the user to make sure to use the correct array on its respective device. I don't think I'll make it automatic, since it's not clear what to do if multiple devices are involved. It's also possible to create arrays whose memory is accessible from other devices.

[matinraayai]

@maleadt @jpsamaroo are SubArrays supposed to be a type of AnyROCArray or AnyCuArray? One thing I noticed is that this function doesn't get dispatched with a type of SubArray, backed by a ROCArray parent.

[maleadt]

Ah yeah, we disabled those subtypes in Adapt.jl (which the Any* types are based on): https://github.com/JuliaGPU/Adapt.jl/blob/d9f852a61ee42258e543f5ff3b28e1f3abbb48bc/src/wrappers.jl#L86-L90=. That's because Adapt.jl used to be only used by CUDA.jl, where CuArray is used to encode views and reshapes (i..e, without using the array wrappers). Either we make it 'mandatory' for users of the Any* union types to handle SubArray/Reinterpret/Reshape themselves, or we should generalize Adapt.jl

[matinraayai]

@maleadt do CuSubArrays currently count as AnyCuArrays? Can I do the same thing CUDA.jl is doing to make AMDGPU support this type hierarchy? Ideally I want it in Adapt.jl because it seems useful but for now I need the workaround for our project.

[maleadt]

There is no CuSubArray, CuArray itself is used to represent contiguous views (it encodes an offset argument for that), reshapes and reinterpretations. That would probably be useful for AMDGPU.jl too (because it simplifies dispatching to vendor libraries), but would be up to @jpsamaroo to decide.

[jpsamaroo]

There is no CuSubArray, CuArray itself is used to represent contiguous views (it encodes an offset argument for that), reshapes and reinterpretations. That would probably be useful for AMDGPU.jl too (because it simplifies dispatching to vendor libraries), but would be up to @jpsamaroo to decide.

Yeah that sounds like the reasonable approach to me.