Support sparse arrays (SparseArrays, Yao)

[mofeing]

MLIR has builtin support for sparse tensors and as far as we know, XLA supports them. We just haven't yet added support for nor tried them.

One case of particular interest for me is that Yao uses sparse arrays for representing some quantum gates. Specifically it uses a LuxurySparseArray type which has been giving me problems. Although it should be able to track the TracedRNumbers and end up with an SparseArray{TracedRNumber} kind of array, this is not what is happening (I'm running into errors many times) and it's also far from optimal (it's better to use a TracedRArray).

Currently, we force the dense representation of parametric Yao gates but we should reconsider the implementation and test it again once we add support for sparse tensors.

[wsmoses]

hm can you add the particular mwe and erring code here from the PR?

Specifically I want to know what is required for us to implement such that ideally we don't need the extension

[wsmoses]

the reason is not because having an extension is bad (I'm all for extensions), but because its necessity implies there's a class of julia code we don't support and having that specified/tracked is useful

[mofeing]

so the first MWE is not a problem with sparse arrays but a interaction problem between Reactant and YaoBlocks. specifically, with the Rz gate (Rx and Ry work).

julia> using Reactant, YaoBlocks

julia> θ = ConcreteRNumber(0.0)
ConcreteRNumber{Float64}(0.0)

julia> f = Reactant.compile((θ,)) do ϕ
               mat(Rz(ϕ))
           end
ERROR: MethodError: no method matching ComplexF64(::Reactant.TracedRNumber{Float64})
The type `ComplexF64` exists, but no method is defined for this combination of argument types when trying to construct it.

Closest candidates are:
  (::Type{Complex{T}} where T<:Real)(::Any, ::Any)
   @ Base complex.jl:14
  (::Type{T})(::T) where T<:Number
   @ Core boot.jl:900
  ComplexF64(::ConcreteRNumber)
   @ Reactant ~/.julia/packages/Reactant/m1CaM/src/ConcreteRArray.jl:33
  ...

Stacktrace:
  [1] mat(::Type{ComplexF64}, R::RotationGate{2, Reactant.TracedRNumber{Float64}, ZGate})
    @ YaoBlocks ~/.julia/packages/YaoBlocks/DuDXk/src/primitive/rotation_gate.jl:87
  [2] mat(x::RotationGate{2, Reactant.TracedRNumber{Float64}, ZGate})
    @ YaoBlocks ~/.julia/packages/YaoBlocks/DuDXk/src/abstract_block.jl:120
...

i had to implement the following methods in the end (which might be interesting to add in a PR)...

function (::Type{Complex{T}})(x::Reactant.TracedRNumber{T}) where {T}
    result = Reactant.MLIR.IR.TensorType((), Reactant.MLIR.IR.Type(Complex{T}))
    res = Reactant.MLIR.IR.result(Reactant.MLIR.Dialects.stablehlo.convert(x.mlir_data; result))
    return Reactant.TracedRNumber{Complex{T}}((), res)
end

(::Type{T})(x::Reactant.TracedRNumber{T}) where {T} = x

Base.convert(::Type{T}, x::Reactant.TracedRNumber{T}) where {T<:Number} = x

... but still gave me an error

ERROR: TypeError: in typeassert, expected ComplexF64, got a value of type Reactant.TracedRNumber{ComplexF64}
Stacktrace:
  [1] setindex!(A::Vector{ComplexF64}, x::Reactant.TracedRNumber{ComplexF64}, i::Int64)
    @ Base ./array.jl:976
  [2] macro expansion
    @ ./broadcast.jl:968 [inlined]
  [3] macro expansion
    @ ./simdloop.jl:77 [inlined]
  [4] copyto!
    @ ./broadcast.jl:967 [inlined]
  [5] copyto!
    @ ./broadcast.jl:920 [inlined]
  [6] copy
    @ ./broadcast.jl:892 [inlined]
  [7] materialize
    @ ./broadcast.jl:867 [inlined]
  [8] broadcast_preserving_zero_d
    @ ./broadcast.jl:856 [inlined]
  [9] *(A::Reactant.TracedRNumber{ComplexF64}, B::Vector{ComplexF64})
    @ Base ./arraymath.jl:21
 [10] *
    @ ~/.julia/juliaup/julia-1.11.1+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/diagonal.jl:253 [inlined]
 [11] *
    @ ./operators.jl:596 [inlined]
 [12] mat(::Type{ComplexF64}, R::RotationGate{2, Reactant.TracedRNumber{Float64}, ZGate})
    @ YaoBlocks ~/.julia/packages/YaoBlocks/DuDXk/src/primitive/rotation_gate.jl:87
 [13] mat
    @ ~/.julia/packages/YaoBlocks/DuDXk/src/abstract_block.jl:120 [inlined]
...

if I change f to use Rx or Ry gates, then it works flawlessly

julia> f = Reactant.compile((θ,)) do ϕ
               mat(Ry(ϕ))
           end
Reactant.Compiler.Thunk{Symbol("###11_reactant#232")}()

julia> f(θ)
2×2 ConcreteRArray{ComplexF64, 2}:
 1.0+0.0im  -0.0+0.0im
 0.0+0.0im   1.0+0.0im

the reason why Rz seems to fail here is that Rz, unlike Rx and Ry, is diagonal and thus Yao uses a LinearAlgebra.Diagonal representation... but it sets the eltype to ComplexF64 https://github.com/QuantumBFS/Yao.jl/blob/67d315f1f4a5d92b0cad0112452cbe231ed47a8b/lib/YaoBlocks/src/primitive/rotation_gate.jl#L85-L88

julia> mat(Rz(θ))
2×2 LinearAlgebra.Diagonal{ComplexF64, Vector{ComplexF64}}:
 1.0-0.0im      ⋅
     ⋅      1.0-0.0im

for this case I don't think the problem is the LinearAlgebra.Diagonal, but that Yao sets the eltype to ComplexF64 instead of TracedRNumber{ComplexF64}

there's a second problem with LuxurySparseArray but it's late. will try to write about it tomorrow.

[mofeing]

another MWE, using controlled rotation gates

julia> θ = ConcreteRArray(π/4)
0-dimensional ConcreteRArray{Float64, 0}:
0.7853981633974483

julia> mat(control(1, 2=>Rx(θ))(2))
ERROR: MethodError: no method matching ConcreteRArray{Float64, 0}(::ConcreteRArray{Float64, 0})
The type `ConcreteRArray{Float64, 0}` exists, but no method is defined for this combination of argument types when trying to construct it.

Closest candidates are:
  (::Type{ConcreteRArray{T, N}} where {T, N})(::Any, ::Any)
   @ Reactant ~/.julia/packages/Reactant/m1CaM/src/ConcreteRArray.jl:6

Stacktrace:
 [1] RotationGate{2, ConcreteRArray{Float64, 0}, XGate}(block::XGate, theta::ConcreteRArray{Float64, 0})
   @ YaoBlocks ~/.julia/packages/YaoBlocks/DuDXk/src/primitive/rotation_gate.jl:22
 [2] RotationGate(block::XGate, theta::ConcreteRArray{Float64, 0})
   @ YaoBlocks ~/.julia/packages/YaoBlocks/DuDXk/src/primitive/rotation_gate.jl:26
 [3] Rx(theta::ConcreteRArray{Float64, 0})
   @ YaoBlocks ~/.julia/packages/YaoBlocks/DuDXk/src/primitive/rotation_gate.jl:45
 [4] top-level scope
   @ REPL[35]:1

this is fixed with

julia> (::Type{T})(x::T) where {T<:ConcreteRArray} = x

julia> mat(control(1, 2=>Rx(θ))(2))
4×4 SparseArrays.SparseMatrixCSC{ComplexF64, Int64} with 6 stored entries:
 1.0+0.0im          ⋅               ⋅              ⋅
     ⋅      0.92388+0.0im           ⋅          0.0-0.382683im
     ⋅              ⋅           1.0+0.0im          ⋅
     ⋅          0.0-0.382683im      ⋅      0.92388+0.0im

for the Rz, I need again another fix

julia> function (::Type{Complex{T}})(x::Reactant.ConcreteRArray{T}) where {T}
           return ConcreteRArray(ComplexF64.(x))
       end

julia> mat(control(1, 2=>Rz(θ))(2))
4×4 LinearAlgebra.Diagonal{ComplexF64, Vector{ComplexF64}}:
 1.0+0.0im          ⋅               ⋅              ⋅
     ⋅      0.92388-0.382683im      ⋅              ⋅
     ⋅              ⋅           1.0+0.0im          ⋅
     ⋅              ⋅               ⋅      0.92388+0.382683im

check out that the resulting eltype is ComplexF64, not ConcreteRNumber{ComplexF64}

---

when compiling, the error is the following:

julia> g(x) = mat(control(1, 2=>Rz(x))(2))
g (generic function with 1 method)

julia> @jit g(θ)
ERROR: MethodError: no method matching ComplexF64(::Reactant.TracedRArray{Float64, 0})
The type `ComplexF64` exists, but no method is defined for this combination of argument types when trying to construct it.

Closest candidates are:
  (::Type{Complex{T}} where T<:Real)(::Any, ::Any)
   @ Base complex.jl:14
  ComplexF64(::ConcreteRNumber)
   @ Reactant ~/.julia/packages/Reactant/m1CaM/src/ConcreteRArray.jl:33
  Complex{T}(::Complex) where T<:Real
   @ Base complex.jl:43
  ...

Stacktrace:
  [1] mat(::Type{ComplexF64}, R::RotationGate{2, Reactant.TracedRArray{Float64, 0}, ZGate})
    @ YaoBlocks ~/.julia/packages/YaoBlocks/DuDXk/src/primitive/rotation_gate.jl:87
  [2] mat(::Type{ComplexF64}, c::ControlBlock{RotationGate{2, Reactant.TracedRArray{Float64, 0}, ZGate}, 1, 1})
    @ YaoBlocks ~/.julia/packages/YaoBlocks/DuDXk/src/composite/control.jl:160
  [3] mat(x::ControlBlock{RotationGate{2, Reactant.TracedRArray{Float64, 0}, ZGate}, 1, 1})
    @ YaoBlocks ~/.julia/packages/YaoBlocks/DuDXk/src/abstract_block.jl:120
  [4] g
    @ ./REPL[45]:1 [inlined]
  [5] (::Tuple{})(none::Reactant.TracedRArray{Float64, 0})
    @ Base.Experimental ./<missing>:0
  [6] (::Reactant.var"#27#37"{Bool, Bool, typeof(g), Tuple{ConcreteRArray{Float64, 0}}, Vector{Union{ReactantCore.MissingTracedValue, Reactant.TracedRArray, Reactant.TracedRNumber}}, Tuple{Reactant.TracedRArray{Float64, 0}}})()
    @ Reactant ~/.julia/packages/Reactant/m1CaM/src/utils.jl:152
  [7] block!(f::Reactant.var"#27#37"{Bool, Bool, typeof(g), Tuple{ConcreteRArray{Float64, 0}}, Vector{Union{ReactantCore.MissingTracedValue, Reactant.TracedRArray, Reactant.TracedRNumber}}, Tuple{Reactant.TracedRArray{Float64, 0}}}, blk::Reactant.MLIR.IR.Block)
    @ Reactant.MLIR.IR ~/.julia/packages/Reactant/m1CaM/src/mlir/IR/Block.jl:201
  [8] make_mlir_fn(f::Function, args::Tuple{ConcreteRArray{Float64, 0}}, kwargs::Tuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, no_args_in_result::Bool, construct_function_without_args::Bool, do_transpose::Bool)
    @ Reactant ~/.julia/packages/Reactant/m1CaM/src/utils.jl:116
  [9] make_mlir_fn
    @ ~/.julia/packages/Reactant/m1CaM/src/utils.jl:36 [inlined]
...

the method to fix ComplexF64(::TracedRArray{Float64,0}) can be implemented but breaks Julia semantics and I believe the problem is not that we are missing a Julia method but that Yao goes too... low-level.

[mofeing]

sorry for the mess, but it's not easy to track the source(s) of the problem. the PR #306 fixes only the rotation gates, but not controlled rotation gates:

julia> θ = ConcreteRNumber(π/4)
ConcreteRNumber{Float64}(0.7853981633974483)

julia> g(x) = mat(Rz(x))
g (generic function with 1 method)

julia> @jit g(θ)
2×2 ConcreteRArray{ComplexF64, 2}:
 0.92388-0.382683im      0.0+0.0im
     0.0+0.0im       0.92388+0.382683im

julia> g(x) = @allowscalar mat(control(1,2=>Rx(x[]))(2))
g (generic function with 1 method)

julia> @jit g(θ)
ERROR: MethodError: no method matching ComplexF64(::Reactant.TracedRNumber{ComplexF64})
The type `ComplexF64` exists, but no method is defined for this combination of argument types when trying to construct it.

Closest candidates are:
  (::Type{Complex{T}} where T<:Real)(::Any, ::Any)
   @ Base complex.jl:14
  (::Type{T})(::T) where T<:Number
   @ Core boot.jl:900
  ComplexF64(::ConcreteRNumber)
   @ Reactant ~/.julia/packages/Reactant/HPtbx/src/ConcreteRArray.jl:46
  ...

Stacktrace:
  [1] convert(::Type{ComplexF64}, x::Reactant.TracedRNumber{ComplexF64})
    @ Base ./number.jl:7
  [2] setindex!(A::Matrix{ComplexF64}, x::Reactant.TracedRNumber{ComplexF64}, i::Int64)
    @ Base ./array.jl:976
  [3] copyto_unaliased!
    @ ./abstractarray.jl:1087 [inlined]
  [4] copyto!(dest::Matrix{ComplexF64}, src::Reactant.TracedRArray{ComplexF64, 2})
    @ Base ./abstractarray.jl:1061
  [5] copyto_axcheck!
    @ ./abstractarray.jl:1167 [inlined]
  [6] Array
    @ ./array.jl:615 [inlined]
  [7] Array
    @ ./boot.jl:603 [inlined]
  [8] cunmat(nbit::Int64, cbits::Tuple{Int64}, cvals::Tuple{Int64}, U0::Reactant.TracedRArray{ComplexF64, 2}, locs::Tuple{Int64})
    @ YaoBlocks ~/.julia/packages/YaoBlocks/DuDXk/src/routines.jl:219
  [9] mat
    @ ~/.julia/packages/YaoBlocks/DuDXk/src/composite/control.jl:160 [inlined]
 [10] mat
    @ ~/.julia/packages/YaoBlocks/DuDXk/src/abstract_block.jl:120 [inlined]
 [11] macro expansion
    @ ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:206 [inlined]
 [12] g
    @ ./REPL[31]:1 [inlined]
 [13] (::Tuple{})(none::Reactant.TracedRNumber{Float64})
    @ Base.Experimental ./<missing>:0
 [14] (::Reactant.var"#32#42"{Bool, Bool, typeof(g), Tuple{ConcreteRNumber{Float64}}, Vector{Union{ReactantCore.MissingTracedValue, Reactant.TracedRArray, Reactant.TracedRNumber}}, Tuple{Reactant.TracedRNumber{Float64}}})()
    @ Reactant ~/.julia/packages/Reactant/HPtbx/src/utils.jl:162
 [15] block!(f::Reactant.var"#32#42"{Bool, Bool, typeof(g), Tuple{ConcreteRNumber{Float64}}, Vector{Union{ReactantCore.MissingTracedValue, Reactant.TracedRArray, Reactant.TracedRNumber}}, Tuple{Reactant.TracedRNumber{Float64}}}, blk::Reactant.MLIR.IR.Block)
    @ Reactant.MLIR.IR ~/.julia/packages/Reactant/HPtbx/src/mlir/IR/Block.jl:201
 [16] make_mlir_fn(f::Function, args::Tuple{ConcreteRNumber{Float64}}, kwargs::Tuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, no_args_in_result::Bool, construct_function_without_args::Bool, do_transpose::Bool)
    @ Reactant ~/.julia/packages/Reactant/HPtbx/src/utils.jl:120
 [17] make_mlir_fn
    @ ~/.julia/packages/Reactant/HPtbx/src/utils.jl:40 [inlined]
...

[mofeing]

in today's meeting, we came to the conclusion that the way Yao implements mat is not very conventional and there are some non-idiomatic implementations that should be fixed there. nevertheless, it has uncovered that eltype on ConcreteRArray and TracedRArray is not completely well-defined and that #287 needs some more development to fix that

the Yao PR is gonna be merged now but should be revisited once #287 is merged