[FRONTEND] Fixed bugs in global symbols resolution of @triton.jit'd f…

…unctions (triton-lang#136)
quintinwang5 · Jul 27, 2021 · 298aead · 298aead
1 parent 94ce6aa
commit 298aead
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/triton/code_gen.py b/python/triton/code_gen.py
@@ -659,9 +659,12 @@ def parse(self):
 
     def __call__(self, *args, generator: CodeGenerator, **meta):
         try:
+            gscope = generator.gscope.copy()
             lscope = generator.lscope.copy()
             values = generator.module.get_values().copy()
+            generator.gscope = sys.modules[self.fn.__module__].__dict__
             ret = generator.visit_FunctionDef(self.parse().body[0], inline=True, arg_values=args)
+            generator.gscope = gscope
             generator.lscope = lscope
             generator.module.set_values(values)
             return ret