Move nargs/isva to CodeInfo #54341

Keno · 2024-05-03T06:26:03Z

This changes the canonical source of truth for va handling from Method to CodeInfo. There are multiple goals for this change:

This addresses a longstanding complaint about the way that CodeInfo-returning generated functions work. Previously, the va-ness or not of the returned CodeInfo always had to match that of the generator. For Cassette-like transforms that generally have one big generator function that is varargs (while then looking up lowered code that is not varargs), this could become quite annoying. It's possible to workaround, but there is really no good reason to tie the two together. As we observed when we implemented OpaqueClosures, the vararg-ness of the signature and the vararg arguments->tuple transformation are mostly independent concepts. With this PR, generated functions can return CodeInfos with whatever combination of nargs/isva is convenient.
This change requires clarifying where the va processing boundary is in inference. irinterp: Move irinterp va processing #54076 was already moving in that direction for irinterp, and this essentially does much of the same for regular inference. As a consequence the constprop cache is now using non-va-cooked signatures, which I think is preferable.
This further decouples codegen from the presence of a Method (which is already not assumed, since the code being generated could be a toplevel thunk, but some codegen features are only available to things that come from Methods). There are a number of upcoming features that will require codegen of things that are not quite method specializations (See design doc linked in Allow CodeInstance in Expr(:invoke) #52797 and things like RFC: Effect Preconditions - or - the future of @inbounds #50641). This helps pave the road for that.
I've previously considered expanding the kinds of vararg signatures that can be described (see e.g. lowering: Plumb through arg destructuring code into all wrapper functions #53851), which also requires a decoupling of the signature and ast notions of vararg. This again lays the groundwork for that, although I have no immediate plans to implement this change.

Impact wise, this adds an internal field, which is not too breaking, but downstream clients vary in how they construct their CodeInfos and the current way they're doing it will likely be incorrect after this change, so they will require a small two-line adjustment. We should perhaps consider pulling out some of the more common patterns into a more stable package, since interface in most of the last few releases, but that's a separate issue.

vtjnash

Seems okay, but isn't this currently an ABI violation, as the ABI for the caller is defined by the flag in Method, whereas this means codegen will emit incorrect code for the callee? We could go the route where that k&r C did and remove isva from the ABI, but I thought that was somewhat universally considered a design mistake now

vtjnash · 2024-05-03T13:59:43Z

src/method.c

@@ -735,6 +741,9 @@ JL_DLLEXPORT jl_code_info_t *jl_code_for_staged(jl_method_instance_t *linfo, siz
                }
                jl_error("The function body AST defined by this @generated function is not pure. This likely means it contains a closure, a comprehension or a generator.");
            }
+            // TODO: This should ideally be in the lambda expression,


This is impossible (as it is defined by value not syntax), so it should not be a TODO

The todo is to change the syntax (as we did for OpaqueClosure).

That seems like a breaking change

