Introduce QuantumToolboxMetalExt

.buildkite/Metal_Ext.yml

@@ -0,0 +1,34 @@
+steps:
+  - label: "Metal Julia {{matrix.version}}"
+    matrix:
+      setup:
+        version:
+          - "1.10" # oldest
+          #- "1"    # latest
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "{{matrix.version}}"
+      - JuliaCI/julia-test#v1:
+          test_args: "--quickfail"
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+          dirs:
+            - src
+            - ext
+    agents:
+      queue: "juliaecosystem"
+      os: "macos"
+      arch: "aarch64"
+    env:
+      GROUP: "Metal_Ext"
+      SECRET_CODECOV_TOKEN: "ZfhQu/IcRLqNyZ//ZNs5sjBPaV76IHfU5gui52Qn+Rp8tOurukqgScuyDt+3HQ4R0hJYBw1/Nqg6jmBsvWSc9NEUx8kGsUJFHfN3no0+b+PFxA8oJkWc9EpyIsjht5ZIjlsFWR3f0DpPqMEle/QyWOPcal63CChXR8oAoR+Fz1Bh8GkokLlnC8F9Ugp9xBlu401GCbyZhvLTZnNIgK5yy9q8HBJnBg1cPOhI81J6JvYpEmcIofEzFV/qkfpTUPclu43WNoFX2DZPzbxilf3fsAd5/+nRkRfkNML8KiN4mnmjHxPPbuY8F5zC/PS5ybXtDpfvaMQc01WApXCkZk0ZAQ==;U2FsdGVkX1+eDT7dqCME5+Ox5i8GvWRTQbwiP/VYjapThDbxXFDeSSIC6Opmon+M8go22Bun3bat6Fzie65ang=="
+    timeout_in_minutes: 60
+    if: | 
+          // Don't run Buildkite if the commit message includes the text [skip ci], [ci skip], or [no ci]
+          // Don't run Buildkite for PR draft
+          // Only run Buildkite when new commits and PR are made to main branch
+          build.message !~ /\[skip ci\]/ &&
+          build.message !~ /\[ci skip\]/ &&
+          build.message !~ /\[no ci\]/   &&
+          !build.pull_request.draft      &&
+          (build.branch =~ /main/ || build.pull_request.base_branch =~ /main/)

.buildkite/pipeline.yml

@@ -2,15 +2,25 @@
 steps:
   - label: ":runner: Dynamically launch pipelines"
     plugins:
-      - staticfloat/forerunner:
+      - staticfloat/forerunner: # CUDA
           watch:
             - ".buildkite/pipeline.yml"
             - ".buildkite/CUDA_Ext.yml"
             - "src/**"
             - "ext/QuantumToolboxCUDAExt.jl"
             - "test/runtests.jl"
-            - "test/cuda_ext.jl"
+            - "test/ext-test/cuda_ext.jl"
             - "Project.toml"
           target: ".buildkite/CUDA_Ext.yml"
+      - staticfloat/forerunner: # Metal
+          watch:
+            - ".buildkite/pipeline.yml"
+            - ".buildkite/Metal_Ext.yml"
+            - "src/**"
+            - "ext/QuantumToolboxMetalExt.jl"
+            - "test/runtests.jl"
+            - "test/ext-test/metal_ext.jl"
+            - "Project.toml"
+          target: ".buildkite/Metal_Ext.yml"
     agents:
       queue: "juliagpu"

.github/workflows/CI.yml

@@ -8,7 +8,8 @@ on:
       - '.github/workflows/CI.yml'
       - 'src/**'
       - 'ext/**'
-      - 'test/**'
+      - 'test/runtests.jl'
+      - 'test/core-test/**'
       - 'Project.toml'
   pull_request:
     branches:

Project.toml

@@ -25,9 +25,11 @@ StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 
 [extensions]
 QuantumToolboxCUDAExt = "CUDA"
+QuantumToolboxMetalExt = "Metal"
 
 [compat]
 ArrayInterface = "6, 7"
@@ -39,6 +41,7 @@ Graphs = "1.7"
 IncompleteLU = "0.2"
 LinearAlgebra = "<0.0.1, 1"
 LinearSolve = "2"
+Metal = "1"
 OrdinaryDiffEqCore = "1"
 OrdinaryDiffEqTsit5 = "1"
 Pkg = "<0.0.1, 1"
@@ -54,6 +57,7 @@ julia = "1.10"
 
 [extras]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]

docs/make.jl

