Use correct types for constexpr arguments (#8248)

A previous commit improved str_to_ty to support constexprs, but its use
in ast_to_ttir did not correctly handle nested tuples: fixing this
properly requires recursion. Instead of invoking str_to_ty, it's easier
to just fix up the types ourselves. Fortunately, we have the arguments
passed into the kernel and can use this to go in and correct the
parameters before code generation sees it.

Author: saagarjha

python/test/unit/language/test_tuple.py

@@ -256,6 +256,19 @@ def m_to_the_n(X, shape: tl.constexpr, strides, m_n):
     torch.testing.assert_close(x, expected_x, rtol=0, atol=0)
 
 
+def test_passing_nested_tuple_with_constexpr(device):
+
+    @triton.jit
+    def test(x):
+        # This creates a new scope, which will force a copy of liveins. It's
+        # important for this to happen as it forces IR flattening/unflattening,
+        # which relies on the types being correct for the roundtrip to succeed.
+        for _ in range(1):
+            tl.static_assert(x[1][0] == 2)
+
+    test[(1, )](((1, ), (tl.constexpr(2), )))
+
+
 def test_passing_tuple_to_make_tensor_descriptor(device, with_allocator):
 
     @triton.jit

python/triton/compiler/code_generator.py

@@ -1589,16 +1589,24 @@ def ret(self, node: ast.Call):
 
 def ast_to_ttir(fn, src, context, options, codegen_fns, module_map, module=None):
     arg_types = [None] * len(fn.arg_names)
-    const_iter = iter(src.constants.items())
-    kc, vc = next(const_iter, (None, None))
-
-    for i, (ks, v) in enumerate(src.signature.items()):
-        idx = fn.arg_names.index(ks)
-        cexpr = None
-        if kc is not None and kc[0] == i:
-            cexpr = vc
-            kc, vc = next(const_iter, (None, None))
-        arg_types[idx] = str_to_ty(v, cexpr)
+
+    for k, v in src.signature.items():
+        idx = fn.arg_names.index(k)
+        arg_types[idx] = str_to_ty(v, None)
+
+    def apply_constexpr_types(argument, indices, value):
+        index = indices.pop()
+        if len(indices) == 0:
+            if isinstance(argument, list):
+                argument[index] = constexpr(value).type
+            else:
+                argument.types[index] = constexpr(value).type
+        else:
+            apply_constexpr_types(argument[index], indices, value)
+
+    for path, value in src.constants.items():
+        apply_constexpr_types(arg_types, list(path)[::-1], value)
+
     prototype = ASTFunction([], arg_types, src.constants, src.attrs)
     file_name, begin_line = get_jit_fn_file_line(fn)
     # query function representation