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/compile/bufferized_ndarray.py

@@ -377,7 +377,7 @@ def asm_repack(self, ctx, lhs, obj):
         """
         Repack the buffer from C context.
         """
-        (self.tns.asm_repack(ctx, lhs.tns, obj.tns),)
+        self.tns.asm_repack(ctx, lhs.tns, obj.tns)
         ctx.exec(
             asm.Block(
                 asm.SetAttr(lhs, "tns", obj.tns),

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):
@@ -68,7 +74,7 @@ class Level(FTyped, ABC):
 
     @property
     @abstractmethod
-    def shape(self):
+    def shape(self) -> tuple:
         """
         Shape of the fibers in the structure.
         """
@@ -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,17 @@ 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))
+
+    def __str__(self):
+        return f"FiberTensorFType({self.lvl})"
 
     @property
     def shape(self):
@@ -209,24 +227,82 @@ def position_type(self):
 
     @property
     def buffer_factory(self):
+        return self.lvl.buffer_factory
+
+    def unfurl(self, ctx, tns, ext, mode, proto):
+        op = None
+        if isinstance(mode, ntn.Update):
+            op = mode.op
+        tns = ctx.resolve(tns).obj
+        obj = self.lvl.get_fields_class(tns, 0, asm.Literal(self.position_type(0)), op)
+        return self.lvl.unfurl(ctx, ntn.Stack(obj, self.lvl), ext, mode, proto)
+
+    def lower_freeze(self, ctx, tns, op):
+        return tns
+
+    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):
         """
-        Returns the ftype of the buffer used for the fibers.
-        This is typically a NumpyBufferFType or similar.
+        Unpack the into asm context.
         """
-        return self.lvl.buffer_factory
+        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 asm_repack(self, ctx, lhs, obj):
+        """
+        Repack the buffer from the context.
+        """
+        ctx.exec(asm.Repack(obj.buf_s))
+        return
 
 
-def tensor(lvl: LevelFType, position_type: type | None = None):
+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

@@ -1,32 +1,48 @@
+import operator
 from abc import ABC
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, NamedTuple
 
 import numpy as np
 
-from ..fiber_tensor import Level, LevelFType
+from ... import finch_assembly as asm
+from ... import finch_notation as ntn
+from ...compile import looplets as lplt
+from ..fiber_tensor import FiberTensorFields, Level, LevelFType
+
+
+class DenseLevelFields(NamedTuple):
+    tns: FiberTensorFields
+    nind: int
+    pos: asm.AssemblyNode
+    op: Any
 
 
 @dataclass(unsafe_hash=True)
 class DenseLevelFType(LevelFType, ABC):
     lvl: Any
     dimension_type: Any = None
+    pos: asm.AssemblyNode | None = None
+    op: Any = None
 
     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]))
 
+    def __str__(self):
+        return f"DenseLevelFType({self.lvl})"
+
     @property
     def ndim(self):
         return 1 + self.lvl.ndim
@@ -63,6 +79,45 @@ def buffer_factory(self):
         """
         return self.lvl.buffer_factory
 
+    def get_fields_class(self, tns, nind, pos, op):
+        return DenseLevelFields(tns, nind, pos, op)
+
+    def unfurl(self, ctx, tns, ext, mode, proto):
+        def child_accessor(ctx, idx):
+            pos_2 = asm.Variable(
+                ctx.freshen(ctx.idx, f"_pos_{self.ndim - 1}"), self.pos
+            )
+            ctx.exec(
+                asm.Assign(
+                    pos_2,
+                    asm.Call(
+                        asm.Literal(operator.add),
+                        [
+                            tns.obj.pos,
+                            asm.Call(
+                                asm.Literal(operator.mul),
+                                [
+                                    tns.obj.tns.stride[tns.obj.nind],
+                                    asm.Variable(ctx.idx.name, ctx.idx.type_),
+                                ],
+                            ),
+                        ],
+                    ),
+                )
+            )
+            return ntn.Stack(
+                self.lvl.get_fields_class(
+                    tns.obj.tns, tns.obj.nind + 1, pos_2, tns.obj.op
+                ),
+                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)
@@ -77,11 +132,15 @@ class DenseLevel(Level):
     _format: DenseLevelFType
     lvl: Any
     dimension: Any
+    pos: asm.AssemblyNode | None = None
 
     @property
-    def shape(self):
+    def shape(self) -> tuple:
         return (self.dimension, *self.lvl.shape)
 
     @property
     def ftype(self):
         return self._format
+
+    def with_pos(self, pos: asm.AssemblyNode) -> "DenseLevel":
+        return DenseLevel(self._format, self.lvl, self.dimension, pos)