Align Numba codegen to the new assembly. E2E Numba test. (#161)

* Align Numba codegen to new assembly. E2E Numba test.

* Delete `override` import

* Remove `shape` from `BufferizedNDArrayFType`

* Remove colon from labels

* Remove dangling shapes

* Cleanup

* Apply review comments

* Custom serializer for BufferizedNDArray

* Keep only `numba_type` for `BufferizedNDArray`

* Remove `construct_from_numba` from `BufferizedNDArray`

* Rename function to `provision_tensors`

* Move `numba_type` to `AssemblyStructFType`

* Moved last serialization part to `AssemblyStructFType`

* Use kind and itemsize instead of char for dtype str

* move the serialization stuff around a bit (just to keep the dependency order consistent)

* pre-commit

---------

Co-authored-by: Willow Ahrens <[email protected]>

Author: mtsokol

src/finchlite/__init__.py

@@ -9,7 +9,6 @@
     ExtentFType,
     NotationContext,
     dimension,
-    extent,
 )
 from .galley import (
     DC,
@@ -20,6 +19,7 @@
 from .interface import (
     EagerTensor,
     LazyTensor,
+    Mode,
     Scalar,
     abs,
     acos,
@@ -54,6 +54,7 @@
     floordiv,
     fuse,
     fused,
+    get_default_scheduler,
     greater,
     greater_equal,
     less,
@@ -82,6 +83,7 @@
     pow,
     prod,
     reduce,
+    set_default_scheduler,
     sin,
     sinh,
     split_dims,
@@ -121,6 +123,7 @@
     "FTyped",
     "FiberTensorFType",
     "LazyTensor",
+    "Mode",
     "NotationContext",
     "NumpyBuffer",
     "NumpyBufferFType",
@@ -160,14 +163,14 @@
     "elementwise",
     "equal",
     "expand_dims",
-    "extent",
     "fill_value",
     "fisinstance",
     "flatten",
     "floordiv",
     "ftype",
     "fuse",
     "fused",
+    "get_default_scheduler",
     "greater",
     "greater_equal",
     "less",
@@ -196,6 +199,7 @@
     "pow",
     "prod",
     "reduce",
+    "set_default_scheduler",
     "shape_type",
     "sin",
     "sinh",

src/finchlite/algebra/__init__.py

@@ -24,6 +24,7 @@
 from .tensor import (
     Tensor,
     TensorFType,
+    TensorPlaceholder,
     element_type,
     fill_value,
     shape_type,
@@ -34,6 +35,7 @@
     "StableNumber",
     "Tensor",
     "TensorFType",
+    "TensorPlaceholder",
     "conjugate",
     "conjugate",
     "element_type",

src/finchlite/algebra/tensor.py

@@ -9,9 +9,9 @@
 
 class TensorFType(FType, ABC):
     @property
-    def ndim(self) -> int:
+    def ndim(self) -> np.intp:
         """Number of dimensions of the tensor."""
-        return len(self.shape_type)
+        return np.intp(len(self.shape_type))
 
     @property
     @abstractmethod
@@ -21,13 +21,13 @@ def fill_value(self) -> Any:
 
     @property
     @abstractmethod
-    def element_type(self):
+    def element_type(self) -> Any:
         """Data type of the tensor elements."""
         ...
 
     @property
     @abstractmethod
-    def shape_type(self) -> tuple:
+    def shape_type(self) -> tuple[type, ...]:
         """Shape type of the tensor. The shape type is a tuple of the index
         types in the tensor. It's the type of each element in tns.shape. It
         should be an actual tuple, rather than a tuple type, so that it can hold
@@ -46,16 +46,10 @@ class Tensor(FTyped, ABC):
     """
 
     @property
-    def ndim(self) -> int:
+    def ndim(self) -> np.intp:
         """Number of dimensions of the tensor."""
         return self.ftype.ndim
 
-    @property
-    @abstractmethod
-    def shape(self):
-        """Shape of the tensor as a tuple."""
-        ...
-
     @property
     @abstractmethod
     def ftype(self) -> TensorFType:
@@ -143,7 +137,7 @@ class NDArrayFType(TensorFType):
     This includes the fill value, element type, and shape type.
     """
 
-    def __init__(self, dtype: np.dtype, ndim: int):
+    def __init__(self, dtype: np.dtype, ndim: np.intp):
         self._dtype = dtype
         self._ndim = ndim
 
@@ -159,7 +153,7 @@ def __repr__(self) -> str:
         return f"NDArrayFType(dtype={repr(self._dtype)}, ndim={self._ndim})"
 
     @property
-    def ndim(self) -> int:
+    def ndim(self) -> np.intp:
         return self._ndim
 
     @property
@@ -176,5 +170,14 @@ def shape_type(self) -> tuple:
 
 
 register_property(
-    np.ndarray, "ftype", "__attr__", lambda x: NDArrayFType(x.dtype, x.ndim)
+    np.ndarray, "ftype", "__attr__", lambda x: NDArrayFType(x.dtype, np.intp(x.ndim))
 )
+
+
+class TensorPlaceholder:
+    def __init__(self, dtype):
+        self._dtype = dtype
+
+    @property
+    def dtype(self):
+        return self._dtype

src/finchlite/autoschedule/_utils.py

@@ -1,19 +1,23 @@
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 
 
 def intersect(x1: Iterable, x2: Iterable) -> tuple:
     return tuple(x for x in x1 if x in x2)
 
 
-def is_subsequence(x1: tuple, x2: tuple) -> bool:
-    return x1 == tuple(x for x in x2 if x in x1)
+def extend_uniqe(x1: Iterable, x2: Iterable) -> tuple:
+    return tuple(x1) + setdiff(x2, x1)
 
 
-def setdiff(x1: tuple, x2: tuple) -> tuple:
+def is_subsequence(x1: Iterable, x2: Iterable) -> bool:
+    return tuple(x1) == tuple(x for x in x2 if x in x1)
+
+
+def setdiff(x1: Iterable, x2: Iterable) -> tuple:
     return tuple([x for x in x1 if x not in x2])
 
 
-def with_subsequence(x1: tuple, x2: tuple) -> tuple:
+def with_subsequence(x1: Sequence, x2: Iterable) -> tuple:
     res = list(x2)
     indices = [idx for idx, val in enumerate(x2) if val in x1]
     for idx, i in enumerate(indices):