Add AMDGPU allocator support

Author: lkdvos

Project.toml

@@ -17,6 +17,7 @@ TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 
 [weakdeps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2"
@@ -26,6 +27,7 @@ JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 
 [extensions]
+TensorOperationsAMDGPUExt = "AMDGPU"
 TensorOperationsBumperExt = "Bumper"
 TensorOperationsChainRulesCoreExt = "ChainRulesCore"
 TensorOperationsMooncakeExt = "Mooncake"
@@ -66,7 +68,6 @@ julia = "1.10"
 
 [extras]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2"

ext/TensorOperationsAMDGPUExt.jl

@@ -0,0 +1,49 @@
+module TensorOperationsAMDGPUExt
+
+using AMDGPU
+using TensorOperations
+using TensorOperations: TensorOperations as TO
+
+#-------------------------------------------------------------------------------------------
+# Allocator
+#-------------------------------------------------------------------------------------------
+
+TO.tensoradd_type(TC, A::AnyRocArray, pA::Index2Tuple, conjA::Bool) =
+    ROCArray{TC, TO.numind(pA)}
+
+function TO.tensoralloc_add(
+        TC, A::AbstractArray, pA::Index2Tuple, conjA::Bool,
+        istemp::Val, allocator::TO.AMDAllocator
+    )
+    ttype = ROCArray{TC, TO.numind(pA)}
+    structure = TO.tensoradd_structure(A, pA, conjA)
+    return TO.tensoralloc(ttype, structure, istemp, allocator)::ttype
+end
+
+function TO.tensoralloc_contract(
+        TC,
+        A::AbstractArray, pA::Index2Tuple, conjA::Bool,
+        B::AbstractArray, pB::Index2Tuple, conjB::Bool,
+        pAB::Index2Tuple,
+        istemp::Val, allocator::TO.AMDAllocator
+    )
+    ttype = ROCArray{TC, TO.numind(pAB)}
+    structure = TO.tensorcontract_structure(A, pA, conjA, B, pB, conjB, pAB)
+    return TO.tensoralloc(ttype, structure, istemp, allocator)::ttype
+end
+
+# NOTE: the general implementation in the `DefaultAllocator` case works just fine, without
+# selecting an explicit memory model
+function TO.tensoralloc(
+        ::Type{<:ROCArray{T, N}}, structure,
+        ::Val{istemp}, allocator::TO.AMDAllocator
+    ) where {T, N}
+    return ROCArray{T, N}(undef, structure)
+end
+
+function TO.tensorfree!(C::ROCArray, ::TO.AMDAllocator)
+    AMDGPU.unsafe_free!(C)
+    return nothing
+end
+
+end

src/implementation/allocator.jl

@@ -30,6 +30,13 @@ parameters `Min`, `Mout`, `Mtemp` can be any of the CUDA.jl memory types, i.e.
 """
 struct CUDAAllocator{Mout, Min, Mtemp} end
 
+"""
+    AMDAllocator()
+
+Allocator that uses the AMD memory manager and will thus allocate `ROCArray` instances.
+"""
+struct AMDAllocator end
+
 """
     ManualAllocator()