[Cherry-pick][RESOLVED] [FRONTEND] Wrap non-jit function call results in constexpr automatically (#8002) (#544)

Summary:
⚠️ **MERGE CONFLICTS DETECTED** ⚠️

This cherry-pick contains merge conflicts that require manual resolution.

Original Commit: b93eefd
Original Author: peterbell10
Original Date: 2025-08-28 22:30:24 +0100

**Action Required:**
1. Check out this branch locally
2. Resolve the merge conflicts in the affected files
3. Commit the resolved changes
4. Update this PR

Original commit message:
```
[FRONTEND] Wrap non-jit function call results in constexpr automatically (#8002)

Fixes
triton-lang/triton#8001 (comment)
```

This PR was automatically cherry-picked from the upstream triton-lang/triton repository.
The conflicts have been committed with conflict markers for easier resolution.

Pull Request resolved: #544

Reviewed By: dshi7

Differential Revision: D86147298

Pulled By: agron911

fbshipit-source-id: 7987bddfa82532da67243016570ad14f91abb2ce

Author: agron911

python/test/gluon/test_frontend.py

@@ -2250,3 +2250,11 @@ def test_infer_layout_for_padded_shared(target):
   }
 }
 """)
+
+
+@filecheck_test
+@gluon.jit
+def test_layout_zeros():
+    # CHECK: #blocked = #ttg.blocked
+    # CHECK: arith.constant dense<0.000000e+00> : tensor<128xf32, #blocked>
+    ttgl.zeros([128], ttgl.float32, layout=ttgl.BlockedLayout([1], [32], [4], [0]))

python/triton/compiler/code_generator.py

@@ -1,4 +1,5 @@
 import ast
+import builtins
 import contextlib
 import copy
 import inspect
@@ -1373,7 +1374,15 @@ def call_Function(self, node, fn, args, kws):
         if fn in self.builtin_namespace.values():
             args = map(_unwrap_if_constexpr, args)
         ret = fn(*args, **kws)
-        return _apply_to_tuple_values(ret, lambda x: x) if _is_namedtuple(type(ret)) else ret
+
+        def wrap_constexpr(x):
+            if _is_triton_value(x):
+                return x
+            return constexpr(x)
+
+        if isinstance(ret, (builtins.tuple, language.tuple)):
+            return _apply_to_tuple_values(ret, wrap_constexpr)
+        return wrap_constexpr(ret)
 
     def call_Method(self, node, fn, fn_self, args, kws):
         if isinstance(fn, JITFunction):

python/triton/language/core.py

@@ -3183,7 +3183,7 @@ def kernel(A, B, C, D, BLOCK: tl.constexpr):
 # -----------------------
 
 
-class static_range:
+class static_range(base_value):
     """
     Iterator that counts upward forever.
 
@@ -3223,7 +3223,7 @@ def __next__(self):
         raise RuntimeError("static_range can only be used in @triton.jit'd functions")
 
 
-class range:
+class range(base_value):
     """
     Iterator that counts upward forever.
 
@@ -3293,7 +3293,7 @@ def __next__(self):
         raise RuntimeError("tl.range can only be used in @triton.jit'd functions")
 
 
-class condition:
+class condition(base_value):
     """
     While loop condition wrapper.