@@ -39,6 +39,7 @@ const PAGES = [
         "Two-time correlation functions" => [],
         "Extensions" => [
             "users_guide/extensions/cuda.md",
+            "users_guide/extensions/metal.md",
         ],
     ],
     "Tutorials" => [

docs/src/users_guide/extensions/metal.md

@@ -0,0 +1,161 @@
+# [Extension for Metal.jl](@id doc:Metal)
+
+## Introduction
+
+This is an extension to support `QuantumObject.data` conversion from standard dense and sparse CPU arrays to GPU ([`Metal.jl`](https://github.com/JuliaGPU/Metal.jl)) arrays.
+
+This extension will be automatically loaded if user imports both `QuantumToolbox` and [`Metal.jl`](https://github.com/JuliaGPU/Metal.jl):
+
+```julia
+using QuantumToolbox
+using Metal
+```
+
+We wrapped several functions in `CUDA` and `CUDA.CUSPARSE` in order to not only converting `QuantumObject.data` into GPU arrays, but also changing the element type and word size (`32` and `64`) since some of the GPUs perform better in `32`-bit. The functions are listed as follows (where input `A` is a [`QuantumObject`](@ref)):
+
+- `cu(A; word_size=64)`: return a new [`QuantumObject`](@ref) with `CUDA` arrays and specified `word_size`.
+- `CuArray(A)`: If `A.data` is a dense array, return a new [`QuantumObject`](@ref) with `CUDA.CuArray`.
+- `CuArray{T}(A)`: If `A.data` is a dense array, return a new [`QuantumObject`](@ref) with `CUDA.CuArray` under element type `T`.
+- `CuSparseVector(A)`: If `A.data` is a sparse vector, return a new [`QuantumObject`](@ref) with `CUDA.CUSPARSE.CuSparseVector`.
+- `CuSparseVector{T}(A)`: If `A.data` is a sparse vector, return a new [`QuantumObject`](@ref) with `CUDA.CUSPARSE.CuSparseVector` under element type `T`.
+- `CuSparseMatrixCSC(A)`: If `A.data` is a sparse matrix, return a new [`QuantumObject`](@ref) with `CUDA.CUSPARSE.CuSparseMatrixCSC`.
+- `CuSparseMatrixCSC{T}(A)`: If `A.data` is a sparse matrix, return a new [`QuantumObject`](@ref) with `CUDA.CUSPARSE.CuSparseMatrixCSC` under element type `T`.
+- `CuSparseMatrixCSR(A)`: If `A.data` is a sparse matrix, return a new [`QuantumObject`](@ref) with `CUDA.CUSPARSE.CuSparseMatrixCSR`.
+- `CuSparseMatrixCSR{T}(A)`: If `A.data` is a sparse matrix, return a new [`QuantumObject`](@ref) with `CUDA.CUSPARSE.CuSparseMatrixCSR` under element type `T`.
+
+We suggest to convert the arrays from CPU to GPU memory by using the function `cu` because it allows different `data`-types of input [`QuantumObject`](@ref).
+
+Here are some examples:
+
+## Converting dense arrays
+
+```julia
+V = fock(2, 0) # CPU dense vector
+```
+
+```
+Quantum Object:   type=Ket   dims=[2]   size=(2,)
+2-element Vector{ComplexF64}:
+ 1.0 + 0.0im
+ 0.0 + 0.0im
+```
+
+```julia
+cu(V)
+```
+
+```
+Quantum Object:   type=Ket   dims=[2]   size=(2,)
+2-element CuArray{ComplexF64, 1, CUDA.DeviceMemory}:
+ 1.0 + 0.0im
+ 0.0 + 0.0im
+```
+
+```julia
+cu(V; word_size = 32)
+```
+
+```
+Quantum Object:   type=Ket   dims=[2]   size=(2,)
+2-element CuArray{ComplexF32, 1, CUDA.DeviceMemory}:
+ 1.0 + 0.0im
+ 0.0 + 0.0im
+```
+
+```julia
+M = Qobj([1 2; 3 4]) # CPU dense matrix
+```
+
+```
+Quantum Object:   type=Operator   dims=[2]   size=(2, 2)   ishermitian=false
+2×2 Matrix{Int64}:
+ 1  2
+ 3  4
+```
+
+```julia
+cu(M)
+```
+
+```
+Quantum Object:   type=Operator   dims=[2]   size=(2, 2)   ishermitian=false
+2×2 CuArray{Int64, 2, CUDA.DeviceMemory}:
+ 1  2
+ 3  4
+```
+
+```julia
+cu(M; word_size = 32)
+```
+
+```
+Quantum Object:   type=Operator   dims=[2]   size=(2, 2)   ishermitian=false
+2×2 CuArray{Int32, 2, CUDA.DeviceMemory}:
+ 1  2
+ 3  4
+```
+
+## Converting sparse arrays
+
+```julia
+V = fock(2, 0; sparse=true) # CPU sparse vector
+```
+
+```
+Quantum Object:   type=Ket   dims=[2]   size=(2,)
+2-element SparseVector{ComplexF64, Int64} with 1 stored entry:
+  [1]  =  1.0+0.0im
+```
+
+```julia
+cu(V)
+```
+
+```
+Quantum Object:   type=Ket   dims=[2]   size=(2,)
+2-element CuSparseVector{ComplexF64, Int32} with 1 stored entry:
+  [1]  =  1.0+0.0im
+```
+
+```julia
+cu(V; word_size = 32)
+```
+
+```
+Quantum Object:   type=Ket   dims=[2]   size=(2,)
+2-element CuSparseVector{ComplexF32, Int32} with 1 stored entry:
+  [1]  =  1.0+0.0im
+```
+
+```julia
+M = sigmax() # CPU sparse matrix
+```
+
+```
+Quantum Object:   type=Operator   dims=[2]   size=(2, 2)   ishermitian=true
+2×2 SparseMatrixCSC{ComplexF64, Int64} with 2 stored entries:
+     ⋅      1.0+0.0im
+ 1.0+0.0im      ⋅    
+```
+
+```julia
+cu(M)
+```
+
+```
+Quantum Object:   type=Operator   dims=[2]   size=(2, 2)   ishermitian=true
+2×2 CuSparseMatrixCSC{ComplexF64, Int32} with 2 stored entries:
+     ⋅      1.0+0.0im
+ 1.0+0.0im      ⋅    
+```
+
+```julia
+cu(M; word_size = 32)
+```
+
+```
+Quantum Object:   type=Operator   dims=[2]   size=(2, 2)   ishermitian=true
+2×2 CuSparseMatrixCSC{ComplexF32, Int32} with 2 stored entries:
+     ⋅      1.0+0.0im
+ 1.0+0.0im      ⋅    
+```

ext/QuantumToolboxMetalExt.jl

@@ -0,0 +1,56 @@
+module QuantumToolboxMetalExt
+
+using QuantumToolbox
+import Metal: mtl, MtlArray
+
+@doc raw"""
+    MtlArray(A::QuantumObject)
+
+If `A.data` is an arbitrary array, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `Metal.MtlArray` for gpu calculations.
+
+Note that this function will always change element type into `32`-bit (`Int32`, `Float32`, and `ComplexF32`).
+"""
+MtlArray(A::QuantumObject{<:AbstractArray{T}}) where {T<:Number} = QuantumObject(MtlArray(A.data), A.type, A.dims)
+MtlArray(A::QuantumObject{<:AbstractArray{T}}) where {T<:Int64} = QuantumObject(MtlArray{Int32}(A.data), A.type, A.dims)
+MtlArray(A::QuantumObject{<:AbstractArray{T}}) where {T<:Float64} =
+    QuantumObject(MtlArray{Float32}(A.data), A.type, A.dims)
+MtlArray(A::QuantumObject{<:AbstractArray{T}}) where {T<:ComplexF64} =
+    QuantumObject(MtlArray{ComplexF32}(A.data), A.type, A.dims)
+
+@doc raw"""
+    MtlArray{T}(A::QuantumObject)
+
+If `A.data` is an arbitrary array, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `Metal.MtlArray` with element type `T` for gpu calculations.
+"""
+MtlArray{T}(A::QuantumObject{<:AbstractArray{Tq}}) where {T,Tq<:Number} =
+    QuantumObject(MtlArray{T}(A.data), A.type, A.dims)
+
+@doc raw"""
+    mtl(A::QuantumObject)
+
+Return a new [`QuantumObject`](@ref) where `A.data` is in the type of `Metal` arrays for gpu calculations.
+
+Note that this function will always change element type into `32`-bit (`Int32`, `Float32`, and `ComplexF32`).
+"""
+mtl(A::QuantumObject{<:AbstractArray{T}}) where {T<:Int64} = QuantumObject(MtlArray{Int32}(A.data), A.type, A.dims)
+mtl(A::QuantumObject{<:AbstractArray{T}}) where {T<:Float64} = QuantumObject(MtlArray{Float32}(A.data), A.type, A.dims)
+mtl(A::QuantumObject{<:AbstractArray{T}}) where {T<:ComplexF64} =
+    QuantumObject(MtlArray{ComplexF32}(A.data), A.type, A.dims)
+
+## TODO: Remove the following part if Metal.jl support `sparse`
+import QuantumToolbox: _spre, _spost, _sprepost
+_spre(A::MtlArray, Id::AbstractMatrix) = kron(Id, A)
+_spost(B::MtlArray, Id::AbstractMatrix) = kron(transpose(B), Id)
+_sprepost(A::MtlArray, B::MtlArray) = kron(transpose(B), A)
+
+## TODO: Remove the following part if Metal.jl support `kron`
+import LinearAlgebra: kron
+LinearAlgebra.kron(A::Diagonal, B::MtlArray{T}) where {T} = MtlArray{T}(LinearAlgebra.kron(mtl(A), B))
+LinearAlgebra.kron(A::MtlArray{T}, B::Diagonal) where {T} = MtlArray{T}(LinearAlgebra.kron(A, mtl(B)))
+LinearAlgebra.kron(A::Transpose{T1,<:MtlArray}, B::MtlArray{T2}) where {T1,T2} = MtlArray(LinearAlgebra.kron(A, B))
+LinearAlgebra.kron(A::MtlArray{T1}, B::Transpose{T2,<:MtlArray}) where {T1,T2} = MtlArray(LinearAlgebra.kron(A, B))
+LinearAlgebra.kron(A::Transpose{T1,<:MtlArray}, B::Transpose{T2,<:MtlArray}) where {T1,T2} =
+    MtlArray(LinearAlgebra.kron(A, B))
+LinearAlgebra.kron(A::MtlArray{T1}, B::MtlArray{T2}) where {T1,T2} = MtlArray(LinearAlgebra.kron(A, B))
+
+end