Potentially both ... and Vararg{ could be recognized syntactically without being too breaking, but I wouldn't want to do it without a bigger motivation than this. Or in 2.0 of course.

Breaking design changes should not have TODO comments though. The syntactic thing is possible, but breaks referential transparency and seems to gain nothing but a slightly more convoluted lowering, so I don't think that seems worthwhile

Keno · 2024-05-05T01:31:17Z

Seems okay, but isn't this currently an ABI violation, as the ABI for the caller is defined by the flag in Method, whereas this means codegen will emit incorrect code for the callee?

As far as I could tell, the only use of the isva flag in ABI determination was for the specsig decision (if it was VA, it'd check if the specsig had a vararg in it and bail if so). This changes that to just look at the specsig unconditionally and only bail if unbound, which should capture the intent, but is perhaps even slightly more consistent. The special codegen for vararg tuples remains, but that was purely a callee consideration and did not affect the ABI.

vtjnash

Ah, okay, I suppose it does seem unlikely for a non-isva method to have reached that code, so it won't matter if it generates a different calling convention now. It didn't realize we had made a guarantee that non-isva methods will always use exactly the same calling convention as an isva method with the same arguments (eg that the calling convention is defined by the caller rather than the callee). But I guess since we do that for a couple releases now, we can make that a guarantee not to later change that again

Keno · 2024-05-05T04:07:43Z

@nanosoldier runtests()

nanosoldier · 2024-05-05T13:30:32Z

The package evaluation job you requested has completed - possible new issues were detected.
The full report is available.

Keno · 2024-05-05T21:07:37Z

Ok, Zygote and Diffractor need adjustment for this (expected). Other than that, there is one interesting failure https://s3.amazonaws.com/julialang-reports/nanosoldier/pkgeval/by_hash/fb6f42d_vs_2cf469d/MutableArithmetics.primary.log. Looks good otherwise though. Will do the Zygote/Diffractor adjustment tonight or tomorrow.

vtjnash · 2024-05-06T00:50:41Z

I am still very skeptical it is worthwhile making all CodeInfo 4 bytes larger and the ABI slightly worse to make it possible for external tools not to correctly normalize and inline their IR

Keno · 2024-05-06T01:01:02Z

That's not the primary motivation. The primary motivation is to allow for specializations whose ABI may not match the signature (which needs additional things outside of this). That said, to address the size concern, I can use a bit in the flags to copy the value from the method during uncompress. I don't understand the ABI concern. Could you elaborate which ABI you think is worse in this PR. It is intended to be largely unchanged.

vtjnash · 2024-05-06T01:25:24Z

src/codegen.cpp

-        if (jl_is_vararg(jl_tparam(sig, jl_nparams(sig) - 1)))
-            return false;
-    }
+    if (jl_vararg_kind(jl_tparam(sig, jl_nparams(sig) - 1)) == JL_VARARG_UNBOUND)


It doesn't look like the other queries in this function will be computed correctly for a bound or numeric var (being based on nparams instead of nargs+isva). Though that seems a somewhat pre-existing problem.

Yes, we can adjust that, but I don't think we ever did the change to actually normalize it bound vararg, so I'm not sure it ever actually shows up here at the moment. Regardless, as you said, separate concern, since the reachability of that isn't dependent on the isva flag.

vtjnash · 2024-05-06T01:32:43Z

I am worried if this could hamstring our ability to generate a fast ABI to something that is a Vararg signature but not concrete (which is all that we currently implement optimizations for right now). For example, optimizing the initial parameters (nargs minus isva) even if we still pass the rest of the Vararg segment as a boxed list (or an optimized list+length). Those are not an optimization we implement now, but it is one that we likely should add later.

Keno · 2024-05-06T01:43:27Z

I don't think this prevents us from changing the ABI based on that (computed from the method if available). Of course this change would require that to have a little extra code to potentially repack the tuple if there's a mismatch but we have most have that already anyway.

vtjnash · 2024-05-06T02:51:20Z

You cannot fix an ABI mistake by "repacking" though. If the ABI states that you pass nargs-isva and then the vararg count then the vararg pointer, you better pass exactly that based on the CodeInstance you are calling, and not some other random registers based on the content of the CodeInfo

Keno · 2024-05-06T02:54:27Z

You cannot fix an ABI mistake by "repacking" though. If the ABI states that you pass nargs-isva and then the vararg count then the vararg pointer, you better pass exactly that based on the CodeInstance you are calling, and not some other random registers based on the content of the CodeInfo

Yes, sure, but when you codegen the codeinfo, you see that same info, so even if the CodeInfo says, that there's only one arg and isva is true, you can implicitly generate a bunch of getfields and repack the tuple. We do that exact same thing now where we pass all the arguments individually and then call tuple on it implicitly based on nargs.

Keno · 2024-05-06T03:14:59Z

Said another way, there's two orthogonal but related questions:

How does the semantic list of arguments get passed between the caller and the callee
How does the semantic list of arguments get split up between the slots of the callee.

Right now 1 is a property of (mi->specTypes, m->isva, ci) and 2 is a property of (m->nargs, m->isva). This changes 2 to be in CodeInfo instead. I also removed the m->isva from 1, because I think it's more consistent to not have it depend on the method, but I'd be fine putting it back if necessary. If one were to put it back that would be m->isva, not code->isva, because it needs to be consistent between the caller (which does not have access to the CodeInfo) and the callee.

However, the two questions are largely separate. Yes, it would be preferable to have the ABI match closely to what the CodeInfo wants for its slots so that there doesn't need to be much copying, but it's not a semantic requirement.

vtjnash · 2024-05-07T03:34:53Z

Alright, yeah, if we put back the nargs and isva flag on Method, it seems much clearer that this is intended solely as a compression trick (avoiding needing to explicitly unpack arguments as that will be done implicitly by the interpreter and codegen instead)

Keno · 2024-05-07T15:03:36Z

Alright, yeah, if we put back the nargs and isva flag on Method

They're still there

Keno · 2024-05-07T17:16:25Z

@vtjnash Are your concerns adequately addressed? If so, I'll rebase this, merge the package adjustments, kick off another pkgeval run, fix any remaining issues and get this in.

vtjnash · 2024-05-07T18:12:41Z

Yeah, without looking at the code changes, the aspects discussed sound good now

Keno · 2024-05-07T19:42:28Z

Reproducer for the MutableArithmetics issue (nothing special, just a typo in the code, but I always find it interesting what code patterns are unique among 5000 packages + Base):

global SIDE_EFFECT::Int
function foo(a, b, c, d, args...)
	# Side effect to force constprop rather than semi-concrete
	global SIDE_EFFECT = a + b + c + d
	return SIDE_EFFECT
end
bar(args...) = foo(4, args...)

code_typed(bar, Tuple{Vararg{Int}})

This changes the canonical source of truth for va handling from `Method` to `CodeInfo`. There are multiple goals for this change: 1. This addresses a longstanding complaint about the way that CodeInfo-returning generated functions work. Previously, the va-ness or not of the returned CodeInfo always had to match that of the generator. For Cassette-like transforms that generally have one big generator function that is varargs (while then looking up lowered code that is not varargs), this could become quite annoying. It's possible to workaround, but there is really no good reason to tie the two together. As we observed when we implemented OpaqueClosures, the vararg-ness of the signature and the `vararg arguments`->`tuple` transformation are mostly independent concepts. With this PR, generated functions can return CodeInfos with whatever combination of nargs/isva is convenient. 2. This change requires clarifying where the va processing boundary is in inference. #54076 was already moving in that direction for irinterp, and this essentially does much of the same for regular inference. As a consequence the constprop cache is now using non-va-cooked signatures, which I think is preferable. 3. This further decouples codegen from the presence of a `Method` (which is already not assumed, since the code being generated could be a toplevel thunk, but some codegen features are only available to things that come from Methods). There are a number of upcoming features that will require codegen of things that are not quite method specializations (See design doc linked in #52797 and things like #50641). This helps pave the road for that. 4. I've previously considered expanding the kinds of vararg signatures that can be described (see e.g. #53851), which also requires a decoupling of the signature and ast notions of vararg. This again lays the groundwork for that, although I have no immediate plans to implement this change. Impact wise, this adds an internal field, which is not too breaking, but downstream clients vary in how they construct their `CodeInfo`s and the current way they're doing it will likely be incorrect after this change, so they will require a small two-line adjustment. We should perhaps consider pulling out some of the more common patterns into a more stable package, since interface in most of the last few releases, but that's a separate issue.

Keno · 2024-05-09T20:00:21Z

@nanosoldier runtests(["PolygonAlgorithms", "CassetteOverlay", "LocalOperators", "LoopFieldCalc", "Blobs", "EarlyStopping", "ResumableFunctions", "IRTools", "SpaceIndices", "LazyStack", "FactorLoadingMatrices", "BufferedStreams", "ZygoteRules", "DifferentiableFlatten", "ExponentialAction", "UnitfulChainRules", "ODEInterface", "GCMAES", "MutableArithmetics", "SimJulia", "Cthulhu", "BangBang", "ParameterHandling", "Optimisers", "ManifoldDiff", "SciMLOperators", "ArrayPadding", "LinearOperators", "Roots", "FinancialToolbox", "DynamicExpressions", "OMEinsum", "DiffImageRotation", "FiniteDiff", "BallroomSkatingSystem", "OneHotArrays", "ForwardDiffPullbacks", "RadonKA", "FMIImport", "StandardizedRestrictedBoltzmannMachines", "CenteredRBMs", "DifferentiableEigen", "Pyehtim", "FoldRNA", "AdvRBMs", "TreeRecipe", "IsingModels", "ComponentArrays", "Diffractor", "PhysicalParticles", "YaoBlocks", "GRAPE", "CCDReduction", "QuantumControl", "QuantumControlBase", "SliceMap", "TensorGames", "GPLikelihoods", "ParameterSchedulers", "BayesianLinearRegressors", "OptimalPortfolios", "HypercubeTransform", "GLPK", "OptimizationOptimisers", "Stheno", "LinearMixingModels", "SignalTemporalLogic", "CoordRefSystems", "AbstractGPs", "BigO", "ValueShapes", "DistributionMeasures", "CapAndHomalg", "Bijectors", "Miter", "MLUtils", "TransmuteDims", "VLBILikelihoods", "BaytesMCMC", "PlanetOrbits", "FourierTools", "SNOW", "OptimizationNLopt", "Backboner", "LuxTestUtils", "DecisionMakingPolicies", "ThermodynamicIntegration", "NonconvexCore", "NonconvexSearch", "PhysicalMeshes", "Polynomials4ML", "DECAES", "NonconvexNOMAD", "NonconvexMetaheuristics", "NonconvexMMA", "ConjugateComputationVI", "PhysicalFFT", "GaussianVariationalInference", "StructuredOptimization", "DiscoDiff", "QAOA", "QuadratureCuba", "QuadratureCubature", "PhysicalFDM", "NonconvexPercival", "UnderwaterAcoustics", "MGVI", "DiffFusion", "DeconvOptim", "OptimizationOptimJL", "PottsGumbelRBMLayers", "TuringBenchmarking", "BetaML", "SphericalFunctions", "FinanceModels", "SuperfluidRotSpec", "Braket", "MarginalLogDensities", "DerivableFunctions", "SurrogatesRandomForest", "WaveOpticsPropagation", "ParametricMCPs", "MaxEntropyGraphs", "FLOWPanel", "RegressionDiscontinuity", "FSimBase", "ReinforcementLearningZoo", "RvSpectML", "Fronts", "LongwaveModePropagator", "CausalityTools", "NonconvexIpopt", "RvLineList", "ControlSystemIdentification", "NonconvexJuniper", "Plots", "NonconvexPavito", "JumpProblemLibrary", "OpenQuantumSystems", "NonconvexBayesian", "Yunir", "ElectrochemicalKinetics", "SubsidenceChron", "GeoArtifacts", "TKTDsimulations", "OptimizationMOI", "SwissVAMyKnife", "TidierPlots", "Chamber", "IterativeLearningControl"])

nanosoldier · 2024-05-10T15:04:23Z

The package evaluation job you requested has completed - possible new issues were detected.
The full report is available.

KristofferC · 2024-05-11T09:34:06Z

The PkgEval failures are known?

Internal error: during type inference of
arrays(DataFrames.DataFrame)
Encountered unexpected error in runtime:
AssertionError(msg="invalid `given_argtypes` for `mi`")
va_process_argtypes at ./compiler/inferenceresult.jl:122
InferenceState at ./compiler/inferencestate.jl:306
InferenceState at

etc?

Keno · 2024-05-11T12:55:41Z

Yes. I fixed CassetteOverlay and the rest depend on ResumableFunctions and will be fixed once they merge my PR.

Krastanov · 2024-05-12T04:12:48Z

@Keno , could you double check the fix you have submitted to ResumableFunctions. I am a bit lost on what it is supposed to mean (it seems to be acting on a variable that does not exist in the given scope)

src/ircode.c

aviatesk · 2024-05-13T09:09:28Z

base/compiler/typeinfer.jl

@@ -468,7 +468,7 @@ function adjust_effects(sv::InferenceState)
        # this frame is known to be safe
        ipo_effects = Effects(ipo_effects; nothrow=true)
    end
-    if is_inaccessiblemem_or_argmemonly(ipo_effects) && all(1:narguments(sv, #=include_va=#true)) do i::Int
+    if is_inaccessiblemem_or_argmemonly(ipo_effects) && all(1:narguments(sv, #=include_va=#true)) do i::UInt


I'd prefer sticking with Int for this sort of operations.

aviatesk · 2024-05-13T09:20:49Z

base/compiler/inferenceresult.jl

+    if length(given_argtypes) == 0 || length(cache_argtypes) == 0
+        return Any[]
+    end


Why do we need this check?

Toplevel inference has no arguments. I think that used to bypass this path.

aviatesk · 2024-05-13T09:58:57Z

base/compiler/abstractinterpretation.jl

+        argt = widenslotwrapper(argtypes[i])
+        if isvarargtype(argt)
+            @assert fieldcount(sig) == i
+            argt = unwrapva(argt)
+        end
+        thentype = elsetype = tmeet(𝕃ᵢ, argt, fieldtype(sig, i))


This check wasn't needed previously, but has it become necessary with this PR? argtypes used by abstract_apply can certainly include Vararg, but I find it curious why the compiler functioned well without this check before.

I suspect it was theoretically reachable, but the change to apply this to the pre-processed signature made it more likely. I didn't bother investigating.

aviatesk · 2024-05-13T10:03:03Z

base/compiler/abstractinterpretation.jl

-    if condargs !== nothing
-        given_argtypes = let condargs=condargs
-            va_process_argtypes(𝕃, given_argtypes, mi) do isva_given_argtypes::Vector{Any}, last::Int
-                # invalidate `Conditional` imposed on varargs
-                for (slotid, i) in condargs
-                    if slotid ≥ last && (1 ≤ i ≤ length(isva_given_argtypes)) # `Conditional` is already widened to vararg-tuple otherwise
-                        isva_given_argtypes[i] = widenconditional(isva_given_argtypes[i])
-                    end
-                end
-            end
-        end
-    else
-        given_argtypes = va_process_argtypes(𝕃, given_argtypes, mi)
-    end


With this change, we might be able to switch the weird interface of [Simple|Widened|Conditional]Argtypes (if you can even call it an 'interface') to something better like AbstractLattice.

Keno requested a review from aviatesk May 3, 2024 06:26

vtjnash reviewed May 3, 2024

View reviewed changes

Keno force-pushed the kf/vasemanticssplit branch from 133e8ce to fb6f42d Compare May 5, 2024 01:28

vtjnash reviewed May 5, 2024

View reviewed changes

Keno force-pushed the kf/vasemanticssplit branch from fb6f42d to ecc57b0 Compare May 5, 2024 06:25

vtjnash reviewed May 6, 2024

View reviewed changes

Keno force-pushed the kf/vasemanticssplit branch from ecc57b0 to a6b2bd8 Compare May 6, 2024 02:27

Keno force-pushed the kf/vasemanticssplit branch from a6b2bd8 to 344616a Compare May 7, 2024 19:45

Keno force-pushed the kf/vasemanticssplit branch from 344616a to d45ef7b Compare May 7, 2024 21:15

Keno merged commit d01d256 into master May 11, 2024
6 of 8 checks passed

Keno deleted the kf/vasemanticssplit branch May 11, 2024 04:45

aviatesk reviewed May 13, 2024

View reviewed changes

DilumAluthge mentioned this pull request May 13, 2024

Follow ups to #54341 (Move nargs/isva to CodeInfo) #54453

Merged

aviatesk added a commit that referenced this pull request May 14, 2024

Followups to #54341 (Move nargs/isva to CodeInfo)

bf43075

aviatesk added a commit that referenced this pull request May 14, 2024

Follow ups to #54341 (Move nargs/isva to CodeInfo) (#54453)

c3f1479

Move nargs/isva to CodeInfo #54341

Move nargs/isva to CodeInfo #54341

Conversation

Keno commented May 3, 2024

vtjnash left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Keno commented May 5, 2024

vtjnash left a comment

Choose a reason for hiding this comment

Keno commented May 5, 2024

nanosoldier commented May 5, 2024

Keno commented May 5, 2024

vtjnash commented May 6, 2024

Keno commented May 6, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

vtjnash commented May 6, 2024

Keno commented May 6, 2024

vtjnash commented May 6, 2024

Keno commented May 6, 2024

Keno commented May 6, 2024

vtjnash commented May 7, 2024

Keno commented May 7, 2024

Keno commented May 7, 2024

vtjnash commented May 7, 2024

Keno commented May 7, 2024

Keno commented May 9, 2024

nanosoldier commented May 10, 2024

KristofferC commented May 11, 2024

Keno commented May 11, 2024

Krastanov commented May 12, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment