cuda extension

QuantumBFS · Oct 7, 2023 · ae6c465 · ae6c465
1 parent e01d6ff
commit ae6c465
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,17 +1,24 @@
 name = "YaoToEinsum"
 uuid = "9b173c7b-dc24-4dc5-a0e1-adab2f7b6ba9"
 authors = ["GiggleLiu <cacate0129@gmail.com> and contributors"]
-version = "0.2.0"
+version = "0.2.1"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922"
 Yao = "5872b779-8223-5990-8dd0-5abbb0748c8c"
 
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+
+[extensions]
+YaoToEinsumCUDAExt = "CUDA"
+
 [compat]
+CUDA = "4, 5"
 OMEinsum = "0.7"
 Yao = "0.8"
-julia = "1"
+julia = "1.9"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

diff --git a/ext/YaoToEinsumCUDAExt.jl b/ext/YaoToEinsumCUDAExt.jl
@@ -0,0 +1,9 @@
+module YaoToEinsumCUDAExt
+using CUDA, YaoToEinsum
+
+export cu
+
+function cu(tnet::TensorNetwork)
+    return TensorNetwork(tnet.code, tnet.tensors .|> CuArray)
+end
+end