diff --git a/Project.toml b/Project.toml index 38005b6..12373fe 100644 --- a/Project.toml +++ b/Project.toml @@ -5,6 +5,7 @@ version = "1.0.0-DEV" [deps] KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" [weakdeps] ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" @@ -12,8 +13,18 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" [extensions] BCKChainRulesCore = "ChainRulesCore" +[compat] +MathOptInterface = "1.42.0" +OpenCL = "0.10.3" +pocl_jll = "7.0.0" + [extras] +OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" +pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd" +MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" +MathOptSetDistances = "3b969827-a86c-476c-9527-bb6f1a8fbad5" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Test"] +test = ["Test", "MathOptInterface", "Random", "MathOptSetDistances", "OpenCL", "pocl_jll"] diff --git a/src/BatchConeKernels.jl b/src/BatchConeKernels.jl index 5dd070b..85dbb71 100644 --- a/src/BatchConeKernels.jl +++ b/src/BatchConeKernels.jl @@ -1,6 +1,11 @@ module BatchConeKernels +using MathOptInterface +const MOI = MathOptInterface + using KernelAbstractions const KA = KernelAbstractions +include("distance_sets.jl") + end diff --git a/src/distance_sets.jl b/src/distance_sets.jl new file mode 100644 index 0000000..dd9ce18 --- /dev/null +++ b/src/distance_sets.jl @@ -0,0 +1,147 @@ +## COV_EXCL_START + +@kernel function kerd_zero(Y, @Const(X)) + i = @index(Global) + acc = zero(eltype(X)) + @inbounds for j in 1:size(X, 1) + acc += X[j, i]^2 + end + @inbounds Y[i] = sqrt(acc) +end + +@kernel function kerd_nonneg(Y, @Const(X)) + i = @index(Global) + acc = zero(eltype(X)) + @inbounds for j in 1:size(X, 1) + vj = X[j, i] + acc += (vj < 0 ? -vj : 0)^2 + end + @inbounds Y[i] = sqrt(acc) +end + +@kernel function kerd_nonpos(Y, @Const(X)) + i = @index(Global) + acc = zero(eltype(X)) + @inbounds for j in 1:size(X, 1) + vj = X[j, i] + acc += (vj > 0 ? vj : 0)^2 + end + @inbounds Y[i] = sqrt(acc) +end + +@kernel function kerd_soc(Y, @Const(X), n) + i = @index(Global) + t = X[1, i] + acc = zero(eltype(X)) + @inbounds for j in 2:n + acc += X[j, i]^2 + end + normx = sqrt(acc) + if normx <= t + @inbounds Y[i] = zero(normx) + elseif normx <= -t + @inbounds Y[i] = sqrt(t^2 + normx^2) + else + @inbounds Y[i] = (normx - t) / sqrt(2.0) + end +end + +@kernel function kerd_rsoc(Y, @Const(X), n) + i = @index(Global) + t = X[1, i] + u = X[2, i] + acc = zero(eltype(X)) + @inbounds for j in 3:n + acc += X[j, i]^2 + end + dot_xs = acc + r1 = max(-t, 0) + r2 = max(-u, 0) + r3 = max(dot_xs - 2 * t * u, 0) + @inbounds Y[i] = sqrt(r1^2 + r2^2 + r3^2) +end + +@kernel function kerd_pow(Y, @Const(X), exponent) + i = @index(Global) + x = X[1, i] + y = X[2, i] + z = X[3, i] + if x < 0 || y < 0 + r1 = max(-x, 0) + r2 = max(-y, 0) + @inbounds Y[i] = sqrt(r1^2 + r2^2) + else + result = abs(z) - x^exponent * y^(1-exponent) + @inbounds Y[i] = max(result, 0) + end +end + +@kernel function kerd_dpow(Y, @Const(X), exponent) + i = @index(Global) + u = X[1, i] + v = X[2, i] + w = X[3, i] + ce = 1 - exponent + r1 = max(-u, 0) + r2 = max(-v, 0) + r3 = max(abs(w) - (u/exponent)^exponent * (v/ce)^ce, 0) + @inbounds Y[i] = sqrt(r1^2 + r2^2 + r3^2) +end + +@kernel function kerd_exp(Y, @Const(X)) + i = @index(Global) + x = X[1, i] + y = X[2, i] + z = X[3, i] + if x <= 0 && abs(y) < 1e-10 && z >= 0 + @inbounds Y[i] = zero(x) + else + r1 = max(-y, 0) + r2 = max(y * exp(x / y) - z, 0) #FIXME: div by zero + @inbounds Y[i] = sqrt(r1^2 + r2^2) + end +end + +@kernel function kerd_dexp(Y, @Const(X)) + i = @index(Global) + u = X[1, i] + v = X[2, i] + w = X[3, i] + if abs(u) < 1e-10 && v >= 0 && w >= 0 + @inbounds Y[i] = zero(u) + else + r1 = max(u, 0) + r2 = max(-u * exp(v / u) - ℯ * w, 0) #FIXME: div by zero + @inbounds Y[i] = sqrt(r1^2 + r2^2) + end +end + +## COV_EXCL_STOP + +@inline function distance_to_set(::Type{C}, Y, X, backend=CPU()) where C <: MOI.Zeros + kerd_zero(backend)(Y, X, ndrange=size(X, 2)) +end +@inline function distance_to_set(::Type{C}, Y, X, backend=CPU()) where C <: MOI.Nonnegatives + kerd_nonneg(backend)(Y, X, ndrange=size(X, 2)) +end +@inline function distance_to_set(::Type{C}, Y, X, backend=CPU()) where C <: MOI.Nonpositives + kerd_nonpos(backend)(Y, X, ndrange=size(X, 2)) +end +@inline function distance_to_set(::Type{C}, Y, X, n, backend=CPU()) where C <: MOI.SecondOrderCone + kerd_soc(backend)(Y, X, n, ndrange=size(X, 2)) +end +@inline function distance_to_set(::Type{C}, Y, X, n, backend=CPU()) where C <: MOI.RotatedSecondOrderCone + kerd_rsoc(backend)(Y, X, n, ndrange=size(X, 2)) +end +@inline function distance_to_set(::Type{C}, Y, X, exponent, backend=CPU()) where C <: MOI.PowerCone + kerd_pow(backend)(Y, X, exponent, ndrange=size(X, 2)) +end +@inline function distance_to_set(::Type{C}, Y, X, exponent, backend=CPU()) where C <: MOI.DualPowerCone + kerd_dpow(backend)(Y, X, exponent, ndrange=size(X, 2)) +end +# @inline function distance_to_set(::Type{C}, Y, X, backend=CPU()) where C <: MOI.ExponentialCone +# kerd_exp(backend)(Y, X, ndrange=size(X, 2)) +# end +# @inline function distance_to_set(::Type{C}, Y, X, backend=CPU()) where C <: MOI.DualExponentialCone +# kerd_dexp(backend)(Y, X, ndrange=size(X, 2)) +# end \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index 2c8f3c1..c37a9b4 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,2 +1,97 @@ using BatchConeKernels +using KernelAbstractions using Test +using MathOptSetDistances +using MathOptInterface +using Random + +const MOI = MathOptInterface +const MOSD = MathOptSetDistances +const BCK = BatchConeKernels + + +using OpenCL, pocl_jll + +@testset "distance" begin + Random.seed!(1234) + n = 5 + batch = 100 + Xcpu = randn(n, batch) + X = OpenCL.CLArray(Xcpu) + Y = OpenCL.zeros(batch) + + # Zero cone + BCK.distance_to_set(MOI.Zeros, Y, X, OpenCLBackend()) + Ycpu = Array(Y) + ref_zeros = [distance_to_set(MOSD.DefaultDistance(), Xcpu[:,i], MOI.Zeros(n)) for i in 1:batch] + @test maximum(abs.(Ycpu .- ref_zeros)) < 1e-6 + + # Nonnegatives cone + BCK.distance_to_set(MOI.Nonnegatives, Y, X, OpenCLBackend()) + Ycpu = Array(Y) + ref_nonneg = [distance_to_set(MOSD.DefaultDistance(), Xcpu[:,i], MOI.Nonnegatives(n)) for i in 1:batch] + @test maximum(abs.(Ycpu .- ref_nonneg)) < 1e-6 + + # Nonpositives cone + BCK.distance_to_set(MOI.Nonpositives, Y, X, OpenCLBackend()) + Ycpu = Array(Y) + ref_nonpos = [distance_to_set(MOSD.DefaultDistance(), Xcpu[:,i], MOI.Nonpositives(n)) for i in 1:batch] + @test maximum(abs.(Ycpu .- ref_nonpos)) < 1e-6 + + # SOC + n_soc = 3 + X_soccpu = randn(n_soc, batch) + X_soc = OpenCL.CLArray(X_soccpu) + Y_soc = OpenCL.zeros(batch) + BCK.distance_to_set(MOI.SecondOrderCone, Y_soc, X_soc, n_soc, OpenCLBackend()) + Y_soccpu = Array(Y_soc) + ref_soc = [distance_to_set(MOSD.DefaultDistance(), X_soccpu[:,i], MOI.SecondOrderCone(n_soc)) for i in 1:batch] + @test maximum(abs.(Y_soccpu .- ref_soc)) < 1e-3 + + # RotatedSecondOrderCone + n_rsoc = 4 + X_rsoccpu = randn(n_rsoc, batch) + X_rsoc = OpenCL.CLArray(X_rsoccpu) + Y_rsoc = OpenCL.zeros(batch) + BCK.distance_to_set(MOI.RotatedSecondOrderCone, Y_rsoc, X_rsoc, n_rsoc, OpenCLBackend()) + Y_rsoccpu = Array(Y_rsoc) + ref_rsoc = [distance_to_set(MOSD.DefaultDistance(), X_rsoccpu[:,i], MOI.RotatedSecondOrderCone(n_rsoc)) for i in 1:batch] + @test maximum(abs.(Y_rsoccpu .- ref_rsoc)) < 1e-3 + + # PowerCone + X_powcpu = randn(3, batch) + X_pow = OpenCL.CLArray(X_powcpu) + Y_pow = OpenCL.zeros(batch) + exponent = 0.4 + BCK.distance_to_set(MOI.PowerCone, Y_pow, X_pow, exponent, OpenCLBackend()) + Y_powcpu = Array(Y_pow) + ref_pow = [distance_to_set(MOSD.DefaultDistance(), X_powcpu[:,i], MOI.PowerCone(exponent)) for i in 1:batch] + @test maximum(abs.(Y_powcpu .- ref_pow)) < 1e-3 + + # DualPowerCone + X_dpowcpu = rand(3, batch) + X_dpow = OpenCL.CLArray(X_dpowcpu) + Y_dpow = OpenCL.zeros(batch) + BCK.distance_to_set(MOI.DualPowerCone, Y_dpow, X_dpow, exponent, OpenCLBackend()) + Y_dpowcpu = Array(Y_dpow) + ref_dpow = [distance_to_set(MOSD.DefaultDistance(), X_dpowcpu[:,i], MOI.DualPowerCone(exponent)) for i in 1:batch] + @test maximum(abs.(Y_dpowcpu .- ref_dpow)) < 1e-3 + + # # Exponential cone + # X_expcpu = randn(3, batch) + # X_exp = OpenCL.CLArray(X_expcpu) + # Y_exp = OpenCL.zeros(batch) + # BCK.distance_to_set(MOI.ExponentialCone, Y_exp, X_exp, OpenCLBackend()) + # Y_expcpu = Array(Y_exp) + # ref_exp = [distance_to_set(MOSD.DefaultDistance(), X_expcpu[:,i], MOI.ExponentialCone()) for i in 1:batch] + # @test maximum(abs.(Y_expcpu .- ref_exp)) < 1e-3 + + # # Dual Exponential cone + # X_dexpcpu = randn(3, batch) + # X_dexp = OpenCL.CLArray(X_dexpcpu) + # Y_dexp = OpenCL.zeros(batch) + # BCK.distance_to_set(MOI.DualExponentialCone, Y_dexp, X_dexp, OpenCLBackend()) + # Y_dexpcpu = Array(Y_dexp) + # ref_dexp = [distance_to_set(MOSD.DefaultDistance(), X_dexpcpu[:,i], MOI.DualExponentialCone()) for i in 1:batch] + # @test maximum(abs.(Y_dexpcpu .- ref_dexp)) < 1e-3 +end