[quake, cc] Improved alloca ops

include/cudaq/Optimizer/Dialect/CC/CCOps.td

@@ -688,14 +688,22 @@ def cc_AllocaOp : CCOp<"alloca", [
 
   let arguments = (ins
     TypeAttr:$elementType,
-    Optional<AnyInteger>:$seqSize
+    Optional<AnySignlessInteger>:$seqSize
   );
   let results = (outs cc_PointerType:$address);
 
-  let assemblyFormat = [{
-    $elementType (`[` $seqSize^ `]`)? `:` functional-type(operands, results)
-      attr-dict
-  }];
+  let hasFolder = 1;
+  let hasCustomAssemblyFormat = 1;
+  let builders = [
+    OpBuilder<(ins "mlir::Type":$elementType, "mlir::Value":$seqSize), [{
+      auto resTy = cudaq::cc::PointerType::get(seqSize ?
+        cudaq::cc::ArrayType::get(elementType) : elementType);
+      return build($_builder, $_state, resTy, elementType, seqSize);
+    }]>,
+    OpBuilder<(ins "mlir::Type":$elementType), [{
+      return build($_builder, $_state, elementType, mlir::Value{});
+    }]>
+  ];
 }
 
 def cc_LoadOp : CCOp<"load",

include/cudaq/Optimizer/Dialect/CC/CCTypes.td

@@ -108,6 +108,12 @@ def cc_ArrayType : CCType<"Array", "array"> {
 
     bool isUnknownSize() const { return getSize() == unknownSize; }
   }];
+
+  let builders = [
+    TypeBuilderWithInferredContext<(ins "mlir::Type":$elementType), [{
+      return Base::get(elementType.getContext(), elementType, unknownSize);
+    }]>,
+  ];
 }
 
 //===----------------------------------------------------------------------===//

include/cudaq/Optimizer/Dialect/Quake/QuakeOps.td

@@ -80,11 +80,14 @@ def quake_AllocaOp : QuakeOp<"alloca", [MemoryEffects<[MemAlloc, MemWrite]>]> {
     }]>,
     OpBuilder<(ins "size_t":$size), [{
       return build($_builder, $_state, $_builder.getType<VeqType>(size), {});
+    }]>,
+    OpBuilder<(ins "mlir::Type":$ty), [{
+      return build($_builder, $_state, ty, {});
     }]>
   ];
 
   let assemblyFormat = [{
-    (`[` $size^ `:` type($size)`]`)? qualified(type($ref_or_vec)) attr-dict
+    qualified(type($ref_or_vec)) (`[` $size^ `:` type($size) `]`)? attr-dict
   }];
 
   let hasCanonicalizer = 1;

lib/Frontend/nvqpp/ConvertDecl.cpp

@@ -253,7 +253,10 @@ bool QuakeBridgeVisitor::VisitVarDecl(clang::VarDecl *x) {
       // this is a qreg<N> q;
       auto qregSizeVal = builder.create<mlir::arith::ConstantIntOp>(
           loc, qregSize, builder.getIntegerType(64));
-      qreg = builder.create<quake::AllocaOp>(loc, qType, qregSizeVal);
+      if (qregSize != 0)
+        qreg = builder.create<quake::AllocaOp>(loc, qType);
+      else
+        qreg = builder.create<quake::AllocaOp>(loc, qType, qregSizeVal);
     }
     symbolTable.insert(name, qreg);
     // allocated_qreg_names.push_back(name);
@@ -267,12 +270,10 @@ bool QuakeBridgeVisitor::VisitVarDecl(clang::VarDecl *x) {
       symbolTable.insert(name, val);
       return pushValue(val);
     }
-    auto qregSizeVal = builder.create<mlir::arith::ConstantIntOp>(
-        loc, 1, builder.getIntegerType(64));
     auto zero = builder.create<mlir::arith::ConstantIntOp>(
         loc, 0, builder.getIntegerType(64));
     auto qregSizeOne = builder.create<quake::AllocaOp>(
-        loc, quake::VeqType::get(builder.getContext(), 1), qregSizeVal);
+        loc, quake::VeqType::get(builder.getContext(), 1));
     Value addressTheQubit =
         builder.create<quake::ExtractRefOp>(loc, qregSizeOne, zero);
     symbolTable.insert(name, addressTheQubit);

lib/Frontend/nvqpp/ConvertExpr.cpp

