Add AMDGPU and JLArrays buffer support

[lkdvos]

Follow-up to #293: wires ROCArray and JLArray into the storage-generic BufferAllocator, via AMDBufferAllocator and JLBufferAllocator.

The JLArray backing is half the point — it makes the foreign-storage paths testable without GPU hardware, which they previously weren't.

Note on the implementation

Neither package can serve temporaries the way the CUDA extension does. AMDGPU.unsafe_wrap queries hipPointerGetAttributes on every call — the per-temporary host round-trip the buffer exists to avoid — and would re-hipHostRegister a sub-range of already-pinned memory for host buffers; JLArrays has no unsafe_wrap at all.

Both instead share the buffer's refcounted DataRef at an offset, as their own reshape does. No API call, and the buffer stays alive by refcount.

One wrinkle: JLArray offsets are in elements, not bytes, so an element type larger than the padding alignment would land on a truncated offset and alias a previous temporary. buffer_arraytype returns nothing for those, falling back to a regular allocation.

Testing

• New JLArray testset in test/allocator.jl (134 tests) runs in ordinary CI.
• Shared BufferAllocator ($AT) testset in test/gpu.jl covers JLArray, CuArray and ROCArray, so buildkite's rocm queue exercises the AMD path.
• Full suite passes locally (CUDA + JLArray); AMD verified only by precompilation, dispatch and inference checks — no ROCm device on hand, so the rocm queue is the real check.

🤖 Generated with Claude Code

[kshyatt]

Not sure about the name AMDBufferAllocator, what about ROCBufferAllocator for consistency with the array type names?

[lkdvos]

I was having this exact debate with myself, but it's inconsistent no matter what: we have AMDAllocator and CUDAAllocator for the other allocator structs, and also CUDABufferAllocator 😢. If you have a strong preference, I can definitely go for CuBufferAllocator and ROCBufferAllocator, but I can't change the AMDAllocator and CUDAAllocator without breaking changes...

(

TensorOperations.jl/src/implementation/allocator.jl

Lines 13 to 38 in 58c6e03

	"""
	CUDAAllocator{Mout,Min,Mtemp}()
	Allocator that uses the CUDA memory manager and will thus allocate `CuArray` instances. The
	parameters `Min`, `Mout`, `Mtemp` can be any of the CUDA.jl memory types, i.e.
	`CUDA.DeviceMemory`, `CUDA.UnifiedMemory` or `CUDA.HostMemory`.
	* `Mout` is used to determine how to deal with output tensors; with `Mout=CUDA.HostMemory`
	or `Mout=CUDA.UnifiedMemory` the CUDA runtime will ensure that the data is also available
	at in the host memory, and can thus be converted back to normal arrays using
	`unsafe_wrap(Array, outputtensor)`. If `Mout=CUDA.DeviceMemory` the data will remain on
	the GPU, untill an explict `Array(outputtensor)` is called.
	* `Min` is used to determine how to deal with input tensors; with `Min=CUDA.HostMemory` the
	CUDA runtime will itself take care of transferring the data to the GPU, otherwise it is
	copied explicitly.
	* `Mtemp` is used to allocate space for temporary tensors; it defaults to
	`CUDA.default_memory` which is `CUDA.DeviceMemory`. Only if many or huge temporary tensors
	are expected could it be useful to choose `CUDA.UnifiedMemory`.
	"""
	struct CUDAAllocator{Mout, Min, Mtemp} end
	"""
	AMDAllocator()
	Allocator that uses the AMD memory manager and will thus allocate `ROCArray` instances.
	"""
	struct AMDAllocator end

)

[kshyatt]

ah ok then I don't feel strongly enough to break stuff, AMDBufferAllocator it is

[lkdvos]

great, will merge once the AMD tests turn green then. (Can you approve the PR so it automerges?)

[codecov]

Codecov Report

❌ Patch coverage is 96.72131% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
ext/TensorOperationsAMDGPUExt.jl	93.10%	2 Missing ⚠️

Files with missing lines	Coverage Δ
ext/TensorOperationsJLArraysExt.jl	100.00% <100.00%> (ø)
src/implementation/allocator.jl	95.68% <ø> (+0.03%)	⬆️
ext/TensorOperationsAMDGPUExt.jl	65.11% <93.10%> (+65.11%)	⬆️

... and 5 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.