Cleanup and comments

Author: kshyatt

ext/TensorKitAdaptExt.jl

@@ -16,10 +16,10 @@ function Adapt.adapt_structure(to, x::DiagonalTensorMap)
     return DiagonalTensorMap(data′, x.domain)
 end
 function Adapt.adapt_structure(::Type{T}, x::BraidingTensor{T′, S, A}) where {T <: Number, T′, S, A}
-    return BraidingTensor(space(x), TensorKit.similarstoragetype(A, T), x.adjoint)
+    return BraidingTensor{T}(space(x), x.adjoint)
 end
-function Adapt.adapt_structure(::Type{TA}, x::BraidingTensor{T, S, A}) where {TA <: DenseArray{<:Number}, T, S, A}
-    return BraidingTensor(space(x), TA, x.adjoint)
+function Adapt.adapt_structure(::Type{TA}, x::BraidingTensor{T, S, A}) where {T′, TA <: DenseArray{T′}, T, S, A}
+    return BraidingTensor{T′, S, TA}(space(x), x.adjoint)
 end
 
 end

ext/TensorKitCUDAExt/TensorKitCUDAExt.jl

@@ -14,7 +14,7 @@ using TensorKit.Factorizations
 using TensorKit.Strided
 using TensorKit.Factorizations: AbstractAlgorithm
 using TensorKit: SectorDict, tensormaptype, scalar, similarstoragetype, AdjointTensorMap, scalartype, project_symmetric_and_check
-import TensorKit: randisometry, rand, randn, _set_subblock!
+import TensorKit: randisometry, rand, randn, fill_braidingsubblock!
 
 using TensorKit: MatrixAlgebraKit
 
@@ -23,15 +23,14 @@ using Random
 include("cutensormap.jl")
 include("truncation.jl")
 
-function TensorKit._set_subblock!(data::TD, val) where {T, TD <: Union{<:CuMatrix{T}, <:StridedViews.StridedView{T, 4, <:CuArray{T}}}}
+function TensorKit.fill_braidingsubblock!(data::TD, val) where {T, TD <: Union{<:CuMatrix{T}, <:StridedViews.StridedView{T, 4, <:CuArray{T}}}}
     @kernel function fill_subblock_kernel!(subblock, val)
         idx = @index(Global, Cartesian)
-        @inbounds subblock[idx[1], idx[2], idx[2], idx[1]] = val
+        idx_val = idx[1] == idx[4] && idx[2] == idx[3] ? val : zero(val)
+        @inbounds subblock[idx] = idx_val
     end
     kernel = fill_subblock_kernel!(KernelAbstractions.get_backend(data))
-    d1 = size(data, 1)
-    d2 = size(data, 2)
-    kernel(data, val; ndrange = (d1, d2))
+    kernel(data, val; ndrange = size(data))
     return data
 end
 

src/tensors/braidingtensor.jl

