Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Association Dialect #173

Open
NeuralCoder3 opened this issue Jan 16, 2023 · 4 comments
Open

Association Dialect #173

NeuralCoder3 opened this issue Jan 16, 2023 · 4 comments
Labels
enhancement New feature or request

Comments

@NeuralCoder3
Copy link
Collaborator

Maybe an association dialect would be helpful.

Often, when we rewrite things, we want to replace something with a semantically equal term on another level.
The destination term can often be written (more or less) conveniently in thorin.
This is especially the case for staged lowering like present in the autodiff and matrix dialect.

These associations are currently implemented using naming conventions and internals (a hack around externals that are erased).

The association dialect would keep track of associations in the thorin code and expose map on the C++ side.
No special care to quoting/reification is needed as we usually only want to associate axioms anyway.

Example:

// in associate
%assoc.associate: [T1:*,T2:*] -> %assoc.association_class -> T1 -> T2 -> %assoc.association

// in matrix
%matrix.association :  %assoc.association_class

.extern associate_matrix() {
  (%assoc.associate %matrix.association %matrix.prod internal_matrix_mapRed_prod, 
  ...)
}

There could also be generalized passes on associations (e.g. replace all terms in a class with their associated counterpart).

This seems like a simple way to cleverly utilize the dialect infrastructure and get away from string matching.

This dialect could be combined or extend the sketched rewrite-mapping dialect from previous discussions.

@NeuralCoder3 NeuralCoder3 added the enhancement New feature or request label Jan 16, 2023
@leissa
Copy link
Member

leissa commented Jan 16, 2023

I thought about this as well. Maybe we can simply add an optional value to an axiom and have a pass that reduces axioms with their associated values later in the compilation pipeline?

We could a add a bit syntactic sugar to directly do sth like this:

.axm %foo.id (x: .Nat) -> .Nat = x;

Sugar for

.ax %foo.id: .Nat -> .Nat = .lm (x: .Nat) -> .Nat = x;

@NeuralCoder3
Copy link
Collaborator Author

Opaque functions as axioms would not be enough in my case.

For matrix, I want to have multiple options

  • rewriting to map-reduce (which is in turn rewritten to affine for which is in turn rewritten to recursion)
  • rewriting using an external function (blas)

For autodiff, I do not want to replace the axiom but rather in certain circumstances get an associated function (the augmented function that contains the pullback). (This is roughly equivalent to rewriting autodiff.ad axiom to axiom_diff.)

@leissa
Copy link
Member

leissa commented Jan 16, 2023

Ah, I see. Another contender is #90. But this still isn't enough for your autodiff example, correct?

@NeuralCoder3
Copy link
Collaborator Author

Yes, #90 is rewrite-mapping dialect I was referring to.
This would (at least in the axiom mapping case) be enough. One would just have to introduce a rewrite-axiom to invoke the mapping. For ad, this could directly be autodiff.ad.
The mapping would then go through thorin.
In this regard, the association dialect is isomorphic to a part of #90.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

No branches or pull requests

2 participants