tablegen: Add StaticIf to select based on static condition

[Pangoraw]

cc @mofeing, we can collaborate to integrate in your fft pr.

I am open to other naming ideas.

[wsmoses]

I think probably we should remove the explicit needs primal and shadow, and instead directly integrate this into the use analysis (lookup SelectIfActive in this case we wouldn’t do the check on if active but check on whatever the string is)

[Pangoraw]

Looking into it, I could not find the custom use analysis for SelectIfActive, it seems that instead SelectIfActive has no analysis:

       SmallVector<bool> cachedArguments(Operation *op,
                                 MGradientUtilsReverse *gutils) const {
         SmallVector<bool> toret(op->getNumOperands(), false);
         for (size_t idx=0; idx<op->getNumOperands(); idx++) {
            bool used = false;
            // Rule (Op ?:$x, ?:$y)
            // Arg 0 : (SelectIfActive ?:$x, (CMul (DiffeRet), ?:$y), ?:$x)
            // Arg 1 : (DiffeRet)
            toret[idx] = used;
         }
         return toret;
       }

And it currently works because either it is used in the Forward diff which seem to rely on the analysis or the primals are used outside of the SelectIfActive node.

It seems that the assert here does not trigger due to the way enzyme-tblgen might be built?

Enzyme/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp

Line 1208 in 1ca3ceb

assert(!usesCustom);

So really, it should base its use on its arguments.

[wsmoses]

Ah yeah we should add support there (actually for both static if and if active). The reason we never hit that branch is because use analysis is auto generated from reverse mode rules and I think we only use select if active in forward.

However here we really want to add the correct diff use analysis, since then we can for example use static if to define rules for reduce and static if it’s a known max min or add etc

[wsmoses]

But also it should do so recursively for the operands of the select if active (depending on if active)

[Pangoraw]

Right, the same analysis need to be used for SelectIfComplex as well (or it should be replaced with an instance of StaticIf)

[wsmoses]

I mean ideally we can define selectifcomplex in tablegen with selectif