[ActivityAnalysis] Remove isConstantValue call in activity analysis

[rmoyard]

Hello!

I am trying to update Catalyst (https://github.com/PennyLaneAI/catalyst) to the latest LLVM, MLIR and Enzyme version. I have noticed since this change was introduced
#1558 that we do not get correct results. I am not familiar with your codebase but I could identify that !Hypothesis->isConstantValue(TR, SI->getPointerOperand()); is the problematic line. I think there is some side effect calling this line, because removing cop2 from the conditional is not sufficient, the call to isConstantValue has to be removed for us to get correct results.

I would appreciate your help to solve this issue :)

[wsmoses]

Do you have an example of a failing test?

Cc @pengmai

[rmoyard]

@wsmoses Sorry in advance even our simplest example produces a lot of code. But you can find the .ll file before and after the changes. https://gist.github.com/rmoyard/dce7bef76d41c689335d42e6a5bff30e It is produced simply by taking the derivative of this quantum function. Let me know if you want more details.

dev = qml.device("lightning.qubit", wires=2)

def f(x):
    qml.RX(x * 4, wires=0)
    return qml.expval(qml.PauliY(0))

@catalyst.qjit(keep_intermediate=True)
def compiled(x: float):
    g = qml.qnode(qml.device("lightning.qubit", wires=1), diff_method="adjoint")(f)
    h = catalyst.grad(g, method="auto")
    return h(x)

We use register gradient for the quantum part and the classical part is done entirely (here x*4) by backprop. The thing that is curious is that the wrong values only happen when we have the classical part.

[rmoyard]

By looking at the activity analysis and just keeping relevant information I noticed some differences for the value %12, it goes from constant value 1 to 0. I think that is a mistake because this part of the code should not enter the computation of the gradient.

%.idx = shl i64 %3, 3 cv=1 ci=1
%8 = tail call ptr @_mlir_memref_to_llvm_alloc(i64 %.idx) #4 cv=0 ci=1
%9 = tail call ptr @_mlir_memref_to_llvm_alloc(i64 72) #4 cv=1 ci=1
%10 = ptrtoint ptr %9 to i64 cv=1 ci=1
%11 = add i64 %10, 63 cv=1 ci=1
%12 = and i64 %11, -64 cv=1 ci=1
%13 = inttoptr i64 %12 to ptr cv=1 ci=1

%.idx = shl i64 %3, 3 cv=1 ci=1
%8 = tail call ptr @_mlir_memref_to_llvm_alloc(i64 %.idx) #4 cv=0 ci=1
%9 = tail call ptr @_mlir_memref_to_llvm_alloc(i64 72) #4 cv=1 ci=1
%10 = ptrtoint ptr %9 to i64 cv=1 ci=1
%11 = add i64 %10, 63 cv=1 ci=1
%12 = and i64 %11, -64 cv=0 ci=1
%13 = inttoptr i64 %12 to ptr cv=0 ci=1

From what I can understand there is no more search UP for %12. (full activity here https://gist.github.com/rmoyard/a8f9f0821d9611734f69aece7380d001)

Could someone tell me how is the isConstant function is working? And could that influence the results?

[wsmoses]

Probably a bit longer than you wanted, but a few months back I did a deep dive on the internals of the analyses that may be helpful: https://www.youtube.com/watch?v=VTyHPKR5dDQ&ab_channel=EnzymeAD

We definitely should add more docs to this (PRs welcome).

We also are hoping to write a lot of this up in a few papers, but we haven't gotten to that yet.

[wsmoses]

I think activity analysis starts here: https://www.youtube.com/watch?v=VTyHPKR5dDQ&t=47m43s&ab_channel=EnzymeAD

[rmoyard]

Hi @wsmoses ! I have added the integration test, let me know if there is a better folder where to put it. You can review the changes