Dense Level #183

mtsokol · 2025-09-24T09:31:32Z

This PR introduces dense level logic, for interpreters and Numba backend codegen.

mtsokol · 2025-09-27T16:25:57Z

Hi @willow-ahrens,
I managed to lower notation to assembly with FiberTensor and Dense levels. Here are some questions:

How can I access strides of a Fiber tensor? Is it even possible? I think it's part of a dense level only but in our assembly strides are associated with the Tensor object.

Matmul lowered with tensor(dense(dense(element))) inputs:

<Outdated>

src/finchlite/tensor/fiber_tensor.py

mtsokol · 2025-10-03T13:53:48Z

Hi @willow-ahrens,
Would you have spare time to give it a first look?

matmul with 2-dim dense format looks now similar to Julia's version:

def matmul(C: FiberTensorFType(DenseLevelFType(DenseLevelFType(ElementLevelFType(fv=0.0)))), A: FiberTensorFType(DenseLevelFType(DenseLevelFType(ElementLevelFType(fv=0.0)))), B: FiberTensorFType(DenseLevelFType(DenseLevelFType(ElementLevelFType(fv=0.0))))) -> FiberTensorFType(DenseLevelFType(DenseLevelFType(ElementLevelFType(fv=0.0)))):
    m: ExtentFType(start=int64, end=int64) = dimension(A, 0)
    n: ExtentFType(start=int64, end=int64) = dimension(B, 1)
    p: ExtentFType(start=int64, end=int64) = dimension(A, 1)
    A_: FiberTensorFType(DenseLevelFType(DenseLevelFType(ElementLevelFType(fv=0.0)))) = A
    A__buf: np_buf_t(float64) = A_.lvl.lvl.lvl.val
    A__buf_slot: np_buf_t(float64) = unpack(A__buf)
    B_: FiberTensorFType(DenseLevelFType(DenseLevelFType(ElementLevelFType(fv=0.0)))) = B
    B__buf: np_buf_t(float64) = B_.lvl.lvl.lvl.val
    B__buf_slot: np_buf_t(float64) = unpack(B__buf)
    C_: FiberTensorFType(DenseLevelFType(DenseLevelFType(ElementLevelFType(fv=0.0)))) = C
    C__buf: np_buf_t(float64) = C_.lvl.lvl.lvl.val
    C__buf_slot: np_buf_t(float64) = unpack(C__buf)
    for i in range(0, length(slot(C__buf_slot, np_buf_t(float64)))):
        store(slot(C__buf_slot, np_buf_t(float64)), i, 0.0)
    for i in range(m.start, m.end):
        i__pos: int64 = add(0, mul(A_.lvl.stride, i))
        i__pos_2: int64 = add(0, mul(C_.lvl.stride, i))
        for k in range(p.start, p.end):
            k__pos: int64 = add(i__pos, mul(A_.lvl.lvl.stride, k))
            k__pos_2: int64 = add(0, mul(B_.lvl.stride, k))
            for j in range(n.start, n.end):
                j__pos: int64 = add(k__pos_2, mul(B_.lvl.lvl.stride, j))
                j__pos_2: int64 = add(i__pos_2, mul(C_.lvl.lvl.stride, j))
                a_ik: float64 = load(slot(A__buf_slot, np_buf_t(float64)), k__pos)
                b_kj: float64 = load(slot(B__buf_slot, np_buf_t(float64)), j__pos)
                c_ij: float64 = mul(a_ik, b_kj)
                store(slot(C__buf_slot, np_buf_t(float64)), j__pos_2, add(load(slot(C__buf_slot, np_buf_t(float64)), j__pos_2), c_ij))
    repack(C__buf_slot)
    matmul_return: FiberTensorFType(DenseLevelFType(DenseLevelFType(ElementLevelFType(fv=0.0)))) = C
    return matmul_return

mtsokol · 2025-10-03T13:54:45Z

src/finchlite/autoschedule/compiler.py

                    ntn.Variable(
                        name,
-                        BufferizedNDArrayFType(
+                        type(val)(


That's something to address as a separate larger task

mtsokol · 2025-10-04T20:17:11Z

src/finchlite/autoschedule/compiler.py

-                            val.ndim,
-                            TupleFType.from_tuple(val.shape_type),
-                        ),
+                        val.from_kwargs(val.to_kwargs()),


Proper handling of the format inference is a separate larger task (to map traits.jl), here I only convert it to a "dict of attributes" where I can override some of them (for example promoted dtype from two inputs).

mtsokol · 2025-10-04T20:18:04Z

src/finchlite/autoschedule/compiler.py

+            levels_to_add = [
+                idx for idx, f in enumerate(result_fields) if f not in fields
+            ]
+            result_rep = result_rep.add_levels(levels_to_add)


As in the comment above - temporary selection of format (including removing/adding levels), this will be moved to a separate module.

mtsokol · 2025-10-04T20:20:12Z

src/finchlite/codegen/numba_backend.py

-    ]
-    return fmt(*args)
+    }
+    return fmt(**kwargs)


formats now have "multiple constructors", whether we construct from Numba or in the facing API

mtsokol · 2025-10-04T20:20:52Z

