Fix issues in lift-array-alloc

[annagrin]

Description

Fix crashes and over-optimizations in lift-array-alloc:

Example:

  qubits = cudaq.qvector(2)
  arr1 = [1]
  t = arr1[0]
  arr2 = [0]
  arr2[0] = t
  b = arr2[0]
  x(qubits[b])

This example generates the following IR:

func.func @test3() {
  %c0_i64 = arith.constant 0 : i64
  %c1_i64 = arith.constant 1 : i64

  // qubits = cudaq.qvector(2)
  %0 = quake.alloca !quake.veq<2>

  // arr1 = [1]
  %1 = cc.alloca !cc.array<i64 x 1>
  %2 = cc.cast %1 : (!cc.ptr<!cc.array<i64 x 1>>) -> !cc.ptr<i64>
  cc.store %c1_i64, %2 : !cc.ptr<i64>

  // t = arr1[0]
  %3 = cc.load %2 : !cc.ptr<i64>

  // arr2 = [0]
  %4 = cc.alloca !cc.array<i64 x 1>
  %5 = cc.cast %4 : (!cc.ptr<!cc.array<i64 x 1>>) -> !cc.ptr<i64>
  cc.store %c0_i64, %5 : !cc.ptr<i64> // Dominates the next store, don't lift

  // arr2[0] = t
  cc.store %3, %5 : !cc.ptr<i64>

  // b = arr2[0]
  %6 = cc.load %5 : !cc.ptr<i64>

  // x(qubits[b])
  %7 = quake.extract_ref %0[%6] : (!quake.veq<2>, i64) -> !quake.ref
  quake.x %7 : (!quake.ref) -> ()
  return
}

And lift-array-alloc misses the second store to arr2 below and incorrectly assumes that arr2[0]==0 instead of arr2[0]==1:

  cc.store %c0_i64, %5 : !cc.ptr<i64> // Dominates the next store, don't lift
  cc.store %3, %5 : !cc.ptr<i64>

The fix detects that the first store dominates the second and skips the optimization.

[schweitzpgi]

I don't see why the replaceOpWithNewOp needs to be changed here. It will

• create the new ExtractValueOp
• replace all the users of the original LoadOp to use the new ExtractValueOp

which is exactly what your change does. Additionally, we don't need or want to aggressively delete Ops. For example, passes may not be successful, and the MLIR infrastructure is designed to allow it to cleanly back out of the changes without leaving the IR broken. We should just let MLIR do the cleanups anywhere that it can. In this case, a LoadOp with 0 users will simply be removed by the next DCE. We don't have to do anything special.

[annagrin]

Not sure why but the original code crashes on uccsd example

[schweitzpgi]

How does any of this get past the check here? This is intended to check that we do not have multiple stores to the same offset of the array.

[schweitzpgi]

I ran a little test, and this check correctly blocks the transformation.

func.func @f() {
    %0 = cc.alloca !cc.array<i64 x 1>
    %1 = arith.constant 1 : i64
    %2 = cc.compute_ptr %0[0] : (!cc.ptr<!cc.array<i64 x 1>>) -> !cc.ptr<i64>
    cc.store %1, %2 : !cc.ptr<i64>
    %3 = arith.constant 3 : i64
    cc.store %3, %2 : !cc.ptr<i64>
    return
}

Running lift-array-alloc does not change this code, which is the correct thing to happen.

[annagrin]

I think the failing example differs on the second store not using a constant to store to a pointer

[schweitzpgi]

The pass should refuse to convert this code because there are 2 stores to the same offset. But I don't see why it wasn't already. But this is a good test to make sure it doesn't regress further.

[annagrin]

second store is not a constant. I will update the fix

[github-actions]

CUDA Quantum Docs Bot: A preview of the documentation can be found here.