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

[FRONTEND] Add logical operations on constexprs (#1033)

Browse Source
This commit is contained in:
Sophia Wisdom
2023-01-04 18:06:32 -08:00
committed by GitHub
parent bc73bbb12c
commit 411bacb2a8
2 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/triton/compiler.py
View File

@@ -734,10 +734,6 @@ class CodeGenerator(ast.NodeVisitor):
assert len(node.values) == 2 assert len(node.values) == 2
lhs = self.visit(node.values[0]) lhs = self.visit(node.values[0])
rhs = self.visit(node.values[1]) rhs = self.visit(node.values[1])
if isinstance(lhs, triton.language.constexpr):
lhs = lhs.value
if isinstance(rhs, triton.language.constexpr):
rhs = rhs.value
fn = { fn = {
ast.And: 'logical_and', ast.And: 'logical_and',

12
python/triton/language/core.py
View File

@@ -403,6 +403,18 @@ class constexpr:
def __neg__(self): def __neg__(self):
return constexpr(-self.value) return constexpr(-self.value)
def __and__(self, other):
return constexpr(self.value & other.value)
def logical_and(self, other):
return constexpr(self.value and other.value)
def __or__(self, other):
return constexpr(self.value | other.value)
def logical_or(self, other):
return constexpr(self.value or other.value)
def __pos__(self): def __pos__(self):
return constexpr(+self.value) return constexpr(+self.value)