gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[PYTHON] Added atomic_add (#94)

Browse Source
This commit is contained in:
Philippe Tillet
2021-04-29 09:13:45 -04:00
committed by Philippe Tillet
parent d7f87929fa
commit 2b75158426
5 changed files with 65 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/triton/code_gen.py
View File

@@ -249,9 +249,13 @@ class CodeGenerator(ast.NodeVisitor):
ast.Is: '__eq__',
ast.IsNot: '__ne__',
}[type(node.ops[0])]
if self.is_triton_object(lhs) or self.is_triton_object(rhs):
if self.is_triton_object(lhs):
return getattr(lhs, fn)(rhs, builder=self.builder)
return getattr(lhs, fn)(rhs)
elif self.is_triton_object(rhs):
fn = fn[:2] + 'r' + fn[2:]
return getattr(rhs, fn)(lhs, builder=self.builder)
else:
return getattr(lhs, fn)(rhs)
def visit_UnaryOp(self, node):
op = self.visit(node.operand)