Support Enum Type construction and comparision

Signed-off-by: Qiqi Xiao <qiqix@nvidia.com>

Author: xiaoqiqi177

changelog.d/enum.md

@@ -0,0 +1,4 @@
+<!--- SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!--- SPDX-License-Identifier: Apache-2.0 -->
+
+- Added comparison and constructor support for Python ``enum.Enum`` inside kernels. Enum members can now be compared with ``==`` / ``!=`` and constructed from a constant value (e.g. ``Color(0)``).

experimental/cuda-lang/test/examples/test_wmma.py

@@ -338,7 +338,7 @@ def kernel(a, b, c, d):
         if tid >= 32:
             return
 
-        if static_eval(config.acc_layout == wmma.Layout.COL_MAJOR):
+        if config.acc_layout == wmma.Layout.COL_MAJOR:
             c_ptr = c.get_element_pointer((tile_col, tile_row))
             acc_ldm = global_m
         else:
@@ -353,14 +353,14 @@ def kernel(a, b, c, d):
         )
 
         for kk in range(0, global_k, wmma_k):
-            if static_eval(config.a_layout == wmma.Layout.COL_MAJOR):
+            if config.a_layout == wmma.Layout.COL_MAJOR:
                 a_ptr = a.get_element_pointer((kk, tile_row))
                 a_ldm = global_m
             else:
                 a_ptr = a.get_element_pointer((tile_row, kk))
                 a_ldm = global_k
 
-            if static_eval(config.b_layout == wmma.Layout.ROW_MAJOR):
+            if config.b_layout == wmma.Layout.ROW_MAJOR:
                 b_ptr = b.get_element_pointer((kk, tile_col))
                 b_ldm = global_n
             else:
@@ -374,7 +374,7 @@ def kernel(a, b, c, d):
                 config.satf,
             )
 
-        if static_eval(config.acc_layout == wmma.Layout.COL_MAJOR):
+        if config.acc_layout == wmma.Layout.COL_MAJOR:
             d_ptr = d.get_element_pointer((tile_col, tile_row))
         else:
             d_ptr = d.get_element_pointer((tile_row, tile_col))

src/cuda/tile/_ir/core_ops.py

@@ -27,7 +27,7 @@
     DataclassInfo, DataclassTy, DataclassValue, BoundMethodValue, BoundMethodTy, InvalidType, \
     ContextManagerTy, ContextManagerLifecycle, LiveCapturedScope, ClosureTy, ClosureValue, \
     RangeIterType, RangeValue, TypeTy, ModuleTy, NONE, SliceType, StringTy, FormattedStringTy, \
-    StringFormat, FormattedStringValue, FormattedPiece, DictTy, DictValue
+    StringFormat, FormattedStringValue, FormattedPiece, DictTy, DictValue, EnumTy
 from cuda.tile._ir.typing_support import type_of_constant_python_value, \
     loose_type_of_constant_python_value, get_dataclass_info, as_third_party_dtype_spec
 from cuda.tile._ir2bytecode import BytecodeContext
@@ -648,6 +648,16 @@ async def getattr_dataclass_impl(object: Var, name: Var):
     else:
         return sym2var(cls_attr, constant_only=True)
 
+
+@impl(getattr, overload=(EnumTy, "name"))
+def getattr_enum_name_impl(object: Var, name: Var):
+    return sym2var(object.get_constant().name)
+
+
+@impl(getattr, overload=(EnumTy, "value"))
+def getattr_enum_value_impl(object: Var, name: Var):
+    return sym2var(object.get_constant().value)
+
 # ===========================================================================================
 
 
@@ -707,6 +717,12 @@ def comparison_string_impl(fn: str, x: Var, y: Var):
     return binop_propagate_constant(fn, x.get_type().value, y.get_type().value, None)
 
 
+@comparison_operator_impl(_registry, EnumTy, EnumTy)
+def comparison_enum_impl(fn: str, x: Var, y: Var):
+    from cuda.tile._ir.arithmetic_ops import binop_propagate_constant
+    return binop_propagate_constant(fn, x.get_constant(), y.get_constant(), None)
+
+
 # ===========================================================================================
 # Print
 # ===========================================================================================
