[lang] Use Scalar/Pointer/Vector types instead of TileTy in cuda.lang

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

Author: gbonik

experimental/cuda-lang/src/cuda/lang/__init__.py

@@ -5,6 +5,7 @@
 from cuda.tile._memory_model import (
     MemoryScope,
     MemoryOrder,
+    MemorySpace
 )
 
 from ._execution import (
@@ -15,11 +16,19 @@
 
 from ._compile import compile_simt
 
-from ._stub import (
+from cuda.tile._datatype import (
     pointer_dtype,
     opaque_pointer_dtype,
     is_pointer_dtype,
     PointerInfo,
+)
+from ._stub.types import (
+    Scalar,
+    Vector,
+    Pointer,
+)
+from ._stub.core_api import (
+    dtype_of,
     warp_size,
     full_mask,
     block_idx,
@@ -46,11 +55,7 @@
     setmaxregister_increase,
     setmaxregister_decrease,
     elect_sync,
-    Constant,
     Array,
-    Pointer,
-    Vector,
-    printf,
     shared_array,
     local_array,
     address_space_cast,
@@ -65,17 +70,33 @@
     syncthreads,
     syncwarp,
     nvvm,
-    libdevice,
+    nanosleep,
+    griddepcontrol_wait,
+    griddepcontrol_launch_dependents,
+    memory_barrier,
+)
+from cuda.tile._stub import (
+    Constant
+)
+
+from cuda.lang._stub import libdevice
+
+from cuda.lang._stub.tensor_map import (
     TensorMapSwizzle,
     TensorMap,
     tensor_map_tiled,
+)
+
+from cuda.lang._stub.tcgen05 import (
     CTAGroup,
     Tcgen05LdStShape,
     tcgen05_alloc,
     tcgen05_dealloc,
     tcgen05_commit,
     tcgen05_ld,
-    nanosleep,
+)
+
+from cuda.lang._stub.mbarrier import (
     MbarrierScope,
     mbarrier_init,
     mbarrier_invalidate,
@@ -87,12 +108,12 @@
     mbarrier_test_wait_parity,
     mbarrier_try_wait,
     mbarrier_try_wait_parity,
+)
+
+from cuda.lang._stub.cluster_launch_control import (
     clusterlaunchcontrol_try_cancel,
     clusterlaunchcontrol_is_canceled,
     clusterlaunchcontrol_get_first_block_idx,
-    griddepcontrol_wait,
-    griddepcontrol_launch_dependents,
-    memory_barrier,
 )
 
 from ._datatype import (
@@ -114,10 +135,10 @@
     uint64,
     mbarrier,
     clusterlaunchcontrol_token,
-    MemorySpace,
 )
 
 __all__ = (
+    "dtype_of",
     "pointer_dtype",
     "opaque_pointer_dtype",
     "is_pointer_dtype",
@@ -172,7 +193,6 @@
     "uint32",
     "uint64",
     "Constant",
-    "printf",
     "shared_array",
     "local_array",
     "address_space_cast",
@@ -184,6 +204,7 @@
     "syncwarp",
     "Array",
     "Pointer",
+    "Scalar",
     "Vector",
     "nvvm",
     "libdevice",

experimental/cuda-lang/src/cuda/lang/_compile.py

@@ -24,7 +24,6 @@
 from cuda.lang._passes.ir2mlir import ir2mlir
 from cuda.lang._passes.flatten_cfg import flatten_cfg
 from cuda.lang._passes.simt_semantics import simt_semantic_analysis
-from cuda.lang._passes.canonicalize_parameters import canonicalize_parameters
 from cuda.lang._passes.handle_dyn_shared_mem import handle_dynamic_shared_memory
 from cuda.lang._passes.hoist_tensor_map import hoist_tensor_maps, HoistedTensorMap
 from cuda.lang.compilation import (
@@ -137,7 +136,6 @@ def get_function_ir(
             function.param_locs,
             ctx
         )
-        canonicalize_parameters(params, builder)
         hir2ir(function, params.aggregate_vars, ctx)
     func_body = ctx.make_block("entry", function.body.loc)
     func_body.params = sum((vars for vars, _ in params.nonconstant_flat_vars), ())

experimental/cuda-lang/src/cuda/lang/_datatype.py

@@ -2,13 +2,8 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import TypeAlias, Union
+from typing import TypeAlias
 
-from cuda.lang._ir.type import (
-    MemorySpace,
-    TileTy,
-)
-from cuda.tile._stub import Tile
 from cuda.tile._datatype import (
     DType,
     bfloat16,
@@ -52,11 +47,11 @@
 )
 
 
-def to_torch_dtype(dtype: DType | TileTy):
-    import torch
+def to_torch_dtype(dtype: DType, /):
+    if not isinstance(dtype, DType):
+        raise TypeError("Expected a DType object")
 
-    if isinstance(dtype, TileTy):
-        dtype = dtype.dtype
+    import torch
 
     dtype_map = {
         bool_: torch.bool,
@@ -87,11 +82,7 @@ def to_torch_dtype(dtype: DType | TileTy):
     raise NotImplementedError(f"No torch dtype mapping for {dtype}")
 
 
-TypeSpec: TypeAlias = Union[DType | TileTy]
-
-
-def is_any_pointer(value):
-    return isinstance(value, Tile) and value.ndim == 0 and is_pointer_dtype(value.dtype)
+TypeSpec: TypeAlias = DType
 
 
 __all__ = [
@@ -102,7 +93,6 @@ def is_any_pointer(value):
     "is_boolean",
     "is_integral",
     "is_signed",
-    "is_any_pointer",
     "is_pointer_dtype",
     "pointer_dtype",
     "opaque_pointer_dtype",
@@ -131,5 +121,4 @@ def is_any_pointer(value):
     "DType",
     "to_torch_dtype",
     "default_int_type",
-    "MemorySpace",
 ]

experimental/cuda-lang/src/cuda/lang/_execution.py

@@ -33,7 +33,7 @@ class kernel(_cext.TileDispatcher):
 
             @cl.kernel
             def kernel():
-                cl.printf("Hello!\\n")
+                print("Hello!")
 
             cl.launch(stream, (1,), (3,), kernel, ())
 

experimental/cuda-lang/src/cuda/lang/_ir/_host_program.py

@@ -14,10 +14,11 @@
     binary_bitwise_tensorlike_raw
 )
 from cuda.tile._ir.core_ops import TypedConst, Assign, loosely_typed_const, strictly_typed_const
+from cuda.lang._ir.type import ScalarTy
 from cuda.tile._ir.ops import (
     AssumeBounded, AssumeDivBy,
 )
-from cuda.tile._ir.typing_support import I32_TY, I64_TY
+from cuda.tile._datatype import int32, int64
 
 HostOpcode = Literal["Const", "KernelArgI32", "KernelArgI64", "Mul", "Add", "RoundUpToPow2"]
 
@@ -50,9 +51,9 @@ def get_host_programs_by_var(kernel_body: ir.Block) -> dict[str, HostProgram]:
     ret = dict()
     for i, p in enumerate(kernel_body.params):
         ty = p.get_type()
-        if ty == I32_TY:
+        if ty == ScalarTy(int32):
             ret[p.name] = HostProgram(opcodes=["KernelArgI32"], op_attrs=[i])
-        elif ty == I64_TY:
+        elif ty == ScalarTy(int64):
             ret[p.name] = HostProgram(opcodes=["KernelArgI64"], op_attrs=[i])
 
     for op in kernel_body:

experimental/cuda-lang/src/cuda/lang/_ir/ir.py

@@ -6,11 +6,6 @@
 from dataclasses import dataclass, field
 import itertools
 from collections import defaultdict
-from typing import Sequence
-
-from typing_extensions import override
-
-from cuda.tile._datatype import DType
 from cuda.tile._ir.ir import (
     Block as TileBlock,
     Builder as TileBuilder,
@@ -22,9 +17,8 @@
     attribute,
     add_operation,
     format_var,
-    AggregateValue, TypingHooks,
+    AggregateValue,
 )
-from cuda.tile._ir.type import TensorLikeTy, TileTy
 
 
 class Builder:
@@ -91,18 +85,13 @@ def to_string(
         return f"{' ' * indent}^{self._name}({params}):\n{ops}"
 
 
-class _LangTypingHooks(TypingHooks):
-    @override
-    def get_tensor_like_type(self, dtype: DType, shape: Sequence[int]) -> TensorLikeTy:
-        return TileTy(dtype, shape)
-
-
 class IRContext(TileIRContext):
     def __init__(self, log_ir_on_error: bool = True):
+        from cuda.lang._ir.type import LangTypingHooks
         self._block_names: dict[int, str] = {}
         self._block_counter: dict[str, itertools.count] = defaultdict(itertools.count)
         super().__init__(log_ir_on_error, tileiras_version=None,
-                         typing_hooks=_LangTypingHooks())
+                         typing_hooks=LangTypingHooks())
 
     def make_block(self, name: str, loc: Loc, params: tuple[Var, ...] = ()) -> Block:
         block = Block(self, loc)