Feat: Tensor, DType, PatternMatcher, etc

caten/__init__.py

@@ -0,0 +1,10 @@
+from . import dtype, helpers, ir, tensor
+from .tensor import ATenSpec, ATen, ATenMath, ATenMovements, ATenNN, ATenLinalg, ATenBase, get_backend, Tensor
+from .runtime import cpu
+
+__all__ = [
+    "dtype",
+    "helpers",
+    "ir",
+    "tensor"
+]

caten/dtype.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+from dataclasses import dataclass, fields
+
+class DTypeMetaClass(type):
+    dcache: dict[tuple, DType] = {}
+    def __call__(cls, *args, **kwargs):
+        if (ret:=DTypeMetaClass.dcache.get(args, None)) is not None: return ret
+        DTypeMetaClass.dcache[args] = ret = super().__call__(*args)
+        return ret
+
+# TODO: Vector/Packed DType
+@dataclass(frozen=True, eq=False)
+class DType:
+    name: str
+    @staticmethod
+    def new(name:str): return DType(name)
+
+## definitions
+float64 = DType.new("float64")
+float32 = DType.new("float32")
+int64 = DType.new("int64")
+int32 = DType.new("int32")
+
+## dtype aliases
+index = int64

caten/ir.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+from abc import ABCMeta, abstractmethod
+from typing import List, Dict, Any
+from dataclasses import dataclass
+from .dtype import DType
+
+@dataclass(frozen=True)
+class ATenAxis():
+    shape: ATenOp
+    stride: ATenOp
+    offset: ATenOp
+    incf: ATenOp
+    def index(self, i: ATenOp):
+        # TODO: Assert i.T.dtype is dtype.index
+        return Mul(self.stride, Add(Mul(i, self.incf), self.offset))
+
+@dataclass(frozen=True)
+class ATenOpType():
+    shape: List[ATenAxis]
+    dtype: DType
+    offset: ATenOp
+
+@dataclass(frozen=True)
+class ATenOp(metaclass=ABCMeta):
+    args: List[AtenOp]
+    T: ATenOpType
+    @classmethod
+    @abstractmethod
+    def from_astexpr(cls):
+        pass
+
+    @abstractmethod
+    def infer_dtype(self):
+        pass
+## == Tensor Graph ============================================================
+class UnaryOps():
+    def verify(self): verify_tensor_op(self, 1)
+class BinaryOps():
+    def verify(self): verify_tensor_op(self, 2)
+class TernaryOps():
+    def verify(self): verify_tensor_op(self, 3)
+### UnaryOps
+class Neg(ATenOp, UnaryOps):
+    """
+    OUT = -X
+    """
+    pass
+
+class Recip(ATenOp, UnaryOps):
+    pass
+
+class Sin(ATenOp, UnaryOps):
+    pass
+
+class Exp2(ATenOp, UnaryOps):
+    pass
+
+class Log2(ATenOp, UnaryOps):
+    pass
+
+class Sqrt(ATenOp, UnaryOps):
+    pass
+
+class Cast(ATenOp, UnaryOps):
+    pass
+
+class Bitcast(ATenOp, UnaryOps):
+    pass
+
+class Not(ATenOp, UnaryOps):
+    """
+    Logical not if the X is a boolean
+    otherwise lognot ~x
+    """
+    pass
+### BinaryOps
+class Add(ATenOp, BinaryOps):
+    """
+    OUT = Add(X, Y)
+    """
+    @classmethod
+    def from_ast_expr(cls):
+        pass
+
+class Mul(ATenOp, BinaryOps):
+    """
+    OUT = Mul(X, Y)
+    """
+    @classmethod
+    def from_ast_expr(cls):
+        pass
+
+class IDiv(ATenOp, BinaryOps):
+    pass
+
+class And(ATenOp, BinaryOps):
+    pass
+
+class Or(ATenOp, BinaryOps):
+    pass
+
+class And(ATenOp, BinaryOps):
+    pass
+
+class Xor(ATenOp, BinaryOps):
+    pass
+
+class Max(ATenOp, BinaryOps):
+    pass
+
+class Mod(ATenOp, BinaryOps):
+    pass
+
+class Neq(ATenOp, BinaryOps):
+    pass
+
+class Lt(ATenOp, BinaryOps):
+    pass
+### TernaryOps
+class Where(ATenOp, TernaryOps):
+    pass
+
+### Allocation
+class Variable(ATenOp):
+    symbol: str
+
+class Allocate(ATenOp):
+    """
+    Allocate(S1, S2, S3, ...)
+    """
+    pass
+
+## == JIT =====================================================================
+class Reduce(ATenOp):
+    """
+    OUT = Reduce(A, B, op=BinaryOps)
+    """
+    op: BinaryOps
+    @classmethod
+    def from_ast_expr(cls):
+        pass
+
+class Store(ATenOp):
+    pass
+## ControlFlow
+class Range(ATenOp):
+    pass
+
+class Loop(ATenOp):
+    pass
+
+class When(ATenOp):
+    pass
+
+class Progn(ATenOp):
+    pass
+## == ScheduleOps ============================================================
+class Polyhedral(ATenOp):
+    pass
+
+def Var():
+    pass
+
+# e.g.:
+# a = T.Var("A[m n]", float32)
+# P.stmt("...")[a]

