AMDGPU updates

Author: Katharine Hyatt

Project.toml

@@ -27,6 +27,7 @@ cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 [sources]
 GPUArrays = {rev = "ksh/more_diag", url = "https://github.com/JuliaGPU/GPUArrays.jl"}
 MatrixAlgebraKit = {rev = "ksh/tk", url = "https://github.com/QuantumKitHub/MatrixAlgebraKit.jl"}
+AMDGPU = {rev = "ksh/diag_norm", url = "https://github.com/JuliaGPU/AMDGPU.jl"}
 cuTENSOR = {subdir = "lib/cutensor", url = "https://github.com/JuliaGPU/CUDA.jl"}
 
 [extensions]
@@ -65,6 +66,7 @@ julia = "1.10"
 
 [extras]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -80,4 +82,4 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 
 [targets]
-test = ["Adapt", "Aqua", "Combinatorics", "CUDA", "cuTENSOR", "GPUArrays", "LinearAlgebra", "TensorOperations", "Test", "TestExtras", "ChainRulesCore", "ChainRulesTestUtils", "FiniteDifferences", "Zygote"]
+test = ["Adapt", "AMDGPU", "Aqua", "Combinatorics", "CUDA", "cuTENSOR", "GPUArrays", "LinearAlgebra", "TensorOperations", "Test", "TestExtras", "ChainRulesCore", "ChainRulesTestUtils", "FiniteDifferences", "Zygote"]

ext/TensorKitAMDGPUExt/TensorKitAMDGPUExt.jl

@@ -1,6 +1,6 @@
 module TensorKitAMDGPUExt
 
-using AMDGPU, AMDGPU.CUBLAS, LinearAlgebra
+using AMDGPU, AMDGPU.rocBLAS, LinearAlgebra
 using AMDGPU: @allowscalar
 import AMDGPU: rand as rocrand, rand! as rocrand!, randn as rocrandn, randn! as rocrandn!
 
@@ -16,7 +16,7 @@ using Random
 
 include("roctensormap.jl")
 
