ExaModelsC: emit several models into one library

ExaModelsC/Project.toml

@@ -6,6 +6,7 @@ authors = ["Sungho Shin <sushin@mit.edu>"]
 [deps]
 ExaModels = "1037b233-b668-4ce9-9b63-f9f681f55dd2"
 JuliaC = "acedd4c2-ced6-4a15-accc-2607eb759ba2"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

ExaModelsC/test/builder_check.py

@@ -0,0 +1,31 @@
+# Drives the schema + builder ABI of a compiled structured model from Python:
+# positional values against the schema's field order, the table as a dict of
+# columns. Inputs mirror the S_* constants in runtests.jl — keep them in sync.
+#
+# usage: python builder_check.py <libpath> <prefix> <n> <outfile>
+import sys
+
+import numpy as np
+
+import cnlpmodels
+
+libpath, prefix, n, outfile = sys.argv[1], sys.argv[2], int(sys.argv[3]), sys.argv[4]
+
+lib = cnlpmodels.load(libpath)
+v0 = np.linspace(0.1, 0.6, n)
+lo = np.full(n, -5.0)
+tab = {
+    "i": np.array([2, 5, 6]),
+    "w": np.array([1.5, 3.0, 0.5]),
+    "s": np.array([2.0, -1.0, 0.0]),
+}
+
+m = cnlpmodels.CModel(lib, n, v0, lo, tab, prefix=prefix)
+x = np.linspace(0.5, 3.0, n)
+
+with open(outfile, "w") as f:
+    f.write("nvar %d\n" % m.nvar)
+    f.write("ncon %d\n" % m.ncon)
+    f.write("obj %.17g\n" % m.obj(x))
+    f.write("grad " + " ".join("%.17g" % v for v in m.grad(x)) + "\n")
+    f.write("cons " + " ".join("%.17g" % v for v in m.cons(x)) + "\n")

ExaModelsC/test/fixtures/RecipeKernels/ext/RecipeKernelsExaModels.jl

@@ -7,4 +7,9 @@ import RecipeKernels
 built with it names a module that cannot be a dependency of anything."
 @inline ramp(i) = 0.25 * i
 
+"""A tuple-returning argument function owned by the EXTENSION, for the argfun
+guard test: it must get past the tuple-return check so the guard itself is what
+refuses it."""
+exttables(n) = (Int(n),)
+
 end # module

ExaModelsC/test/fixtures/RecipeKernels/src/RecipeKernels.jl

@@ -16,7 +16,7 @@ recompiles, whereas a closure carries a gensym that is not.
 """
 module RecipeKernels
 
-export alternating, offsets
+export alternating, offsets, doubled_args, parsed_args, table2d, datant
 
 "Alternating starting point — the shape a per-index start generator takes."
 @inline alternating(i) = isodd(i) ? -1.2 : 1.0
@@ -25,4 +25,30 @@ export alternating, offsets
 form and has to run once the size is known."
 offsets(n) = [2 * div(i - 1, 2) for i in 1:2:max(n - 2, 0)]
 
+# Argument functions for the argfun surface: named, package-owned, returning
+# the argument TUPLE the core is instantiated with — the contract
+# `compile_library` requires so the generated library can call them by name.
+
+"The integer-kind argument function: `P_new(n)` hands `n` to it."
+doubled_args(n::Integer) = (2 * Int(n),)
+
+"The string-kind argument function: `P_new_str(s)` hands the string to it —
+standing in for a case-file path parsed on the far side of the boundary."
+parsed_args(s::AbstractString) = (parse(Int, s),)
+
+"A 2-D table of heterogeneous tuples over an open size — the COPS-shaped
+deferred data (a Matrix of (Int,Int,Int,F64,F64,F64) rows over `1:n × 1:3`)
+that exercises product iterators and per-element instantiate typing."
+table2d(n) = [(i, j, i + j, 0.5, 1.5, 2.5) for i in 1:n, j in 1:3]
+
+"A NamedTuple-returning data function — the shape a modelling library's
+`*_data(n)` takes, with fields PROJECTED out of the deferred call
+(`d.v_start` for starts, `d.tab` as an iterator): the projections are
+`ArgIndexed` nodes over the one `ArgNode1`, which is the structure the
+gasoil-class cores carry."
+datant(n) = (
+    v_start = fill(0.5, 3 * n),
+    tab = [(i, j, 0.5 + i) for i in 1:n, j in 1:3],
+)
+
 end # module

src/argument.jl

@@ -181,17 +181,17 @@ julia> ExaModels.instantiate(x, (nh = 2,)) === x     # identity, no arg dependen
 true
 ```
 """