@@ -1718,11 +1718,15 @@ bool QuakeBridgeVisitor::VisitCXXConstructExpr(clang::CXXConstructExpr *x) {
           auto sizeVal = popValue();
           if (isa<quake::VeqType>(sizeVal.getType()))
             return pushValue(sizeVal);
+          if (regTy.hasSpecifiedSize())
+            return pushValue(builder.create<quake::AllocaOp>(loc, regTy));
           return pushValue(
               builder.create<quake::AllocaOp>(loc, regTy, sizeVal));
         }
         auto qregSizeVal = builder.create<arith::ConstantIntOp>(
             loc, regTy.getSize(), builder.getIntegerType(64));
+        if (regTy.hasSpecifiedSize())
+          return pushValue(builder.create<quake::AllocaOp>(loc, regTy));
         return pushValue(
             builder.create<quake::AllocaOp>(loc, regTy, qregSizeVal));
       }

lib/Optimizer/Dialect/CC/CCOps.cpp

@@ -1,10 +1,10 @@
-/*************************************************************** -*- C++ -*- ***
+/****************************************************************-*- C++ -*-****
  * Copyright (c) 2022 - 2023 NVIDIA Corporation & Affiliates.                  *
  * All rights reserved.                                                        *
  *                                                                             *
  * This source code and the accompanying materials are made available under    *
  * the terms of the Apache License 2.0 which accompanies this distribution.    *
- *******************************************************************************/
+ ******************************************************************************/
 
 #include "cudaq/Optimizer/Dialect/CC/CCOps.h"
 #include "cudaq/Optimizer/Builder/Factory.h"
@@ -45,6 +45,62 @@ cudaq::cc::AddressOfOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// AllocaOp
+//===----------------------------------------------------------------------===//
+
+void cudaq::cc::AllocaOp::print(OpAsmPrinter &p) {
+  p << ' ' << getElementType();
+  if (auto size = getSeqSize())
+    p << '[' << size << " : " << size.getType() << ']';
+}
+
+ParseResult cudaq::cc::AllocaOp::parse(OpAsmParser &parser,
+                                       OperationState &result) {
+  Type eleTy;
+  if (parser.parseType(eleTy))
+    return failure();
+  result.addAttribute("elementType", TypeAttr::get(eleTy));
+  Type resTy;
+  if (succeeded(parser.parseOptionalLSquare())) {
+    OpAsmParser::UnresolvedOperand operand;
+    Type operTy;
+    if (parser.parseOperand(operand) || parser.parseColonType(operTy) ||
+        parser.parseRSquare() ||
+        parser.resolveOperand(operand, operTy, result.operands))
+      return failure();
+    resTy = cc::PointerType::get(cc::ArrayType::get(eleTy));
+  } else {
+    resTy = cc::PointerType::get(eleTy);
+  }
+  if (!resTy || parser.parseOptionalAttrDict(result.attributes) ||
+      parser.addTypeToList(resTy, result.types))
+    return failure();
+  return success();
+}
+
+OpFoldResult cudaq::cc::AllocaOp::fold(ArrayRef<Attribute> params) {
+  if (params.size() == 1) {
+    // If allocating a contiguous block of elements and the size of the block is
+    // a constant, fold the size into the cc.array type and allocate a constant
+    // sized block.
+    if (auto intAttr = dyn_cast_or_null<IntegerAttr>(params[0])) {
+      auto size = intAttr.getInt();
+      if (size > 0) {
+        auto resTy = cast<cc::ArrayType>(
+            cast<cc::PointerType>(getType()).getElementType());
+        auto arrTy = cc::ArrayType::get(resTy.getContext(),
+                                        resTy.getElementType(), size);
+        getOperation()->setAttr("elementType", TypeAttr::get(arrTy));
+        getResult().setType(cc::PointerType::get(arrTy));
+        getOperation()->eraseOperand(0);
+        return getResult();
+      }
+    }
+  }
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//

lib/Optimizer/Dialect/Quake/QuakeOps.cpp

@@ -49,35 +49,40 @@ Value quake::createConstantAlloca(PatternRewriter &builder, Location loc,
       loc, quake::VeqType::getUnsized(builder.getContext()), newAlloca);
 }
 
