Implement function inlining. #7713

dnwpark · 2024-09-04T15:20:09Z

Allows for the ir ast trees to be directly embedded into the tree of a calling expression.

This is accomplished by:

Retrieving the function body and recompiling ir in the ir compiler
- Use a new context detached from the inlining context
- Added a function compile_function_inline to do only the necessary steps, excluding things materializing sets.
Substituting any irast.Parameter nodes from the function with a new irast.InlinedParameterExpr
- This is necessary because irast.Parameter is used for runtime parameters when compiling queries.
Compiling the arguments in a subrelation in the pg compiler
- Include the rvar and path value in the appropriate relations.

related #7692

Note: The volatility check for inlined function bodies was removed. This is because functions which use globals are immutable, since they expect a globals json to be passed in. This is not true of their inlined bodies, which become stable.

msullivan · 2024-09-05T00:00:01Z

I think it would be useful to try a 2-argument function as part of your early tests

msullivan · 2024-09-09T21:42:40Z

edb/edgeql/compiler/stmtctx.py

    env = context.Environment(
        schema=schema,
        options=options,
+        path_scope=(
+            inlining_context.path_scope
+            if inlining_context is not None else
+            None
+        ),
        alias_result_view_name=options.result_view_name,
+        scope_tree_nodes=(
+            inlining_context.env.scope_tree_nodes
+            if inlining_context is not None else
+            None
+        ),
+    )
+    ctx = context.ContextLevel(
+        None,
+        context.ContextSwitchMode.NEW,
+        env=env,
+        aliases=(
+            inlining_context.aliases
+            if inlining_context is not None else
+            None
+        ),
+        scope_id_ctr=(
+            inlining_context.scope_id_ctr
+            if inlining_context is not None else
+            None
+        ),


We should use the existing context, with appropriate new scopes applied to it.
Doing it in a completely separate context is not going to work at all in anything but the most simple situations.

And especially the same environment. (Or at least, an almost identical environment, since we might need to change options or something, though that is a little fraught too.)

msullivan · 2024-09-09T23:03:57Z

edb/ir/ast.py

+class InlinedParameter(ImmutableBase):
+
+    # Parameter is used for both function parameters and runtime parameters.
+    # If a function is inlined in compile_FunctionCall, this is used to
+    # specify the function's parameters.
+
+    # During function compilation, arguments are converted from names to either
+    # int for positional argument or str for named argument.
+    arg_key: int | str
+
+    required: bool
+    is_global: bool


The more typical way to do something like this in our compiler would probably be to substitute in irast.Set objects that have the same path_id as as the arguments but have an irast.RefExpr as their expr, to indicate that they are bound somewhere.

They should then get picked up by their path id when doing the backend compilation, without needing to add an extra map.

msullivan

Overall this looks pretty good.

I think we need to share more of env and ctx when compiling the function body, though.

I think there are some things that are being done mostly to workaround not sharing enough state, and a bunch more will come up when doing DML or when writing other sorts of tests.

msullivan · 2024-09-19T22:54:56Z

edb/edgeql/compiler/func.py

+            context=sd.CommandContext(
+                # Probably not correct. Need to store modaliases while compiling
+                # functions?
+                modaliases=ctx.modaliases,


modaliases shouldn't be inherited by the function body

though it also doesn't matter, since functions are normalized into only having fully qualified names

msullivan · 2024-09-20T18:05:26Z

edb/edgeql/compiler/func.py

-                optional=arg.param_typemod == ft.TypeModifier.OptionalType,
+                optional=(
+                    arg.param_typemod == ft.TypeModifier.OptionalType
+                    or arg.is_default


Hm.... Is adding arg.is_default fully necessary here? I think in some cases it might result in worse codegen

It was a consequence of how has_inlined_defaults interacted with the rest of the inlining. The simplest solution seems to be to just disable inlining defaults when inlining the function body.

msullivan · 2024-09-20T18:30:37Z

edb/edgeql/compiler/func.py

+            else:
+                return arg


What is this case for?

Default values get added to inlined_args as a irast.Set instead of CallArg. In that case, the default value is inlined directly. This should work because defaults are always constants.

msullivan · 2024-09-20T19:52:25Z

edb/edgeql/compiler/func.py

+class ArgumentInliner(ast.NodeTransformer):
+
+    mapped_args: dict[irast.PathId, irast.PathId]


I realize now that it should be possible (and probably nicer) to put the inlined arguments in the context somehow while compiling the inlined function, and inserting them directly when compiling Parameters, rather then substituting them in after the fact.

I decided to not go down that path since the backend compiler does stuff with args too. Probably possible though.

msullivan · 2024-09-20T19:53:44Z

edb/edgeql/compiler/inference/volatility.py

-    if ir.body is not None:
-        body_volatility = infer_volatility(ir.body, env=env)
-        if body_volatility > ir.volatility:
-            raise errors.QueryError(
-                f'inline function body expression is {body_volatility} '
-                f'while the function is declared as {ir.volatility}',
-                span=ir.span,
-            )


I would like to understand while this is happening better. (Though it's not super critical)

msullivan · 2024-09-20T20:21:39Z

edb/pgsql/compiler/relgen.py

+        if pathctx.maybe_get_path_var(
+            stmt, prefix_path_id, aspect=pgce.PathAspect.VALUE, env=ctx.env
+        ) is None:
+            assert isinstance(source_rvar.query, pgast.Query)
+            # When inlining functions, the prefix path will be the argument
+            # path id, while ir_source.path_id will be the "correct" path id.
+            # Remap the argument path id here to ensure the path bond is
+            # able to find the source rvar when joining.
+            prefix_path_id = pathctx.map_path_id(
+                prefix_path_id, source_rvar.query.view_path_id_map
+            )
+


This doesn't feel like it /should/ be necessary; it seems like it should be fixable on the edgeql->IR side.
Think about it a little more, and then update the comment to say it is a HACK we should revisit if it doesn't seem easy to fix.

The most straightforward solution here was to just update the paths of all pointers in the edgeql->ir side instead of trying to rely on the path mapping.

msullivan · 2024-09-20T20:23:57Z

edb/pgsql/compiler/relgen.py

+            with newctx.subrel() as arg_ctx:
+                args = _compile_call_args(ir_set, ctx=arg_ctx)
+                arg_rvar = relctx.rvar_for_rel(arg_ctx.rel, ctx=ctx)
+                for ir_arg in expr.args.values():
+                    arg_path_id = ir_arg.expr.path_id
+                    relctx.include_rvar(
+                        newctx.rel,
+                        arg_rvar,
+                        arg_path_id,
+                        ctx=newctx,
+                    )
+                    if arg_scope_stmt := relctx.maybe_get_scope_stmt(
+                        arg_path_id, ctx=newctx
+                    ):
+                        # The rvar is joined to newctx.rel, but other sets may
+                        # look for it in the scope statement. Make sure it's
+                        # available.
+                        pathctx.put_path_value_rvar(
+                            arg_scope_stmt,
+                            arg_path_id,
+                            arg_rvar,
+                        )
+
+            if expr.inline_arg_path_ids:
+                for param_path_id, arg_path_id in (
+                    expr.inline_arg_path_ids.items()
+                ):
+                    pathctx.put_path_id_map(
+                        arg_ctx.rel,
+                        param_path_id,
+                        arg_path_id
+                    )


This probably wants to be its own function, because otherwise it drowns out the common case of just calling _compile_call_args

msullivan · 2024-09-20T20:39:32Z

edb/schema/functions.py

+    # Infer cardinality, multiplicity, and volatility
+    inf_ctx = inference.make_ctx(env=ctx.env)
+    inference.infer_cardinality(
+        ir_set, scope_tree=ctx.path_scope, ctx=inf_ctx
+    )
+    inference.infer_multiplicity(
+        ir_set, scope_tree=ctx.path_scope, ctx=inf_ctx
+    )
+    inference.infer_volatility(ir_set, env=ctx.env)


I think we shouldn't need to do this? It'll all get done on the inlined body, right?

Oh, I think the reason why it doesn't work without inferring them here is because not enough of ctx/env are being shared

msullivan · 2024-09-20T21:08:49Z

edb/edgeql/compiler/context.py

+        set_types: Optional[Dict[irast.Set, s_types.Type]] = None,
+        scope_tree_nodes: Optional[
+            MutableMapping[int, irast.ScopeTreeNode]
+        ] = None,
+        type_rewrites: Optional[dict[
+            Tuple[s_types.Type, bool],
+            irast.Set | None | Literal[True]
+        ]] = None,


This is missing a lot of fields that need to be preserved in Environment for correct operation in some case or another, especially once DML is supported.
Definitely needed:

materialized_sets

dml_exprs

dml_stmts

dml_rewrites

schema_refs

I think possibly needed, or needed if things are done a little differently:

pointer_derivation_map (for card inference?)

pointer_specified_info (for card inference?)

query_globals

Plus, I think having all of the caches shared will help get more cache hits, even if it isn't necessary.
And I think many of the other things wouldn't hurt to have be shared.

So I think instead of constructing a fully new env, we should probably make a straight copy of the env and just update the few fields we care about changing.

It might /just/ be options we need to change, though? In which case it might actually just be best to update the options field and then restore it after.

msullivan · 2024-09-20T21:15:58Z

edb/schema/functions.py

+    ctx = stmtctx.init_context(
+        schema=schema,
+        options=get_compiler_options(
+            schema,
+            context,
+            func_name=func_name,
+            params=params,
+            track_schema_ref_exprs=track_schema_ref_exprs,
+            inlining_context=inlining_context,
+        ),
+        inlining_context=inlining_context,
+    )


Do we really need to go through init_context here? I think we probably should be trying to share the context mostly. (I think it probably would work if we create a detached() context?)

I tried to see if I could get a detached context working but I think the inlined context needs to have more than what is available in a detached context.

msullivan · 2024-09-24T19:56:58Z

tests/test_edgeql_functions.py

+    async def test_edgeql_functions_inline_object_04(self):
+        await self.con.execute('''
+            create type Bar {
+                create required property a -> int64;
+            };
+            insert Bar{a := 1};
+            insert Bar{a := 2};
+            insert Bar{a := 3};
+            create function foo(x: Bar) -> int64 {
+                set is_inlined := true;
+                using (x.a);
+            };
+        ''')


Could we do a test like

create function foo() -> int64 { set is_inlined := true; using (count(Bar)); };

and then query

select (Bar, foo())

and make sure that we get 3 back as the result from foo().
I want to make sure that paths in the inlined functions don't end up getting captured by factoring with the enclosing query.

msullivan

This looks pretty good, though I'm worried about the path factoring issue I asked for a test for, which we'll need to talk about

msullivan · 2024-09-24T22:10:41Z

edb/edgeql/compiler/__init__.py

@@ -201,6 +202,7 @@ def compile_ast_to_ir(
    *,
    script_info: Optional[irast.ScriptInfo] = None,
    options: Optional[CompilerOptions] = None,
+    inlining_context: Optional[context.ContextLevel] = None,


You don't use these entrypoints anymore, right? Can we drop the inlining_context from them?

msullivan · 2024-09-24T23:57:00Z

Great work!

dnwpark force-pushed the function-inline branch from 6dbdee8 to fe66209 Compare September 4, 2024 21:24

dnwpark force-pushed the function-inline branch 6 times, most recently from 72c96d7 to 0a659e2 Compare September 6, 2024 21:41

dnwpark changed the base branch from master to ast-visitor-irast September 6, 2024 23:29

dnwpark force-pushed the ast-visitor-irast branch from 11c9d17 to 7e61c76 Compare September 9, 2024 17:11

dnwpark force-pushed the function-inline branch from 0a659e2 to 000c220 Compare September 9, 2024 17:11

msullivan reviewed Sep 9, 2024

View reviewed changes

dnwpark force-pushed the ast-visitor-irast branch from 7e61c76 to 77b42b2 Compare September 10, 2024 16:06

dnwpark force-pushed the function-inline branch 2 times, most recently from 3be4bf1 to 52c32b1 Compare September 10, 2024 16:29

Base automatically changed from ast-visitor-irast to master September 10, 2024 16:30

dnwpark force-pushed the function-inline branch 8 times, most recently from 67f672b to bf311de Compare September 17, 2024 15:35

dnwpark force-pushed the function-inline branch 5 times, most recently from 1783407 to 3598bec Compare September 19, 2024 16:53

dnwpark force-pushed the function-inline branch 2 times, most recently from 37ee25c to b85e67a Compare September 19, 2024 20:02

msullivan reviewed Sep 20, 2024

View reviewed changes

dnwpark force-pushed the function-inline branch from b85e67a to 97908c6 Compare September 23, 2024 21:23

dnwpark added 15 commits September 24, 2024 13:42

Add inlining functions.

9fd98ba

Pass existing context into inlined function compile.

84215b8

Add inlining default positional arguments.

9b766c6

Pack variadic arguments into an array.

719e983

Add param and arg ids to path id map.

66cfc6f

Compile inline functions without finalizing.

211915d

Ensure inlined functions always check for null args.

3ef7941

Fix type intersections not finding the correct rvar.

a1b08a0

Remove check for inlined function body volatility.

17237a8

Add tests.

9592273

Change init_context to copy almost all fields.

f5672a9

Factor out compiling inlined function args.

2ce748a

Add TODO to get rid of implicit globals params.

686f45c

Disable inlined defaults when inlining the function body.

fe48d68

Update pointer paths in frontend instead of relying on path mapping.

45dbd62

dnwpark force-pushed the function-inline branch from 3e5230d to 45dbd62 Compare September 24, 2024 17:43

msullivan reviewed Sep 24, 2024

View reviewed changes

Use detached context instead copying context.

a82fb74

msullivan approved these changes Sep 24, 2024

View reviewed changes

Remove unused inlining_context argument from compile_ast_to_ir.

a35e431

Remove inlined_ctx from compile_ast_fragment_to_ir.

d19da8a

dnwpark force-pushed the function-inline branch from 8af3eb0 to d19da8a Compare September 25, 2024 00:53

dnwpark merged commit 131a406 into master Sep 25, 2024
25 checks passed

dnwpark deleted the function-inline branch September 25, 2024 01:17

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement function inlining. #7713

Implement function inlining. #7713

dnwpark commented Sep 4, 2024 •

edited

Loading

msullivan commented Sep 5, 2024

msullivan Sep 9, 2024

msullivan Sep 9, 2024

msullivan Sep 9, 2024

msullivan left a comment

msullivan Sep 19, 2024

msullivan Sep 19, 2024

msullivan Sep 20, 2024

dnwpark Sep 24, 2024

msullivan Sep 20, 2024

dnwpark Sep 20, 2024

msullivan Sep 20, 2024

dnwpark Sep 23, 2024

msullivan Sep 20, 2024

msullivan Sep 20, 2024

dnwpark Sep 24, 2024

msullivan Sep 20, 2024

msullivan Sep 20, 2024

msullivan Sep 20, 2024

msullivan Sep 20, 2024

msullivan Sep 20, 2024

dnwpark Sep 23, 2024

msullivan Sep 24, 2024

msullivan left a comment

msullivan Sep 24, 2024

msullivan commented Sep 24, 2024

		class ArgumentInliner(ast.NodeTransformer):

		mapped_args: dict[irast.PathId, irast.PathId]

Implement function inlining. #7713

Implement function inlining. #7713

Conversation

dnwpark commented Sep 4, 2024 • edited Loading

msullivan commented Sep 5, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

msullivan 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

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

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

msullivan left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

msullivan commented Sep 24, 2024

dnwpark commented Sep 4, 2024 •

edited

Loading