fix: speed up sparse preparation for GPU arrays (#818)

* fix: speed up sparse preparation for GPU arrays

* Remove conditional GPU tests

Author: gdalle

.buildkite/pipeline.yml

@@ -1,6 +1,5 @@
 steps:
   - label: "DI GPU tests"
-    if: build.pull_request.labels includes "gpu"
     plugins:
       - JuliaCI/julia#v1:
           version: "1"

DifferentiationInterface/CHANGELOG.md

@@ -7,10 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.7.2]
+
 ### Feat
 
 - Backend switching for Mooncake ([#768])
 
+### Fixed
+
+- Speed up sparse preparation for GPU arrays ([#818])
+
 ## [0.7.1]
 
 ### Feat
@@ -42,12 +48,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Allocate Enzyme shadow memory during preparation ([#782])
 
-[unreleased]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.1...main
+[unreleased]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.2...main
+[0.7.2]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.1...DifferentiationInterface-v0.7.2
 [0.7.1]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.0...DifferentiationInterface-v0.7.1
 [0.7.0]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.54...DifferentiationInterface-v0.7.0
 [0.6.54]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.53...DifferentiationInterface-v0.6.54
 [0.6.53]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.52...DifferentiationInterface-v0.6.53
 
+[#818]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/818
 [#812]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/812
 [#810]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/810
 [#809]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/809

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.7.1"
+version = "0.7.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -42,7 +42,9 @@ DifferentiationInterfaceGPUArraysCoreExt = "GPUArraysCore"
 DifferentiationInterfaceGTPSAExt = "GTPSA"
 DifferentiationInterfaceMooncakeExt = "Mooncake"
 DifferentiationInterfacePolyesterForwardDiffExt = [
-    "PolyesterForwardDiff", "ForwardDiff", "DiffResults"
+    "PolyesterForwardDiff",
+    "ForwardDiff",
+    "DiffResults",
 ]
 DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"]
 DifferentiationInterfaceSparseArraysExt = "SparseArrays"

DifferentiationInterface/ext/DifferentiationInterfaceGPUArraysCoreExt/DifferentiationInterfaceGPUArraysCoreExt.jl

@@ -13,9 +13,7 @@ end
 function DI.multibasis(a::AbstractGPUArray{T}, inds) where {T}
     b = similar(a)
     fill!(b, zero(T))
-    for i in inds
-        @allowscalar b[i] = one(T)
-    end
+    view(b, inds) .= one(T)
     return b
 end