-@inline instantiate(x, a...) = x
-@inline instantiate(::ArgSource{K}, a...) where {K} = a[K]
-@inline instantiate(n::ArgIndexed{I, J}, a...) where {I, J} =
+@inline instantiate(x, a::Vararg{Any,N}) where {N} = x
+@inline instantiate(::ArgSource{K}, a::Vararg{Any,N}) where {K, N} = a[K]
+@inline instantiate(n::ArgIndexed{I, J}, a::Vararg{Any,N}) where {I, J, N} =
     _arg_access(instantiate(getfield(n, :inner), a...), J)
-@inline instantiate(n::ArgNode1, a...) =
+@inline instantiate(n::ArgNode1, a::Vararg{Any,N}) where {N} =
     getfield(n, :f)(instantiate(getfield(n, :inner), a...))
-@inline instantiate(n::ArgNode2, a...) = getfield(n, :f)(
+@inline instantiate(n::ArgNode2, a::Vararg{Any,N}) where {N} = getfield(n, :f)(
     instantiate(getfield(n, :inner1), a...),
     instantiate(getfield(n, :inner2), a...),
 )
-@inline instantiate(n::ArgCall, a...) =
+@inline instantiate(n::ArgCall, a::Vararg{Any,N}) where {N} =
     getfield(n, :f)(map(x -> instantiate(x, a...), getfield(n, :args))...)
 
 @inline _arg_access(x, j::Symbol) = getproperty(x, j)
@@ -204,8 +204,34 @@ true
 # that looks fully instantiated and is not.  Mapping costs nothing observable:
 # an immutable tuple rebuilds `===` to itself, and each element is passed
 # through by identity when it has no dependency of its own.
-@inline instantiate(t::Tuple, a...) = map(x -> instantiate(x, a...), t)
-@inline instantiate(t::NamedTuple, a...) = map(x -> instantiate(x, a...), t)
+# `Vararg{Any,N}` forces specialization on the arguments' concrete types:
+# a vararg that is only splatted through is otherwise left unspecialized
+# (Julia's passthrough heuristic), and the `map` below then carries a dynamic
+# call that `juliac --trim=safe` cannot resolve. Harmless under the JIT,
+# load-bearing for AOT — same reason on the `ExaCore` method in nlp.jl.
+# Unrolled as well as specialized: `Vararg{Any,N}` is necessary — without it
+# the vararg is not specialized at all — but on a core the size of an AC
+# OPF's the `map` closure is still an unresolved call under `--trim=safe`.
+# The recursion gives each element a direct `instantiate` call.
+#
+# THE PARAMETER ROLES ARE LOAD-BEARING, both of them. The ELEMENTS stay
+# structural (`t::Tuple`, recursed via `first`/`Base.tail`): splatting them
+# into a bare vararg (`_f(a, x, xs...)`) puts the heterogeneous element
+# types into a pass-through-only vararg, which Julia leaves unspecialized —
+# the per-element types are lost, and a core whose objectives project
+# fields out of a deferred data call (gasoil's shape) widens at the
+# instantiate boundary: 6 verifier errors, reproduced and bisected to this
+# hunk alone (6 → 0 on the swap, real COPS gasoil as the instrument). The
+# ARGUMENTS ride the specialized `Vararg{Any,N}`. Inverting either role
+# reintroduces the widening; the ExaModelsC suite's models are too small
+# and too homogeneous to see it — the COPS CI pin on this branch is the
+# regression gate for this class.
+@inline _instantiate_each(::Tuple{}, a::Vararg{Any,N}) where {N} = ()
+@inline _instantiate_each(t::Tuple, a::Vararg{Any,N}) where {N} =
+    (instantiate(first(t), a...), _instantiate_each(Base.tail(t), a...)...)
+@inline instantiate(t::Tuple, a::Vararg{Any,N}) where {N} = _instantiate_each(t, a...)
+@inline instantiate(t::NamedTuple{K}, a::Vararg{Any,N}) where {K, N} =
+    NamedTuple{K}(_instantiate_each(Tuple(t), a...))
 
 """
     _anyarg(xs...)

src/graph.jl

@@ -229,6 +229,54 @@ struct Node2{F,I1,I2} <: AbstractNode
     inner2::I2
 end
 
+# ── Re-indexing a built subexpression ─────────────────────────────────────────
+#
+# `@add_expr` stores a body that is spliced in wherever it is referenced, so the
+# stored form has to be re-indexable.  A closure did that by being called again
+# with a new index; a node does it by substituting the `DataSource` it was built
+# against.
+#
+# Storing a node rather than a closure is what makes subexpressions usable in a
+# recipe.  A closure captures the `Variable`s it mentions, so when a size is a
+# placeholder the `ArgSource` ends up inside the closure's *type*, where
+# `instantiate` cannot reach it — and Julia emits no constructor for a closure,
+# so it cannot be rebuilt field-wise either.  Nodes have constructors and
+# dispatch, and the operation lives in the type parameter, so this walk is
+# static: no reflection, no `@generated`, nothing `--trim=safe` will refuse.
+@inline _reindex(::DataSource, i) = i
+@inline _reindex(n::DataIndexed{I, J}, i) where {I, J} =
+    _reindexed_access(_reindex(getfield(n, :inner), i), J)
+# Once the source has been replaced by a concrete element — an integer, or the
+# tuple a multi-dimensional or product-iterator generator destructures — the
+# lookup is no longer symbolic and has to be performed, not rebuilt.  Leaving it
+# as a `DataIndexed` around a raw tuple is what broke the multi-dim cases.
+@inline _reindexed_access(inner::AbstractNode, J) = DataIndexed(inner, J)
+@inline _reindexed_access(inner, J::Symbol) = getproperty(inner, J)
+@inline _reindexed_access(inner, J) = getindex(inner, J)
+@inline _reindex(n::Var, i) = Var(_reindex(getfield(n, :i), i))
+@inline _reindex(n::ParameterNode, i) = ParameterNode(_reindex(getfield(n, :i), i))
+@inline function _reindex(n::Node1{F}, i) where {F}
+    # Julia synthesises no partially-parameterised constructor, so the child
+    # types are given explicitly; they are known here, so this stays static.
+    inner = _reindex(getfield(n, :inner), i)
+    # A fully-concrete subterm is COMPUTED, not wrapped — the closure this
+    # storage replaces evaluated such terms eagerly, and a node with only
+    # Real children has no evaluation method (register.jl's one-sided Real
+    # methods tie for two: `s[2]` substituting into `i - 1` was ambiguous at
+    # the adjoint call). `F.instance` exists because ops are named functions.
+    inner isa Real && return F.instance(inner)
+    return Node1{F, typeof(inner)}(inner)
+end
+@inline function _reindex(n::Node2{F}, i) where {F}
+    a = _reindex(getfield(n, :inner1), i)
+    b = _reindex(getfield(n, :inner2), i)
+    a isa Real && b isa Real && return F.instance(a, b)
+    return Node2{F, typeof(a), typeof(b)}(a, b)
+end
+# Leaves with no data dependence — constants, variable/parameter sources, plain
+# numbers — are their own re-indexing.
+@inline _reindex(x, i) = x
+
 struct FirstFixed{F}
     inner::F
 end
@@ -542,30 +590,30 @@ end
 # (`Constant`, `Null`, `VarSource`, `DataSource`, …) fall through to the generic
 # identity in argument.jl.
 
-@inline function instantiate(n::Var{I}, a...) where {I}
+@inline function instantiate(n::Var{I}, a::Vararg{Any,N}) where {I, N}
     i = instantiate(n.i, a...)
     return Var{typeof(i)}(i)
 end
-@inline function instantiate(n::ParameterNode{I}, a...) where {I}
+@inline function instantiate(n::ParameterNode{I}, a::Vararg{Any,N}) where {I, N}
     i = instantiate(n.i, a...)
     return ParameterNode{typeof(i)}(i)
 end
-@inline function instantiate(n::Node1{F,I}, a...) where {F,I}
+@inline function instantiate(n::Node1{F,I}, a::Vararg{Any,N}) where {F,I, N}
     i = instantiate(n.inner, a...)
     return Node1{F,typeof(i)}(i)
 end
-@inline function instantiate(n::Node2{F,I1,I2}, a...) where {F,I1,I2}
+@inline function instantiate(n::Node2{F,I1,I2}, a::Vararg{Any,N}) where {F,I1,I2, N}
     i1 = instantiate(n.inner1, a...)
     i2 = instantiate(n.inner2, a...)
     return Node2{F,typeof(i1),typeof(i2)}(i1, i2)
 end
-@inline instantiate(n::SumNode, a...) = SumNode(instantiate(n.inners, a...))
-@inline instantiate(n::ProdNode, a...) = ProdNode(instantiate(n.inners, a...))
+@inline instantiate(n::SumNode, a::Vararg{Any,N}) where {N} = SumNode(instantiate(n.inners, a...))
+@inline instantiate(n::ProdNode, a::Vararg{Any,N}) where {N} = ProdNode(instantiate(n.inners, a...))
 # `DataIndexed` overrides `getproperty` to keep building access paths, so its
 # own field has to be read with `getfield`.
-@inline instantiate(n::DataIndexed{I,J}, a...) where {I,J} =
+@inline instantiate(n::DataIndexed{I,J}, a::Vararg{Any,N}) where {I,J, N} =
     DataIndexed(instantiate(getfield(n, :inner), a...), J)
-@inline instantiate(p::Pair, a...) = instantiate(p.first, a...) => instantiate(p.second, a...)
+@inline instantiate(p::Pair, a::Vararg{Any,N}) where {N} = instantiate(p.first, a...) => instantiate(p.second, a...)
 # An `ArgLeaf` does not survive instantiation: it *becomes* the scalar, leaving
 # a graph indistinguishable from one built with concrete sizes.
-@inline instantiate(n::ArgLeaf, a...) = instantiate(getfield(n, :a), a...)
+@inline instantiate(n::ArgLeaf, a::Vararg{Any,N}) where {N} = instantiate(getfield(n, :a), a...)

src/nlp.jl

@@ -58,7 +58,7 @@ end
 Base.show(io::IO, s::Expression) = _show_expression(io, s)
 function _show_expression(io::IO, s::Expression)
     expr = try
-        _expr_string(s.f(DataSource()))
+        _expr_string(s.f)
     catch
         "(?)"
     end
@@ -579,7 +579,9 @@ is known statically and destructuring stays inferable.
 # rather than re-derived: it is the float type the core was created with, and
 # instantiating changes sizes, never the element type.
 
-function instantiate(c::ExaCore{T}, a...) where {T}
+# `Vararg{Any,N}`: see the note on the `Tuple` method in argument.jl — forces
+# specialization so the field-by-field mapping stays static under `--trim`.
+function instantiate(c::ExaCore{T}, a::Vararg{Any,N}) where {T, N}
     return ExaCore{T}(
         c.name,
         c.backend,
@@ -613,24 +615,24 @@ function instantiate(c::ExaCore{T}, a...) where {T}
     )
 end
 
-instantiate(v::Variable, a...) =
+instantiate(v::Variable, a::Vararg{Any,N}) where {N} =
     Variable(instantiate(v.size, a...), instantiate(v.length, a...), instantiate(v.offset, a...),
              v.name, instantiate(v.tag, a...))
-instantiate(p::Parameter, a...) =
+instantiate(p::Parameter, a::Vararg{Any,N}) where {N} =
     Parameter(instantiate(p.size, a...), instantiate(p.length, a...), instantiate(p.offset, a...),
               instantiate(p.tag, a...))
-instantiate(e::Expression, a...) =
+instantiate(e::Expression, a::Vararg{Any,N}) where {N} =
     Expression(instantiate(e.size, a...), instantiate(e.length, a...), instantiate(e.f, a...),
                instantiate(e.iter, a...), instantiate(e.tag, a...))
-instantiate(o::Objective, a...) = Objective(instantiate(o.f, a...), instantiate(o.itr, a...))
-instantiate(c::Constraint, a...) =
+instantiate(o::Objective, a::Vararg{Any,N}) where {N} = Objective(instantiate(o.f, a...), instantiate(o.itr, a...))
+instantiate(c::Constraint, a::Vararg{Any,N}) where {N} =
     Constraint(instantiate(c.f, a...), instantiate(c.itr, a...), instantiate(c.offset, a...),
                instantiate(c.size, a...), instantiate(c.tag, a...))
-instantiate(c::ConstraintAugmentation, a...) =
+instantiate(c::ConstraintAugmentation, a::Vararg{Any,N}) where {N} =
     ConstraintAugmentation(instantiate(c.f, a...), instantiate(c.itr, a...),
                            instantiate(c.oa, a...), instantiate(c.dims, a...),
                            instantiate(c.tag, a...))
-instantiate(f::SIMDFunction, a...) =
+instantiate(f::SIMDFunction, a::Vararg{Any,N}) where {N} =
     SIMDFunction(instantiate(f.f, a...), f.comp1, f.comp2,
                  instantiate(f.o0, a...), instantiate(f.o1, a...), instantiate(f.o2, a...),
                  f.o1step, f.o2step)
@@ -924,24 +926,24 @@ end
 @inline function Base.getindex(s::Expression, i::I) where {I <: Integer}
     _bound_check(s.size, i)
     idx = i - _start(s.size[1]) + 1
-    return s.f(s.iter[idx])
+    return _reindex(s.f, s.iter[idx])
 end
 @inline function Base.getindex(s::Expression, i)
     # Symbolic index case - the symbolic index IS the iterator element
     # No adjustment needed; the index is used directly in expression building
-    return s.f(i)
+    return _reindex(s.f, i)
 end
 @inline function Base.getindex(s::Expression, is::Vararg{I, N}) where {I <: Integer, N}
     @assert(length(is) == length(s.size), "Expression index dimension error")
     _bound_check(s.size, is)
     idx = idxx(is .- (_start.(s.size) .- 1), _length.(s.size))
-    return s.f(s.iter[idx])
+    return _reindex(s.f, s.iter[idx])
 end
 @inline function Base.getindex(s::Expression, is...)
     # Symbolic indices case - the symbolic indices ARE the iterator elements
     # No adjustment needed; the indices are used directly in expression building
     @assert(length(is) == length(s.size), "Expression index dimension error")
-    return s.f(is)
+    return _reindex(s.f, is)
 end
 
 @inline function Base.getindex(p::P, i) where {P<:Parameter}
@@ -1048,11 +1050,11 @@ end
 
 # `start = (f(i) for i in 1:arg.N)` — the body is untouched (it runs per element
 # once the iterator is concrete); only what it iterates is deferred.
-@inline instantiate(g::Base.Generator, a...) = Base.Generator(g.f, instantiate(g.iter, a...))
+@inline instantiate(g::Base.Generator, a::Vararg{Any,N}) where {N} = Base.Generator(g.f, instantiate(g.iter, a...))
 
 # A product iterator is arg-dependent when any of the ranges it crosses is.
 @inline _anyarg(p::Base.Iterators.ProductIterator, xs...) = _anyarg(p.iterators..., xs...)
-@inline instantiate(p::Base.Iterators.ProductIterator, a...) =
+@inline instantiate(p::Base.Iterators.ProductIterator, a::Vararg{Any,N}) where {N} =
     Base.Iterators.product(instantiate(p.iterators, a...)...)
 
 # `append!` mutates its accumulator and returns it.  That is exactly right while
@@ -1744,7 +1746,9 @@ c, s = add_expr(c, x[i, k]^2 for (i, k) in itr)
     gen = _adapt_gen(gen)
     n = length(gen.iter)
 
-    ex = Expression(ns, n, gen.f, collect(gen.iter), tag)
+    # Store the body as a node, built once against a `DataSource`, rather than
+    # as the generator's closure — see `_reindex` in graph.jl.
+    ex = Expression(ns, n, gen.f(DataSource()), collect(gen.iter), tag)
     return (ExaCore(c; refs = add_refs(c.refs, name, ex)), ex)
 end
 
@@ -2116,6 +2120,15 @@ function multipliers(result::SolverCore.AbstractExecutionStats, y::Constraint)
 end
 
 _adapt_gen(gen) = Base.Generator(gen.f, collect(gen.iter))
+# A placeholder iterable resolves at instantiation, so what to do with it can
+# only be decided then. The unconditional deferred `collect` was wrong on the
+# GPU: `collect(::CuArray)` is a host Vector, so a converted argument was
+# silently pulled back to the CPU and the kernel met a non-bitstype — invisible
+# on the CPU, where `collect` of a Vector is just a copy.
+_maybe_collect(x) = collect(x)
+_maybe_collect(x::Union{AbstractArray, AbstractRange}) = x
+@inline _adapt_gen(gen::Base.Generator{I}) where {I<:AbstractArgNode} =
+    Base.Generator(gen.f, ArgCall(_maybe_collect, (gen.iter,)))
 # `for i in 1:nh, j in 1:nc` over symbolic ranges: `collect` on the product
 # needs `axes`, and `axes` needs concrete lengths, so it cannot run yet.  Defer
 # the collect itself — it happens at instantiation, on real ranges.

test/ArgumentTest/ArgumentTest.jl

@@ -396,6 +396,31 @@ function runtests()
             @test repr(1:arg.N) == "(1:arg.N)"
             @test repr(length(arg.v)) == "length(arg.v)"
         end
+
+        @testset "a converted iterable stays on its backend" begin
+            # The deferred collect in _adapt_gen must not migrate an argument
+            # back to the host: collect(::CuArray) is a Vector, so a recipe
+            # instantiated with device-converted arguments produced a model
+            # whose kernels met a non-bitstype host array. Invisible on the
+            # CPU, where collect of a Vector is just a copy — which is why
+            # this asserts RESIDENCE, not values.
+            for backend in Main.BACKENDS
+                backend === nothing && continue
+                c, d = ExaModels.ExaCore(nargs = Val(1))
+                ExaModels.@add_var(c, x, length(d.rows))
+                ExaModels.@add_obj(c, (x[r.i] - r.w)^2 for r in d.rows)
+                rows = [(i = i, w = Float64(i)) for i = 1:5]
+                dev = ExaModels.convert_array(rows, backend)
+                m = ExaModels.ExaModel(c, (; rows = dev); backend = backend)
+                @testset "$backend" begin
+                    # Residence via typeof equality. Only a device backend
+                    # discriminates — on CPU() the converted array IS a Vector,
+                    # so the pre-fix collect was a same-type copy there; the
+                    # CUDA leg is the one this test exists for.
+                    @test typeof(m.objs[1].itr) == typeof(dev)
+                end
+            end
+        end
     end
 end