caten/runtime/cpu.py

@@ -0,0 +1,13 @@
+import caten as C
+
+class CPUTensor(C.ATenBase):
+    def allocate(self):
+        pass
+
+    def free(self):
+        pass
+
+    def compile(self):
+        pass
+
+C.ATenBase.register("CPU", CPUTensor)

caten/tensor.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+from abc import ABCMeta, abstractmethod
+from typing import Any, Optional, Tuple, Union, ClassVar
+import os
+import caten.ir as ir
+# [TODO]
+# - Tensor => Fused Tensor Graph Construction
+# - Tensor Kernel Construction
+# - Then, start working on auto scheduler
+## Backend Abstraction
+DEVICE_TO_TENSOR = {}
+def get_backend(): return os.environ.get("BACKEND", "CPU")
+##
+
+class ATenSpec:
+    """
+    Usage: C.Tensor[float32, "M", "N"] -> TensorSpec(M N)
+    """
+    def __init__(self, shape: Tuple[Any, ...], dtype: Any = None):
+        self.shape = shape
+        self.dtype = dtype
+    def __repr__(self) -> str: return f"TensorSpec({self.shape}, {self.dtype})"
+
+class ATen:
+    op: ATenOp # ATen is just a wrapper for ATenOp
+    @classmethod
+    def from_shape(cls, shape: List[ATenOp]):
+        return ir.Allocate(shape) # TODO
+
+    def apply(self, op: Callable, *args: List, **kwargs) -> ATen: return ATen(op=op(*args, **kwargs))
+
+    def __class_getitem__(cls, item: Union[Any, Tuple[Any, ...]]) -> TensorSpec:
+        # Usage: C.Tensor[10, 10] -> TensorSpec((10, 10))
+        # TODO
+        pass
+
+    def polyhedral(self):
+        pass
+
+class ATenMath():
+    pass
+
+class ATenMovements():
+    pass
+
+class ATenNN():
+    pass
+
+class ATenLinalg():
+    pass
+
+class ATenBase(ATen, ATenMath, ATenNN, ATenMovements, ATenLinalg, metaclass=ABCMeta):
+    ## == AbstractionLayer
+    @staticmethod
+    def register(device_id: str, cls: ClassVar):
+        DEVICE_TO_TENSOR[device_id] = cls
+
+    @abstractmethod
+    def allocate(self):
+        pass
+
+    @abstractmethod
+    def free(self):
+        pass
+
+    @abstractmethod
+    def compile(self):
+        pass
+
+class Tensor(ATenBase):
+    def __new__(cls, *args, **kwargs):
+        impl = DEVICE_TO_TENSOR.get(get_backend())
+        if impl is None: raise ValueError(f"Unknown BACKEND={get_backend()}")
+        return impl(*args, **kwargs)
+## For-Style Graph Construction
+def kernel(get_kernel: bool = False) -> Callable:
+    def decorator(func: Callable) -> Callable:
+        pass
+    return decorator
+
+# how to generate polyhedral model from tensor ops?
+# rangeify -> range/when ==> polyhedral model
+# with C.range(10, 10):
+# with C.when(10, 10)
+class range():
+    pass
+
+class when():
+    pass

test/test_kernel.py

@@ -0,0 +1,4 @@
+import caten as C
+
+def test_tensor():
+    print(C.Tensor())