[FRONTEND] Fix argument annotation after host side TensorDescriptor (triton-lang#6509)

`val_paths.index(attr_path)` computes the index into the python values
after flattening tuples. Instead, we can use `cursor` which corresponds
to the block argument index in the ttir.

Author: peterbell10

python/test/unit/cuda/test_tensor_descriptor.py

@@ -1402,3 +1402,17 @@ def test_host_tensor_descriptor_matmul(num_stages, num_ctas, BLOCK_M, BLOCK_N, B
         # TODO: The use of stmatrix for Blackwell is currently not supported.
         # Only a subset of TMEM and stmatrix layout pairs are compatible, for example 16x256bx2 and m8n8x4.
         assert "stmatrix.sync.aligned.m8n8.x4.shared.b16" in kernel.asm["ptx"]
+
+
+@requires_tma
+def test_specialization_after_host_tensordesc():
+
+    @triton.jit
+    def kernel(a, b):
+        pass
+
+    device = "cuda"
+    A = torch.randn(1024, device=device)
+    desc = TensorDescriptor.from_tensor(A, [128])
+    h = kernel.warmup(desc, 16, grid=(1, ))
+    assert ", %arg3: i32 {tt.divisibility = 16 : i32}" in h.asm["ttir"]

python/triton/compiler/code_generator.py

@@ -253,16 +253,16 @@ def make_template(ty):
         vals = make_template(self.arg_types)
         is_val = lambda path, _: path not in self.constants and _ is not None
         val_paths = list(find_paths_if(self.arg_types, is_val))
-        # > set attributes
-        for attr_path, attr_specs in self.attrs.items():
-            for attr_name, attr_val in attr_specs:
-                if attr_path in val_paths:
-                    fn.set_arg_attr(val_paths.index(attr_path), attr_name, attr_val)
         # > add IR values to the template
         cursor = 0
         handles = [fn.args(i) for i in range(fn.get_num_args())]
         for path in val_paths:
             ty = get_iterable_path(self.arg_types, path)
+            # > set attributes
+            attr_specs = self.attrs.get(path, [])
+            for attr_name, attr_val in attr_specs:
+                fn.set_arg_attr(cursor, attr_name, attr_val)
+            # > build frontend value
             val, cursor = ty._unflatten_ir(handles, cursor)
             set_iterable_path(vals, path, val)
         # > add constexpr values to the template