Use TensorLikeTy instead of TileTy in binary elementwise ops

This should be a non-functional change. This is the first step toward
separating the type systems of cuda.tile and cuda.lang, while still
sharing some common operation implementations.

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

Author: gbonik

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

@@ -7,8 +7,11 @@
 from typing import Literal
 
 from cuda.lang._ir import ir
-from cuda.tile._ir.ops import loosely_typed_const, raw_binary_arithmetic, strictly_typed_const, \
-    unary, _UNARY_BOOL_INT, raw_binary_bitwise, TypedConst, AssumeBounded, AssumeDivBy, Assign
+from cuda.tile._ir.ops import (
+    loosely_typed_const, binary_arithmetic_tensorlike_raw,
+    strictly_typed_const, unary, _UNARY_BOOL_INT, TypedConst, AssumeBounded, AssumeDivBy, Assign,
+    binary_bitwise_tensorlike_raw
+)
 from cuda.tile._ir.typing_support import I32_TY, I64_TY
 
 HostOpcode = Literal["Const", "KernelArgI32", "KernelArgI64", "Mul", "Add", "RoundUpToPow2"]
@@ -70,10 +73,10 @@ def host_program_to_ir(program: HostProgram, kernel_params: tuple[ir.Var, ...])
             case "KernelArgI64": stack.append(kernel_params[next(attrs)])
             case "Mul":
                 b = stack.pop()
-                stack[-1] = raw_binary_arithmetic("mul", stack[-1], b)
+                stack[-1] = binary_arithmetic_tensorlike_raw("mul", stack[-1], b)
             case "Add":
                 b = stack.pop()
-                stack[-1] = raw_binary_arithmetic("add", stack[-1], b)
+                stack[-1] = binary_arithmetic_tensorlike_raw("add", stack[-1], b)
             case "RoundUpToPow2":
                 stack[-1] = _round_up_ir(stack[-1], next(attrs))
             case _:
@@ -86,7 +89,7 @@ def host_program_to_ir(program: HostProgram, kernel_params: tuple[ir.Var, ...])
 def _round_up_ir(value: ir.Var, alignment: int) -> ir.Var:
     value_ty = value.get_type()
     mask = strictly_typed_const(alignment - 1, value_ty)
-    value_plus_mask = raw_binary_arithmetic("add", value, mask)
+    value_plus_mask = binary_arithmetic_tensorlike_raw("add", value, mask)
     neg_mask = unary('neg', _UNARY_BOOL_INT, mask)
-    rounded = raw_binary_bitwise('and_', value_plus_mask, neg_mask)
+    rounded = binary_bitwise_tensorlike_raw('and_', value_plus_mask, neg_mask)
     return rounded

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

@@ -6,7 +6,11 @@
 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,
@@ -18,8 +22,9 @@
     attribute,
     add_operation,
     format_var,
-    AggregateValue,
+    AggregateValue, TypingHooks,
 )
+from cuda.tile._ir.type import TensorLikeTy, TileTy
 
 
 class Builder:
@@ -86,11 +91,18 @@ 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):
         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)
+        super().__init__(log_ir_on_error, tileiras_version=None,
+                         typing_hooks=_LangTypingHooks())
 
     def make_block(self, name: str, loc: Loc, params: tuple[Var, ...] = ()) -> Block:
         block = Block(self, loc)

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

@@ -24,19 +24,16 @@
     require_tile_type, require_constant_bool, require_constant_pointer_info,
     require_scalar_pointer_type,
 )
