FLOYao GPU Support

[p-luo]

Currently, Yao can support GPU computations, via https://github.com/QuantumBFS/CuYao.jl. Does FLOYao have the same capabilities, and if not, can it be implemented easily?

[jlbosse]

Currently FLOYao does not support GPU computations. But I just played around with it by changing

FLOYao.jl/src/majorana_reg.jl

Lines 33 to 35 in 8b840d0

	struct MajoranaReg{T<:Real} <: AbstractRegister{2}
	state::Matrix{T}
	end

to struct MajoranaReg{T} state::CUDA.CuArray{T,2,CUDA.DeviceMemory} end to see how hard it would be to add CUDA support.

It seems that most of the instruct!s from https://github.com/QuantumBFS/FLOYao.jl/blob/master/src/instruct.jl work out of the box because they only do GPU friendly operations. What needs fixing or extra logic on GPUs is

• FLOYao.jl/src/instruct.jl

  Lines 66 to 79 in 8b840d0

  	function Yao.instruct!(reg::MajoranaReg, gate::AbstractMatrix, locs)
  	# unless the gate is a tensor product of Z's (but _not_ using Yaos kron
  	# block a gate acting on non-consecutive qubits will never be a
  	# FLO gate
  	areconsecutive(locs) || throw(NonFLOException("$gate on $locs is not a FLO gate"))
  	W = qubit2majoranaevolution(gate, locs)
  	matlocs = (2*(locs[1]-1)+1:2(locs[end]))
  	@warn """Calling manual instruct!($reg, $gate, $locs).
  	You can greatly speed up your FLO gates by exactly implementing unsafe_apply!()
  	and instruct!() for them. See FLOYao/src/instruct.jl and FLOYao/src/apply_composite.jl
  	for how to do that.
  	"""
  	reg.state[matlocs,:] .= W * reg.state[matlocs,:]
  	end

• FLOYao.jl/src/expect.jl

  Lines 29 to 43 in 8b840d0

  	function majorana_expect(block::AbstractMatrix, locs, reg::MajoranaReg)
  	fullH = zero(reg.state)
  	matlocs = (2*(locs[1]-1)+1:2(locs[end]))
  	fullH[matlocs,matlocs] .= block
  	expval = sum(1:nqubits(reg), init=zero(eltype(reg.state))) do i
  	reg.state[:,2i] ⋅ (fullH * reg.state[:,2i-1]) - reg.state[:,2i-1] ⋅ (fullH * reg.state[:,2i])
  	end
  	offset_expval = sum(1:nqubits(reg), init=zero(eltype(reg.state))) do i
  	- reg.state[:,2i] ⋅ (fullH * reg.state[:,2i]) - reg.state[:,2i-1] ⋅ (fullH * reg.state[:,2i-i])
  	end |> imag
  	return (- expval - offset_expval) / 4
  	end

  needs minimal fixes to get block onto the GPU and the same is true in a couple more places where we do matrix matrix multiplications where one matrix is built on the CPU.
• The sparse matrix utils around

  FLOYao.jl/src/utils.jl

  Lines 358 to 401 in 8b840d0

  	# -------------------------------------------
  	# Utilities for fast sparse matrix operations
  	# -------------------------------------------
  	"""
  	fast_add!(A::AbstractMatrix, B::SparseMatrixCSC)
  	Fast implementation of `A .+= B` for sparse `B`.
  	"""
  	function fast_add!(A::AbstractMatrix, B::SparseMatrixCSC)
  	@assert size(A, 1) == size(B, 1) && size(A, 2) == size(B, 2) "Dimension mismatch"
  	for j = 1:size(B, 2)
  	for k in SparseArrays.nzrange(B, j)
  	@inbounds A[B.rowval[k],j] += B.nzval[k]
  	end
  	end
  	return A
  	end
  	function fast_add!(A::AbstractMatrix, B::AbstractMatrix)
  	return A .+= B
  	end
  	"""
  	fast_overlap(y::AbstractVecOrMat, A::SparseMatrixCSC, x::AbstractVecOrMat)
  	Fast implementation of `tr(y' * A * x)` for sparse `A`.
  	"""
  	function fast_overlap(y::AbstractVecOrMat{T1}, A::SparseMatrixCSC{T2}, x::AbstractVecOrMat{T3}) where {T1,T2,T3}
  	@assert size(x, 1) == size(A, 2) && size(y, 1) == size(A, 1) && size(x, 2) == size(y, 2) "Dimension mismatch"
  	g = zero(promote_type(T1, T2, T3))
  	@inbounds for j = 1:size(A, 2)
  	for k in SparseArrays.nzrange(A, j)
  	i = A.rowval[k]
  	for m = 1:size(y, 2)
  	g += conj(y[i, m]) * A.nzval[k] * x[j, m]
  	end
  	end
  	end
  	return g
  	end
  	function fast_overlap(y::AbstractVecOrMat{T1}, A::AbstractMatrix{T2}, x::AbstractVecOrMat{T3}) where {T1,T2,T3}
  	return y ⋅ (A * x)
  	end

  need to be made GPU friendly.

Currently I don't have the capacity do add CUDA support, but PR's are welcome! If you decide to give it a go it would probably be best to broadly follow the strategy that CuYao currently follows to have it as a package extension. Of course I am also happy to assist, you can reach me e.g. on the julialang slack as
JanLukas Bosse