Skip to content

Conversation

@TheRealMichaelWang
Copy link
Collaborator

Added a Einsum to Indirect-Einsum (Insum) Lowerer. Additionally the following were also added:

  • A GetAttribute EinsumExpr AST node. It gets an attribute/property of a tensor (i.e. the element or coordinate array of a COO tensor)
  • A sparse tensor that inherits from eager tensor, with the corresponding FType. It is implemented in basic COO form.

The Insum Lowerer itself essentially does the following in order to convert an Einsum with a single sparse tensor to an Insum

  • Build a tensor T whose elements store the number of non-zero elements in each corresponding fiber of the sparse tensor that is reduced
  • Check if each reduced fiber in the sparse tensor is all non-zero or not by checking T. We then initialize the initial reduction values of our output tensor based on whether the corresponding reduced fiber is all non-zero or not.
    • If all elements in the reduced fiber of the sparse tensor are non-zero we can safely use the default initial value for the einsum's reduction operation
    • If there exist any zeros in the reduced fiber of the sparse tensor, we do the following.
      • We rewrite the arg/pointwise expression of the Einsum where we assume that all the accesses to the sparse tensor's elements are zero. We call this AST EinsumExpr, "ZeroArg".
      • We initialize the initial reduction value of our output tensor to the reduction of the default initial reduction value and the "ZeroArg" we calculated in the previous step.
      • In doing so we make an assumption about the reduction operator. For clarity let the reduction operator be called f, let the default initial value for f be called i, and let the array of elements being reduced called s. The core assumption we make is that if s has any zero elements, reduce(f, s) = reduce(f(i, 0), reduce(f, nonzero(s))). Also note that 0 is not necessarily the only value used in f(i, 0), but we still make a similar assumption.

TheRealMichaelWang and others added 30 commits October 9, 2025 10:12
* Still needs to be tested
* Currently all are failing
…in.Plan

* Fixed errors handling parameters in EinsumLowerer
* Properly handled einsums in return value
* Return values will now be encoded by ein.Produce nodes in the Einsum's bodies tuple
* Reverted to old printing that supports ein.Produce
* Fixed misc issues related to reverting changes
* Sparse tensor implementation is COO
* Stores an array of coordinates of non-zero elements and an array of non-zero elements
* Get Attribute is a general purpose node that can be used to retreive the coordinate and element arrays from a sparse tensor
* We create a count tensor T, where each element in T is the number of non-zero elements in each corresponding reduced fiber of the sparse tensor
* We initialize the initial reduction values based on whether each reduced fiber is all non-zero or not
TheRealMichaelWang and others added 29 commits October 21, 2025 11:41
…plan

* Renamed lower_to_pointwise_op to compile_expr and lower_to_pointwise to compile_operand for clarity
… unused rename_einsum method. Updated logic for processing Aggregate cases to improve clarity and efficiency.
…gate queries, improving flexibility in handling various query formats.
… inlining its logic directly into the compile_plan method. This change simplifies the code and enhances clarity in handling Aggregate cases.
…e case structure, improving code clarity and reducing redundancy. This change enhances the handling of various query formats, including Reformat and Reorder cases.
…solidating logic into a single case structure, improving code clarity and reducing redundancy. This update also removes outdated code, enhancing maintainability.
* Imported logic AST nodes will be refered to with lgc. prefix prependded for clarity in einsumlowerer
* Will add back in another optimization pass, sometime in the future
… and doesn't rely on building a vector of statements/bodies
Copy link
Member

@willow-ahrens willow-ahrens left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm adding my review to remove the request, we've shifted our focus temporarily to the right-hand-side insum PR. Feel free to re-request once the other PR has been merged into main and integrated into this one and this one is ready.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants