update cuda deps, move to KernelAbstractions

Author: lpawela

.github/workflows/CI.yml

@@ -1,7 +1,9 @@
 name: CI
 on:
-  - push
-  - pull_request
+  push:
+    branches: [master]
+  pull_request:
+    types: [opened, synchronize, reopened]
 jobs:
   test:
     name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
@@ -10,11 +12,10 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.7'
-          - '1.8'
+          - '1.11'
+          - '1.12'
         os:
           - ubuntu-latest
-          - macOS-latest
         arch:
           - x64
     steps:

Project.toml

@@ -1,25 +1,31 @@
 name = "MatrixEnsembles"
 uuid = "c7015dd7-3fb7-4a4c-827e-526313618491"
 authors = ["Łukasz Pawela <lukasz.pawela@gmail.com>"]
-version = "0.1.4"
+version = "0.2.0"
 
 [deps]
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+
+[extensions]
+MatrixEnsemblesCUDAExt = "CUDA"
+
 [compat]
+CUDA = "5"
 DocStringExtensions = "0.8.3, 0.9"
-CUDA = "4"
-Requires = "1.1"
+KernelAbstractions = "0.9"
 StatsBase = "0.33.2, 0.34"
-julia = "1.4, 1.5"
+julia = "1.9"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 
 [targets]
-test = ["Test"]
+test = ["Test", "CUDA"]

README.md

@@ -4,7 +4,7 @@
 
 ## Sampling random matrices on the GPU
 
-We have introduced an experimental implementation of sampling of random matrices and random quantum objects on the GPU. In order to use this feature, the `CUDA` package is required. To import `MatrixEnsembles` with GPU support use
+We have introduced an experimental implementation of sampling of random matrices and random quantum objects on the GPU using `KernelAbstractions.jl`. In order to use this feature, the `CUDA` package is required to load the necessary extension. To import `MatrixEnsembles` with GPU support use
 ```julia
 using CUDA, MatrixEnsembles
 ```

docs/Project.toml

@@ -0,0 +1,9 @@
+[deps]
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+MatrixEnsembles = "c7015dd7-3fb7-4a4c-827e-526313618491"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+
+[compat]
+Documenter = "1"
+MatrixEnsembles = "0.1, 0.2"
+CUDA = "5"

docs/make.jl

