[MLIR] Specialize active callbacks to their own function

[erick-xanadu]

Context: Enzyme allows one to specify custom gradients for specific functions. In order to specify custom gradients for callbacks, callbacks need to be specialized to their own specific functions. E.g., instead of having the following code:

func.call @callback(%identifier, %argc, %retc, %0, %1, ..., %m, %n)

And be unable to register custom gradients for @callback. Specialize callbacks to their identifiers like so:

func.func @active_callback_123(%arg0, %arg1, ..., %argm, %argn) {
  %identifier = llvm.constant ...
  %argc = llvm.constant ...
  %retc = llvm.constant ...
  llvm.call @callback(%identifier, %argc, %retc, %arg0, %arg1, ... %argm, %argn)
  return
}

  // ...
  func.call @active_callback_123(%0, %1, ... %m, %n)
  // ..

And now we can register a custom gradient for callback_123 and any other callback_456.

Description of the Change:

• Iterate over all ActiveCallbackOps and create the specialize function. Each ActiveCallbackOp will be annotated with the specialized function.
• During the lowering of ActiveCallbackOps to LLVM-IR replace with a call to the specialized function.

[sc-60494]

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 98.04%. Comparing base (7248c12) to head (00cf68c).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #735   +/-   ##
=======================================
  Coverage   98.04%   98.04%           
=======================================
  Files          69       69           
  Lines        9536     9538    +2     
  Branches      762      763    +1     
=======================================
+ Hits         9350     9352    +2     
  Misses        151      151           
  Partials       35       35           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[dime10]

Thanks Erick, looks good! Regarding reusing the transformation from results into arguments and memrefs into struct pointers, is that something that would be applicable to this PR?

[dime10]

Should this go directly before gradient lowering?

[erick-xanadu]

Why?

[dime10]

I think I was tying the callbacks to gradients because of the registration with enzyme. I was also thinking what is the lowest point in the pipeline it is first needed.

Is there a reason it should be the first thing in the pipeline?

[dime10]

Unrelated to this PR but I cannot tell at all what this pass is for or related to (from its name) 😅

[erick-xanadu]

Do you have a different name that you prefer? Would outline-active-callback-pass make more sense?

[dime10]

Oh I was referring to annotate-function

[dime10]

An inactive callback has no results right?

[dime10]

Ah maybe I got this wrong because I though debug.callback -> inactive_callback, but I guess the active_callback also calls the inactive callback.

[erick-xanadu]

I am not sure if you were saying that I should get rid of the results in the InactiveCallbackOp. I just did. The reason why I didn't get rid of it is because I was still exploring different ways to implement this. To be honest, I have a better way in mind now 😅 but I will prioritize getting this PR in and subsequent PRs will focus on cleaning up.

[dime10]

but I will prioritize getting this PR in and subsequent PRs will focus on cleaning up

Generally I will say that I don't much like the practice of splitting up PRs along a "rough first implementation" / "cleanup" line, because it makes it almost impossible to track the deficiencies of the first PR and make sure they all got addressed in the second PR.
A split along functionality is much more sensible imo.

This PR is fine but something to keep in mind for the future.

[erick-xanadu]

I agree with you, but it is difficult to reconcile the "please close tickets before the end of iteration" with "no rough draft implementation PRs" when:

1. we did not have time to design the whole epic in advance to forsee everything in each story
2. stories build upon previous stories that may not be used until later.

I do not know if it would be possible to have tickets that do not correspond to merging of PRs but instead to just get the implementation. And a final ticket for merging the implementation.

That to me, seems appropriate for large features like this one. It would satisfy the opinion of:

1. Small PRs are better (as it help for reviews)
2. No draft PRs (as everything is already finished except for merging)

But would put a lower priority on "close tickets before the end of the story" and "each ticket must be a PR"

[erick-xanadu]

@dime10

Thanks Erick, looks good! Regarding reusing the transformation from results into arguments and memrefs into struct pointers, is that something that would be applicable to this PR?

No. But maybe a future PR might change this. I'm thinking whether using memrefs directly in the specialization and later undergoing the pointer-to-struct ABI transform would be good in terms of readability / invariants. I think so, but I will implement it in a different PR.

[erick-xanadu]

Closing in favour of #782