Repository navigation
Can we use blas_contract!'s AD rules for contract!, even for fermions? #562
Description
Motivation:
we already have custom rules for blas_contract!, which at least in the Enzyme case save us a lot of compilation overhead as Enzyme does not have to walk the entire interior call graph of the function for all supported type combinations. For tensors without a plain array view (every non-Trivial sector), TO.tensorcontract! goes through TensorKit.contract!, which has no custom Enzyme rule. Enzyme then has to analyze contract!'s whole call graph for every contraction specialization. The only part of that call graph which has a custom Enzyme rule is blas_contract!. Avoiding doing all this call graph walking could save us ~30% walltime compiling the rule.
I can't find the comment at this moment but I think @Jutho commented when we were first developing the Mooncake and Enzyme extensions and asked whether this would be possible, and the conclusion was "no" because of fermions. But is this true?
For symmetric tensors, this should be valid to do because every other call in contract! is inactive.
I got a 🤖 to test that the tensor derivatives are the same for blas_contract! and contract! for:
| space | checks passed | branches covered | rule invocations | time |
|---|---|---|---|---|
Trivial |
25115 / 25115 | all 4 | 172 | 25 min |
Irrep[ℤ₂] |
27823 / 27823 | all 4 | 172 | 32 min |
Irrep[CU₁] |
35082 / 35082 | all 4 | 172 | 67 min |
FermionParity ⊠ Irrep[SU₂] ⊠ Irrep[U₁] (VfHubbard) |
90748 / 90748 | all 4 | 172 | 76 min |
Irrep[SU₂] (VRepSU₂) |
41005 / 41005 | all 4 | 172 | 97 min |
Irrep[SU₃] (SU3Irrep, irreps 1, 3, 3̄, 8) |
50810 / 50810 | all 4 | 172 | 95 min |
FermionSpin (VfRepSU₂) |
67369 / 67369 | all 4 | 172 | 110 min |
Irrep[A₄] ⊠ GroupElement[ℤ₄, 2], IsingBimodule ⊠ Irrep[A₄] |
skipped: non-symmetric braiding |
With a bunch of index swaps. But I'm not sure this is comprehensive enough to convince us that we can do this, so if anyone has other symmetries I should check, please do chime in.
Activity
kshyatt commentedon Oct 1, 2026
I also tried just aggressively marking everything besides
blas_contract!incontract!asinactive, and it didn't help :(lkdvos commentedon Oct 1, 2026
I think there are some things that are mixed here, or I might not fully understand the question:
TensorOperationscontract pullback rules, which we can't because these are missing the additional twist that is inserted on contracted dual legs. This has nothing to do with whatTensorKit.blas_contract!vsTensorKit.contract!compute.blas_contract!instead ofcontract!is more about runtime, where I'm assuming that exchangingAandBon the forward pass to avoid permutations should have the same effect on the backwards pass, but I definitely did not think too much about whether this is actually true, and whether intercepting one level higher would effectively end up with the same result because the pullbacks end up callingcontractwhich should again reorder if necessary.contract!call effectively compiles 4blas_contract!calls because of the manual union split we have there, where each of these gets different partitions and argument types, therefore requiring 4 compiled methods and also 4 compiled enzyme versions.What I am not sure of however is if we actually solve this by intercepting one layer higher, or simply defer the problem. At least in the forwards computation, we compile 4
blas_contract!methods percontract!, but there are also 4contract!signatures that end up calling these same 4blas_contract!, so it ends up being a little more work up front that might then be reused later. I don't know enough about Enzyme rules to know if this same reasoning works there though[edit] As a sidenote, I recently came across the
Base.inferencebarrierfunction that can be used to avoid compiling the functions that end up not being called, not sure if that could be relevant/of help herekshyatt commentedon Oct 1, 2026
Oh nice the inference barrier might indeed help. Sorry that this is kind of all over the place, it was a result of 7am debugging. That's what I get for citing a half remembered comment...
kshyatt commentedon Oct 1, 2026
It's not just this, I don't think, but it does contribute a significant amount. It's not just 4 compiled Enzyme versions btw, I think it has to compile for every allowable type combination going into
blas_contract!although my understanding of that part of the internals is not so good.