-from cuda.lang._ir.type import (
-    LocalArrayContextManagerTy, ContextManagerState, TensorMapTy,
-    dtype_to_tensor_map_type, ArrayValue, PointerInfoTy
-)
+from cuda.tile._ir.type import TensorLikeTy
 from cuda.tile._ir.ops import (
-    binary_arithmetic,
+    binary_arithmetic_tensorlike,
+    binary_arithmetic_tensorlike_raw,
     loosely_typed_const,
     tile_impl_registry,
     bind_method,
     build_tuple,
     strictly_typed_const,
     astype,
-    raw_binary_arithmetic,
     Return,
     return_,
     Assign,
@@ -76,6 +73,8 @@
 from .. import _stub as stub
 
 from .type import (
+    LocalArrayContextManagerTy, ContextManagerState, TensorMapTy,
+    dtype_to_tensor_map_type, ArrayValue, PointerInfoTy,
     MemorySpace,
     Type,
     make_vector_ty,
@@ -285,8 +284,8 @@ def _array_linear_offset(array: Var, indices: tuple[Var, ...]) -> Var:
     for index, stride in zip(indices, array_val.strides, strict=True):
         index = astype(index, datatype.uint64)
         stride = astype(stride, datatype.uint64)
-        scaled = raw_binary_arithmetic("mul", index, stride)
-        offset = raw_binary_arithmetic("add", offset, scaled)
+        scaled = binary_arithmetic_tensorlike_raw("mul", index, stride)
+        offset = binary_arithmetic_tensorlike_raw("add", offset, scaled)
     return offset
 
 
@@ -538,7 +537,7 @@ def _is_pointer_type(ty):
     return isinstance(ty, TileTy) and is_pointer_dtype(ty.dtype)
 
 
-@impl(operator.add, overload=(TileTy, TileTy))
+@impl(operator.add, overload=(TensorLikeTy, TensorLikeTy))
 async def add_impl(x: Var, y: Var) -> Var:
     xty, yty = x.get_type(), y.get_type()
     if _is_pointer_type(yty):
@@ -554,7 +553,7 @@ async def add_impl(x: Var, y: Var) -> Var:
         return await call_function(operator.add, x, y)
 
 
-@impl(operator.sub, overload=(TileTy, TileTy))
+@impl(operator.sub, overload=(TensorLikeTy, TensorLikeTy))
 async def sub_impl(x: Var, y: Var) -> Var:
     xty, yty = x.get_type(), y.get_type()
     if _is_pointer_type(xty):
@@ -563,7 +562,7 @@ async def sub_impl(x: Var, y: Var) -> Var:
             raise TileTypeError(f"Expected integer pointer offset, got {offset_dtype}")
         y = astype(y, datatype.int64)
         c0 = loosely_typed_const(0)
-        offset = binary_arithmetic('sub', c0, y)
+        offset = binary_arithmetic_tensorlike('sub', c0, y)
         return _pointer_with_offset(x, offset)
     if _is_pointer_type(yty):
         raise TileTypeError('It is invalid to subtract a pointer from an integer')
@@ -868,7 +867,7 @@ def shared_array_impl(shape: Var, dtype: Var, dynamic: Var, alignment: Var) -> O
 
         if size is None or total_size is None:
             total_size = None
-            total_size_var = raw_binary_arithmetic("mul", total_size_var, size_var)
+            total_size_var = binary_arithmetic_tensorlike_raw("mul", total_size_var, size_var)
         else:
             total_size *= size
             total_size_var = strictly_typed_const(total_size, size_ty)

src/cuda/tile/_compile.py

@@ -27,12 +27,14 @@
 from cuda.tile._bytecode.version import BytecodeVersion
 from cuda.tile._cext import get_compute_capability, TileContext, default_tile_context
 from cuda.tile._compiler_options import CompilerOptions
+from cuda.tile._datatype import DType
 from cuda.tile._exception import (
     TileCompilerError,
     TileCompilerExecutionError,
     TileCompilerTimeoutError, FunctionDesc, Loc
 )
 from cuda.tile._ir import ir, hir
+from cuda.tile._ir.ir import TypingHooks
 from cuda.tile._ir.ops import loosely_typed_const, flatten_block_parameters, tile_impl_registry
 from cuda.tile._ir.type import TileTy, ArrayTy, ListTy
 from cuda.tile._passes.ast2hir import get_function_hir
@@ -234,6 +236,11 @@ def unique_path_from_func_desc(base_dir: str, desc: FunctionDesc, suffix: str, m
         yield f
 
 
+class _TileTypingHooks(TypingHooks):
+    def get_tensor_like_type(self, dtype: DType, shape: Sequence[int]) -> TileTy:
+        return TileTy(dtype, shape)
+
+
 class _IrKeeper:
     def __init__(self,
                  ann_func: AnnotatedFunction,
@@ -260,7 +267,8 @@ def get_final_ir(self, signature_index: int) -> ir.Block:
             sig = self.signatures[signature_index]
             param_names = tuple(self.ann_func.pysig.parameters.keys())
             ir_ctx = ir.IRContext(log_ir_on_error=self._log_cutile_ir,
-                                  tileiras_version=self.bytecode_version)
+                                  tileiras_version=self.bytecode_version,
+                                  typing_hooks=_TileTypingHooks())
             with ir.Builder(ir_ctx, self._func_hir.body.loc) as ir_builder:
                 params = _create_kernel_parameters(sig.parameters,
                                                    self.ann_func.constant_parameter_mask,

src/cuda/tile/_ir/ir.py

@@ -18,19 +18,26 @@
 )
 
 from cuda.tile._ir.aggregate_value import AggregateValue
-from cuda.tile._ir.type import Type, InvalidType, LooselyTypedScalar
+from cuda.tile._ir.type import Type, InvalidType, LooselyTypedScalar, TensorLikeTy
 from cuda.tile._exception import (
     TileTypeError, Loc, TileInternalError
 )
+from cuda.tile._datatype import DType
 from cuda.tile._bytecode.version import BytecodeVersion
 
 
 if TYPE_CHECKING:
     from cuda.tile._ir2bytecode import BytecodeContext
 
 
+class TypingHooks:
+    def get_tensor_like_type(self, dtype: DType, shape: Sequence[int]) -> TensorLikeTy:
+        raise NotImplementedError()
+
+
 class IRContext:
-    def __init__(self, log_ir_on_error: bool, tileiras_version: BytecodeVersion):
+    def __init__(self, log_ir_on_error: bool, tileiras_version: BytecodeVersion,
+                 typing_hooks: TypingHooks):
         self._all_vars: Dict[str, str] = {}
         self._counter_by_name: Dict[str, Iterator[int]] = defaultdict(itertools.count)
         self._temp_counter = itertools.count()
@@ -41,6 +48,7 @@ def __init__(self, log_ir_on_error: bool, tileiras_version: BytecodeVersion):
         self._aggregate_values: Dict[str, Any] = dict()
         self.tileiras_version: BytecodeVersion = tileiras_version
         self._function_specialization_id_counter = itertools.count()
+        self.typing_hooks = typing_hooks
 
     def next_function_specialization_id(self) -> str:
         # Monotonic counter used as a unique id when creating concrete FunctionDescs

src/cuda/tile/_ir/op_impl.py

@@ -452,11 +452,15 @@ def require_constant_axis_order(var: Var, rank: int) -> Tuple[int, ...]:
     return tuple(normalize_axis(x, rank, var) for x in value)
 
 
-def require_tile_type(var: Var) -> TileTy:
+def ensure_tile(var: Var) -> Var[TileTy]:
     ty = var.get_type()
     if not isinstance(ty, TileTy):
         raise _make_type_error(f"Expected a tile, but given value has type {ty}", var)
-    return ty
+    return var
+
+
+def require_tile_type(var: Var) -> TileTy:
+    return ensure_tile(var).get_type()
 
 
 def require_tile_or_tile_tuple_type(var: Var) -> TileTy | TupleTy: