Association Dialect

[NeuralCoder3]

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.

[leissa]

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]

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]

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

[NeuralCoder3]

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.