-const ROCDiagonalTensorMap{T, S} = DiagonalTensorMap{T, S, ROCVector{T, AMDGPU.DeviceMemory}}
+const ROCDiagonalTensorMap{T, S} = DiagonalTensorMap{T, S, ROCVector{T, AMDGPU.Mem.HIPBuffer}}
 
 """
     ROCDiagonalTensorMap{T}(undef, domain::S) where {T,S<:IndexSpace}
@@ -44,6 +44,10 @@ function ROCDiagonalTensorMap(data::ROCVector{T}, V::S) where {T, S}
     return ROCDiagonalTensorMap{T, S}(data, V)
 end
 
+function ROCDiagonalTensorMap(data::Vector{T}, V::S) where {T, S}
+    return ROCDiagonalTensorMap{T, S}(ROCVector{T}(data), V)
+end
+
 function TensorKit.Factorizations.MAK.initialize_output(::typeof(svd_full!), t::ROCDiagonalTensorMap, alg::DiagonalAlgorithm)
     V_cod = fuse(codomain(t))
     V_dom = fuse(domain(t))

ext/TensorKitAMDGPUExt/roctensormap.jl

@@ -1,11 +1,11 @@
-const ROCTensorMap{T,S,N₁,N₂} = TensorMap{T,S,N₁,N₂, ROCVector{T,AMDGPU.DeviceMemory}}
+const ROCTensorMap{T,S,N₁,N₂} = TensorMap{T,S,N₁,N₂, ROCVector{T,AMDGPU.Mem.HIPBuffer}}
 const ROCTensor{T, S, N} = ROCTensorMap{T, S, N, 0}
 
 const AdjointROCTensorMap{T,S,N₁,N₂} = AdjointTensorMap{T,S,N₁,N₂,ROCTensorMap{T,S,N₁,N₂}}
 
 function TensorKit.tensormaptype(S::Type{<:IndexSpace}, N₁, N₂, TorA::Type{<:StridedROCArray})
     if TorA <: ROCArray
-        return TensorMap{eltype(TorA),S,N₁,N₂,ROCVector{eltype(TorA), AMDGPU.DeviceMemory}}
+        return TensorMap{eltype(TorA),S,N₁,N₂,ROCVector{eltype(TorA), AMDGPU.Mem.HIPBuffer}}
     else
         throw(ArgumentError("argument $TorA should specify a scalar type (`<:Number`) or a storage type `<:ROCVector{<:Number}`"))
     end
@@ -68,6 +68,10 @@ function ROCTensorMap(data::AbstractDict{<:Sector,<:ROCMatrix}, codom::TensorSpa
     return ROCTensorMap(data, codom ← dom)
 end
 
+function ROCTensorMap(ts::TensorMap{T, S, N₁, N₂, A}) where {T, S, N₁, N₂, A}
+    return ROCTensorMap{T, S, N₁, N₂}(ROCArray(ts.data), ts.space)
+end
+
 for (fname, felt) in ((:zeros, :zero), (:ones, :one))
     @eval begin
         function AMDGPU.$fname(codomain::TensorSpace{S},
@@ -87,7 +91,7 @@ for (fname, felt) in ((:zeros, :zero), (:ones, :one))
     end
 end
 
-for randfun in (:curand, :curandn)
+for randfun in (:rocrand, :rocrandn)
     randfun! = Symbol(randfun, :!)
     @eval begin
         # converting `codomain` and `domain` into `HomSpace`
@@ -171,7 +175,7 @@ vi_scalartype(::Type{<:ROCTensorMap{T}}) where {T} = T
 vi_scalartype(::Type{<:ROCArray{T}}) where {T} = T
 
 function TensorKit.similarstoragetype(TT::Type{<:ROCTensorMap{TTT,S,N₁,N₂}}, ::Type{T}) where {TTT,T,S,N₁,N₂}
-    return ROCVector{T, AMDGPU.DeviceMemory}
+    return ROCVector{T, AMDGPU.Mem.HIPBuffer}
 end
 
 function Base.convert(TT::Type{ROCTensorMap{T,S,N₁,N₂}},

src/TensorKit.jl

@@ -80,7 +80,7 @@ export left_orth, right_orth, left_null, right_null,
     eigh_full!, eigh_full, eigh_trunc!, eigh_trunc, eig_full!, eig_full, eig_trunc!,
     eig_trunc,
     eigh_vals!, eigh_vals, eig_vals!, eig_vals,
-    isposdef, isposdef!, ishermitian, isisometry, isunitary, sylvester, rank, cond
+    isposdef, isposdef!, ishermitian, isisometric, isunitary, sylvester, rank, cond
 
 export braid, braid!, permute, permute!, transpose, transpose!, twist, twist!, repartition,
     repartition!

src/factorizations/adjoint.jl

@@ -29,11 +29,11 @@ function MAK.right_null!(t::AdjointTensorMap, N, alg::AbstractAlgorithm)
     return N
 end
 
-function MAK.is_left_isometry(t::AdjointTensorMap; kwargs...)
-    return is_right_isometry(adjoint(t); kwargs...)
+function MAK.is_left_isometric(t::AdjointTensorMap; kwargs...)
+    return is_right_isometric(adjoint(t); kwargs...)
 end
-function MAK.is_right_isometry(t::AdjointTensorMap; kwargs...)
-    return is_left_isometry(adjoint(t); kwargs...)
+function MAK.is_right_isometric(t::AdjointTensorMap; kwargs...)
+    return is_left_isometric(adjoint(t); kwargs...)
 end
 
 # 2-arg functions

src/factorizations/factorizations.jl

@@ -30,9 +30,9 @@ include("pullbacks.jl")
 
 TensorKit.one!(A::AbstractMatrix) = MatrixAlgebraKit.one!(A)
 
-function MatrixAlgebraKit.isisometry(t::AbstractTensorMap, (p₁, p₂)::Index2Tuple)
+function MatrixAlgebraKit.isisometric(t::AbstractTensorMap, (p₁, p₂)::Index2Tuple)
     t = permute(t, (p₁, p₂); copy = false)
-    return isisometry(t)
+    return isisometric(t)
 end
 
 #------------------------------#
@@ -84,14 +84,14 @@ function LinearAlgebra.isposdef!(t::AbstractTensorMap)
 end
 
 # TODO: tolerances are per-block, not global or weighted - does that matter?
-function MatrixAlgebraKit.is_left_isometry(t::AbstractTensorMap; kwargs...)
+function MatrixAlgebraKit.is_left_isometric(t::AbstractTensorMap; kwargs...)
     domain(t) ≾ codomain(t) || return false
-    f((c, b)) = MatrixAlgebraKit.is_left_isometry(b; kwargs...)
+    f((c, b)) = MatrixAlgebraKit.is_left_isometric(b; kwargs...)
     return all(f, blocks(t))
 end
-function MatrixAlgebraKit.is_right_isometry(t::AbstractTensorMap; kwargs...)
+function MatrixAlgebraKit.is_right_isometric(t::AbstractTensorMap; kwargs...)
     domain(t) ≿ codomain(t) || return false
-    f((c, b)) = MatrixAlgebraKit.is_right_isometry(b; kwargs...)
+    f((c, b)) = MatrixAlgebraKit.is_right_isometric(b; kwargs...)
     return all(f, blocks(t))
 end