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] Fixed contains_return_op behavior (#1317)

Browse Source
This commit is contained in:
Philippe Tillet
2023-03-10 23:58:28 -08:00
committed by GitHub
parent 3fe3adbcde
commit 5a786cf778
1 changed files with 14 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
python/triton/compiler.py
View File

@@ -183,9 +183,23 @@ class CodeGenerator(ast.NodeVisitor):
break
return stmts and isinstance(stmt, ast.Return)
# TODO: should be its own AST visitor
def contains_return_op(self, node):
if isinstance(node, ast.Return):
return True
elif isinstance(node, ast.Assign):
return self.contains_return_op(node.value)
elif isinstance(node, ast.Module):
pred = lambda s: self.contains_return_op(s)
return any(pred(s) for s in node.body)
elif isinstance(node, ast.FunctionDef):
pred = lambda s: self.contains_return_op(s)
return any(pred(s) for s in node.body)
elif isinstance(node, ast.Call):
fn = self.visit(node.func)
if isinstance(fn, triton.JITFunction):
return self.contains_return_op(fn.parse())
return False
elif isinstance(node, ast.If):
pred = lambda s: self.contains_return_op(s)
ret = any(pred(s) for s in node.body)