[Frontend] Fix a few things in the frontend (#7060)

Fix assigning to tuples of other nodes, such as `a.x, b.y = unpack_me()`

Fix `@constexpr_function` so that the functions can still be called from
Python.

Add `__eq__` to `constexpr_type` so that the parser can reconcile types
of liveouts.

Author: Mogball

python/test/unit/language/test_frontend.py

@@ -179,6 +179,11 @@ def __init__(self, a, b):
     def create(a):
         return AggregateWithConstexpr(a, tl.constexpr(42))
 
+    @triton.jit
+    def modify(self, a):
+        self.a = a
+        return self
+
 
 @triton.jit
 def add_rhs_constexpr(agg):
@@ -196,3 +201,59 @@ def test_aggregate_with_constexpr():
     # CHECK: tt.func private @"test_frontend.add_rhs_constexpr__test_frontend.AggregateWithConstexpr<i32S4S, constexpr[42]>
     # CHECK: %cst = arith.constant dense<42> : tensor<4xi32>
     # CHECK: arith.addi %arg0, %cst : tensor<4xi32>
+
+
+@tl.constexpr_function
+def constexpr_function(x):
+    return x + 1
+
+
+@filecheck_test
+@triton.jit
+def test_constexpr_function_from_jit():
+    # CHECK-LABEL: test_constexpr_function
+    x: tl.constexpr = constexpr_function(7)
+    # CHECK: make_range {end = 8 : i32, start = 0 : i32}
+    tl.arange(0, x)
+
+
+def test_constexpr_function_from_python():
+    assert constexpr_function(7) == 8
+
+
+@triton.jit
+def swap(pair):
+    return pair.second, pair.first
+
+
+@filecheck_test
+@triton.jit
+def test_assign_tuple_attrs():
+    # CHECK-LABEL: test_assign_tuple_attrs
+    p = Pair(tl.arange(0, 4), tl.arange(4, 8))
+    # CHECK: [[P:%.*]]:2 = tt.call @{{.*}}swap
+    p.first, p.second = swap(p)
+    # CHECK: call @{{.*}}anchor{{.*}}([[P]]#0)
+    # CHECK: call @{{.*}}anchor{{.*}}([[P]]#1)
+    anchor(p.first)
+    anchor(p.second)
+
+
+@filecheck_test
+@triton.jit
+def test_reassign_aggregate_with_constexpr():
+    # CHECK-LABEL: test_reassign_aggregate_with_constexpr
+    agg = AggregateWithConstexpr.create(tl.arange(0, 4))
+    var = 1
+    # CHECK: [[AGG:%.*]] = scf.if {{.*}} -> (tensor<4xi32>)
+    # CHECK:   [[VALUE:%.*]] = tt.call {{.*}}modify
+    # CHECK:   yield [[VALUE]]
+    # CHECK: else
+    # CHECK:   [[VALUE:%.*]] = tt.call {{.*}}modify
+    # CHECK:   yield [[VALUE]]
+    if var == 0:
+        agg = agg.modify(tl.arange(4, 8))
+    else:
+        agg = agg.modify(tl.arange(8, 12))
+    # CHECK: call @{{.*}}anchor{{.*}}([[AGG]])
+    anchor(agg)

python/triton/compiler/code_generator.py

@@ -571,8 +571,8 @@ def assignTarget(self, target, value):
         if isinstance(target, ast.Subscript):
             return self.visit_Subscript_Store(target, value)
         if isinstance(target, ast.Tuple):
-            for i, name in enumerate(target.elts):
-                self.set_value(self.visit(name), value.values[i])
+            for i, target in enumerate(target.elts):
+                self.assignTarget(target, value.values[i])
             return
         if isinstance(target, ast.Attribute):
             base = self.visit(target.value)

python/triton/language/core.py

@@ -178,6 +178,9 @@ class constexpr_type(base_type):
     def __init__(self, value):
         self.value = value
 
+    def __eq__(self, other):
+        return self.value == other.value
+
     def __repr__(self) -> str:
         return f"constexpr[{self.value}]"
 
@@ -338,7 +341,7 @@ def constexpr_function(f):
     @wraps(f)
     def wrapper(*args, **kwargs):
         # de-constexpr arguments and discard the _builder keyword argument:
-        args = [getattr(x, "value", x) for x in args]
+        args = [_unwrap_if_constexpr(x) for x in args]
         kwargs = {k: getattr(v, "value", v) for (k, v) in kwargs.items() if k != "_builder"}
 
         # call the raw Python function f: