Cone distance kernels

Project.toml

@@ -5,15 +5,26 @@ 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"
 
 [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"]

src/BatchConeKernels.jl

@@ -1,6 +1,11 @@
 module BatchConeKernels
 
+using MathOptInterface
+const MOI = MathOptInterface
+
 using KernelAbstractions
 const KA = KernelAbstractions
 
+include("distance_sets.jl")
+
 end

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

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