github/ROCm
1
1
Fork 0
mirror of https://github.com/ROCm/ROCm.git synced 2026-04-05 03:01:17 -04:00
Code Issues Packages Projects Releases Wiki Activity

[FRONTEND] Fix contains_return_op when analyzing functions in another module (#1365)

Browse Source
This commit is contained in:
Philippe Tillet
2023-03-18 15:02:45 -07:00
committed by GitHub
parent 6d61a5ca23
commit 2f035c0611
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/triton/compiler.py
View File

@@ -150,10 +150,10 @@ class CodeGenerator(ast.NodeVisitor):
self.global_uses[name] = value
return value
lookup_order = local_lookup, self.gscope.get, self.builtin_namespace.get
absent_marker = object()
def name_lookup(name: str) -> Any:
lookup_order = local_lookup, self.gscope.get, self.builtin_namespace.get
absent = absent_marker
for lookup_function in lookup_order:
value = lookup_function(name, absent)
@@ -201,7 +201,11 @@ class CodeGenerator(ast.NodeVisitor):
elif isinstance(node, ast.Call):
fn = self.visit(node.func)
if isinstance(fn, triton.JITFunction):
return self.contains_return_op(fn.parse())
old_gscope = self.gscope
self.gscope = sys.modules[fn.fn.__module__].__dict__
ret = self.contains_return_op(fn.parse())
self.gscope = old_gscope
return ret
return False
elif isinstance(node, ast.If):
pred = lambda s: self.contains_return_op(s)