@@ -2,14 +2,14 @@ using Documenter, MatrixEnsembles
 
 format = Documenter.HTML(edit_link = "master",
                          prettyurls = get(ENV, "CI", nothing) == "true",
+                         assets = ["assets/favicon.ico"]
 )
 
 makedocs(
     clean = true,
     format = format,
     sitename = "MatrixEnsembles.jl",
     authors = "Łukasz Pawela",
-    assets = ["assets/favicon.ico"],
     pages = [
         "Home" => "index.md",
         "Manual" => Any[
@@ -20,8 +20,8 @@ makedocs(
     ]
 )
 
-deploydocs(
-    deps = Deps.pip("pygments", "mkdocs", "python-markdown-math"),
-    target = "build",
-    repo = "github.com/iitis/MatrixEnsembles.jl.git"
-)
+# deploydocs(
+#     deps = Deps.pip("pygments", "mkdocs", "python-markdown-math"),
+#     target = "build",
+#     repo = "github.com/iitis/MatrixEnsembles.jl.git"
+# )

docs/src/index.md

@@ -20,7 +20,8 @@ states,
 * manipulating them with quantum channels
 * calculating functionals on these objects, *i.e. trace norm, diamond norm, entropy, fidelity*,
 * application of random matrix theory in quantum
-information processing.
+information processing,
+* random matrix sampling on GPU.
 
 ## [References](@id refs)
 

docs/src/man/random.md

@@ -159,6 +159,24 @@ o*o'
 ```
 For convenience we provide the following type aliases `const COE = CircularEnsemble{1}`, `const CUE = CircularEnsemble{2}`, `const CSE = CircularEnsemble{4}`.
 
+## GPU Sampling
+
+`MatrixEnsembles.jl` supports sampling random matrices directly on the GPU using `CUDA.jl`. This functionality is provided via a package extension that is loaded automatically when `CUDA` is imported.
+
+To use GPU sampling:
+```julia
+using CUDA
+using MatrixEnsembles
+
+# Create an ensemble (e.g., Circular Unitary Ensemble)
+c = CUE(1024)
+
+# Sample directly to a CuArray
+u = MatrixEnsembles.curand(c)
+```
+
+The `curand` function mimics `rand` but uses the `CUDA.default_rng()` and optimized kernels (implemented with `KernelAbstractions.jl`) to generate matrices directly on the device, minimizing host-device data transfer.
+
 ## [References](@id refs_rand)
 
 [1] B. Collins, I. Nechita, *Random matrix techniques in quantum information theory*, Journal of Mathematical Physics, 2016;57(1):015215.

ext/MatrixEnsemblesCUDAExt.jl

@@ -0,0 +1,17 @@
+module MatrixEnsemblesCUDAExt
+
+using MatrixEnsembles
+using CUDA
+using Random
+
+import MatrixEnsembles: curand
+
+function MatrixEnsembles.curand(d::QIContinuousMatrixDistribution)
+    rand(CUDA.default_rng(), d)
+end
+
+function __init__()
+    CUDA.allowscalar(false)
+end
+
+end

src/MatrixEnsembles.jl

@@ -1,36 +1,27 @@
 module MatrixEnsembles
 using LinearAlgebra
+using KernelAbstractions
 import Base: rand, size
 using Random: GLOBAL_RNG, AbstractRNG
-using Requires
-
 
 export rand, size, QIContinuousMatrixDistribution
 export curand
+
 abstract type QIContinuousMatrixDistribution; end
 
 rand(c::QIContinuousMatrixDistribution) = rand(GLOBAL_RNG, c)
 
+"""
+    curand(dist)
+
+Generate a random sample from `dist` using the default CUDA RNG.
+Requires `CUDA.jl` to be loaded.
+"""
+function curand end
+
 include("ginibre.jl")
 include("circular.jl")
 include("wigner.jl")
 include("wishart.jl")
 
-function __init__()
-    @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba" begin
-        if CUDA.functional()
-            const CuArray = CUDA.CuArray
-            const CuVector = CUDA.CuVector
-            const CuMatrix = CUDA.CuMatrix
-            const CuSVD = CUDA.CUSOLVER.CuSVD
-            const CuQR = CUDA.CUSOLVER.CuQR
-
-            CUDA.allowscalar(false)
-            include("cuda/circular.jl") 
-            include("cuda/ginibre.jl")
-            include("cuda/wigner.jl")
-            include("cuda/wishart.jl")
-        end
-    end
-end
 end

src/circular.jl

@@ -15,13 +15,50 @@ const COE = CircularEnsemble{1}
 const CUE = CircularEnsemble{2}
 const CSE = CircularEnsemble{4}
 
+@kernel function identity_kernel!(A)
+    I, J = @index(Global, NTuple)
+    A[I, J] = I == J ? 1 : 0
+end
+
+@kernel function symplectic_kernel!(A)
+    I, J = @index(Global, NTuple)
+    # Check if we are in the same 2x2 block diagonal
+    ki = (I - 1) ÷ 2
+    kj = (J - 1) ÷ 2
+    
+    if ki == kj
+        r = (I - 1) % 2
+        c = (J - 1) % 2
+        
+        if r == 0 && c == 1
+            A[I, J] = -1
+        elseif r == 1 && c == 0
+            A[I, J] = 1
+        else
+            A[I, J] = 0
+        end
+    else
+        A[I, J] = 0
+    end
+end
+
 function _qr_fix!(z::AbstractMatrix)
     q, r = qr!(z)
     d = diag(r)
     ph = d./abs.(d)
     idim = size(r, 1)
-    q = Matrix(q)[:, 1:idim]
-    q = transpose(ph) .* q
+    
+    # Generic densification: Use kernel to create identity
+    m = size(q, 1)
+    dest = similar(z, m, idim)
+    
+    backend = KernelAbstractions.get_backend(dest)
+    kernel = identity_kernel!(backend)
+    kernel(dest, ndrange=size(dest))
+    
+    q_dense = q * dest
+    
+    transpose(ph) .* q_dense
 end
 
 function _qr_fix(z::AbstractMatrix)
@@ -44,7 +81,12 @@ end
 function rand(rng::AbstractRNG, c::CSE)
     z = rand(rng, c.g)
     u = _qr_fix!(z)
-    ur = cat([[0 -1; 1 0] for _=1:c.d÷2]..., dims=[1,2])
+    
+    ur = similar(z, eltype(z), c.d, c.d)
+    backend = KernelAbstractions.get_backend(ur)
+    kernel = symplectic_kernel!(backend)
+    kernel(ur, ndrange=size(ur))
+    
     ur*u*ur'*transpose(u)
 end