From f1848759fab6a63f1e743468cce798e2e58c87bc Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Thu, 1 May 2025 11:11:16 +0100 Subject: [PATCH 1/9] HAX get merging working in a simple use case (breaking others) --- mypyc/codegen/emitmodule.py | 4 ++-- mypyc/irbuild/env_class.py | 6 +++++- mypyc/irbuild/function.py | 2 +- mypyc/irbuild/generator.py | 22 +++++++++++++--------- 4 files changed, 21 insertions(+), 13 deletions(-) diff --git a/mypyc/codegen/emitmodule.py b/mypyc/codegen/emitmodule.py index 8474be62579d..19a82b1f0f5b 100644 --- a/mypyc/codegen/emitmodule.py +++ b/mypyc/codegen/emitmodule.py @@ -246,8 +246,8 @@ def compile_scc_to_ir( # Insert refcount handling. insert_ref_count_opcodes(fn) - if fn in env_user_functions: - insert_spills(fn, env_user_functions[fn]) + #if fn in env_user_functions: + # insert_spills(fn, env_user_functions[fn]) # Switch to lower abstraction level IR. lower_ir(fn, compiler_options) diff --git a/mypyc/irbuild/env_class.py b/mypyc/irbuild/env_class.py index ab786fe71dda..bc631e7d6b2e 100644 --- a/mypyc/irbuild/env_class.py +++ b/mypyc/irbuild/env_class.py @@ -47,6 +47,7 @@ class is generated, the function environment has not yet been ) env_class.attributes[SELF_NAME] = RInstance(env_class) if builder.fn_info.is_nested: + assert False # If the function is nested, its environment class must contain an environment # attribute pointing to its encapsulating functions' environment class. env_class.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_infos[-2].env_class) @@ -58,7 +59,7 @@ class is generated, the function environment has not yet been def finalize_env_class(builder: IRBuilder) -> None: """Generate, instantiate, and set up the environment of an environment class.""" - instantiate_env_class(builder) + #instantiate_env_class(builder) # Iterate through the function arguments and replace local definitions (using registers) # that were previously added to the environment with references to the function's @@ -76,6 +77,7 @@ def instantiate_env_class(builder: IRBuilder) -> Value: ) if builder.fn_info.is_nested: + assert False builder.fn_info.callable_class._curr_env_reg = curr_env_reg builder.add( SetAttr( @@ -125,6 +127,7 @@ def load_outer_env( Returns the register where the environment class was loaded. """ + assert False env = builder.add(GetAttr(base, ENV_ATTR_NAME, builder.fn_info.fitem.line)) assert isinstance(env.type, RInstance), f"{env} must be of type RInstance" @@ -238,6 +241,7 @@ def setup_func_for_recursive_call(builder: IRBuilder, fdef: FuncDef, base: Impli prev_env = builder.fn_infos[-2].env_class prev_env.attributes[fdef.name] = builder.type_to_rtype(fdef.type) + assert False if isinstance(base, GeneratorClass): # If we are dealing with a generator class, then we need to first get the register # holding the current environment class, and load the previous environment class from diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py index cb9a1a3dc4a3..c1557e410b2d 100644 --- a/mypyc/irbuild/function.py +++ b/mypyc/irbuild/function.py @@ -243,7 +243,7 @@ def c() -> None: # are free in their nested functions. Generator functions need an environment class to # store a variable denoting the next instruction to be executed when the __next__ function # is called, along with all the variables inside the function itself. - if contains_nested or is_generator: + if contains_nested: # or is_generator: setup_env_class(builder) if is_nested or in_non_ext: diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index 74c8d27a6324..f40821fad434 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -64,8 +64,10 @@ def gen_generator_func( setup_generator_class(builder) load_env_registers(builder) gen_arg_defaults(builder) + g = instantiate_generator_class(builder) + builder.fn_info._curr_env_reg = g finalize_env_class(builder) - builder.add(Return(instantiate_generator_class(builder))) + builder.add(Return(g)) args, _, blocks, ret_type, fn_info = builder.leave() func_ir, func_reg = gen_func_ir(args, blocks, fn_info) @@ -126,18 +128,19 @@ def instantiate_generator_class(builder: IRBuilder) -> Value: # generator class gets instantiated from the callable class' '__call__' method, and hence # we use the callable class' environment register. Otherwise, we use the original # function's environment register. - if builder.fn_info.is_nested: - curr_env_reg = builder.fn_info.callable_class.curr_env_reg - else: - curr_env_reg = builder.fn_info.curr_env_reg + #if builder.fn_info.is_nested: + # curr_env_reg = builder.fn_info.callable_class.curr_env_reg + #else: + # curr_env_reg = builder.fn_info.curr_env_reg # Set the generator class' environment attribute to point at the environment class # defined in the current scope. - builder.add(SetAttr(generator_reg, ENV_ATTR_NAME, curr_env_reg, fitem.line)) + #builder.add(SetAttr(generator_reg, ENV_ATTR_NAME, curr_env_reg, fitem.line)) # Set the generator class' environment class' NEXT_LABEL_ATTR_NAME attribute to 0. zero = Integer(0) - builder.add(SetAttr(curr_env_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line)) + builder.add(SetAttr(generator_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line)) + #builder.add(SetAttr(curr_env_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line)) return generator_reg @@ -145,7 +148,8 @@ def setup_generator_class(builder: IRBuilder) -> ClassIR: name = f"{builder.fn_info.namespaced_name()}_gen" generator_class_ir = ClassIR(name, builder.module_name, is_generated=True) - generator_class_ir.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_info.env_class) + builder.fn_info.env_class = generator_class_ir + #generator_class_ir.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_info.env_class) generator_class_ir.mro = [generator_class_ir] builder.classes.append(generator_class_ir) @@ -392,7 +396,7 @@ def setup_env_for_generator_class(builder: IRBuilder) -> None: cls.send_arg_reg = exc_arg cls.self_reg = builder.read(self_target, fitem.line) - cls.curr_env_reg = load_outer_env(builder, cls.self_reg, builder.symtables[-1]) + cls.curr_env_reg = cls.self_reg # load_outer_env(builder, cls.self_reg, builder.symtables[-1]) # Define a variable representing the label to go to the next time # the '__next__' function of the generator is called, and add it From 8eab5a964cec5147fd1fe8ce0c95aebbbef16ab8 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 30 May 2025 14:00:23 +0100 Subject: [PATCH 2/9] WIP work on fixing non-merged cases --- mypyc/irbuild/context.py | 5 +++++ mypyc/irbuild/env_class.py | 7 ++----- mypyc/irbuild/function.py | 2 +- mypyc/irbuild/generator.py | 7 ++++--- 4 files changed, 12 insertions(+), 9 deletions(-) diff --git a/mypyc/irbuild/context.py b/mypyc/irbuild/context.py index a740f0b821d9..8d35c0ce2599 100644 --- a/mypyc/irbuild/context.py +++ b/mypyc/irbuild/context.py @@ -95,6 +95,11 @@ def curr_env_reg(self) -> Value: assert self._curr_env_reg is not None return self._curr_env_reg + def can_merge_generator_and_env_classes(self) -> bool: + # In simple cases we can place the environment into the generator class, + # instead of having two separate classes. + return self.is_generator and not self.is_nested and not self.contains_nested + class ImplicitClass: """Contains information regarding implicitly generated classes. diff --git a/mypyc/irbuild/env_class.py b/mypyc/irbuild/env_class.py index bc631e7d6b2e..b0909f86686a 100644 --- a/mypyc/irbuild/env_class.py +++ b/mypyc/irbuild/env_class.py @@ -47,7 +47,6 @@ class is generated, the function environment has not yet been ) env_class.attributes[SELF_NAME] = RInstance(env_class) if builder.fn_info.is_nested: - assert False # If the function is nested, its environment class must contain an environment # attribute pointing to its encapsulating functions' environment class. env_class.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_infos[-2].env_class) @@ -59,7 +58,8 @@ class is generated, the function environment has not yet been def finalize_env_class(builder: IRBuilder) -> None: """Generate, instantiate, and set up the environment of an environment class.""" - #instantiate_env_class(builder) + if not builder.fn_info.can_merge_generator_and_env_classes(): + instantiate_env_class(builder) # Iterate through the function arguments and replace local definitions (using registers) # that were previously added to the environment with references to the function's @@ -77,7 +77,6 @@ def instantiate_env_class(builder: IRBuilder) -> Value: ) if builder.fn_info.is_nested: - assert False builder.fn_info.callable_class._curr_env_reg = curr_env_reg builder.add( SetAttr( @@ -127,7 +126,6 @@ def load_outer_env( Returns the register where the environment class was loaded. """ - assert False env = builder.add(GetAttr(base, ENV_ATTR_NAME, builder.fn_info.fitem.line)) assert isinstance(env.type, RInstance), f"{env} must be of type RInstance" @@ -241,7 +239,6 @@ def setup_func_for_recursive_call(builder: IRBuilder, fdef: FuncDef, base: Impli prev_env = builder.fn_infos[-2].env_class prev_env.attributes[fdef.name] = builder.type_to_rtype(fdef.type) - assert False if isinstance(base, GeneratorClass): # If we are dealing with a generator class, then we need to first get the register # holding the current environment class, and load the previous environment class from diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py index c1557e410b2d..26257d7c3b88 100644 --- a/mypyc/irbuild/function.py +++ b/mypyc/irbuild/function.py @@ -243,7 +243,7 @@ def c() -> None: # are free in their nested functions. Generator functions need an environment class to # store a variable denoting the next instruction to be executed when the __next__ function # is called, along with all the variables inside the function itself. - if contains_nested: # or is_generator: + if contains_nested or (is_generator and not builder.fn_info.can_merge_generator_and_env_classes()): setup_env_class(builder) if is_nested or in_non_ext: diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index f40821fad434..5fa05c1673b0 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -64,10 +64,11 @@ def gen_generator_func( setup_generator_class(builder) load_env_registers(builder) gen_arg_defaults(builder) - g = instantiate_generator_class(builder) - builder.fn_info._curr_env_reg = g + gen = instantiate_generator_class(builder) + if builder.fn_info.can_merge_generator_and_env_classes(): + builder.fn_info._curr_env_reg = gen finalize_env_class(builder) - builder.add(Return(g)) + builder.add(Return(gen)) args, _, blocks, ret_type, fn_info = builder.leave() func_ir, func_reg = gen_func_ir(args, blocks, fn_info) From 7a8a67b02f3450fd2b208be2ce3a7b5d89f09d98 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 30 May 2025 14:05:07 +0100 Subject: [PATCH 3/9] Finish adding back support for non-merged use cases --- mypyc/irbuild/generator.py | 48 ++++++++++++++++++++++---------------- 1 file changed, 28 insertions(+), 20 deletions(-) diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index 5fa05c1673b0..974e4c04d4ca 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -125,23 +125,26 @@ def instantiate_generator_class(builder: IRBuilder) -> Value: fitem = builder.fn_info.fitem generator_reg = builder.add(Call(builder.fn_info.generator_class.ir.ctor, [], fitem.line)) - # Get the current environment register. If the current function is nested, then the - # generator class gets instantiated from the callable class' '__call__' method, and hence - # we use the callable class' environment register. Otherwise, we use the original - # function's environment register. - #if builder.fn_info.is_nested: - # curr_env_reg = builder.fn_info.callable_class.curr_env_reg - #else: - # curr_env_reg = builder.fn_info.curr_env_reg - - # Set the generator class' environment attribute to point at the environment class - # defined in the current scope. - #builder.add(SetAttr(generator_reg, ENV_ATTR_NAME, curr_env_reg, fitem.line)) - - # Set the generator class' environment class' NEXT_LABEL_ATTR_NAME attribute to 0. - zero = Integer(0) - builder.add(SetAttr(generator_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line)) - #builder.add(SetAttr(curr_env_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line)) + if builder.fn_info.can_merge_generator_and_env_classes(): + zero = Integer(0) + builder.add(SetAttr(generator_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line)) + else: + # Get the current environment register. If the current function is nested, then the + # generator class gets instantiated from the callable class' '__call__' method, and hence + # we use the callable class' environment register. Otherwise, we use the original + # function's environment register. + if builder.fn_info.is_nested: + curr_env_reg = builder.fn_info.callable_class.curr_env_reg + else: + curr_env_reg = builder.fn_info.curr_env_reg + + # Set the generator class' environment attribute to point at the environment class + # defined in the current scope. + builder.add(SetAttr(generator_reg, ENV_ATTR_NAME, curr_env_reg, fitem.line)) + + # Set the generator class' environment class' NEXT_LABEL_ATTR_NAME attribute to 0. + zero = Integer(0) + builder.add(SetAttr(curr_env_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line)) return generator_reg @@ -149,8 +152,10 @@ def setup_generator_class(builder: IRBuilder) -> ClassIR: name = f"{builder.fn_info.namespaced_name()}_gen" generator_class_ir = ClassIR(name, builder.module_name, is_generated=True) - builder.fn_info.env_class = generator_class_ir - #generator_class_ir.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_info.env_class) + if builder.fn_info.can_merge_generator_and_env_classes(): + builder.fn_info.env_class = generator_class_ir + else: + generator_class_ir.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_info.env_class) generator_class_ir.mro = [generator_class_ir] builder.classes.append(generator_class_ir) @@ -397,7 +402,10 @@ def setup_env_for_generator_class(builder: IRBuilder) -> None: cls.send_arg_reg = exc_arg cls.self_reg = builder.read(self_target, fitem.line) - cls.curr_env_reg = cls.self_reg # load_outer_env(builder, cls.self_reg, builder.symtables[-1]) + if builder.fn_info.can_merge_generator_and_env_classes(): + cls.curr_env_reg = cls.self_reg + else: + cls.curr_env_reg = load_outer_env(builder, cls.self_reg, builder.symtables[-1]) # Define a variable representing the label to go to the next time # the '__next__' function of the generator is called, and add it From 2f0d3243c00cca1fc540aada024d248ce5e7ec88 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 30 May 2025 14:26:54 +0100 Subject: [PATCH 4/9] WIP --- mypyc/codegen/emitmodule.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mypyc/codegen/emitmodule.py b/mypyc/codegen/emitmodule.py index 19a82b1f0f5b..8474be62579d 100644 --- a/mypyc/codegen/emitmodule.py +++ b/mypyc/codegen/emitmodule.py @@ -246,8 +246,8 @@ def compile_scc_to_ir( # Insert refcount handling. insert_ref_count_opcodes(fn) - #if fn in env_user_functions: - # insert_spills(fn, env_user_functions[fn]) + if fn in env_user_functions: + insert_spills(fn, env_user_functions[fn]) # Switch to lower abstraction level IR. lower_ir(fn, compiler_options) From 423e7a26ed323b7b7cf68ba10938eb061e0d7a68 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 30 May 2025 14:36:31 +0100 Subject: [PATCH 5/9] Fix spilling --- mypyc/transform/spill.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/mypyc/transform/spill.py b/mypyc/transform/spill.py index 3c014ca2c0da..21b2147d02b0 100644 --- a/mypyc/transform/spill.py +++ b/mypyc/transform/spill.py @@ -28,18 +28,20 @@ def insert_spills(ir: FuncIR, env: ClassIR) -> None: # TODO: Actually for now, no Registers at all -- we keep the manual spills entry_live = {op for op in entry_live if not isinstance(op, Register)} - ir.blocks = spill_regs(ir.blocks, env, entry_live, live) + ir.blocks = spill_regs(ir.blocks, env, entry_live, live, ir.arg_regs[0]) def spill_regs( - blocks: list[BasicBlock], env: ClassIR, to_spill: set[Value], live: AnalysisResult[Value] + blocks: list[BasicBlock], env: ClassIR, to_spill: set[Value], live: AnalysisResult[Value], + self_reg: Register ) -> list[BasicBlock]: for op in blocks[0].ops: if isinstance(op, GetAttr) and op.attr == "__mypyc_env__": env_reg = op break else: - raise AssertionError("could not find __mypyc_env__") + # Environment has been merged into generator object + env_reg = self_reg spill_locs = {} for i, val in enumerate(to_spill): From 636df11f625cc91edb31c92364d1afdb26459b81 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 30 May 2025 15:03:02 +0100 Subject: [PATCH 6/9] Fix nested generator --- mypyc/irbuild/generator.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index 974e4c04d4ca..abf04f9f30ed 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -64,10 +64,13 @@ def gen_generator_func( setup_generator_class(builder) load_env_registers(builder) gen_arg_defaults(builder) - gen = instantiate_generator_class(builder) if builder.fn_info.can_merge_generator_and_env_classes(): + gen = instantiate_generator_class(builder) builder.fn_info._curr_env_reg = gen - finalize_env_class(builder) + finalize_env_class(builder) + else: + finalize_env_class(builder) + gen = instantiate_generator_class(builder) builder.add(Return(gen)) args, _, blocks, ret_type, fn_info = builder.leave() From 0cc15570c5c91e676f3ddcbf8582ba985c87fe44 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Mon, 2 Jun 2025 14:58:09 +0100 Subject: [PATCH 7/9] Update comments --- mypyc/irbuild/generator.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index abf04f9f30ed..9dea0ee5f7c2 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -129,6 +129,7 @@ def instantiate_generator_class(builder: IRBuilder) -> Value: generator_reg = builder.add(Call(builder.fn_info.generator_class.ir.ctor, [], fitem.line)) if builder.fn_info.can_merge_generator_and_env_classes(): + # Set the generator instance to the initial state (zero). zero = Integer(0) builder.add(SetAttr(generator_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line)) else: @@ -145,7 +146,7 @@ def instantiate_generator_class(builder: IRBuilder) -> Value: # defined in the current scope. builder.add(SetAttr(generator_reg, ENV_ATTR_NAME, curr_env_reg, fitem.line)) - # Set the generator class' environment class' NEXT_LABEL_ATTR_NAME attribute to 0. + # Set the generator instance's environment to the initial state (zero). zero = Integer(0) builder.add(SetAttr(curr_env_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line)) return generator_reg From 1799fcd2c5505e49825733194c05de8baab1cd85 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Mon, 2 Jun 2025 14:58:28 +0100 Subject: [PATCH 8/9] Lint --- mypyc/irbuild/function.py | 4 +++- mypyc/transform/spill.py | 7 +++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py index 26257d7c3b88..dbebc350bb6c 100644 --- a/mypyc/irbuild/function.py +++ b/mypyc/irbuild/function.py @@ -243,7 +243,9 @@ def c() -> None: # are free in their nested functions. Generator functions need an environment class to # store a variable denoting the next instruction to be executed when the __next__ function # is called, along with all the variables inside the function itself. - if contains_nested or (is_generator and not builder.fn_info.can_merge_generator_and_env_classes()): + if contains_nested or ( + is_generator and not builder.fn_info.can_merge_generator_and_env_classes() + ): setup_env_class(builder) if is_nested or in_non_ext: diff --git a/mypyc/transform/spill.py b/mypyc/transform/spill.py index 21b2147d02b0..1ea7aae297b3 100644 --- a/mypyc/transform/spill.py +++ b/mypyc/transform/spill.py @@ -32,8 +32,11 @@ def insert_spills(ir: FuncIR, env: ClassIR) -> None: def spill_regs( - blocks: list[BasicBlock], env: ClassIR, to_spill: set[Value], live: AnalysisResult[Value], - self_reg: Register + blocks: list[BasicBlock], + env: ClassIR, + to_spill: set[Value], + live: AnalysisResult[Value], + self_reg: Register, ) -> list[BasicBlock]: for op in blocks[0].ops: if isinstance(op, GetAttr) and op.attr == "__mypyc_env__": From e3570340f48120259e7d94060b204e8d075bae39 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Mon, 2 Jun 2025 15:04:36 +0100 Subject: [PATCH 9/9] Fix self check --- mypyc/transform/spill.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mypyc/transform/spill.py b/mypyc/transform/spill.py index 1ea7aae297b3..d92dd661e7eb 100644 --- a/mypyc/transform/spill.py +++ b/mypyc/transform/spill.py @@ -38,6 +38,7 @@ def spill_regs( live: AnalysisResult[Value], self_reg: Register, ) -> list[BasicBlock]: + env_reg: Value for op in blocks[0].ops: if isinstance(op, GetAttr) and op.attr == "__mypyc_env__": env_reg = op