WIP Dense level

Author: mtsokol

src/finchlite/__init__.py

@@ -108,6 +108,9 @@
     DenseLevelFType,
     ElementLevelFType,
     FiberTensorFType,
+    dense,
+    element,
+    fiber_tensor,
 )
 
 __all__ = [
@@ -158,11 +161,14 @@
     "cos",
     "cosh",
     "defer",
+    "dense",
     "dimension",
+    "element",
     "element_type",
     "elementwise",
     "equal",
     "expand_dims",
+    "fiber_tensor",
     "fill_value",
     "fisinstance",
     "flatten",

src/finchlite/codegen/numpy_buffer.py

@@ -76,6 +76,10 @@ def __str__(self):
         arr_str = str(self.arr).replace("\n", "")
         return f"np_buf({arr_str})"
 
+    def __repr__(self):
+        arr_repr = repr(self.arr).replace("\n", "")
+        return f"NumpyBuffer({arr_repr})"
+
 
 class NumpyBufferFType(CBufferFType, NumbaBufferFType, CStackFType):
     """

src/finchlite/interface/lazy.py

@@ -332,24 +332,24 @@ def __ne__(self, other):
 register_property(LazyTensor, "asarray", "__attr__", lambda x: x)
 
 
-def asarray(arg: Any, format="bufferized") -> Any:
+def asarray(arg: Any, format=None) -> Any:
     """
     Convert given argument and return wrapper type instance.
     If input argument is already array type, return unchanged.
 
     Args:
         arg: The object to be converted.
+        format: The format for the result array.
 
     Returns:
         The array type result of the given object.
     """
-    if format != "bufferized":
-        raise Exception(f"Only bufferized format is now supported, got: {format}")
+    if format is None:
+        if hasattr(arg, "asarray"):
+            return arg.asarray()
+        return query_property(arg, "asarray", "__attr__")
 
-    if hasattr(arg, "asarray"):
-        return arg.asarray()
-
-    return query_property(arg, "asarray", "__attr__")
+    return format(arg)
 
 
 def defer(arr) -> LazyTensor:

src/finchlite/tensor/__init__.py

@@ -1,4 +1,4 @@
-from .fiber_tensor import FiberTensor, FiberTensorFType, Level, LevelFType, tensor
+from .fiber_tensor import FiberTensor, FiberTensorFType, Level, LevelFType, fiber_tensor
 from .level import (
     DenseLevel,
     DenseLevelFType,
@@ -19,5 +19,5 @@
     "LevelFType",
     "dense",
     "element",
-    "tensor",
+    "fiber_tensor",
 ]

src/finchlite/tensor/fiber_tensor.py

@@ -1,9 +1,15 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Generic, TypeVar
+from typing import Generic, NamedTuple, TypeVar
 
-from finchlite.algebra import Tensor, TensorFType
-from finchlite.symbolic import FType, FTyped
+import numpy as np
+
+from .. import finch_assembly as asm
+from .. import finch_notation as ntn
+from ..algebra import Tensor, register_property
+from ..codegen.numpy_buffer import NumpyBuffer
+from ..compile.lower import FinchTensorFType
+from ..symbolic import FType, FTyped
 
 
 class LevelFType(FType, ABC):
@@ -161,8 +167,14 @@ def buffer_factory(self):
         return self.lvl.buffer_factory
 
 
+class FiberTensorFields(NamedTuple):
+    stride: tuple[asm.Variable, ...]
+    buf: asm.Variable
+    buf_s: asm.Slot
+
+
 @dataclass(unsafe_hash=True)
-class FiberTensorFType(TensorFType):
+class FiberTensorFType(FinchTensorFType):
     """
     An abstract base class representing the ftype of a fiber tensor.
 
@@ -177,11 +189,14 @@ def __post_init__(self):
         if self._position_type is None:
             self._position_type = self.lvl.position_type
 
-    def __call__(self, shape):
+    def __call__(self, *, shape=None, val=None):
         """
         Creates an instance of a FiberTensor with the given arguments.
         """
-        return FiberTensor(self.lvl(shape), self.lvl.position_type(1))
+        if shape is None:
+            shape = val.shape
+            val = NumpyBuffer(val.reshape[-1])
+        return FiberTensor(self.lvl(shape, val), self.lvl.position_type(1))
 
     @property
     def shape(self):
@@ -209,24 +224,82 @@ def position_type(self):
 
     @property
     def buffer_factory(self):
-        """
-        Returns the ftype of the buffer used for the fibers.
-        This is typically a NumpyBufferFType or similar.
-        """
         return self.lvl.buffer_factory
 
-
-def tensor(lvl: LevelFType, position_type: type | None = None):
+    def unfurl(self, ctx, tns, ext, mode, proto):
+        op = None
+        if isinstance(mode, ntn.Update):
+            op = mode.op
+        tns = ctx.resolve(tns).obj
+        acc_t = self.lvl.unfurl(
+            ctx,
+        )
+
+        # acc_t = BufferizedNDArrayAccessorFType(self, 0, self.buf_t.length_type, op)
+        # obj = BufferizedNDArrayAccessorFields(
+        #     tns, 0, asm.Literal(self.buf_t.length_type(0)), op
+        # )
+        return acc_t.unfurl(ctx, ntn.Stack(obj, acc_t), ext, mode, proto)
+
+    def lower_freeze(self, ctx, tns, op):
+        raise NotImplementedError
+
+    def lower_thaw(self, ctx, tns, op):
+        raise NotImplementedError
+
+    def lower_unwrap(self, ctx, obj):
+        raise NotImplementedError
+
+    def lower_increment(self, ctx, obj, val):
+        raise NotImplementedError
+
+    def lower_declare(self, ctx, tns, init, op, shape):
+        i_var = asm.Variable("i", self.buffer_factory.length_type)
+        body = asm.Store(
+            tns.obj.buf_s,
+            i_var,
+            asm.Literal(init.val),
+        )
+        ctx.exec(
+            asm.ForLoop(i_var, asm.Literal(np.intp(0)), asm.Length(tns.obj.buf_s), body)
+        )
+        return
+
+    def asm_unpack(self, ctx, var_n, val):
+        """
+        Unpack the into asm context.
+        """
+        stride = []
+        shape_type = self.shape_type
+        for i in range(self.ndim):
+            stride_i = asm.Variable(f"{var_n}_stride_{i}", shape_type[i])
+            stride.append(stride_i)
+            stride_e = asm.GetAttr(val, asm.Literal("strides"))
+            stride_i_e = asm.GetAttr(stride_e, asm.Literal(f"element_{i}"))
+            ctx.exec(asm.Assign(stride_i, stride_i_e))
+        buf = asm.Variable(f"{var_n}_buf", self.buffer_factory)
+        buf_e = asm.GetAttr(val, asm.Literal("buf"))
+        ctx.exec(asm.Assign(buf, buf_e))
+        buf_s = asm.Slot(f"{var_n}_buf_slot", self.buffer_factory)
+        ctx.exec(asm.Unpack(buf_s, buf))
+
+        return FiberTensorFields(tuple(stride), buf, buf_s)
+
+
+def fiber_tensor(lvl: LevelFType, position_type: type | None = None):
     """
     Creates a FiberTensorFType with the given level ftype and position type.
 
     Args:
         lvl: The level ftype to be used for the tensor.
-        pos_type: The type of positions within the tensor. Defaults to None.
+        position_type: The type of positions within the tensor. Defaults to None.
 
     Returns:
         An instance of FiberTensorFType.
     """
     # mypy does not understand that dataclasses generate __hash__ and __eq__
     # https://github.com/python/mypy/issues/19799
     return FiberTensorFType(lvl, position_type)  # type: ignore[abstract]
+
+
+register_property(FiberTensor, "asarray", "__attr__", lambda x: x)

src/finchlite/tensor/level/dense_level.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 
+from ...compile import looplets as lplt
 from ..fiber_tensor import Level, LevelFType
 
 
@@ -16,15 +17,15 @@ def __post_init__(self):
         if self.dimension_type is None:
             self.dimension_type = np.intp
 
-    def __call__(self, shape):
+    def __call__(self, shape, val=None):
         """
         Creates an instance of DenseLevel with the given ftype.
         Args:
             shape: The shape to be used for the level. (mandatory)
         Returns:
             An instance of DenseLevel.
         """
-        lvl = self.lvl(shape=shape[1:])
+        lvl = self.lvl(shape[1:], val)
         return DenseLevel(self, lvl, self.dimension_type(shape[0]))
 
     @property
@@ -63,6 +64,16 @@ def buffer_factory(self):
         """
         return self.lvl.buffer_factory
 
+    def unfurl(self):
+        def child_accessor(ctx, idx):
+            self.lvl
+
+        return lplt.Lookup(
+            body=lambda ctx, idx: lplt.Leaf(
+                body=lambda ctx: child_accessor(ctx, idx),
+            )
+        )
+
 
 def dense(lvl, dimension_type=None):
     return DenseLevelFType(lvl, dimension_type=dimension_type)

src/finchlite/tensor/level/element_level.py

@@ -4,6 +4,7 @@
 import numpy as np
 
 from ...codegen import NumpyBufferFType
+from ...compile import looplets as lplt
 from ...symbolic import FType, ftype
 from ..fiber_tensor import Level, LevelFType
 
@@ -26,17 +27,18 @@ def __post_init__(self):
         self._element_type = self.val_format.element_type
         self._fill_value = self._element_type(self._fill_value)
 
-    def __call__(self, shape=()):
+    def __call__(self, shape=(), val=None):
         """
         Creates an instance of ElementLevel with the given ftype.
         Args:
-            fmt: The ftype to be used for the level.
+            shape: Should be always `()`, used for validation.
+            val: The value to store in the ElementLevel instance.
         Returns:
             An instance of ElementLevel.
         """
         if len(shape) != 0:
             raise ValueError("ElementLevelFType must be called with an empty shape.")
-        return ElementLevel(self)
+        return ElementLevel(self, val)
 
     @property
     def ndim(self):
@@ -62,6 +64,16 @@ def shape_type(self):
     def buffer_factory(self):
         return self._buffer_factory
 
+    def unfurl(self):
+        def child_accessor(ctx, idx):
+            pass
+
+        return lplt.Lookup(
+            body=lambda ctx, idx: lplt.Leaf(
+                body=lambda ctx: child_accessor(ctx, idx),
+            )
+        )
+
 
 def element(
     fill_value=None,
@@ -78,6 +90,7 @@ def element(
         element_type: The type of elements stored in the level.
         position_type: The type of positions within the level.
         buffer_factory: The factory used to create buffers for the level.
+        val_format: Format of the value stored in the level.
 
     Returns:
         An instance of ElementLevelFType.

src/finchlite/tensor/level/sparse_list_level.py

@@ -4,14 +4,18 @@
     DenseLevelFType,
     ElementLevelFType,
     FiberTensorFType,
+    NumpyBuffer,
     NumpyBufferFType,
+    dense,
+    element,
+    fiber_tensor,
 )
 
 
 def test_fiber_tensor_attributes():
     fmt = FiberTensorFType(DenseLevelFType(ElementLevelFType(0.0)))
     shape = (3,)
-    a = fmt(shape)
+    a = fmt(shape=shape)
 
     # Check shape attribute
     assert a.shape == shape
@@ -33,3 +37,11 @@ def test_fiber_tensor_attributes():
 
     # Check buffer_format exists
     assert a.buffer_factory == NumpyBufferFType
+
+
+def test_fiber_tensor():
+    fmt = fiber_tensor(
+        dense(dense(element(np.int64(0), np.int64, np.intp, NumpyBufferFType)))
+    )
+
+    fmt(shape=(3, 4), val=NumpyBuffer(np.arange(12)))