@@ -787,6 +803,9 @@ def _expand_var(var: Var | str, format_spec: str | None = None,
             leaf_vars.append(var)
         elif isinstance(ty, DTypeSpec):
             format_parts.append(str(ty.dtype))
+        elif isinstance(ty, EnumTy):
+            member = var.get_constant()
+            format_parts.append(f"{ty.enum_ty.__name__}.{member.name}")
         else:
             raise TileTypeError(f"Can't print value of type {ty}")
 

src/cuda/tile/_passes/hir2ir.py

@@ -5,13 +5,16 @@
 import sys
 from contextlib import contextmanager
 import dataclasses
+
+from enum import Enum
 from typing import Sequence, Mapping, Callable
 
 from .ast2hir import get_function_hir
 from .. import TileTypeError
 from .._coroutine_util import resume_after, run_coroutine
 from .._datatype import PointerInfo
-from .._exception import Loc, FunctionDesc, TileInternalError, TileError, TileRecursionError
+from .._exception import Loc, FunctionDesc, TileInternalError, TileError, TileRecursionError, \
+    TileValueError
 from .._execution import is_stub
 from .._ir import hir, ir
 from .._ir.ir import Var, IRContext
@@ -302,6 +305,17 @@ async def call(callee_var: Var, args, kwargs) -> Var | None:
     elif isinstance(callee_ty, TypeTy) and callee_ty.ty is PointerInfo:
         arg_list = _bind_args(_POINTER_INFO_SIGNATURE, "PointerInfo", args, kwargs)
         return await _call_builtin(PointerInfo, arg_list, builder)
+    elif isinstance(callee_ty, TypeTy) and issubclass(callee_ty.ty, Enum):
+        if len(args) != 1 or kwargs:
+            raise TileTypeError("Enum constructor takes exactly one positional argument")
+        arg = args[0]
+        if not arg.is_constant():
+            raise TileTypeError("Enum constructor argument must be a constant")
+        val = arg.get_constant()
+        try:
+            return loosely_typed_const(callee_ty.ty(val))
+        except ValueError:
+            raise TileValueError(f"{val!r} is not a valid {callee_ty.ty.__name__}")
     else:
         raise TileTypeError(f"Cannot call an object of type {callee_ty}")
 

test/test_enum.py

@@ -0,0 +1,189 @@
+# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from enum import Enum, IntEnum
+
+import pytest
+import torch
+
+import cuda.tile as ct
+from cuda.tile import TileTypeError, TileValueError
+
+
+class Color(Enum):
+    RED = 0
+    GREEN = 1
+    BLUE = 2
+
+
+class Status(Enum):
+    OK = "ok"
+    ERROR = "error"
+
+
+class Weight(Enum):
+    LIGHT = 0.5
+    HEAVY = 2.0
+
+
+class Priority(IntEnum):
+    LOW = 0
+    MEDIUM = 1
+    HIGH = 2
+
+
+def test_comparison_eq():
+    @ct.kernel
+    def kernel(out):
+        x = Color.RED
+        if x == Color.RED:
+            ct.scatter(out, (), 1)
+        else:
+            ct.scatter(out, (), -1)
+
+    out = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (out,))
+    assert out.item() == 1
+
+
+def test_comparison_not_equal():
+    @ct.kernel
+    def kernel(out):
+        x = Color.RED
+        if x != Color.GREEN:
+            ct.scatter(out, (), 1)
+        else:
+            ct.scatter(out, (), -1)
+
+    out = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (out,))
+    assert out.item() == 1
+
+
+def test_construction_from_known_int():
+    @ct.kernel
+    def kernel(out):
+        i = 0
+        x = Color(i)
+        if x == Color.RED:
+            ct.scatter(out, (), 10)
+        elif x == Color.GREEN:
+            ct.scatter(out, (), 20)
+        else:
+            ct.scatter(out, (), 30)
+
+    out = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (out,))
+    assert out.item() == 10
+
+
+def test_construction_from_string_value():
+    @ct.kernel
+    def kernel(out):
+        x = Status("ok")
+        if x == Status.OK:
+            ct.scatter(out, (), 1)
+        else:
+            ct.scatter(out, (), 0)
+
+    out = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (out,))
+    assert out.item() == 1
+
+
+def test_construction_from_float_value():
+    @ct.kernel
+    def kernel(out):
+        x = Weight(0.5)
+        if x == Weight.LIGHT:
+            ct.scatter(out, (), 1)
+        else:
+            ct.scatter(out, (), 0)
+
+    out = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (out,))
+    assert out.item() == 1
+
+
+def test_intenum_ordering():
+    @ct.kernel
+    def kernel(out):
+        if Priority.LOW < Priority.HIGH:
+            ct.scatter(out, (), 1)
+        else:
+            ct.scatter(out, (), -1)
+
+    out = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (out,))
+    assert out.item() == 1
+
+
+# ===========================================================================
+# Error cases
+# ===========================================================================
+
+def test_construction_from_runtime_value_raises():
+    @ct.kernel
+    def kernel(x, out):
+        bid = ct.bid(0)
+        _ = Color(bid)
+        ct.scatter(out, (), 0)
+
+    x = torch.zeros(1, device="cuda")
+    out = torch.zeros((), dtype=torch.int32, device="cuda")
+    with pytest.raises(TileTypeError):
+        ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, out))
+
+
+@pytest.mark.parametrize("invalid_value", ["foo", 99])
+def test_construction_from_invalid_type_or_value_raises(invalid_value):
+    @ct.kernel
+    def kernel(out):
+        _ = Color(invalid_value)
+        ct.scatter(out, (), 0)
+
+    out = torch.zeros((), dtype=torch.int32, device="cuda")
+    with pytest.raises(TileValueError):
+        ct.launch(torch.cuda.current_stream(), (1,), kernel, (out,))
+
+
+@pytest.mark.parametrize("enum_value", [Color.BLUE, Status.ERROR,  Weight.HEAVY])
+def test_name_attribute(enum_value):
+    name = enum_value.name
+
+    @ct.kernel
+    def kernel(out):
+        if enum_value.name == name:
+            ct.scatter(out, (), 1)
+
+    out = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (out,))
+    assert out.item() == 1
+
+
+@pytest.mark.parametrize("enum_value", [Color.BLUE, Status.ERROR,  Weight.HEAVY])
+def test_value_attribute(enum_value):
+    value = enum_value.value
+
+    @ct.kernel
+    def kernel(out):
+        if enum_value.value == value:
+            ct.scatter(out, (), 1)
+
+    out = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (out,))
+    assert out.item() == 1
+
+
+def test_enum_ordering_raises():
+    @ct.kernel
+    def kernel(out):
+        if Color.RED < Color.GREEN:
+            ct.scatter(out, (), 1)
+        else:
+            ct.scatter(out, (), -1)
+
+    out = torch.zeros((), dtype=torch.int32, device="cuda")
+    with pytest.raises(TileTypeError):
+        ct.launch(torch.cuda.current_stream(), (1,), kernel, (out,))

