Skip to content

Commit dbbdd81

Browse files
RXGottliebRXGottlieb
authored
Update for Julia v1.12 (#17)
* Update for Julia v1.12 - Explicitly use `Float64(pi)` in some places in written kernels - Switch to using MultiFloats in `quadrant` to avoid use of `BigFloat` on GPU - Improve sparsity detection for more complex expressions - Move `@register_symbolic` to global scope for `SCMC_sigmoid` * Less restrictive version for MultiFloats * Efficiency improvements - Make efficiency improvements to `kgen` and `create_kernel!` - Add minor utilities to support `kgen` and `create_kernel!` improvements - Remove some unused sections of code - Remove some extraneous notes * Remove old function call * Add better precompilation with PrecompileTools - Add option to skip compilation of `kgen`-generated kernels. If `compile=false`, kgen returns a symbol instead of the function itself. * Correct sparsity ID * Correct case in `cos`
1 parent 5f9c86a commit dbbdd81

8 files changed

Lines changed: 273 additions & 324 deletions

File tree

‎Project.toml‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,15 @@
11
name = "SourceCodeMcCormick"
22
uuid = "a7283dc5-4ecf-47fb-a95b-1412723fc960"
33
authors = ["Robert Gottlieb <Robert.x.gottlieb@uconn.edu>"]
4-
version = "0.5.1"
4+
version = "0.5.2"
55

66
[deps]
77
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
88
Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
99
DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
1010
Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
1111
IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
12+
MultiFloats = "bdf0d083-296b-4888-a5b6-7498122e68a5"
1213
PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
1314
Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
1415
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -20,6 +21,7 @@ CUDA = "5"
2021
DocStringExtensions = "0.8 - 0.9"
2122
Graphs = "1"
2223
IfElse = "0.1.0 - 0.1.1"
24+
MultiFloats = "3.1"
2325
PrecompileTools = "~1"
2426
Reexport = "~1"
2527
StaticArrays = "~1"

‎src/SourceCodeMcCormick.jl‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ using DocStringExtensions
99
using Graphs
1010
using CUDA
1111
using StaticArrays: @MVector
12+
using MultiFloats
13+
using PrecompileTools: @setup_workload, @compile_workload
1214
import Dates
1315
import SymbolicUtils: BasicSymbolic, exprtype, SYM, TERM, ADD, MUL, POW, DIV
1416

@@ -48,8 +50,27 @@ include(joinpath(@__DIR__, "relaxation", "relaxation.jl"))
4850
include(joinpath(@__DIR__, "transform", "transform.jl"))
4951
include(joinpath(@__DIR__, "grad", "grad.jl"))
5052
include(joinpath(@__DIR__, "kernel_writer", "kernel_write.jl"))
51-
include(joinpath(@__DIR__, "precompile.jl"))
52-
_precompile_()
53+
54+
@setup_workload begin
55+
@variables x, y
56+
@compile_workload begin
57+
kgen(1 + x + y^2 + x*y, overwrite=true, compile=false)
58+
kgen(1 +
59+
(-x) +
60+
exp(x) +
61+
log(x) +
62+
(1/x) +
63+
abs(x) +
64+
2*y +
65+
(1/(1+exp(-x))) +
66+
x^3 +
67+
x^4 +
68+
x^3.5 +
69+
cos(x) +
70+
x*y, overwrite=true, compile=false)
71+
kgen(x^3)
72+
end
73+
end
5374

5475
export McCormickIntervalTransform, IntervalTransform
5576

‎src/kernel_writer/kernel_write.jl‎

Lines changed: 147 additions & 273 deletions
Large diffs are not rendered by default.

‎src/kernel_writer/math_kernels.jl‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1365,6 +1365,7 @@ end
13651365

13661366
# Sigmoid function
13671367
# max threads: 640
1368+
@register_symbolic SCMC_sigmoid(x) # Register as symbolic so that we can use it later
13681369
function SCMC_sigmoid_kernel(OUT::CuDeviceMatrix, x::CuDeviceMatrix)
13691370
idx = threadIdx().x + (blockIdx().x - Int32(1)) * blockDim().x
13701371
stride = blockDim().x * gridDim().x
@@ -4670,21 +4671,21 @@ function SCMC_cos_kernel(OUT::CuDeviceMatrix, x::CuDeviceMatrix)
46704671
kL = Base.ceil(-0.5 - x[idx,3]/(2.0*pi))
46714672
xL1 = x[idx,3] + 2.0*pi*kL
46724673
xU1 = x[idx,4] + 2.0*pi*kL
4673-
if (xL1 < -pi) || (xL1 > pi)
4674+
if (xL1 < -pi) || (xL1 > Float64(pi))
46744675
eps_min = NaN
46754676
eps_max = NaN
46764677
elseif xL1 <= 0.0
46774678
if xU1 <= 0.0
46784679
eps_min = x[idx,3]
46794680
eps_max = x[idx,4]
4680-
elseif xU1 >= pi
4681+
elseif xU1 >= Float64(pi)
46814682
eps_min = pi - 2.0*pi*kL
46824683
eps_max = -2.0*pi*kL
46834684
else
46844685
eps_min = (cos(xL1) <= cos(xU1)) ? x[idx,3] : x[idx,4]
46854686
eps_max = -2.0*pi*kL
46864687
end
4687-
elseif xU1 <= pi
4688+
elseif xU1 <= Float64(pi)
46884689
eps_min = x[idx,4]
46894690
eps_max = x[idx,3]
46904691
elseif xU1 >= 2.0*pi
@@ -5449,9 +5450,16 @@ function cos_newton_or_golden_section(x0::Float64, xL::Float64, xU::Float64, env
54495450
return xk
54505451
end
54515452

5452-
# Directly from IntervalArithmetic.jl
5453+
# Similar to IntervalArithmetic.jl, but not using `rem2pi`
54535454
function quadrant(x::Float64)
5454-
x_mod2pi = rem2pi(x, RoundNearest)
5455+
bigx = MultiFloats.Float64x2(x)
5456+
bigpi = MultiFloats._MF{Float64,2}((3.141592653589793, 1.2246467991473532e-16))
5457+
rem = Float64(floor(bigx/bigpi))
5458+
if iseven(rem)
5459+
x_mod2pi = Float64(bigx - rem*bigpi)
5460+
else
5461+
x_mod2pi = Float64(bigx - (rem+1)*bigpi)
5462+
end
54555463

54565464
x_mod2pi < -(pi/2.0) && return (Int32(2), x_mod2pi)
54575465
x_mod2pi < 0 && return (Int32(3), x_mod2pi)

‎src/kernel_writer/string_math_kernels.jl‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10291,21 +10291,21 @@ function SCMC_cos_kernel(OUT::String, v1::String, varlist::Vector{String}, spars
1029110291
write(buffer, " kL = Base.ceil(-0.5 - $v1_lo/(2.0*pi))\n")
1029210292
write(buffer, " xL1 = $v1_lo + 2.0*pi*kL\n")
1029310293
write(buffer, " xU1 = $v1_hi + 2.0*pi*kL\n")
10294-
write(buffer, " if (xL1 < -pi) || (xL1 > pi)\n")
10294+
write(buffer, " if (xL1 < -pi) || (xL1 > Float64(pi))\n")
1029510295
write(buffer, " eps_min = NaN\n")
1029610296
write(buffer, " eps_max = NaN\n")
1029710297
write(buffer, " elseif xL1 <= 0.0\n")
1029810298
write(buffer, " if xU1 <= 0.0\n")
1029910299
write(buffer, " eps_min = $v1_lo\n")
1030010300
write(buffer, " eps_max = $v1_hi\n")
10301-
write(buffer, " elseif xU1 >= pi\n")
10301+
write(buffer, " elseif xU1 >= Float64(pi)\n")
1030210302
write(buffer, " eps_min = pi - 2.0*pi*kL\n")
1030310303
write(buffer, " eps_max = -2.0*pi*kL\n")
1030410304
write(buffer, " else\n")
1030510305
write(buffer, " eps_min = (cos(xL1) <= cos(xU1)) ? $v1_lo : $v1_hi\n")
1030610306
write(buffer, " eps_max = -2.0*pi*kL\n")
1030710307
write(buffer, " end\n")
10308-
write(buffer, " elseif xU1 <= pi\n")
10308+
write(buffer, " elseif xU1 <= Float64(pi)\n")
1030910309
write(buffer, " eps_min = $v1_hi\n")
1031010310
write(buffer, " eps_max = $v1_lo\n")
1031110311
write(buffer, " elseif xU1 >= 2.0*pi\n")
@@ -10572,7 +10572,7 @@ function SCMC_cos_kernel(OUT::String, v1::String, varlist::Vector{String}, spars
1057210572
write(buffer, " if $v1_hi - $v1_lo > 3.141592653589793\n")
1057310573
write(buffer, " $OUT_lo = -1.0\n")
1057410574
write(buffer, " $OUT_hi = 1.0\n")
10575-
write(buffer, " elseif lo_quadrant==2 && hi_quadrant==3\n")
10575+
write(buffer, " elseif lo_quadrant==2 || lo_quadrant==3\n")
1057610576
write(buffer, " $OUT_lo = cos(lo)\n")
1057710577
write(buffer, " $OUT_hi = cos(hi)\n")
1057810578
write(buffer, " else\n")
@@ -10601,21 +10601,21 @@ function SCMC_cos_kernel(OUT::String, v1::String, varlist::Vector{String}, spars
1060110601
write(buffer, " kL = Base.ceil(-0.5 - $v1_lo/(2.0*pi))\n")
1060210602
write(buffer, " xL1 = $v1_lo + 2.0*pi*kL\n")
1060310603
write(buffer, " xU1 = $v1_hi + 2.0*pi*kL\n")
10604-
write(buffer, " if (xL1 < -pi) || (xL1 > pi)\n")
10604+
write(buffer, " if (xL1 < -pi) || (xL1 > Float64(pi))\n")
1060510605
write(buffer, " eps_min = NaN\n")
1060610606
write(buffer, " eps_max = NaN\n")
1060710607
write(buffer, " elseif xL1 <= 0.0\n")
1060810608
write(buffer, " if xU1 <= 0.0\n")
1060910609
write(buffer, " eps_min = $v1_lo\n")
1061010610
write(buffer, " eps_max = $v1_hi\n")
10611-
write(buffer, " elseif xU1 >= pi\n")
10611+
write(buffer, " elseif xU1 >= Float64(pi)\n")
1061210612
write(buffer, " eps_min = pi - 2.0*pi*kL\n")
1061310613
write(buffer, " eps_max = -2.0*pi*kL\n")
1061410614
write(buffer, " else\n")
1061510615
write(buffer, " eps_min = (cos(xL1) <= cos(xU1)) ? $v1_lo : $v1_hi\n")
1061610616
write(buffer, " eps_max = -2.0*pi*kL\n")
1061710617
write(buffer, " end\n")
10618-
write(buffer, " elseif xU1 <= pi\n")
10618+
write(buffer, " elseif xU1 <= Float64(pi)\n")
1061910619
write(buffer, " eps_min = $v1_hi\n")
1062010620
write(buffer, " eps_max = $v1_lo\n")
1062110621
write(buffer, " elseif xU1 >= 2.0*pi\n")

‎src/precompile.jl‎

Lines changed: 0 additions & 25 deletions
This file was deleted.

‎src/transform/utilities.jl‎

Lines changed: 55 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -236,18 +236,22 @@ julia> pull_vars(func)
236236
z
237237
```
238238
"""
239-
pull_vars(term::BasicSymbolic) = pull_vars(Num(term))
240-
function pull_vars(term::Num)
239+
pull_vars(term::BasicSymbolic; get_names::Bool=false) = pull_vars(Num(term), get_names=get_names)
240+
function pull_vars(term::Num; get_names::Bool=false)
241241
vars = Num[]
242242
strings = String[]
243243
if ~(typeof(term.val) <: Real)
244244
vars, strings = _pull_vars(term.val, vars, strings)
245245
vars = vars[sort_vars(strings)]
246246
end
247-
return vars
247+
if get_names
248+
return get_name.(vars)
249+
else
250+
return vars
251+
end
248252
end
249253

250-
function pull_vars(terms::Vector{Num})
254+
function pull_vars(terms::Vector{Num}; get_names::Bool=false)
251255
vars = Num[]
252256
strings = String[]
253257
for term in terms
@@ -258,20 +262,28 @@ function pull_vars(terms::Vector{Num})
258262
if ~isempty(vars)
259263
vars = vars[sort_vars(strings)]
260264
end
261-
return vars
265+
if get_names
266+
return get_name.(vars)
267+
else
268+
return vars
269+
end
262270
end
263271

264-
function pull_vars(eqn::Equation)
272+
function pull_vars(eqn::Equation; get_names::Bool=false)
265273
vars = Num[]
266274
strings = String[]
267275
if ~(typeof(eqn.rhs) <: Real)
268276
vars, strings = _pull_vars(eqn.rhs, vars, strings)
269277
vars = vars[sort_vars(strings)]
270278
end
271-
return vars
279+
if get_names
280+
return get_name.(vars)
281+
else
282+
return vars
283+
end
272284
end
273285

274-
function pull_vars(eqns::Vector{Equation})
286+
function pull_vars(eqns::Vector{Equation}; get_names::Bool=false)
275287
vars = Num[]
276288
strings = String[]
277289
for eqn in eqns
@@ -282,9 +294,13 @@ function pull_vars(eqns::Vector{Equation})
282294
if ~isempty(vars)
283295
vars = vars[sort_vars(strings)]
284296
end
285-
return vars
297+
if get_names
298+
return get_name.(vars)
299+
else
300+
return vars
301+
end
286302
end
287-
function pull_vars(eqn::T) where T<:Real
303+
function pull_vars(eqn::T; get_names::Bool=false) where T<:Real
288304
return Num[]
289305
end
290306

@@ -536,6 +552,35 @@ function extract(eqs::Vector{Equation}, ID::Int=length(eqs))
536552
return final_expr
537553
end
538554

555+
556+
"""
557+
shorten(::Vector{Equation}, ::Int)
558+
559+
Given a set of symbolic equations, and a specific element index,
560+
return a Vector{Equation} that only contains elements needed to
561+
evaluate the chosen element.
562+
```
563+
"""
564+
function shorten(eqs::Vector{Equation}, ID::Int)
565+
indices = Int[]
566+
function delve!(idx, indices, LHS, RHS)
567+
if idx in indices
568+
return nothing
569+
else
570+
for var in RHS[idx]
571+
var_idx = findfirst(==(var), LHS)
572+
if ~isnothing(var_idx) && (var_idx != idx)
573+
delve!(var_idx, indices, LHS, RHS)
574+
end
575+
end
576+
push!(indices, idx)
577+
return nothing
578+
end
579+
end
580+
delve!(ID, indices, Symbol.(getfield.(eqs, :lhs)), pull_vars.(getfield.(eqs, :rhs), get_names=true))
581+
return eqs[indices]
582+
end
583+
539584
"""
540585
convex_evaluator(::Num)
541586
convex_evaluator(::Equation)

‎src/transform/write.jl‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,30 @@ function eqn_edges(a::Vector{Equation})
118118
end
119119
return edgelist, vars
120120
end
121+
function eqn_edges(a::Vector{Equation}, vars::Vector{Symbol})
122+
# Create the list of edges
123+
edgelist = Edge{Int}[]
124+
125+
# Create a mapping dictionary
126+
varid = Dict(vars .=> collect(1:length(vars)))
127+
128+
# Identify LHS variables
129+
LHS_id = [varid[x] for x in Symbol.(getfield.(a, :lhs))]
130+
131+
# Identify RHS variables
132+
RHS_id = [[varid[x] for x in pull_vars(RHS, get_names=true)] for RHS in getfield.(a, :rhs)]
133+
134+
# Create edges of RHS -> LHS
135+
for i in eachindex(LHS_id)
136+
for j in eachindex(RHS_id[i])
137+
if RHS_id[i][j] == LHS_id[i]
138+
continue
139+
end
140+
push!(edgelist, Edge(RHS_id[i][j], LHS_id[i]))
141+
end
142+
end
143+
return edgelist
144+
end
121145

122146
# A new topological sort that tries to minimize the number of temporary vectors
123147
# that need to be preallocated
@@ -128,7 +152,7 @@ function topological_sort(g::SimpleDiGraph; order::Vector{Int64}=Int64[])
128152
for j in g.badjlist[i][sortperm(-lengths)]
129153
recursive_add(g, j, order)
130154
end
131-
if ~in(i, order)
155+
if ~in(i, order) && ~isempty(g.badjlist[i])
132156
push!(order, i)
133157
end
134158
end

0 commit comments

Comments
 (0)