Make CuYao an extension

Makefile

@@ -0,0 +1,37 @@
+JL = julia --project
+
+default: init test
+
+init:
+	$(JL) -e 'using Pkg; Pkg.develop([Pkg.PackageSpec(path = joinpath("lib", pkg)) for pkg in ["YaoArrayRegister", "YaoBlocks", "YaoSym", "YaoPlots", "YaoAPI"]]); Pkg.precompile()'
+init-docs:
+	$(JL) -e 'using Pkg; Pkg.activate("docs"); Pkg.develop([Pkg.PackageSpec(path = joinpath("lib", pkg)) for pkg in ["YaoArrayRegister", "YaoBlocks", "YaoSym", "YaoPlots", "YaoAPI"]]); Pkg.precompile()'
+
+update:
+	$(JL) -e 'using Pkg; Pkg.update(); Pkg.precompile()'
+update-docs:
+	$(JL) -e 'using Pkg; Pkg.activate("docs"); Pkg.update(); Pkg.precompile()'
+
+test-Yao:
+	$(JL) -e 'using Pkg; Pkg.test("Yao")'
+test-CuYao:
+	$(JL) -e 'using Pkg; Pkg.activate("ext/CuYao/test"); Pkg.develop(path="."); Pkg.update()'
+	$(JL) -e 'using Pkg; Pkg.activate("ext/CuYao/test"); include("ext/CuYao/test/runtests.jl")'
+
+test-%:
+	$(JL) -e 'using Pkg; Pkg.test("$*")'
+
+test:
+	$(JL) -e 'using Pkg; Pkg.test(["YaoAPI", "YaoArrayRegister", "YaoBlocks", "YaoSym", "YaoPlots", "Yao"])'
+
+coverage:
+	$(JL) -e 'using Pkg; Pkg.test(["YaoAPI", "YaoArrayRegister", "YaoBlocks", "YaoSym", "YaoPlots", "Yao"]; coverage=true)'
+
+servedocs:
+	$(JL) -e 'using Pkg; Pkg.activate("docs"); using LiveServer; servedocs(;skip_dirs=["docs/src/assets", "docs/src/generated"])'
+
+clean:
+	rm -rf docs/build
+	find . -name "*.cov" -type f -print0 | xargs -0 /bin/rm -f
+
+.PHONY: init test

Project.toml

@@ -13,8 +13,15 @@ YaoBlocks = "418bc28f-b43b-5e0b-a6e7-61bbc1a2c1df"
 YaoPlots = "32cfe2d9-419e-45f2-8191-2267705d8dbc"
 YaoSym = "3b27209a-d3d6-11e9-3c0f-41eb92b2cb9d"
 
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+
+[extensions]
+CuYao = "CUDA"
+
 [compat]
 BitBasis = "0.8, 0.9"
+CUDA = "4, 5"
 LinearAlgebra = "1"
 LuxurySparse = "0.7"
 Reexport = "1"

docs/Project.toml

@@ -7,6 +7,7 @@ FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
+LiveServer = "16fef848-5104-11e9-1b77-fb7a48bbb589"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"

docs/make.jl

