[Frontend] Allow @_aggregate types to contain constexprs (#7024)

Constexpr fields don't have an IR representation and are carried through
the type. Also ensure that they are mangled into the type name.

This PR achieves this by adding a `constexpt_type` that contains the
constexpr value itself and makes this the type of `constexpr`. Shoutout
to @peterbell10 for the suggestion

Author: Mogball

python/test/unit/language/test_frontend.py

@@ -141,3 +141,35 @@ def test_call_in_loop():
     # CHECK:   call @accumulate
     for i in range(10):
         acc = accumulate(acc, i)
+
+
+@tl.core._aggregate
+class AggregateWithConstexpr:
+    a: tl.tensor
+    b: tl.constexpr
+
+    def __init__(self, a, b):
+        self.a = a
+        self.b = b
+
+    @staticmethod
+    def create(a):
+        return AggregateWithConstexpr(a, tl.constexpr(42))
+
+
+@triton.jit
+def add_rhs_constexpr(agg):
+    _ = agg.a + agg.b
+
+
+@filecheck_test
+@triton.jit
+def test_aggregate_with_constexpr():
+    # CHECK-LABEL: test_aggregate_with_constexpr
+    # CHECK: tt.call @"add_rhs_constexpr__test_frontend.AggregateWithConstexpr<i32S4S, constexpr[42]>
+    agg = AggregateWithConstexpr.create(tl.arange(0, 4))
+    add_rhs_constexpr(agg)
+
+    # CHECK: tt.func private @"add_rhs_constexpr__test_frontend.AggregateWithConstexpr<i32S4S, constexpr[42]>
+    # CHECK: %cst = arith.constant dense<42> : tensor<4xi32>
+    # CHECK: arith.addi %arg0, %cst : tensor<4xi32>

python/triton/language/core.py

@@ -139,7 +139,59 @@ class const:
     pass
 
 
-class constexpr:
+class base_value:
+    """Base class of values that exist in the triton IR (i.e. not constexprs).
+    """
+    type: base_type
+
+    def _flatten_ir(self, handles: List[ir.value]) -> None:
+        """Flatten frontend value into a sequence of mlir handles, which are appended
+        to the output list
+        """
+        raise NotImplementedError
+
+
+class base_type:
+
+    def __eq__(self, other):
+        raise NotImplementedError("Types must implement __eq__")
+
+    def __ne__(self, other):
+        return not (self == other)
+
+    def _unflatten_ir(self, handles: List[ir.value], cursor: int) -> Tuple[base_value, int]:
+        """Build a frontend value with the current dtype, wrapping a list of existing handles.
+        cursor is the index of the first handle relevant to this value, and the function
+        should return the updated cursor position after any handles consumed by the created value.
+        """
+        raise NotImplementedError
+
+    def mangle(self) -> str:
+        raise NotImplementedError(f"NYI: Type mangling for type {self.__class__}")
+
+    def _flatten_ir_types(self, builder: ir.builder, out: List[ir.type]) -> None:
+        raise NotImplementedError
+
+
+class constexpr_type(base_type):
+
+    def __init__(self, value):
+        self.value = value
+
+    def __repr__(self) -> str:
+        return f"constexpr[{self.value}]"
+
+    def mangle(self) -> str:
+        return repr(self)
+
+    def _flatten_ir_types(self, builder: ir.builder, out: List[ir.type]) -> None:
+        return
+
+    def _unflatten_ir(self, handles: List[ir.value], cursor: int) -> Tuple[base_value, int]:
+        return constexpr(self.value), cursor
+
+
+class constexpr(base_value):
     """
     This class is used to store a value that is known at compile-time.
     """
@@ -149,11 +201,14 @@ def __init__(self, value):
             self.value = value.value
         else:
             self.value = value
-        self.type = constexpr
+        self.type = constexpr_type(value)
 
     def __repr__(self) -> str:
         return f"constexpr[{self.value}]"
 
+    def _flatten_ir(self, handles: List[ir.value]) -> None:
+        return
+
     def __index__(self):
         return self.value
 
@@ -322,40 +377,6 @@ def check_bit_width(value, shift_value):
             )
 
 
-class base_value:
-    """Base class of values that exist in the triton IR (i.e. not constexprs).
-    """
-    type: base_type
-
-    def _flatten_ir(self, handles: List[ir.value]) -> None:
-        """Flatten frontend value into a sequence of mlir handles, which are appended
-        to the output list
-        """
-        raise NotImplementedError
-
-
-class base_type:
-
-    def __eq__(self, other):
-        raise NotImplementedError("Types must implement __eq__")
-
-    def __ne__(self, other):
-        return not (self == other)
-
-    def _unflatten_ir(self, handles: List[ir.value], cursor: int) -> Tuple[base_value, int]:
-        """Build a frontend value with the current dtype, wrapping a list of existing handles.
-        cursor is the index of the first handle relevant to this value, and the function
-        should return the updated cursor position after any handles consumed by the created value.
-        """
-        raise NotImplementedError
-
-    def mangle(self) -> str:
-        raise NotImplementedError(f"NYI: Type mangling for type {self.__class__}")
-
-    def _flatten_ir_types(self, builder: ir.builder, out: List[ir.type]) -> None:
-        raise NotImplementedError
-
-
 # -----------------------
 # dtype
 # -----------------------