src/finchlite/compile/bufferized_ndarray.py

    def __init__(
        self,
-        arr: np.ndarray | NumpyBuffer,
+        val: np.ndarray | NumpyBuffer,


Standarization of names, to match FiberTensor

mtsokol · 2025-10-04T20:22:02Z

src/finchlite/compile/bufferized_ndarray.py

+        shape: tuple[np.integer, ...] | None = None,
+        strides: tuple[np.integer, ...] | None = None,


For simulating multiple constructors.

mtsokol · 2025-10-04T21:03:10Z

@willow-ahrens So right now E2E Numba matmul runs with bufferized and fiber_tensor(dense(... formats.

One part here might look like a bit of a mess, namely from_kwargs and to_kwargs but it's a temporary solution for figuring out output format for mapjoin and aggregates. I can override certain parameters this way, e.g. dtype.

willow-ahrens

Overall, this looks good. I'd like it to pass tests, work for C codegen, and implement the freeze, thaw, etc. I'm more than happy to take a pass at this this week. Let's sync up next tuesday on next steps.

willow-ahrens · 2025-11-11T15:54:11Z

src/finchlite/tensor/fiber_tensor.py

    """

-    @property
-    @abstractmethod


I think the level needs to know this stuff.

Ah, right - it should still be there, done.

willow-ahrens · 2025-11-11T15:55:10Z

src/finchlite/tensor/fiber_tensor.py

+        return tns
+
+    def lower_thaw(self, ctx, tns, op):
+        raise NotImplementedError


We'll need an implementation of these methods. I can work on it.

willow-ahrens · 2025-11-11T15:56:26Z

src/finchlite/tensor/level/dense_level.py

+        val_lvl = asm.GetAttr(val, asm.Literal("lvl"))
+        return self.lvl_t.asm_unpack(ctx, var_n, val_lvl)
+
+    def get_fields_class(self, tns, buf_s, nind, pos, op):


Is this also temporary?

In the def unfurl in FiberTensorFType I need to have Fields class instance so this is a proposition of an abstraction that levels would implement. So for the current form the answer is no, that won't be temporary.

mtsokol · 2025-11-15T21:31:02Z

@willow-ahrens Thank you for taking a look through these changes! I know that this PR is a nightmare - with these temporary hacks for suitable reps and rather vague shape considering that looplets and sparse level are still to come.

Now all tests (and mypy) pass. In test_codegen.py the test_e2e_numba uses dense(dense(element)) format.

mtsokol force-pushed the ms/dense-level branch 2 times, most recently from e02cd84 to 3551c2a Compare September 27, 2025 14:42

mtsokol self-assigned this Sep 27, 2025

mtsokol added the enhancement New feature or request label Sep 27, 2025

mtsokol changed the title ~~[WIP] Dense Level~~ Dense Level Sep 27, 2025

mtsokol commented Sep 30, 2025

View reviewed changes

src/finchlite/tensor/fiber_tensor.py Show resolved Hide resolved

mtsokol commented Sep 30, 2025

View reviewed changes

src/finchlite/tensor/fiber_tensor.py Show resolved Hide resolved

mtsokol added 3 commits October 2, 2025 15:56

WIP Dense level

2c9fb85

Remove unused method

b45493b

Remove BufferizedNDArray mentions from autoschedule.compiler

ca5db23

mtsokol force-pushed the ms/dense-level branch from 27b7a86 to ca5db23 Compare October 2, 2025 13:56

Call different parts of lowering recursively

4f45e9a

mtsokol commented Oct 3, 2025

View reviewed changes

Finish Dense E2E Numba matmul

7f09812

mtsokol marked this pull request as ready for review October 4, 2025 20:13

mtsokol commented Oct 4, 2025

View reviewed changes

small fixes

06ef043

willow-ahrens self-requested a review October 18, 2025 01:53

willow-ahrens and others added 3 commits October 28, 2025 17:18

Merge branch 'main' into ms/dense-level

2f88deb

Merge branch 'main' into ms/dense-level

7530d0c

pre-commit

9746a1c

willow-ahrens requested changes Nov 11, 2025

View reviewed changes

Merge branch 'main' into ms/dense-level

4ad31e1

mtsokol added 3 commits November 15, 2025 21:54

Make all tests pass

f62eda6

Move back LevelFType abstract functions

34b3414

Move back LevelFType abstract functions

9e62bb8

mtsokol and others added 4 commits November 16, 2025 11:10

Merge branch 'main' into ms/dense-level

eefd2ee

Unify val naming across Tensor types

8e451e0

Merge branch 'main' into ms/dense-level

ff366f6

Implement freeze and thaw for Dense level

2aae134

		shape: tuple[np.integer, ...] \| None = None,
		strides: tuple[np.integer, ...] \| None = None,

Dense Level #183

Are you sure you want to change the base?

Dense Level #183

Uh oh!

Conversation

mtsokol commented Sep 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mtsokol commented Sep 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mtsokol commented Oct 3, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mtsokol commented Oct 4, 2025

Uh oh!

willow-ahrens left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mtsokol commented Nov 15, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

mtsokol commented Sep 24, 2025 •

edited

Loading

mtsokol commented Sep 27, 2025 •

edited

Loading