test/test_print.py

@@ -6,6 +6,7 @@
 import ctypes
 import sys
 import traceback
+from enum import Enum
 import torch
 import numpy as np
 import pytest
@@ -239,6 +240,19 @@ def kernel_print_dtype(x, TILE: ct.Constant[int]):
     print(f"dtype = {tx.dtype}")
 
 
+class _PrintColor(Enum):
+    RED = 0
+    GREEN = 1
+
+
+@ct.kernel(opt_level=_OPT_LEVEL)
+def kernel_print_enum(x, TILE: ct.Constant[int]):
+    ct.print(_PrintColor.GREEN)
+    print(_PrintColor.GREEN)
+    ct.print(f"color = {_PrintColor.RED}")
+    print(f"color = {_PrintColor.RED}")
+
+
 _KERNELS_MAP_ = {
     "kernel_printf": kernel_printf,
     "kernel_print": kernel_print,
@@ -257,6 +271,7 @@ def kernel_print_dtype(x, TILE: ct.Constant[int]):
     "kernel_print_tuple_w_tile": kernel_print_tuple_w_tile,
     "kernel_print_tuple_tile_shape": kernel_print_tuple_tile_shape,
     "kernel_print_dtype": kernel_print_dtype,
+    "kernel_print_enum": kernel_print_enum,
 }
 
 
@@ -471,6 +486,14 @@ def test_ct_print_dtype():
     assert actual_outs[3] == "dtype = float32"  # print(f"dtype = {tx.dtype}")
 
 
+def test_ct_print_enum():
+    actual_outs = _run_kernel(kernel_print_enum, (8,), "float32", 8)
+    assert actual_outs[0] == "_PrintColor.GREEN"          # ct.print(_PrintColor.GREEN)
+    assert actual_outs[1] == "_PrintColor.GREEN"          # print(_PrintColor.GREEN)
+    assert actual_outs[2] == "color = _PrintColor.RED"   # ct.print(f"color = {_PrintColor.RED}")
+    assert actual_outs[3] == "color = _PrintColor.RED"   # print(f"color = {_PrintColor.RED}")
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1 and sys.argv[1] == "start_kernel_runner":
         _kernel_runner_main()