From 06ef30a37a7127b6ebe50978bf88153a0fa9cba6 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Wed, 14 Oct 2026 15:15:15 +0200
Subject: [PATCH] fix more complex boundaries for compile/global

---
 kernels/src/kernels/layer/layer.py |  2 +-
 kernels/tests/test_func.py         | 32 +++++++++++++++++++++++++-----
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py
index 4b8d7d4a..08cf5566 100644
--- a/kernels/src/kernels/layer/layer.py
+++ b/kernels/src/kernels/layer/layer.py
@@ -440,7 +440,7 @@ def get(self):
     def set(self, value):
         import torch
 
-        if torch.compiler.is_compiling():
+        if self.uses_global_var or torch.compiler.is_compiling():
             previous = self.global_var
             self.global_var = value
             self.uses_global_var = True
diff --git a/kernels/tests/test_func.py b/kernels/tests/test_func.py
index 891583d8..a7cabfc2 100644
--- a/kernels/tests/test_func.py
+++ b/kernels/tests/test_func.py
@@ -61,13 +61,17 @@ def forward(self, x):
 # To check nested modules
 @use_kernelized_func(surprise_me)
 class Outer(nn.Module):
-    def __init__(self):
+    def __init__(self, num_inner=1):
         super().__init__()
-        self.inner = Inner()
+        self.inner = nn.ModuleList([Inner() for _ in range(num_inner)])
 
     def forward(self, x):
-        # The second call also verifies that Inner restored the outer context.
-        return surprise_me(self.inner(surprise_me(x)))
+        x = surprise_me(x)
+        for inner in self.inner:
+            x = inner(x)
+
+        # Also verifies that Inner restored the outer context
+        return surprise_me(x)
 
 
 def test_decorator():
@@ -292,7 +296,8 @@ def test_kernel_func_nested_dispatch(compile):
     model = Outer()
 
     _bind_forward(model._kernel_funcs["surprise_me"], AddOne.forward)
-    _bind_forward(model.inner._kernel_funcs["double_me"], TimesThree.forward)
+    for inner in model.inner:
+        _bind_forward(inner._kernel_funcs["double_me"], TimesThree.forward)
 
     if compile:
         # We only need to know whether it's safe around get/set so the backend is not relevant
@@ -302,3 +307,20 @@ def test_kernel_func_nested_dispatch(compile):
 
     # Outer (+1) -> Inner (*3) -> Outer (+1): 1 -> 2 -> 6 -> 7
     torch.testing.assert_close(model(x), torch.tensor(7.0))
+
+
+@pytest.mark.parametrize(("num_inner", "expected"), [(1, 7.0), (2, 19.0)])
+def test_kernel_func_nested_compile_context(monkeypatch, num_inner, expected):
+    model = Outer(num_inner)
+
+    _bind_forward(model._kernel_funcs["surprise_me"], AddOne.forward)
+    for inner in model.inner:
+        _bind_forward(inner._kernel_funcs["double_me"], TimesThree.forward)
+
+    # Mimic different compiled boundaries
+    compiling = iter([True, True, False])
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: next(compiling))
+
+    # Same as in `test_kernel_func_nested_dispatch` (just *3 once or twice)
+    x = torch.tensor(1.0)
+    torch.testing.assert_close(model(x), torch.tensor(expected))