@@ -3,7 +3,8 @@
 #====================================================================#
 """
     struct BraidingTensor{T,S<:IndexSpace,A<:DenseVector{T}} <: AbstractTensorMap{T, S, 2, 2}
-    BraidingTensor(V1::S, V2::S, ::Type{A}, adjoint::Bool=false) where {S<:IndexSpace, A <: DenseVector{<:Number}}
+    BraidingTensor(V1::S, V2::S, adjoint::Bool=false) where {S<:IndexSpace}
+    BraidingTensor{T, S, A}(V1::S, V2::S, adjoint::Bool=false) where {T, S, A}
 
 Specific subtype of [`AbstractTensorMap`](@ref) for representing the braiding tensor that
 braids the first input over the second input; its inverse can be obtained as the adjoint.
@@ -27,16 +28,11 @@ struct BraidingTensor{T, S, A} <: AbstractTensorMap{T, S, 2, 2}
     end
 end
 function BraidingTensor{T}(V1::S, V2::S, adjoint::Bool = false) where {T, S <: IndexSpace}
-    return BraidingTensor{T, S, Vector{T}}(V1, V2, adjoint)
+    return BraidingTensor{T, S, similarstoragetype{T}}(V1, V2, adjoint)
 end
 function BraidingTensor(V1::S, V2::S, adjoint::Bool = false) where {S <: IndexSpace}
     T = BraidingStyle(sectortype(S)) isa SymmetricBraiding ? Float64 : ComplexF64
-    return BraidingTensor{T, S, Vector{T}}(V1, V2, adjoint)
-end
-# necessary due to HomSpace ctor below
-function BraidingTensor(V1::S, V2::S, A, adjoint::Bool = false) where {S <: IndexSpace}
-    T = eltype(A)
-    return BraidingTensor{T, S, A}(V1, V2, adjoint)
+    return BraidingTensor{T}(V1, V2, adjoint)
 end
 function BraidingTensor(V1::IndexSpace, V2::IndexSpace, adjoint::Bool = false)
     return BraidingTensor(promote(V1, V2)..., adjoint)
@@ -46,10 +42,10 @@ function BraidingTensor(V::HomSpace, adjoint::Bool = false)
         throw(SpaceMismatch("Cannot define a braiding on $V"))
     return BraidingTensor(V[2], V[1], adjoint)
 end
-function BraidingTensor(V::HomSpace, A, adjoint::Bool = false)
+function BraidingTensor{T, S, A}(V::HomSpace, adjoint::Bool = false) where {T, S, A}
     domain(V) == reverse(codomain(V)) ||
         throw(SpaceMismatch("Cannot define a braiding on $V"))
-    return BraidingTensor(V[2], V[1], A, adjoint)
+    return BraidingTensor{T, S, A}(V[2], V[1], adjoint)
 end
 function BraidingTensor{T}(V::HomSpace, adjoint::Bool = false) where {T}
     domain(V) == reverse(codomain(V)) ||
@@ -93,7 +89,8 @@ function _braiding_factor(f₁, f₂, inv::Bool = false)
     return r
 end
 
-function _set_subblock!(data, val)
+# generates scalar indexing errors on GPU
+function fill_braidingsubblock!(data, val)
     f(I) = ((I[1] == I[4]) & (I[2] == I[3])) * val
     return data .= f.(CartesianIndices(data))
 end
@@ -117,10 +114,8 @@ end
     n2 = d[3] * d[4]
     data_parent = storagetype(b)(undef, prod(d))
     data = sreshape(StridedView(data_parent), d)
-    fill!(data, zero(eltype(b)))
-
     r = _braiding_factor(f₁, f₂, b.adjoint)
-    !isnothing(r) && _set_subblock!(data, r)
+    !isnothing(r) && fill_braidingsubblock!(data, r)
     return data
 end
 
@@ -132,28 +127,33 @@ Base.convert(::Type{TensorMap}, b::BraidingTensor) = TensorMap(b)
 
 Base.complex(b::BraidingTensor{<:Complex}) = b
 function Base.complex(b::BraidingTensor{T, S, A}) where {T, S, A}
-    Ac = similarstoragetype(A, complex(T))
-    return BraidingTensor(space(b), Ac, b.adjoint)
+    Tc = complex(T)
+    Ac = similarstoragetype(A, Tc)
+    return BraidingTensor{Tc, S, Ac}(space(b), b.adjoint)
 end
 
-function _trivial_subblock!(data, b::BraidingTensor)
+# Trivial
+function fill_braidingblock!(data, b::BraidingTensor, s::Trivial)
     V1, V2 = codomain(b)
     d1, d2 = dim(V1), dim(V2)
     subblock = sreshape(StridedView(data), (d1, d2, d2, d1))
-    _set_subblock!(subblock, one(eltype(b)))
+    fill_braidingsubblock!(subblock, one(eltype(b)))
     return data
 end
 
-function _nontrivial_subblock!(data, b::BraidingTensor, s::Sector)
+# Nontrivial
+function fill_braidingblock!(data, b::BraidingTensor, s::Sector)
     base_offset = first(blockstructure(b)[s][2]) - 1
 
     for ((f₁, f₂), (sz, str, off)) in pairs(subblockstructure(space(b)))
         (f₁.coupled == f₂.coupled == s) || continue
         r = _braiding_factor(f₁, f₂, b.adjoint)
-        isnothing(r) && continue
         # change offset to account for single block
         subblock = StridedView(data, sz, str, off - base_offset)
-        _set_subblock!(subblock, r)
+        # without the zero-value, the non-trivial block is not set
+        # correctly in the GPU case
+        val = isnothing(r) ? zero(eltype(data)) : r
+        fill_braidingsubblock!(subblock, val)
     end
     return data
 end
@@ -169,13 +169,8 @@ function block(b::BraidingTensor, s::Sector)
     data = reshape(storagetype(b)(undef, m * n), (m, n))
 
     m * n == 0  && return data # s ∉ blocksectors(b)
-    fill!(data, zero(eltype(b)))
 
-    if sectortype(b) === Trivial
-        return _trivial_subblock!(data, b)
-    else
-        return _nontrivial_subblock!(data, b, s)
-    end
+    return fill_braidingblock!(data, b, s)
 end
 
 # Index manipulations

test/cuda/planar.jl

@@ -15,13 +15,13 @@ for V in spacelist
     @timedtestset "Braiding tensor + CUDA with symmetry: $Istr" verbose = true begin
         W = V[1] ⊗ V[2] ← V[2] ⊗ V[1]
         T = isreal(sectortype(W)) ? Float64 : ComplexF64
-        t1 = @constinferred BraidingTensor(W, CuVector{T, CUDA.DeviceMemory})
+        t1 = @constinferred BraidingTensor{T, spacetype(V[2]), CuVector{T, CUDA.DeviceMemory}}(W)
         @test space(t1) == W
         @test codomain(t1) == codomain(W)
         @test domain(t1) == domain(W)
         @test scalartype(t1) == (isreal(sectortype(W)) ? Float64 : ComplexF64)
         @test storagetype(t1) == CuVector{scalartype(t1), CUDA.DeviceMemory}
-        t2 = @constinferred BraidingTensor(W, CuVector{ComplexF64, CUDA.DeviceMemory})
+        t2 = @constinferred BraidingTensor{ComplexF64, spacetype(V[2]), CuVector{ComplexF64, CUDA.DeviceMemory}}(W)
         @test scalartype(t2) == ComplexF64
         @test storagetype(t2) == CuVector{ComplexF64, CUDA.DeviceMemory}
         t3 = @testinferred adapt(storagetype(t2), t1)