[ALL] Merge master (#447)

python/triton/code_gen.py

@@ -358,9 +358,6 @@ class CodeGenerator(ast.NodeVisitor):
         for stmt in node.orelse:
             ast.NodeVisitor.generic_visit(self, stmt)
 
-    def visit_Str(self, node):
-        return ast.literal_eval(node)
-
     def visit_Subscript(self, node):
         assert node.ctx.__class__.__name__ == "Load"
         lhs = self.visit(node.value)
@@ -441,9 +438,6 @@ class CodeGenerator(ast.NodeVisitor):
     def visit_Index(self, node):
         return self.visit(node.value)
 
-    def visit_NameConstant(self, node):
-        return node.value
-
     def visit_keyword(self, node):
         return {node.arg: self.visit(node.value)}
 
@@ -460,10 +454,23 @@ class CodeGenerator(ast.NodeVisitor):
         if hasattr(fn, '__self__') and self.is_triton_object(fn.__self__) or \
                 sys.modules[fn.__module__] is triton.language.core:
             return fn(*args, _builder=self.builder, **kws)
+        if fn in self.builtins.values():
+            args = [arg.value if isinstance(arg, triton.language.constexpr) else arg
+                    for arg in args]
         return fn(*args, **kws)
 
-    def visit_Num(self, node):
-        return triton.language.constexpr(node.n)
+    def visit_Constant(self, node):
+        return triton.language.constexpr(node.value)
+
+    if sys.version_info < (3, 8):
+        def visit_NameConstant(self, node):
+            return triton.language.constexpr(node.value)
+
+        def visit_Num(self, node):
+            return triton.language.constexpr(node.n)
+
+        def visit_Str(self, node):
+            return triton.language.constexpr(ast.literal_eval(node))
 
     def visit_Attribute(self, node):
         lhs = self.visit(node.value)