[Frontend] Treat JITFunction as constexpr type (#6988)

This teaches a few more places in the frontend to treat JITFunction as
constexprs. This allows, for example, passing lists of functions as
constexprs.

Author: Mogball

python/test/unit/language/test_frontend.py

@@ -185,3 +185,26 @@ def test_aggregate_initializers():
     # CHECK: call @"anchor{{.*}}"([[RANGE]])
     value.modify(tl.arange(4, 8))
     anchor(value)
+
+
+@triton.jit
+def forward(arg):
+    return arg
+
+
+@triton.jit
+def list_of_functions_constexpr(arg, fns: tl.constexpr):
+    for i in tl.static_range(len(fns)):
+        fns[i](arg)
+
+
+@filecheck_test
+@triton.jit
+def test_list_of_functions():
+    # CHECK-LABEL: test_list_of_functions
+    # CHECK: call @"list_of_functions_constexpr{{.*}}cJITFunction(test_frontend:anchor){{.*}}cJITFunction(test_frontend:forward)"
+
+    # CHECK-LABEL: tt.func private @"list_of_functions_constexpr
+    # CHECK-NEXT: call @anchor
+    # CHECK-NEXT: call @forward
+    list_of_functions_constexpr(tl.arange(0, 4), [anchor, forward])

python/triton/compiler/code_generator.py

@@ -48,7 +48,7 @@ def _is_triton_tensor(o: Any) -> bool:
 
 
 def _is_constexpr(o: Any) -> bool:
-    return o is None or isinstance(o, (constexpr, language.core.dtype))
+    return o is None or isinstance(o, (constexpr, language.core.dtype, JITFunction))
 
 
 def _is_non_scalar_tensor(o: Any) -> bool:

python/triton/language/core.py

@@ -1576,6 +1576,7 @@ def type(self):
     aggregate_value.__name__ = cls.__name__
     aggregate_value.__module__ = cls.__module__
     aggregate_value.__qualname__ = cls.__qualname__
+    aggregate_value.__doc__ = cls.__doc__
 
     return aggregate_value
 

python/triton/runtime/jit.py

@@ -711,6 +711,11 @@ def cache_key(self):
             self.used_global_vals = dict(sorted(dependencies_finder.used_global_vals.items()))
         return self.hash
 
+    @property
+    def type(self):
+        from triton.language.core import constexpr
+        return constexpr
+
     def warmup(self, *args, grid, **kwargs):
         return self.run(grid=grid, warmup=True, *map(MockTensor.wrap_dtype, args), **kwargs)