Support BlockArrays in KernelAbstractions bcasting with new extension

Author: kshyatt

Project.toml

@@ -17,16 +17,19 @@ VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 
 [weakdeps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 
 [extensions]
 BlockTensorKitAdaptExt = "Adapt"
+BlockTensorKitGPUArraysExt = "GPUArrays"
 
 [compat]
 Adapt = "4"
 Aqua = "0.8"
 BlockArrays = "1"
 Combinatorics = "1"
 Compat = "4.13"
+GPUArrays = "11.4.1"
 JLArrays = "0.3"
 LinearAlgebra = "1"
 MatrixAlgebraKit = "0.6"
@@ -45,6 +48,7 @@ julia = "1.10"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
@@ -55,4 +59,4 @@ TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
 test = ["Test", "TestExtras", "Random", "Combinatorics", "SafeTestsets", "Aqua", "Adapt", "JLArrays"]
 
 [sources]
-Strided = {url = "https://github.com/QuantumKitHub/Strided.jl", rev = "ksh/jlarrays"}
+Strided = {path = "/home/kshyatt/.julia/dev/Strided"}

ext/BlockTensorKitGPUArraysExt.jl

@@ -0,0 +1,11 @@
+module BlockTensorKitGPUArraysExt
+
+using BlockTensorKit, BlockArrays, GPUArrays, Strided
+using Strided: StridedViews
+using GPUArrays: KernelAbstractions
+
+function KernelAbstractions.get_backend(BA::BlockArrays.BlockArray{T, N, A}) where {T, N, A<:AbstractArray{<:StridedView{T, N, <:AnyGPUArray}}}
+    return KernelAbstractions.get_backend(first(BA.blocks))
+end
+
+end

src/tensors/blocktensor.jl

@@ -114,7 +114,7 @@ function BlockTensorMap(t::AbstractTensorMap, space::TensorMapSumSpace)
     TT = tensormaptype(spacetype(t), numout(t), numin(t), storagetype(t))
     tdst = BlockTensorMap{TT}(undef, space)
     for (f₁, f₂) in fusiontrees(tdst)
-        tdst[f₁, f₂] .= t[f₁, f₂]
+        copy!(tdst[f₁, f₂], t[f₁, f₂])
     end
     return tdst
 end

test/abstracttensor/blocktensor.jl

@@ -5,7 +5,7 @@ using BlockTensorKit
 using Random
 using Combinatorics
 using Adapt
-using Strided, JLArrays
+using JLArrays
 
 Vtr = (
     SumSpace(ℂ^3),
@@ -87,9 +87,7 @@ end
         jl_bt1 = rand(JLVector{T}, W)
         TT = TensorKit.TensorMap{T, spacetype(t1′), numout(t1′), numin(t1′), JLVector{T}}
         jl_bt1′ = @constinferred convert(TT, jl_bt1)
-        JLArrays.@allowscalar begin
-            jl_bt1″ = @inferred BlockTensorMap(jl_bt1′, W)
-        end # still need some logic for copying to a BlockArray of StridedViews
+        jl_bt1″ = @inferred BlockTensorMap(jl_bt1′, W)
         @test jl_bt1 ≈ jl_bt1″
     end
 end