Preserve constexprs returned from functions

It is very easy to convert a constexpr into a tensor but difficult to go
in the reverse direction. The current behavior makes it very difficult
to support certain patterns without ceremony. While technically a
compatibility break, I expect this to be rarely hit because many
operations (like assignment) decay constexprs to tensors anyway.

Author: saagarjha

python/test/unit/language/test_frontend.py

@@ -652,3 +652,18 @@ def max_kernel(a: tl.constexpr, b: tl.constexpr):
 
     with pytest.raises(CompilationError):
         run_parser(max_kernel, args=(1.0, -0.0))
+
+
+@pytest.mark.interpreter
+def test_constexpr_return():
+
+    @triton.jit
+    def get_constexpr_value():
+        return tl.constexpr(42)
+
+    @triton.jit
+    def test():
+        x: tl.constexpr = get_constexpr_value()
+        tl.static_assert(x == 42)
+
+    test[(1, )]()

python/triton/compiler/code_generator.py

@@ -550,7 +550,7 @@ def visit_Return(self, node):
         def decay(value):
             if isinstance(value, language.tuple):
                 return _apply_to_tuple_values(value, decay)
-            elif isinstance(value, (language.constexpr, int, float)):
+            elif isinstance(value, (int, float)):
                 return self.semantic.to_tensor(value)
             return value