Move MemorySpace to tile

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

Author: gbonik

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

@@ -19,7 +19,6 @@
 from cuda.tile._exception import TileCompilerExecutionError
 from cuda.lang._logging import get_log_flags
 from cuda.lang._ir import ir, hir
-from cuda.lang._ir.type import MemorySpace
 from cuda.lang._passes.ast2hir import get_function_hir
 from cuda.lang._passes.ir2mlir import ir2mlir
 from cuda.lang._passes.flatten_cfg import flatten_cfg
@@ -103,8 +102,7 @@ def get_function_ir(
             constant_mask,
             parameter_names,
             function.param_locs,
-            ctx,
-            array_memory_space=MemorySpace.GENERIC
+            ctx
         )
         canonicalize_parameters(params, builder)
         with cuda_lang_impl_registry.as_current():

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

@@ -158,8 +158,6 @@ def satisfies_pointer_constraint(value, constraint: OpaquePointerSpec):
     if isinstance(pointer_ty, TilePointerTy):
         if constraint == opaque_ptr:
             return True
-        if pointer_ty.memory_space is None:
-            return False
         return pointer_ty.memory_space.value == constraint.value
 
     return pointer_ty.memory_space.value == constraint.value

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

@@ -2,7 +2,6 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass
-from enum import Enum
 
 from cuda.lang._ir.ir import LocalArrayContextManagerValue
 from cuda.lang._enums import TensorMapSwizzle
@@ -22,24 +21,14 @@
     EnumTy,
     make_tile_ty,
     ContextManagerTy,
-    ContextManagerState,
+    ContextManagerState, MemorySpace,
 )
 import cuda.tile._datatype as datatype
 from cuda.tile._datatype import DType
 from cuda.tile._ir.ir import Var, AggregateValue
 from cuda.lang._exception import TileTypeError
 
 
-class MemorySpace(Enum):
-    GENERIC = 0
-    GLOBAL = 1
-    SHARED = 3
-    CONSTANT = 4
-    LOCAL = 5
-    TENSOR = 6
-    SHARED_CLUSTER = 7
-
-
 def _is_power_of_2(value: int) -> bool:
     assert isinstance(value, int)
     return value > 0 and value & (value - 1) == 0

src/cuda/tile/_compile.py

@@ -124,8 +124,7 @@ def _create_kernel_parameters(parameter_constraints: Sequence[ParameterConstrain
                               constant_parameter_mask: Sequence[bool],
                               parameter_names: Sequence[str],
                               parameter_locations: Sequence[Loc],
-                              ir_ctx: ir.IRContext,
-                              array_memory_space=None) -> _KernelParameters:
+                              ir_ctx: ir.IRContext) -> _KernelParameters:
     aggregate_vars = []
     nonconstant_flat_vars = []
     for pos, (constraint, is_const, name, loc) in enumerate(
@@ -140,10 +139,10 @@ def _create_kernel_parameters(parameter_constraints: Sequence[ParameterConstrain
             if isinstance(constraint, ScalarConstraint):
                 ty = TileTy(constraint.dtype, ())
             elif isinstance(constraint, ArrayConstraint):
-                ty = _get_array_ty(constraint, array_memory_space)
+                ty = _get_array_ty(constraint)
             elif isinstance(constraint, ListConstraint):
                 assert isinstance(constraint.element, ArrayConstraint)
-                array_ty = _get_array_ty(constraint.element, array_memory_space)
+                array_ty = _get_array_ty(constraint.element)
                 ty = ListTy(array_ty)
             else:
                 raise TypeError(f"Unexpected parameter descriptor type"
@@ -157,7 +156,7 @@ def _create_kernel_parameters(parameter_constraints: Sequence[ParameterConstrain
     return _KernelParameters(aggregate_vars, nonconstant_flat_vars)
 
 
-def _get_array_ty(param: ArrayConstraint, memory_space):
+def _get_array_ty(param: ArrayConstraint):
     for static_stride, bound in zip(param.stride_constant, param.stride_lower_bound_incl,
                                     strict=True):
         if static_stride is not None:
@@ -169,8 +168,7 @@ def _get_array_ty(param: ArrayConstraint, memory_space):
     return ArrayTy(make_tile_ty(param.dtype, ()),
                    shape=(None,) * param.ndim,
                    strides=param.stride_constant,
-                   index_dtype=param.index_dtype,
-                   memory_space=memory_space)
+                   index_dtype=param.index_dtype)
 
 
 def _log_mlir(bytecode_buf):

src/cuda/tile/_ir/type.py

@@ -290,10 +290,24 @@ def size_to_bytecode(s: Optional[int]) -> int:
 # ============== Pointer Type ===============
 
 
+class MemorySpace(enum.Enum):
+    GENERIC = 0
+    GLOBAL = 1
+    SHARED = 3
+    CONSTANT = 4
+    LOCAL = 5
+    TENSOR = 6
+    SHARED_CLUSTER = 7
+
+
 @dataclass(frozen=True)
 class PointerTy(Type):
     pointee_type: Type
-    memory_space: Any = None
+    memory_space: MemorySpace = MemorySpace.GENERIC
+
+    def __str__(self):
+        memspc = "" if self.memory_space == MemorySpace.GENERIC else f", {self.memory_space}"
+        return f"Pointer[{self.pointee_type}{memspc}]"
 
 
 # ============== Tile Type ===============
@@ -344,7 +358,7 @@ def __init__(self,
                  shape: Tuple[Optional[int], ...],
                  strides: Tuple[Optional[int], ...],
                  index_dtype=None,
-                 memory_space=None):
+                 memory_space: MemorySpace = MemorySpace.GENERIC):
         from .._datatype import int32
         assert isinstance(element_type, Type)
         self.element_type = element_type
@@ -403,7 +417,8 @@ def __str__(self):
         strides_str = ('?' if x is None else str(x) for x in self.strides)
         strides_str = "(" + ','.join(strides_str) + ")"
         indexty_str = "" if self.index_dtype == int32 else f",index_dtype={self.index_dtype}]"
-        return f"Array[{type_str},{shape_str}:{strides_str}{indexty_str}]"
+        memspc = "" if self.memory_space == MemorySpace.GENERIC else f", {self.memory_space}"
+        return f"Array[{type_str},{shape_str}:{strides_str}{indexty_str}{memspc}]"
 
 
 # ============== PartitionView Type ===============