Clean up constant creation logic

- Allow only non-aggregate types in loosely_typed_constant().
- Aggregate types must go through sym2var(), which now gets an optional
  `constant_only` flag. This makes sure we also support dataclasses
  and other aggregate types.
- This enables frozen dataclasses as global constants.
- Get rid of get_constant_value(), unify with typeof_pyval().

Signed-off-by: Greg Bonik <gbonik@nvidia.com>

Author: gbonik

changelog.d/global-dataclass.md

@@ -0,0 +1 @@
+- Allow frozen dataclass instances to be used as globals in device code.

src/cuda/tile/_ir/ops.py

@@ -59,7 +59,8 @@
 )
 from .scope import Scope, JumpInfo, ControlFlowInfo
 from .typing_support import (
-    typeof_pyval, loose_type_of_pyval, get_constant_value, get_dataclass_info,
+    get_dataclass_info, as_third_party_dtype_spec, type_of_constant_python_value,
+    loose_type_of_constant_python_value,
 )
 from .type import (
     PartitionViewTy, StridedViewTy, GatherScatterViewTy, TupleTy, TileTy, NoneType,
@@ -72,7 +73,7 @@
 )
 from cuda.tile._datatype import (
     DType, is_integral, is_float, is_signed, is_boolean, is_pointer_dtype, PointerInfo,
-    opaque_pointer_dtype, pointer_dtype,
+    opaque_pointer_dtype, pointer_dtype
 )
 from cuda.tile._ir2bytecode import (
     BytecodeContext, typeid,
@@ -651,21 +652,29 @@ def loosely_typed_const(value: Any,
                         ty: Optional[Type] = None,
                         loose_ty: Optional[Type] = None,
                         name: str | None = None) -> Var:
+    builder = Builder.get_current()
     if ty is None:
-        if isinstance(value, tuple):
-            return build_tuple(tuple(loosely_typed_const(item) for item in value))
-        ty = typeof_pyval(value)
-    ret = strictly_typed_const(value, ty, name=name)
+        ty = type_of_constant_python_value(value, builder.ir_ctx.typing_hooks)
+    assert not ty.is_aggregate(), "Use sym2var(value, constant_only=True) instead"
+
+    # Normalize third party dtype spec objects (e.g. torch.float32 -> ct.float32)
+    if isinstance(ty, DTypeSpec):
+        value = ty.dtype
+
+    ret = _strictly_typed_const_inner(builder, value, ty, name=name)
     if loose_ty is None:
-        loose_ty = loose_type_of_pyval(value)
+        loose_ty = loose_type_of_constant_python_value(value, builder.ir_ctx.typing_hooks)
     ret.set_loose_type(loose_ty)
     return ret
 
 
 def strictly_typed_const(value: Any, ty: Type, name: str | None = None) -> Var:
-    builder = Builder.get_current()
-    result = None if name is None else builder.ir_ctx.make_var(name, builder.loc)
+    return _strictly_typed_const_inner(Builder.get_current(), value, ty, name)
 
+
+def _strictly_typed_const_inner(builder: Builder,
+                                value: Any, ty: Type, name: str | None = None) -> Var:
+    result = None if name is None else builder.ir_ctx.make_var(name, builder.loc)
     ret = builder.add_operation(TypedConst, ty, dict(value=value), result=result)
     if not isinstance(ty, TileTy) or ty.ndim == 0:
         # We currently don't have a way to represent an N-dimensional tile constant
@@ -1893,7 +1902,7 @@ def getattr_tile_dtype_impl(object: Var, name: Var):
 
 @impl(getattr, overload=(TileTy, "shape"))
 def getattr_tile_shape_impl(object: Var, name: Var):
-    return loosely_typed_const(object.get_type().shape)
+    return sym2var(object.get_type().shape, constant_only=True)
 
 
 @impl(getattr, overload=(TileTy, "ndim"))
@@ -1924,12 +1933,12 @@ def getattr_tiled_view_dtype_impl(object: Var, name: Var):
 
 @impl(getattr, overload=(TiledViewTy, "tile_shape"))
 def getattr_tiled_view_tile_shape_impl(object: Var, name: Var):
-    return loosely_typed_const(object.get_type().tile_shape)
+    return sym2var(object.get_type().tile_shape, constant_only=True)
 
 
 @impl(getattr, overload=(TiledViewTy, "traversal_steps"))
 def getattr_tiled_view_traversal_steps_impl(object: Var, name: Var):
-    return loosely_typed_const(object.get_type().traversal_steps)
+    return sym2var(object.get_type().traversal_steps, constant_only=True)
 
 
 @impl(getattr, overload=(TiledViewTy, "num_tiles"))
@@ -1981,7 +1990,7 @@ def getattr_module_impl(object: Var, name: Var):
     ty = object.get_type()
     attr_name = require_constant_str(name)
     try:
-        return loosely_typed_const(getattr(ty.py_mod, attr_name))
+        return sym2var(getattr(ty.py_mod, attr_name), constant_only=True)
     except AttributeError:
         raise TileTypeError(f"Module '{ty.py_mod.__name__}' has no attribute '{attr_name}'")
 
@@ -1991,7 +2000,7 @@ def getattr_type_impl(object: Var, name: Var):
     ty = object.get_type()
     attr_name = require_constant_str(name)
     try:
-        return loosely_typed_const(getattr(ty.ty, attr_name))
+        return sym2var(getattr(ty.ty, attr_name), constant_only=True)
     except AttributeError:
         raise TileTypeError(f"'{ty.ty.__name__}' object has no attribute '{attr_name}'")
 
@@ -2023,7 +2032,7 @@ async def getattr_dataclass_impl(object: Var, name: Var):
         getter = loosely_typed_const(cls_attr.fget)
         return await call(getter, (object,), {})
     else:
-        return loosely_typed_const(cls_attr)
+        return sym2var(cls_attr, constant_only=True)
 
 
 # ===========================================================================================
@@ -2058,7 +2067,11 @@ def assign(value: Var, res: Var) -> None:
 @impl(hir_stubs.identity)
 def identity_impl(x: Var) -> Var:
     if x.is_constant():
-        return loosely_typed_const(x.get_constant(), x.get_type(), x.get_loose_type())
+        ty = x.get_type()
+        if ty.is_aggregate():
+            return make_aggregate(x.get_aggregate(), ty, x.get_loose_type())
+        else:
+            return loosely_typed_const(x.get_constant(), x.get_type(), x.get_loose_type())
     else:
         return x
 
@@ -5268,8 +5281,7 @@ def load_var_impl(name):
         return ret
     elif rn.index >= 0:
         val = scope.func_hir.frozen_global_values[rn.index]
-        val = get_constant_value(val)
-        return loosely_typed_const(val)
+        return sym2var(val, constant_only=True)
     else:
         raise TileSyntaxError(f"Undefined variable {name} used")
 
@@ -5450,30 +5462,38 @@ async def static_foreach_impl(body: hir.Block, items: Var):
         await dispatch_hir_block(body)
 
 
-def sym2var(x: Any) -> Var:
+def sym2var(x: Any, constant_only: bool = False) -> Var:
     # TODO: verify we don't have a stale closure
 
     if isinstance(x, Symbol):
+        if constant_only:
+            raise TileTypeError("Cannot create a constant from a symbolic value")
         return x._var
 
     if isinstance(x, tuple):
-        return build_tuple(tuple(sym2var(item) for item in x))
+        return build_tuple(tuple(sym2var(item, constant_only=constant_only) for item in x))
 
     cls = type(x)
     if dataclasses.is_dataclass(cls):
         info = get_dataclass_info(cls)
-        field_vars = tuple(sym2var(getattr(x, f.name))
+        field_vars = tuple(sym2var(getattr(x, f.name), constant_only=constant_only)
                            for f in dataclasses.fields(cls))
         return build_dataclass_instance(field_vars, info)
 
     if isinstance(x, MethodType):
-        self_var = sym2var(x.__self__)
+        self_var = sym2var(x.__self__, constant_only=constant_only)
         if not isinstance(x.__func__, FunctionType | BuiltinFunctionType):
             raise TileTypeError(f"Object of type {type(x).__name__}"
                                 f" cannot be used as a function for binding a method")
         return bind_method(self_var, x.__func__)
 
-    x = get_constant_value(x)
+    # Transform a third party typed scalar (e.g., np.int16(5)) into a strictly typed constant
+    dtype_spec = as_third_party_dtype_spec(type(x))
+    if dtype_spec is not None:
+        pyval = datatype.numeric_dtype_category(dtype_spec.dtype).pytype(x)
+        ty = Builder.get_current().ir_ctx.typing_hooks.get_tensor_like_type(dtype_spec.dtype, ())
+        return strictly_typed_const(pyval, ty)
+
     return loosely_typed_const(x)
 
 

src/cuda/tile/_ir/typing_support.py

@@ -4,16 +4,17 @@
 import inspect
 import operator
 import dataclasses
-from enum import Enum, IntEnum
+from enum import Enum
 from functools import lru_cache
 from types import ModuleType, FunctionType
 from typing import Any, Callable, Mapping, Union
 
 from cuda.tile import _datatype as datatype, DType
 from cuda.tile._exception import TileTypeError, TileValueError
+from .ir import TypingHooks
 from .type import DataclassInfo, PointerInfoTy
 
-from .type import Type, TupleTy, DTypeConstructor, DTypeSpec, NONE, StringTy, \
+from .type import Type, DTypeConstructor, DTypeSpec, NONE, StringTy, \
     ELLIPSIS, SLICE, ModuleTy, FunctionTy, EnumTy, TypeTy, LooselyTypedScalar, \
     TileTy
 
@@ -52,6 +53,10 @@ def is_dtype(x: Any):
     return isinstance(x, DType) or _safe_get(_dtype_registry, x) is not None
 
 
+def as_third_party_dtype_spec(x: Any) -> DTypeSpec | None:
+    return _safe_get(_dtype_registry, x)
+
+
 def _is_dtype_allowed_as_constructor(dtype: DType) -> bool:
     # Only allow byte aligned numeric dtypes as constructors
     return datatype.is_numeric(dtype) and (dtype.bitwidth % 8 == 0)
@@ -158,19 +163,15 @@ def dtype_of_constant_scalar(val: bool | int | float) -> DType:
         raise TypeError(f'Python value {val} of type {type(val)} is not supported.')
 
 
-def typeof_pyval(val) -> Type:
+def type_of_constant_python_value(val, typing_hooks: TypingHooks) -> Type:
     if val is None:
         return NONE
-    if (t := _safe_get(_dtype_registry, type(val))):
-        return TileTy(t.dtype)
     if isinstance(val, bool | int | float):
-        return TileTy(dtype_of_constant_scalar(val))
+        return typing_hooks.get_tensor_like_type(dtype_of_constant_scalar(val), ())
     if isinstance(val, Enum):
         return EnumTy(type(val))
     if isinstance(val, str):
         return StringTy(val)
-    if isinstance(val, tuple):
-        return TupleTy(tuple(typeof_pyval(v) for v in val))
     if val is Ellipsis:
         return ELLIPSIS
     if isinstance(val, slice):
@@ -186,40 +187,23 @@ def typeof_pyval(val) -> Type:
             return DTypeConstructor(val)
         else:
             return DTypeSpec(val)
-    if (t := _safe_get(_dtype_registry, val)) is not None:
+    if (t := as_third_party_dtype_spec(val)) is not None:
         return t
     if isinstance(val, datatype.PointerInfo):
         return PointerInfoTy(val)
-
     if isinstance(val, type):
         return TypeTy(val)
 
-    # TODO: should we add dlpack?
-    raise TypeError(f'Python value {val} of type {type(val)} is not supported.')
+    ty = type(val)
+    prefix = "" if ty.__module__ == "builtins" else f"{ty.__module__}."
+    raise TileTypeError(f"Cannot create constant from value of type {prefix}{ty.__qualname__}.")
 
 
-def loose_type_of_pyval(value: Any) -> Type:
+def loose_type_of_constant_python_value(value: Any, typing_hooks: TypingHooks) -> Type:
     if isinstance(value, bool | int | float):
         return LooselyTypedScalar(value)
-    elif isinstance(value, tuple):
-        return TupleTy(tuple(loose_type_of_pyval(x) for x in value))
     else:
-        return typeof_pyval(value)
-
-
-_SUPPORTED_CONST_TYPES = (int, float, bool, str, ModuleType, FunctionType, type, Enum, IntEnum)
-
-
-def get_constant_value(val: Any) -> Any:
-    if val is None or isinstance(val, _SUPPORTED_CONST_TYPES) or is_supported_builtin_func(val):
-        return val
-    if is_dtype(val):
-        return to_dtype(val)
-    if isinstance(val, tuple):
-        return tuple(get_constant_value(x) for x in val)
-    typ = type(val)
-    prefix = "" if typ.__module__ == "builtins" else f"{typ.__module__}."
-    raise TileTypeError(f"Cannot create constant from value of type {prefix}{typ.__qualname__}.")
+        return type_of_constant_python_value(value, typing_hooks)
 
 
 @lru_cache

src/cuda/tile/_passes/hir2ir.py

@@ -18,7 +18,7 @@
 from .._ir.ir import Var, IRContext
 from .._ir.op_impl import ImplRegistry
 from .._ir.ops import loosely_typed_const, end_branch, return_, continue_, \
-    break_, store_var, build_dataclass_instance, build_tuple, dtype_constructor
+    break_, store_var, build_dataclass_instance, build_tuple, dtype_constructor, sym2var
 from .._ir.scope import Scope, LocalScope, IntMap
 from .._ir.type import FunctionTy, BoundMethodTy, DTypeConstructor, ClosureTy, \
     ClosureDefaultPlaceholder, StringFormat, TypeTy, TupleTy, BoundMethodValue, TupleValue, \
@@ -390,7 +390,7 @@ def _resolve_operand(x: hir.Operand, scope: Scope) \
     elif isinstance(x, hir.Block | hir.Function | hir.StaticEvalExpression | StringFormat):
         return x
     else:
-        return loosely_typed_const(x)
+        return sym2var(x, constant_only=True)
 
 
 def _bind_args(sig: inspect.Signature, func_name: str, args, kwargs,
@@ -411,6 +411,6 @@ def _bind_args(sig: inspect.Signature, func_name: str, args, kwargs,
                 assert closure_defaults is not None
                 default = closure_defaults[param.default.default_value_index]
             else:
-                default = loosely_typed_const(param.default)
+                default = sym2var(param.default, constant_only=True)
             ret.append(default)
     return ret

test/test_dataclass.py

@@ -49,6 +49,18 @@ def kern(x):
     assert x.tolist() == [2, 7, 30, 40, 5]
 
 
+def test_dataclass_global_capture():
+    fb = FooBar(2, 7)
+
+    @ct.kernel
+    def kern(x):
+        ct.scatter(x, (), fb.foo)
+
+    x = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kern, (x,))
+    assert x.item() == 2
+
+
 def test_dataclass_with_field_named_self():
     @dataclass(frozen=True)
     class Selfish:

test/test_ir_types.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 import copy
+import functools
 import pickle
 
 import pytest
@@ -23,7 +24,7 @@
     IntegerInfo, opaque_pointer_dtype, pointer_dtype, PointerInfo,
 )
 from cuda.tile._ir.ops_utils import promote_dtypes, check_implicit_cast
-from cuda.tile._ir.typing_support import to_dtype, typeof_pyval
+from cuda.tile._ir.typing_support import to_dtype
 import torch
 import numpy as np
 
@@ -338,12 +339,12 @@ def test_torch_dtype_support():
     assert to_dtype(torch.float8_e8m0fnu) == float8_e8m0fnu
 
 
-def test_typeof_pyval():
-    tp = typeof_pyval
+def test_type_of_constant_python_value():
+    from cuda.tile._ir.typing_support import type_of_constant_python_value
+    from cuda.tile._compile import _TileTypingHooks
+    tp = functools.partial(type_of_constant_python_value, typing_hooks=_TileTypingHooks())
     assert tp(1) == TileTy(int32)
     assert tp(1.) == TileTy(float32)
-    assert tp(np.int16(1)) == TileTy(int16)
-    assert tp(np.float64(1.0)) == TileTy(float64)
     assert tp(True) == TileTy(bool_)
     assert tp(None) == NONE
 

test/test_tuple.py

@@ -265,3 +265,16 @@ def kernel(x, M: ct.Constant[int], N: ct.Constant[int]):
     assert x.item() == 1
     ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, 4, 9))
     assert x.item() == -1
+
+
+def test_tuple_global_capture():
+    tup = (100, (101, 102), (103, (104, 105)))
+
+    @ct.kernel
+    def kernel(x):
+        ct.scatter(x, 0, tup[0])
+        ct.scatter(x, 1, tup[2][1][1])
+
+    x = torch.zeros(2, dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (x,))
+    assert x.tolist() == [100, 105]

test/test_typeinfer.py

@@ -470,3 +470,17 @@ def kernel():
     with pytest.raises(TileTypeError,
                        match="Tiles are immutable: item assignment is not supported"):
         ct.launch(torch.cuda.current_stream(), (1,), kernel, ())
+
+
+def test_typeof_numpy_scalar():
+    import numpy as np
+    val = np.int16(4)
+
+    @ct.kernel
+    def kernel():
+        val_dtype = val.dtype
+        ct.static_assert(val_dtype == ct.int16)
+        x = val
+        ct.static_assert(x == 4)
+
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, ())