Merge pull request #1496 from AayushSabharwal/as/better-build-function

ChrisRackauckas · web-flow · commit 4bdefb3697e3 · 2025-03-22T17:00:09.000-01:00
fix: fix `iip_config` with RGF
diff --git a/src/arrays.jl b/src/arrays.jl
@@ -1005,7 +1005,7 @@ function _array_toexpr(x, st)
         [
             Assignment(outsym, term(zeros, Float64, term(map, length, shape(x)))),
             Assignment(Symbol("%$outsym"), inplace_expr(x, outsym))
-        ], outsym, true)
+        ], outsym, false)
 
     toexpr(ex, st)
 end
diff --git a/src/build_function.jl b/src/build_function.jl
@@ -113,10 +113,10 @@ function _build_function(target::JuliaTarget, op, args...;
                          cse = false,
                          nanmath = true,
                          kwargs...)
-
+    op = _recursive_unwrap(op)
     states.rewrites[:nanmath] = nanmath
     dargs = map((x) -> destructure_arg(x[2], !checkbounds, default_arg_name(x[1])), enumerate(collect(args)))
-    fun = Func(dargs, [], unwrap(op))
+    fun = Func(dargs, [], op)
     if wrap_code !== nothing
         fun = wrap_code(fun)
     end
@@ -135,7 +135,9 @@ function _build_function(target::JuliaTarget, op, args...;
     end
 end
 
-const UNIMPLEMENTED_EXPR = :(function (args...); $throw_missing_specialization(length(args)); end)
+function get_unimplemented_expr(dargs)
+    Func(dargs, [], term(throw_missing_specialization, length(dargs)))
+end
 
 SymbolicUtils.Code.get_rewrites(x::Arr) = SymbolicUtils.Code.get_rewrites(unwrap(x))
 
@@ -151,21 +153,21 @@ function _build_function(target::JuliaTarget, op::Union{Arr, ArrayOp, SymbolicUt
                          wrap_code = (identity, identity),
                          iip_config = (true, true),
                          kwargs...)
-
+    op = _recursive_unwrap(op)
     dargs = map((x) -> destructure_arg(x[2], !checkbounds, default_arg_name(x[1])), enumerate(collect(args)))
     states.rewrites[:nanmath] = nanmath
     if iip_config[1]
-        oop_expr = wrap_code[1](Func(dargs, [], unwrap(op)))
+        oop_expr = wrap_code[1](Func(dargs, [], op))
     else
-        oop_expr = UNIMPLEMENTED_EXPR
+        oop_expr = get_unimplemented_expr(dargs)
     end
 
     outsym = DEFAULT_OUTSYM
-    body = inplace_expr(unwrap(op), outsym)
+    body = inplace_expr(op, outsym)
     if iip_config[2]
         iip_expr = wrap_code[2](Func(vcat(outsym, dargs), [], body))
     else
-        iip_expr = UNIMPLEMENTED_EXPR
+        iip_expr = get_unimplemented_expr([outsym; dargs])
     end
 
     if cse
@@ -196,10 +198,7 @@ function _build_and_inject_function(mod::Module, ex)
     elseif ex.head == :(->)
         return _build_and_inject_function(mod, Expr(:function, ex.args...))
     end
-    # XXX: Workaround to specify the module as both the cache module AND context module.
-    # Currently, the @RuntimeGeneratedFunction macro only sets the context module.
-    module_tag = getproperty(mod, RuntimeGeneratedFunctions._tagname)
-    RuntimeGeneratedFunctions.RuntimeGeneratedFunction(module_tag, module_tag, ex; opaque_closures=false)
+    RuntimeGeneratedFunction(mod, mod, ex)
 end
 
 toexpr(n::Num, st) = toexpr(value(n), st)
@@ -305,7 +304,10 @@ function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
                        iip_config = (true, true),
                        nanmath = true,
                        parallel=nothing, cse = false, kwargs...)
-
+    if rhss isa SubArray
+        rhss = copy(rhss)
+    end
+    rhss = _recursive_unwrap(rhss)
     states.rewrites[:nanmath] = nanmath
     # We cannot switch to ShardedForm because it deadlocks with
     # RuntimeGeneratedFunctions
@@ -323,7 +325,7 @@ function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
             oop_expr = wrap_code[1](oop_expr)
         end
     else
-        oop_expr = UNIMPLEMENTED_EXPR
+        oop_expr = get_unimplemented_expr(dargs)
     end
 
 
@@ -340,7 +342,7 @@ function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
             iip_expr = wrap_code[2](iip_expr)
         end
     else
-        iip_expr = UNIMPLEMENTED_EXPR
+        iip_expr = get_unimplemented_expr([DEFAULT_OUTSYM; dargs])
     end
 
     if cse
@@ -466,7 +468,7 @@ function _make_sparse_array(arr, similarto)
         return term(setparent, nzmap(Returns(true), arr), newarr)
     else
         newarr = _make_array(arr.nzval, Vector{symtype(eltype(arr))})
-        return Let([Assignment(:__reference, term(copy, nzmap(Returns(true), arr)))], term(set_nzval, :__reference, newarr), true)
+        return Let([Assignment(:__reference, term(copy, nzmap(Returns(true), arr)))], term(set_nzval, :__reference, newarr), false)
     end
 end
 
@@ -539,10 +541,13 @@ function set_array(s::ShardedForm, closed_args, out, outputidxs, rhss, checkboun
 end
 
 function _set_array(out, outputidxs, rhss::AbstractSparseArray, checkbounds, skipzeros)
-    Let([Assignment(Symbol("%$out"), _set_array(LiteralExpr(:($out.nzval)), nothing, rhss.nzval, checkbounds, skipzeros))], out)
+    Let([Assignment(Symbol("%$out"), _set_array(LiteralExpr(:($out.nzval)), nothing, rhss.nzval, checkbounds, skipzeros))], out, false)
 end
 
 function _set_array(out, outputidxs, rhss::AbstractArray, checkbounds, skipzeros)
+    if parent(rhss) !== rhss
+        return _set_array(out, outputidxs, parent(rhss), checkbounds, skipzeros)
+    end
     if outputidxs === nothing
         outputidxs = collect(eachindex(rhss))
     end
diff --git a/src/variable.jl b/src/variable.jl
@@ -536,7 +536,10 @@ function SymbolicIndexingInterface.symbolic_evaluate(ex::Union{Num, Arr, Symboli
 end
 
 function _recursive_unwrap(val)
-    if symbolic_type(val) == NotSymbolic() && val isa AbstractArray
+    if symbolic_type(val) == NotSymbolic() && val isa Union{AbstractArray, Tuple}
+        if parent(val) !== val
+            return Setfield.@set val.parent = _recursive_unwrap(parent(val))
+        end
         return _recursive_unwrap.(val)
     else
         return unwrap(val)
diff --git a/test/build_function.jl b/test/build_function.jl
@@ -316,3 +316,32 @@ end
     fn = build_function(f(x), DestructuredArgs([f]; create_bindings = false), x; expression = Val{false})
     @test fn([isodd], 3)
 end
+
+@testset "iip_config with RGF" begin
+    @variables a b
+    oop, iip = build_function([a + b, a - b], a, b; iip_config = (false, false), expression = Val{false})
+    @test_throws ArgumentError oop(1, 2)
+    @test_throws ArgumentError iip(ones(2), 1, 2)
+
+    @variables a[1:2]
+    oop, iip = build_function(a .* 2, a; iip_config = (false, false), expression = Val{false})
+    @test_throws ArgumentError oop(ones(2))
+    @test_throws ArgumentError iip(ones(2), ones(2))
+end
+
+@testset "unwrapping/CSE in array of symbolics codegen" begin
+    @variables a b
+    oop, _ = build_function([a^2 + b^2, a^2 + b^2], a, b; expression = Val{true}, cse = true)
+
+    function find_create_array(expr)
+        while expr isa Expr && (!Meta.isexpr(expr, :call) || expr.args[1] != SymbolicUtils.Code.create_array)
+            expr = expr.args[end]
+        end
+        return expr
+    end
+
+    expr = find_create_array(oop)
+    # CSE works, we just need to test that it's happening and OOP is the easiest way to do it
+    @test Meta.isexpr(expr, :call) && expr.args[1] == SymbolicUtils.Code.create_array &&
+          expr.args[end] isa Symbol && expr.args[end-1] isa Symbol
+end
diff --git a/test/build_function_tests/intermediate-exprs-inplace.jl b/test/build_function_tests/intermediate-exprs-inplace.jl
@@ -1,6 +1,8 @@
 :(function (ˍ₋out, u)
       begin
-          ˍ₋out_input_1 = let _out = (zeros)(Float64, (map)(length, (Base.OneTo(5), Base.OneTo(5)))), var"%_out" = for var"%jj′" = (zip)(1:5, (Symbolics.reset_to_one)(1:5))
+          ˍ₋out_input_1 = begin
+                  _out = (zeros)(Float64, (map)(length, (Base.OneTo(5), Base.OneTo(5))))
+                  var"%_out" = for var"%jj′" = (zip)(1:5, (Symbolics.reset_to_one)(1:5))
                           begin
                               j = var"%jj′"[1]
                               j′ = var"%jj′"[2]
diff --git a/test/build_function_tests/intermediate-exprs-outplace.jl b/test/build_function_tests/intermediate-exprs-outplace.jl
@@ -1,6 +1,10 @@
 :(function (u,)
-      let _out = (zeros)(Float64, (map)(length, (Base.OneTo(5), Base.OneTo(5)))), var"%_out" = begin
-                  _out_input_1 = let _out = (zeros)(Float64, (map)(length, (Base.OneTo(5), Base.OneTo(5)))), var"%_out" = for var"%jj′" = (zip)(1:5, (Symbolics.reset_to_one)(1:5))
+      begin
+          _out = (zeros)(Float64, (map)(length, (Base.OneTo(5), Base.OneTo(5))))
+          var"%_out" = begin
+                  _out_input_1 = begin
+                          _out = (zeros)(Float64, (map)(length, (Base.OneTo(5), Base.OneTo(5))))
+                          var"%_out" = for var"%jj′" = (zip)(1:5, (Symbolics.reset_to_one)(1:5))
                                   begin
                                       j = var"%jj′"[1]
                                       j′ = var"%jj′"[2]
diff --git a/test/build_function_tests/manual-limits-outplace.jl b/test/build_function_tests/manual-limits-outplace.jl
@@ -1,5 +1,7 @@
 :(function (u,)
-      let _out = (zeros)(Float64, (map)(length, (Base.OneTo(5), Base.OneTo(5)))), var"%_out" = for var"%jj′" = (zip)(1:5, (Symbolics.reset_to_one)(1:5))
+      begin
+          _out = (zeros)(Float64, (map)(length, (Base.OneTo(5), Base.OneTo(5))))
+          var"%_out" = for var"%jj′" = (zip)(1:5, (Symbolics.reset_to_one)(1:5))
                   begin
                       j = var"%jj′"[1]
                       j′ = var"%jj′"[2]
diff --git a/test/build_function_tests/stencil-broadcast-outplace.jl b/test/build_function_tests/stencil-broadcast-outplace.jl
@@ -1,5 +1,7 @@
 :(function (x,)
-      let _out = (zeros)(Float64, (map)(length, (1:6, 1:6))), var"%_out" = begin
+      begin
+          _out = (zeros)(Float64, (map)(length, (1:6, 1:6)))
+          var"%_out" = begin
                   _out_2_input_1 = (broadcast)(+, x, (adjoint)(x))
                   _out_1 = (view)(_out, 1:6, 1:6)
                   var"%_out_1" = (Symbolics.broadcast_assign!)(_out_1, 0)
diff --git a/test/build_function_tests/stencil-extents-outplace.jl b/test/build_function_tests/stencil-extents-outplace.jl
@@ -1,5 +1,7 @@
 :(function (x,)
-      let _out = (zeros)(Float64, (map)(length, (1:5, 1:5))), var"%_out" = begin
+      begin
+          _out = (zeros)(Float64, (map)(length, (1:5, 1:5)))
+          var"%_out" = begin
                   _out_1 = (view)(_out, 1:5, 1:5)
                   var"%_out_1" = (Symbolics.broadcast_assign!)(_out_1, 0)
                   _out_2 = (view)(_out, 2:4, 2:4)
diff --git a/test/build_function_tests/stencil-transpose-arrayop-outplace.jl b/test/build_function_tests/stencil-transpose-arrayop-outplace.jl
@@ -1,5 +1,7 @@
 :(function (x,)
-      let _out = (zeros)(Float64, (map)(length, (1:6, 1:6))), var"%_out" = begin
+      begin
+          _out = (zeros)(Float64, (map)(length, (1:6, 1:6)))
+          var"%_out" = begin
                   _out_1 = (view)(_out, 2:5, 2:5)
                   var"%_out_1" = for var"%jj′" = (zip)(1:4, (Symbolics.reset_to_one)(1:4))
                           begin
diff --git a/test/build_function_tests/transpose-outplace.jl b/test/build_function_tests/transpose-outplace.jl
@@ -1,5 +1,7 @@
 :(function (x,)
-      let _out = (zeros)(Float64, (map)(length, (Base.OneTo(4), Base.OneTo(4)))), var"%_out" = for var"%jj′" = (zip)(1:4, (Symbolics.reset_to_one)(1:4))
+      begin
+          _out = (zeros)(Float64, (map)(length, (Base.OneTo(4), Base.OneTo(4))))
+          var"%_out" = for var"%jj′" = (zip)(1:4, (Symbolics.reset_to_one)(1:4))
                   begin
                       j = var"%jj′"[1]
                       j′ = var"%jj′"[2]

Original file line number	Diff line number	Diff line change
`@@ -1005,7 +1005,7 @@ function _array_toexpr(x, st)`
`1005`	`1005`	`[`
`1006`	`1006`	`Assignment(outsym, term(zeros, Float64, term(map, length, shape(x)))),`
`1007`	`1007`	`Assignment(Symbol("%$outsym"), inplace_expr(x, outsym))`
`1008`		`- ], outsym, true)`
	`1008`	`+ ], outsym, false)`
`1009`	`1009`
`1010`	`1010`	`toexpr(ex, st)`
`1011`	`1011`	`end`