more cuarray hacks

euro-hpc-pl · Apr 3, 2024 · 98d8208 · 98d8208
1 parent 089163c
commit 98d8208
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/src/base.jl b/src/base.jl
@@ -68,10 +68,10 @@ const MatOrCentral{T,N} = Union{AbstractMatrix{T},CentralTensor{T,N}}
 function dense_central(ten::CentralTensor)
     # @cast V[(u1, u2), (d1, d2)] :=
     #     ten.e11[u1, d1] * ten.e21[u2, d1] * ten.e12[u1, d2] * ten.e22[u2, d2]
-    a11 = reshape(ten.e11, size(ten.e11, 1), :, size(ten.e11, 2))
-    a21 = reshape(ten.e21, :, size(ten.e21, 1), size(ten.e21, 2))
-    a12 = reshape(ten.e12, size(ten.e12, 1), 1, 1, size(ten.e12, 2))
-    a22 = reshape(ten.e22, 1, size(ten.e22, 1), 1, size(ten.e22, 2))
+    a11 = reshape(CuArray(ten.e11), size(ten.e11, 1), :, size(ten.e11, 2))
+    a21 = reshape(CuArray(ten.e21), :, size(ten.e21, 1), size(ten.e21, 2))
+    a12 = reshape(CuArray(ten.e12), size(ten.e12, 1), 1, 1, size(ten.e12, 2))
+    a22 = reshape(CuArray(ten.e22), 1, size(ten.e22, 1), 1, size(ten.e22, 2))
     V = @__dot__(a11 * a21 * a12 * a22)
     V = reshape(V, size(V, 1) * size(V, 2), size(V, 3) * size(V, 4))
     V ./ maximum(V)

diff --git a/src/contractions/central.jl b/src/contractions/central.jl
@@ -84,10 +84,10 @@ function batched_mul!(
     if sl1 * sl2 * sr1 * sr2 < sinter
         # @cast E[(l1, l2), (r1, r2)] :=
             # M.e11[l1, r1] * M.e21[l2, r1] * M.e12[l1, r2] * M.e22[l2, r2]
-        a11 = reshape(M.e11, size(M.e11, 1), :, size(M.e11, 2))
-        a21 = reshape(M.e21, :, size(M.e21, 1), size(M.e21, 2))
-        a12 = reshape(M.e12, size(M.e12, 1), 1, 1, size(M.e12, 2))
-        a22 = reshape(M.e22, 1, size(M.e22, 1), 1, size(M.e22, 2))
+        a11 = reshape(CuArray(M.e11), size(M.e11, 1), :, size(M.e11, 2))
+        a21 = reshape(CuArray(M.e21), :, size(M.e21, 1), size(M.e21, 2))
+        a12 = reshape(CuArray(M.e12), size(M.e12, 1), 1, 1, size(M.e12, 2))
+        a22 = reshape(CuArray(M.e22), 1, size(M.e22, 1), 1, size(M.e22, 2))
         E = @__dot__(a11 * a21 * a12 * a22)
         E = reshape(E, size(E, 1) * size(E, 2), size(E, 3) * size(E, 4))
         E = reshape(E, (sl1 * sl2, sr1 * sr2, 1))