-void quake::AllocaOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
-                                                  MLIRContext *context) {
-  patterns.add<FuseConstantToAllocaPattern>(context);
-}
-
 LogicalResult quake::AllocaOp::verify() {
-  auto resultType = dyn_cast<VeqType>(getResult().getType());
-  if (auto size = getSize()) {
-    std::int64_t argSize = 0;
-    if (auto cnt = dyn_cast_or_null<arith::ConstantOp>(size.getDefiningOp())) {
-      argSize = cnt.getValue().cast<IntegerAttr>().getInt();
-      // TODO: This is a questionable check. We could have a very large unsigned
-      // value that appears to be negative because of two's complement. On the
-      // other hand, allocating 2^64 - 1 qubits isn't going to go well.
-      if (argSize < 0)
-        return emitOpError("expected a non-negative integer size.");
+  // Result must be RefType or VeqType by construction.
+  if (auto resTy = dyn_cast<VeqType>(getResult().getType())) {
+    if (resTy.hasSpecifiedSize()) {
+      if (getSize())
+        return emitOpError("unexpected size operand");
+    } else {
+      if (auto size = getSize()) {
+        if (auto cnt =
+                dyn_cast_or_null<arith::ConstantOp>(size.getDefiningOp())) {
+          std::int64_t argSize = cnt.getValue().cast<IntegerAttr>().getInt();
+          // TODO: This is a questionable check. We could have a very large
+          // unsigned value that appears to be negative because of two's
+          // complement. On the other hand, allocating 2^64 - 1 qubits isn't
+          // going to go well.
+          if (argSize < 0)
+            return emitOpError("expected a non-negative integer size.");
+        }
+      } else {
+        return emitOpError("size operand required");
+      }
     }
-    if (!resultType)
-      return emitOpError(
-          "must return a vector of qubits since a size was provided.");
-    if (resultType.hasSpecifiedSize() &&
-        (static_cast<std::size_t>(argSize) != resultType.getSize()))
-      return emitOpError("expected operand size to match VeqType size.");
-  } else if (resultType && !resultType.hasSpecifiedSize()) {
-    return emitOpError("must return a veq with known size.");
   }
   return success();
 }
 
+void quake::AllocaOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+                                                  MLIRContext *context) {
+  // Use a canonicalization pattern as folding the constant into the veq type
+  // changes the type. Uses may still expect a veq with unspecified size.
+  // Folding is strictly reductive and doesn't allow the creation of ops.
+  patterns.add<FuseConstantToAllocaPattern>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // ExtractRef
 //===----------------------------------------------------------------------===//

python/tests/compiler/adjoint.py

@@ -280,7 +280,7 @@ def test_sample_adjoint_qreg():
 # CHECK-SAME:      %[[VAL_0:.*]]: i32) {
 # CHECK:           %[[VAL_1:.*]] = arith.constant 1 : index
 # CHECK:           %[[VAL_2:.*]] = arith.constant 0 : index
-# CHECK:           %[[VAL_3:.*]] = quake.alloca[%[[VAL_0]] : i32] !quake.veq<?>
+# CHECK:           %[[VAL_3:.*]] = quake.alloca !quake.veq<?>[%[[VAL_0]] : i32]
 # CHECK:           %[[VAL_4:.*]] = quake.vec_size %[[VAL_3]] : (!quake.veq<?>) -> i64
 # CHECK:           %[[VAL_5:.*]] = arith.index_cast %[[VAL_4]] : i64 to index
 # CHECK:           %[[VAL_6:.*]] = cc.loop while ((%[[VAL_7:.*]] = %[[VAL_2]]) -> (index)) {

python/tests/compiler/qalloc.py

@@ -83,7 +83,7 @@ def test_kernel_qalloc_quake_val():
 
 # CHECK-LABEL:   func.func @__nvqpp__mlirgen____nvqppBuilderKernel_{{.*}}(
 # CHECK-SAME:      %[[VAL_0:.*]]: i32) {
-# CHECK:           %[[VAL_1:.*]] = quake.alloca[%[[VAL_0]] : i32] !quake.veq<?>
+# CHECK:           %[[VAL_1:.*]] = quake.alloca !quake.veq<?>[%[[VAL_0]] : i32]
 # CHECK:           return
 # CHECK:         }
 

test/AST-Quake/adjoint-1.cpp

@@ -33,8 +33,7 @@ struct ep {
 };
 
 // CHECK-LABEL:   func.func @__nvqpp__mlirgen__ep()
-// CHECK:           %[[VAL_2:.*]] = arith.constant 3 : i64
-// CHECK:           %[[VAL_3:.*]] = quake.alloca[%[[VAL_2]] : i64] !quake.veq<3>
+// CHECK:           %[[VAL_3:.*]] = quake.alloca !quake.veq<3>
 // CHECK:           %[[VAL_4:.*]] = quake.relax_size %[[VAL_3]] : (!quake.veq<3>) -> !quake.veq<?>
 // CHECK:           quake.apply<adj> @__nvqpp__mlirgen__k %[[VAL_4]] : (!quake.veq<?>) -> ()
 // CHECK:           return

test/AST-Quake/adjoint-3.cpp

@@ -48,7 +48,7 @@ struct QernelZero {
 // CHECK-LABEL:   func.func @__nvqpp__mlirgen__run_circuit
 // CHECK-SAME:        (%{{.*}}: i32, %{{.*}}: i32, %{{.*}}: f64)
 // CHECK:           %[[VAL_5:.*]] = memref.alloca() : memref<f64>
-// CHECK:           %[[VAL_10:.*]] = quake.alloca[%{{.*}} : i64] !quake.veq<?>
+// CHECK:           %[[VAL_10:.*]] = quake.alloca !quake.veq<?>[%{{.*}} : i64]
 // CHECK:           %[[VAL_16:.*]] = memref.load %[[VAL_5]][] : memref<f64>
 // CHECK:           call @__nvqpp__mlirgen__statePrep_A{{.*}}(%[[VAL_10]], %[[VAL_16]]) : (!quake.veq<?>, f64) -> ()
 // CHECK:           cc.scope {

test/AST-Quake/callable-1.cpp

@@ -48,7 +48,7 @@ int main() {
 // CHECK-SAME:        (%[[VAL_0:.*]]: !cc.lambda<(!quake.veq<?>) -> ()>) attributes {{{.*}}"cudaq-entrypoint"{{.*}}} {
 // CHECK:           %[[VAL_1:.*]] = arith.constant 2 : i32
 // CHECK:           %[[VAL_2:.*]] = arith.extsi %[[VAL_1]] : i32 to i64
-// CHECK:           %[[VAL_3:.*]] = quake.alloca[%[[VAL_2]] : i64] !quake.veq<?>
+// CHECK:           %[[VAL_3:.*]] = quake.alloca !quake.veq<?>[%[[VAL_2]] : i64]
 // CHECK:           call @__nvqpp__mlirgen__Z4mainE3$_0(%[[VAL_3]]) : (!quake.veq<?>) -> ()
 // CHECK:           return
 // CHECK:         }

test/AST-Quake/compute_action-2.cpp

@@ -72,7 +72,7 @@ struct ctrlHeisenberg {
 // CHECK:           %[[VAL_2:.*]] = quake.alloca !quake.ref
 // CHECK:           %[[VAL_3:.*]] = memref.load %[[VAL_1]][] : memref<i32>
 // CHECK:           %[[VAL_4:.*]] = arith.extsi %[[VAL_3]] : i32 to i64
-// CHECK:           %[[VAL_5:.*]] = quake.alloca[%[[VAL_4]] : i64] !quake.veq<?>
+// CHECK:           %[[VAL_5:.*]] = quake.alloca !quake.veq<?>[%[[VAL_4]] : i64]
 // CHECK:           %[[VAL_6:.*]] = quake.concat %[[VAL_2]] : (!quake.ref) -> !quake.veq<?>
 // CHECK:           call @__nvqpp__mlirgen__function_magic_func.{{.*}}.ctrl(%[[VAL_6]], %[[VAL_5]]) : (!quake.veq<?>, !quake.veq<?>) -> ()
 // CHECK:           return

test/AST-Quake/ctrl_vector.cpp

@@ -26,9 +26,9 @@ struct lower_ctrl_as_qreg {
 // CHECK-LABEL:   func.func @__nvqpp__mlirgen__lower_ctrl_as_qreg
 // CHECK-SAME: () attributes {{{.*}}"cudaq-entrypoint"{{.*}}} {
 // CHECK:  %[[VAL_0:.*]] = arith.constant 4 : i32
-// CHECK:  %[[VAL_2:.*]] = quake.alloca[%{{.*}} : i64] !quake.veq<?>
+// CHECK:  %[[VAL_2:.*]] = quake.alloca !quake.veq<?>[%{{.*}} : i64]
 // CHECK:  %[[VAL_3:.*]] = arith.constant 2 : i32
-// CHECK:  %[[VAL_5:.*]] = quake.alloca[%{{.*}} : i64] !quake.veq<?>
+// CHECK:  %[[VAL_5:.*]] = quake.alloca !quake.veq<?>[%{{.*}} : i64]
 // CHECK:  %[[VAL_6:.*]] = arith.constant 0 : i32
 // CHECK:  %[[VAL_8:.*]] = quake.extract_ref %[[VAL_5]][%{{.*}}] : (!quake.veq<?>, i64) -> !quake.ref
 // CHECK:  quake.h [%[[VAL_2]]] %[[VAL_8]] : (!quake.veq<?>, !quake.ref) -> ()
@@ -60,7 +60,7 @@ struct test_two_control_call {
 // CHECK:             }
 // CHECK:           } : !cc.lambda<(!quake.ref) -> ()>
 // CHECK:           %[[VAL_4:.*]] = arith.constant 4 : i64
-// CHECK:           %[[VAL_5:.*]] = quake.alloca[%[[VAL_4]] : i64] !quake.veq<4>
+// CHECK:           %[[VAL_5:.*]] = quake.alloca !quake.veq<4>
 // CHECK:           %[[VAL_6:.*]] = quake.alloca !quake.ref
 // CHECK:           quake.apply @__nvqpp__mlirgen__{{.*}}test_two_control_call{{.*}}[%[[VAL_5]]] %[[VAL_6]] : (!quake.veq<4>, !quake.ref) -> ()
 // CHECK:           %[[VAL_7:.*]] = quake.mz %[[VAL_6]] : (!quake.ref) -> i1

test/AST-Quake/dealloc.cpp

@@ -21,11 +21,11 @@ __qpu__ void magic_func(int N) {
 
 // CHECK-LABEL:   func.func @__nvqpp__mlirgen__function_magic_func
 // CHECK:           cc.scope {
-// CHECK:             %[[VAL_4:.*]] = quake.alloca[%{{.*}} : i64] !quake.veq<?>
+// CHECK:             %[[VAL_4:.*]] = quake.alloca !quake.veq<?>[%{{.*}} : i64]
 // CHECK:             quake.dealloc %[[VAL_4]] : !quake.veq<?>
 // CHECK:             cc.continue
 // CHECK:           }
-// CHECK:           %[[VAL_10:.*]] = quake.alloca[%{{.*}} : i64] !quake.veq<?>
+// CHECK:           %[[VAL_10:.*]] = quake.alloca !quake.veq<?>[%{{.*}} : i64]
 // CHECK:           quake.dealloc %[[VAL_10]] : !quake.veq<?>
 // CHECK:           return
 // CHECK:         }

test/AST-Quake/if.cpp

@@ -28,7 +28,7 @@ struct kernel {
 // CHECK:           memref.store %[[VAL_0]], %[[VAL_1]][] : memref<i1>
 // CHECK:           %[[VAL_2:.*]] = arith.constant 2 : i32
 // CHECK:           %[[VAL_3:.*]] = arith.extsi %[[VAL_2]] : i32 to i64
-// CHECK:           %[[VAL_4:.*]] = quake.alloca[%[[VAL_3]] : i64] !quake.veq<?>
+// CHECK:           %[[VAL_4:.*]] = quake.alloca !quake.veq<?>[%[[VAL_3]] : i64]
 // CHECK:           %[[VAL_5:.*]] = memref.load %[[VAL_1]][] : memref<i1>
 // CHECK:           cc.if(%[[VAL_5]]) {
 // CHECK:             cc.scope {
@@ -63,7 +63,7 @@ struct kernel_else {
 // CHECK:           memref.store %[[VAL_0]], %[[VAL_1]][] : memref<i1>
 // CHECK:           %[[VAL_2:.*]] = arith.constant 2 : i32
 // CHECK:           %[[VAL_3:.*]] = arith.extsi %[[VAL_2]] : i32 to i64
-// CHECK:           %[[VAL_4:.*]] = quake.alloca[%[[VAL_3]] : i64] !quake.veq<?>
+// CHECK:           %[[VAL_4:.*]] = quake.alloca !quake.veq<?>[%[[VAL_3]] : i64]
 // CHECK:           %[[VAL_5:.*]] = memref.load %[[VAL_1]][] : memref<i1>
 // CHECK:           cc.if(%[[VAL_5]]) {
 // CHECK:             cc.scope {

test/AST-Quake/lambda_instance.cpp

@@ -22,7 +22,7 @@ struct test0 {
 };
 
 // CHECK-LABEL:   func.func @__nvqpp__mlirgen__test0
-// CHECK:           %[[VAL_2:.*]] = quake.alloca[%{{.*}} : i64] !quake.veq<?>
+// CHECK:           %[[VAL_2:.*]] = quake.alloca !quake.veq<?>[%{{.*}} : i64]
 // CHECK:           %[[VAL_3:.*]] = cc.create_lambda {
 // CHECK:           } : !cc.lambda<(!quake.ref) -> ()>
 // CHECK:           %[[VAL_12:.*]] = quake.extract_ref %[[VAL_2]][%{{.*}}] : (!quake.veq<?>, i64) -> !quake.ref
@@ -46,15 +46,13 @@ struct test1 {
 };
 
 // CHECK-LABEL:   func.func @__nvqpp__mlirgen__test1
-// CHECK-SAME:        () attributes {"cudaq-entrypoint", "cudaq-kernel"} {
-// CHECK:           %[[VAL_3:.*]] = quake.alloca[%{{.*}} : i64] !quake.veq<2>
+// CHECK:           %[[VAL_3:.*]] = quake.alloca !quake.veq<2>
 // CHECK:           %[[VAL_4:.*]] = cc.create_lambda {
 // CHECK:           } : !cc.lambda<(!quake.ref) -> ()>
 // CHECK:           %[[VAL_13:.*]] = quake.extract_ref %[[VAL_3]][%{{.*}}] : (!quake.veq<2>, i64) -> !quake.ref
 // CHECK:           %[[VAL_16:.*]] = quake.extract_ref %[[VAL_3]][%{{.*}}] : (!quake.veq<2>, i64) -> !quake.ref
 // CHECK:           quake.apply @__nvqpp__mlirgen__{{.*}}5test1{{.*}}[%[[VAL_13]]] %[[VAL_16]] : (!quake.ref, !quake.ref) -> ()
 // CHECK:           return
-// CHECK:         }
 
 // CHECK-LABEL:   func.func @__nvqpp__mlirgen__
 // CHECK-SAME:        5test1
@@ -84,7 +82,7 @@ struct test2b {
 };
 
 // CHECK-LABEL:   func.func @__nvqpp__mlirgen__test2b
-// CHECK:           %[[VAL_2:.*]] = quake.alloca[%{{.*}} : i64] !quake.veq<?>
+// CHECK:           %[[VAL_2:.*]] = quake.alloca !quake.veq<?>[%{{.*}} : i64]
 // CHECK:           %[[VAL_4:.*]] = cc.create_lambda {
 // CHECK:           } : !cc.lambda<(!quake.ref) -> ()>
 // CHECK:           call @__nvqpp__mlirgen__instance_test2a{{.*}}(%{{.*}}, %[[VAL_2]]) : (!cc.lambda<(!quake.ref) -> ()>, !quake.veq<?>) -> ()
@@ -130,7 +128,7 @@ struct test2c {
 };
 
 // CHECK-LABEL:   func.func @__nvqpp__mlirgen__test2c
-// CHECK:           %[[VAL_2:.*]] = quake.alloca[%{{.*}} : i64] !quake.veq<?>
+// CHECK:           %[[VAL_2:.*]] = quake.alloca !quake.veq<?>[%{{.*}} : i64]
 // CHECK:           %[[VAL_3:.*]] = cc.create_lambda {
 // CHECK:           } : !cc.lambda<(!quake.ref) -> ()>
 // CHECK:           call @__nvqpp__mlirgen__instance_test2a_c{{.*}}(%{{.*}}, %[[VAL_2]]) : (!cc.lambda<(!quake.ref) -> ()>, !quake.veq<?>) -> ()
@@ -193,7 +191,7 @@ struct test3 {
 // CHECK-SAME:        () attributes {"cudaq-entrypoint", "cudaq-kernel"} {
 // CHECK:           %[[VAL_0:.*]] = arith.constant 2 : i32
 // CHECK:           %[[VAL_1:.*]] = arith.extsi %[[VAL_0]] : i32 to i64
-// CHECK:           %[[VAL_2:.*]] = quake.alloca[%[[VAL_1]] : i64] !quake.veq<?>
+// CHECK:           %[[VAL_2:.*]] = quake.alloca !quake.veq<?>[%[[VAL_1]] : i64]
 // CHECK:           %[[VAL_3:.*]] = cc.create_lambda {
 // CHECK:           ^bb0(%[[VAL_4:.*]]: !quake.ref):
 // CHECK:             cc.scope {