Skip to content

Commit 9753b3f

Browse files
lkdvosclaude
andcommitted
Add GPU overrides for _ind_union in CUDA and AMDGPU extensions
Mirrors the existing `_ind_intersect` overrides: collect GPU index vectors to CPU before computing the union, since `union` does not work on GPU arrays. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent f22622c commit 9753b3f

2 files changed

Lines changed: 6 additions & 2 deletions

File tree

ext/MatrixAlgebraKitAMDGPUExt/MatrixAlgebraKitAMDGPUExt.jl

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -161,9 +161,11 @@ function MatrixAlgebraKit._mul_herm!(C::StridedROCMatrix{T}, A::StridedROCMatrix
161161
return C
162162
end
163163

164-
# TODO: intersect doesn't work on GPU
164+
# TODO: intersect/union don't work on GPU
165165
MatrixAlgebraKit._ind_intersect(A::ROCVector{Int}, B::ROCVector{Int}) =
166166
MatrixAlgebraKit._ind_intersect(collect(A), collect(B))
167+
MatrixAlgebraKit._ind_union(A::ROCVector{Int}, B::ROCVector{Int}) =
168+
MatrixAlgebraKit._ind_union(collect(A), collect(B))
167169

168170
function _sylvester(A::AnyROCMatrix, B::AnyROCMatrix, C::AnyROCMatrix)
169171
hX = sylvester(collect(A), collect(B), collect(C))

ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -165,9 +165,11 @@ function MatrixAlgebraKit._mul_herm!(C::StridedCuMatrix{T}, A::StridedCuMatrix{T
165165
return C
166166
end
167167

168-
# TODO: intersect doesn't work on GPU
168+
# TODO: intersect/union don't work on GPU
169169
MatrixAlgebraKit._ind_intersect(A::CuVector{Int}, B::CuVector{Int}) =
170170
MatrixAlgebraKit._ind_intersect(collect(A), collect(B))
171+
MatrixAlgebraKit._ind_union(A::CuVector{Int}, B::CuVector{Int}) =
172+
MatrixAlgebraKit._ind_union(collect(A), collect(B))
171173

172174
function _sylvester(A::AnyCuMatrix, B::AnyCuMatrix, C::AnyCuMatrix)
173175
# https://github.com/JuliaGPU/CUDA.jl/issues/3021

0 commit comments

Comments
 (0)