@@ -58,6 +58,7 @@ const PAGES = [
         "man/registers.md",
         "man/blocks.md",
         "man/symbolic.md",
+        "man/cuda.md",
         "man/plot.md",
         "man/automatic_differentiation.md",
         "man/simplification.md",

docs/src/man/cuda.md

@@ -0,0 +1,71 @@
+```@meta
+CurrentModule = Yao
+```
+
+# CUDA extension - CuYao
+
+## Tutorial
+`CuYao` is a CUDA extension of Yao, which allows you to run Yao circuits on GPU. The usage of `CuYao` is similar to `Yao`, but with some extra APIs to upload and download registers between CPU and GPU:
+- `cu(reg)` to upload a registe `reg` to GPU, and
+- `cpu(cureg)` to download a register `cureg` from GPU to CPU.
+
+```julia
+julia> using Yao, CUDA
+
+# create a register on GPU
+julia> cureg = rand_state(9; nbatch=1000) |> cu;   # or `curand_state(9; nbatch=1000)`.
+
+# run a circuit on GPU
+julia> cureg |> put(9, 2=>Z);
+
+# measure the register on GPU
+julia> measure!(cureg)
+1000-element CuArray{DitStr{2, 9, Int64}, 1, CUDA.Mem.DeviceBuffer}:
+ 110110100 ₍₂₎
+ 000100001 ₍₂₎
+ 111111001 ₍₂₎
+             ⋮
+ 010001101 ₍₂₎
+ 000100110 ₍₂₎
+
+# download the register to CPU
+julia> reg = cureg |> cpu;
+```
+
+
+## Features
+Supported gates:
+
+- general U(N) gate
+- general U(1) gate
+- X, Y, Z gate
+- T, S gate
+- SWAP gate
+- control gates
+
+Supported register operations:
+
+- measure!, measure_reset!, measure_remove!, select
+- append_qudits!, append_qubits!
+- insert_qudit!, insert_qubits!
+- focus!, relax!
+- join
+- density_matrix
+- fidelity
+- expect
+
+Autodiff:
+- autodiff is supported when the only parameterized gates are rotation gates in a circuit.
+
+## API
+```@docs
+cpu
+curand_state
+cuzero_state
+cuproduct_state
+cuuniform_state
+cughz_state
+```
+
+!!! note
+    the `cu` function is not documented in this module, but it is used to upload a register to GPU.

docs/src/man/plot.md

@@ -11,7 +11,59 @@ end
 
 # Quantum Circuit Visualization
 
-Quantum circuit visualization support for Yao.
+`YaoPlots` is the Quantum circuit visualization component for Yao.
+
+
+## Examples
+#### Example 1: Visualize a QBIR define in Yao
+
+```@example plot
+using Yao.EasyBuild, YaoPlots
+
+# show a qft circuit
+vizcircuit(qft_circuit(5))
+```
+
+If you are using a Pluto/Jupyter notebook, Atom/VSCode editor, you should see the above image in your plotting panel.
+
+#### Example 2: Visualize a single qubit
+```@example plot
+using YaoPlots, Yao
+
+reg = zero_state(1) |> Rx(π/8) |> Rx(π/8)
+rho = density_matrix(ghz_state(2), 1)
+
+bloch_sphere("|ψ⟩"=>reg, "ρ"=>rho; show_projection_lines=true)
+```
+
+See more [examples](lib/YaoPlots/examples/circuits.jl).
+
+### Adjusting the plot attributes
+
+Various attributes of the visualizations can be altered. 
+The plot can be modified, if we change the following attributes
+
+- `YaoPlots.CircuitStyles.linecolor[]` for line color, default value being `"#000000"` (black color)
+- `YaoPlots.CircuitStyles.gate_bgcolor[]` for background color of square blocks, the default value being `"#FFFFFF"` (white color)
+- `YaoPlots.CircuitStyles.textcolor[]` for text color, default value being `"#000000`
+- `YaoPlots.CircuitStyles.lw[]` for line width, default value being `1` (pt)
+- `YaoPlots.CircuitStyles.textsize[]` for text size, default value being `16` (pt)
+- `YaoPlots.CircuitStyles.paramtextsize[]` for parameter text size, for parameterized gates, default value being `10` (pt)
+
+For example,
+
+```@example plot
+using YaoPlots, Yao
+YaoPlots.CircuitStyles.linecolor[] = "pink" 
+YaoPlots.CircuitStyles.gate_bgcolor[] = "yellow" 
+YaoPlots.CircuitStyles.textcolor[] = "#000080" # the navy blue color
+YaoPlots.CircuitStyles.fontfamily[] = "JuliaMono"
+YaoPlots.CircuitStyles.lw[] = 2.5
+YaoPlots.CircuitStyles.textsize[] = 13
+YaoPlots.CircuitStyles.paramtextsize[] = 8
+
+vizcircuit(chain(3, put(1=>X), repeat(3, H), put(2=>Y), repeat(3, Rx(π/2))))
+```
 
 ## Circuit Visualization
 ```@docs

ext/CuYao.jl

@@ -0,0 +1,3 @@
+module CuYao
+include("CuYao/src/CuYao.jl")
+end

ext/CuYao/.gitignore

@@ -0,0 +1,23 @@
+*.jl.cov
+*.jl.*.cov
+*.jl.mem
+
+docs/build/
+docs/site/
+docs/src/tutorial/RegisterBasics.md
+docs/src/tutorial/BlockBasics.md
+docs/src/tutorial/BinaryBasics.md
+docs/src/tutorial/QCBM.md
+
+*.ipynb_checkpoints
+**/*.ipynb_checkpoints
+**/**/*.ipynb_checkpoints
+
+_*.dat
+*.swp
+__pycache__/
+.vscode/
+
+Manifest.toml
+
+_local/

ext/CuYao/README.md

@@ -0,0 +1,8 @@
+## How to run tests
+
+Open a terminal, switch to Yao folder and run the following commands:
+```bash
+make init-CuYao
+make test-CuYao
+```
+

ext/CuYao/benchmarks/TestCuda.jl

@@ -0,0 +1,54 @@
+using CUDA
+using LinearAlgebra
+using BenchmarkTools
+
+function ms!(X::CuArray, s::Number)
+    function kernel(X, s)
+        i = (blockIdx().x-1) * blockDim().x + threadIdx().x
+        @inbounds X[i] *= s
+        return
+    end
+    @cuda blocks=length(X) kernel(X, s)
+    X
+end
+
+function ms1!(X::CuArray, a, b, c)
+    k = f(a)
+    function kernel(X, a, b, c)
+        i = (blockIdx().x-1) * blockDim().x + threadIdx().x
+        k(X, i)
+        k(X, i)
+        k(X, i)
+        #@inbounds X[i] *= a
+        #@inbounds X[i] *= b
+        #@inbounds X[i] *= c
+        return
+    end
+    @cuda blocks=length(X)÷256 threads=256 kernel(X, a, b, c)
+    X
+end
+
+@inline function f(s)
+    @inline function kernel(X, i)
+        X[i]*=s
+    end
+end
+function ms2!(X::CuArray, s::Number)
+    function kernel(X, s)
+        i = (blockIdx().x-1) * blockDim().x + threadIdx().x
+        @inbounds X[i] *= s
+        return
+    end
+    @cuda blocks=length(X)÷256 threads=256 kernel(X, s)
+    X
+end
+
+a = randn(1<<10)
+cua = cu(a)
+@benchmark ms1!($cua, 1.001, 1.001, 1.001)
+@benchmark ms2!(ms2!(ms2!($cua, 1.001), 1.001), 1.001)
+bss = 1:length(cua)
+@benchmark f(1.001).(cua, bss)
+
+@benchmark rmul!($a, 1.01)
+@benchmark ms!($cua, 1.0)

ext/CuYao/benchmarks/gates.jl

@@ -0,0 +1,19 @@
+using Yao, CuYao, CUDA
+using BenchmarkTools
+
+reg = rand_state(12; nbatch=1000)
+creg = reg |> cu
+@benchmark CUDA.@sync creg |> put(12, 3=>Z)
+@benchmark CUDA.@sync creg |> put(12, 3=>X)
+@benchmark reg |> put(12, 3=>Z)
+@benchmark CUDA.@sync creg |> control(12, 6, 3=>X)
+@benchmark reg |> control(12, 6, 3=>X)
+@benchmark CUDA.@sync creg |> put(12, 3=>rot(X, 0.3))
+@benchmark reg |> put(12, 3=>rot(X, 0.3))
+
+reg = rand_state(20)
+creg = reg |> cu
+g = swap(20, 7, 2)
+g = put(20, (7, 2)=>matblock(rand_unitary(4)))
+@benchmark reg |> g
+@benchmark CUDA.@sync creg |> g