Indirect Einsum Backend #181
Conversation
TheRealMichaelWang
commented
Sep 23, 2025
TheRealMichaelWang
commented
- Added Einsum IR
- Logic IR lowered into einsum plan which consists of einsum expressions
- Each einsum node supports pointwise expression
- Added SparseTensor (using COO)
* Added EinsumExtractor barebones and Einsum and Einprod
- Implemented EInsumTransformer transofrm which transforms aggregates and map joins into einsums - Not comprehensive
* Added barebones pointwise AST
* Updated Einsum to fit PointwiseIR
…on that consumes optimized logic IR
* Added support for recursive descent parsing of MapJoins
… logic ir statements seperate from value logic ir statements * Seperates into compile_plan and lower_to_pointwise_op * Added support for handling aggregates (use builtin init_value only for each reduce function, non standard init values aren't supported)
…erators in einsums
…eduler * Fixed circular import issues
src/finchlite/autoschedule/einsum.py
Outdated
|
|
||
| reduceOp: Callable #technically a reduce operation, much akin to the one in aggregate | ||
|
|
||
| input_fields: tuple[Field, ...] |
There was a problem hiding this comment.
Delete input_fields, it can be computed from pointwiseExpr and it's more work to keep them in sync
| reduceOp: Callable #technically a reduce operation, much akin to the one in aggregate | ||
|
|
||
| input_fields: tuple[Field, ...] | ||
| output_fields: tuple[Field, ...] |
| """ | ||
|
|
||
| alias: str | ||
| idxs: tuple[Field, ...] # (Field('i'), Field('j')) |
src/finchlite/interface/eager.py
Outdated
| if isinstance(arg, SparseTensor): | ||
| return DefaultLogicOptimizer(EinsumScheduler(EinsumCompiler())) | ||
|
|
||
| return None |
There was a problem hiding this comment.
This concept should be a separate PR
src/finchlite/interface/eager.py
Outdated
| if isinstance(arg, lazy.LazyTensor): | ||
| return lazy.permute_dims(arg, axis=axis) | ||
| return compute(lazy.permute_dims(arg, axis=axis)) | ||
| return compute(lazy.permute_dims(arg, axis=axis), ctx=get_eager_scheduler(arg)) |
There was a problem hiding this comment.
Ideally (hopefully) we don't need to change eager.py
|
Example 1) Elementwise multiplication # Dense Einsum
C[i] = A[i] * B[i]
# Assuming A is sparse, B and C are dense
# A = {Acrd, Aval} (nonzero coordinate and value array)
C[Acrd[p]] = Aval[p] * B[Acrd[p]]Example 2) Matmul # Dense Matmul Einsum
C[i,j] = A[i,k] * B[k,j]
# Assuming A is sparse, B and C are dense
# A = {AcrdI, AcrdK, Aval} (nonzero coordinate and value COO)
C[AcrdI[p],j] = Aval[p] * B[AcrdK[p],j]Example 3) SDDMM # Dense Sampled Matmul Einsum
C[i,j] = D[i,j] * A[i,k] * B[k,j]
# Assuming D is sparse, A,B and C are dense
# D = {DcrdI, DcrdJ, Dval} (nonzero coordinate and value COO)
C[DcrdI[p],DcrdJ[p]] = Dval[p] * A[DcrdI[p],k] * B[k,DcrdJ[p]]Example 4) SDDMM but with sparse output # Dense Sampled Matmul Einsum
C[i,j] = D[i,j] * A[i,k] * B[k,j]
# Assuming D is sparse, A and B are dense
# output C is SPARSE!
# D = {DcrdI, DcrdJ, Dval} (nonzero coordinate and value COO)
C[p] = Dval[p] * A[DcrdI[p],k] * B[k,DcrdJ[p]] |
|
@TheRealMichaelWang I would like to use the Einsum AST for my own purposes. Do you think you could make a PR with the Einsum language as a separate subtree, and the pass that lowers FinchLogic to Einsum? (i.e. make a folder called "einsum_notation" and add all the nodes to it)? |
|
A general lowering example, which hints at an algorithm to do this generally: Notice that many of the steps above used rewriting simplification to determine whether we can annihilate and to simplify the two cases of A[i, j] = 0 or A[i, j] != 0. In one case, we replace the sparse tensor by a mask that says where the zeros are. In another case, we replace A[i, j] with val[p], and we replace indices with i=>idx[p] and j=>jdx[p]. For each lowering example, you can derive it from these two cases, and then simplify into the final algorithm. |
* Removed redunant output fields from Einsum * Added extract COO einsum IR statement
* Still need to add support for true indirect pointwise access
