diff --git a/lib/Optimizer/Transforms/LiftArrayAlloc.cpp b/lib/Optimizer/Transforms/LiftArrayAlloc.cpp index 5a1d0032148..0f7647b5796 100644 --- a/lib/Optimizer/Transforms/LiftArrayAlloc.cpp +++ b/lib/Optimizer/Transforms/LiftArrayAlloc.cpp @@ -40,8 +40,10 @@ class AllocaPattern : public OpRewritePattern { return failure(); LLVM_DEBUG(llvm::dbgs() << "Candidate was found\n"); - auto eleTy = alloc.getElementType(); - auto arrTy = cast(eleTy); + auto allocTy = alloc.getElementType(); + auto arrTy = cast(allocTy); + auto eleTy = arrTy.getElementType(); + SmallVector values; // Every element of `stores` must be a cc::StoreOp with a ConstantOp as the @@ -89,6 +91,8 @@ class AllocaPattern : public OpRewritePattern { cannotEraseAlloc = isLive = true; } else { for (auto *useuser : user->getUsers()) { + if (!useuser) + continue; if (auto load = dyn_cast(useuser)) { rewriter.setInsertionPointAfter(useuser); LLVM_DEBUG(llvm::dbgs() << "replaced load\n"); @@ -160,14 +164,13 @@ class AllocaPattern : public OpRewritePattern { if (!u) return nullptr; if (auto store = dyn_cast(u)) { - if (op.getOperation() == store.getPtrvalue().getDefiningOp() && - isa_and_present( - store.getValue().getDefiningOp())) { + if (op.getOperation() == store.getPtrvalue().getDefiningOp()) { if (theStore) { LLVM_DEBUG(llvm::dbgs() << "more than 1 store to element of array\n"); return nullptr; } + LLVM_DEBUG(llvm::dbgs() << "found store: " << store << "\n"); theStore = u; } continue; @@ -182,7 +185,13 @@ class AllocaPattern : public OpRewritePattern { } return nullptr; } - return theStore; + return theStore && + isa_and_present( + dyn_cast(theStore) + .getValue() + .getDefiningOp()) + ? theStore + : nullptr; }; auto unsizedArrTy = cudaq::cc::ArrayType::get(arrEleTy); diff --git a/targettests/execution/state_preparation_vector_sizes.cpp b/targettests/execution/state_preparation_vector_sizes.cpp index f6fea7a5697..ed2eff1e5a6 100644 --- a/targettests/execution/state_preparation_vector_sizes.cpp +++ b/targettests/execution/state_preparation_vector_sizes.cpp @@ -21,9 +21,6 @@ #include #include -#include -#include - __qpu__ void test(std::vector inState) { cudaq::qvector q1 = inState; } diff --git a/test/Quake/lift_array.qke b/test/Quake/lift_array.qke index 00574f8e453..c93f3631a65 100644 --- a/test/Quake/lift_array.qke +++ b/test/Quake/lift_array.qke @@ -125,3 +125,73 @@ func.func @test2() -> !quake.veq<2> { // GLOBAL-DAG: cc.global constant private @__nvqpp__mlirgen__function_test_complex_constant_array._Z27test_complex_constant_arrayv.rodata_{{[0-9]+}} (dense<[(0.707106769,0.000000e+00), (0.707106769,0.000000e+00), (0.000000e+00,0.000000e+00), (0.000000e+00,0.000000e+00)]> : tensor<4xcomplex>) : !cc.array x 4> // GLOBAL-DAG: cc.global constant private @__nvqpp__mlirgen__function_custom_h_generator_1._Z20custom_h_generator_1v.rodata_{{[0-9]+}} (dense<[(0.70710678118654757,0.000000e+00), (0.70710678118654757,0.000000e+00), (0.70710678118654757,0.000000e+00), (-0.70710678118654757,0.000000e+00)]> : tensor<4xcomplex>) : !cc.array x 4> // GLOBAL-DAG: cc.global constant private @test2.rodata_{{[0-9]+}} (dense<[1.000000e+00, 2.000000e+00, 6.000000e+00, 9.000000e+00]>" : tensor<4xf64>) : !cc.array + +func.func @test_two_stores() { + %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 + %2 = cc.cast %1 : (!cc.ptr>) -> !cc.ptr + cc.store %c1_i64, %2 : !cc.ptr + + // t = arr1[0] + %3 = cc.load %2 : !cc.ptr + + // arr2 = [0] + %4 = cc.alloca !cc.array + %5 = cc.cast %4 : (!cc.ptr>) -> !cc.ptr + cc.store %c0_i64, %5 : !cc.ptr // Dominates the next store, don't lift + + // arr2[0] = t + cc.store %3, %5 : !cc.ptr + + // b = arr2[0] + %6 = cc.load %5 : !cc.ptr + + // x(qubits[b]) + %7 = quake.extract_ref %0[%6] : (!quake.veq<2>, i64) -> !quake.ref + quake.x %7 : (!quake.ref) -> () + return +} + +// CHECK-LABEL: func.func @test_two_stores() { +// CHECK: %[[VAL_0:.*]] = arith.constant 0 : i64 +// CHECK: %[[VAL_1:.*]] = quake.alloca !quake.veq<2> +// CHECK: %[[VAL_2:.*]] = cc.const_array [1] : !cc.array +// CHECK: %[[VAL_3:.*]] = cc.extract_value %[[VAL_2]][0] : (!cc.array) -> i64 +// CHECK: %[[VAL_4:.*]] = cc.alloca !cc.array +// CHECK: %[[VAL_5:.*]] = cc.cast %[[VAL_4]] : (!cc.ptr>) -> !cc.ptr +// CHECK: cc.store %[[VAL_0]], %[[VAL_5]] : !cc.ptr +// CHECK: cc.store %[[VAL_3]], %[[VAL_5]] : !cc.ptr +// CHECK: %[[VAL_6:.*]] = cc.load %[[VAL_5]] : !cc.ptr +// CHECK: %[[VAL_7:.*]] = quake.extract_ref %[[VAL_1]][%[[VAL_6]]] : (!quake.veq<2>, i64) -> !quake.ref +// CHECK: quake.x %[[VAL_7]] : (!quake.ref) -> () +// CHECK: return +// CHECK: } + +func.func @test_complex_array() { + %cst = complex.constant [0.000000e+00 : f32, 1.000000e+00 : f32] : complex + %cst_0 = complex.constant [1.000000e+00 : f32, 0.000000e+00 : f32] : complex + %0 = cc.alloca !cc.array x 2> + %1 = cc.cast %0 : (!cc.ptr x 2>>) -> !cc.ptr> + cc.store %cst_0, %1 : !cc.ptr> + %2 = cc.compute_ptr %0[1] : (!cc.ptr x 2>>) -> !cc.ptr> + cc.store %cst, %2 : !cc.ptr> + %3 = quake.alloca !quake.veq<1> + %4 = quake.init_state %3, %1 : (!quake.veq<1>, !cc.ptr>) -> !quake.veq<1> + return +} + +// CHECK-LABEL: func.func @test_complex_array() { +// CHECK: %[[VAL_0:.*]] = cc.const_array {{\[}}[1.000000e+00 : f32, 0.000000e+00 : f32], [0.000000e+00 : f32, 1.000000e+00 : f32]{{\]}} : !cc.array x 2> +// CHECK: %[[VAL_1:.*]] = cc.alloca !cc.array x 2> +// CHECK: cc.store %[[VAL_0]], %[[VAL_1]] : !cc.ptr x 2>> +// CHECK: %[[VAL_2:.*]] = cc.cast %[[VAL_1]] : (!cc.ptr x 2>>) -> !cc.ptr> +// CHECK: %[[VAL_3:.*]] = quake.alloca !quake.veq<1> +// CHECK: %[[VAL_4:.*]] = quake.init_state %[[VAL_3]], %[[VAL_2]] : (!quake.veq<1>, !cc.ptr>) -> !quake.veq<1> +// CHECK: return +// CHECK: }