[core] Update quake.apply operation.

Add a verifier and diagnostic message tests.

Fix bugs in apply-op-specialization and canonicalization. Add tests
and handling for !quake.struq while we're here.

Signed-off-by: Eric Schweitz <eschweitz@nvidia.com>

Author: schweitzpgi

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

@@ -397,27 +397,29 @@ def quake_ApplyOp : QuakeOp<"apply",
     Variadic<cc_CallableType>:$indirect_callee, // must be 0 or 1 element
     UnitAttr:$is_adj,
     Variadic<AnyQType>:$controls,
-    Variadic<AnyType>:$args
+    Variadic<AnyType>:$actuals
   );
   let results = (outs Variadic<AnyType>);
 
   let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
+
   let builders = [
     OpBuilder<(ins "mlir::TypeRange":$retTy,
                    "mlir::SymbolRefAttr":$callee,
                    "mlir::UnitAttr":$is_adj,
                    "mlir::ValueRange":$controls,
                    "mlir::ValueRange":$args), [{
-      return build($_builder, $_state, retTy, callee, mlir::ValueRange{},
-                   is_adj, controls, args);
+      return build($_builder, $_state, retTy, callee, {}, is_adj, controls,
+                   args);
     }]>,
     OpBuilder<(ins "mlir::TypeRange":$retTy,
                    "mlir::SymbolRefAttr":$callee,
                    "bool":$is_adj,
                    "mlir::ValueRange":$controls,
                    "mlir::ValueRange":$args), [{
-      return build($_builder, $_state, retTy, callee, mlir::ValueRange{},
-                   is_adj, controls, args);
+      return build($_builder, $_state, retTy, callee, {}, is_adj, controls,
+                   args);
     }]>,
     OpBuilder<(ins "mlir::TypeRange":$retTy,
                    "mlir::Value":$callable,
@@ -446,7 +448,7 @@ def quake_ApplyOp : QuakeOp<"apply",
     operand_range getArgOperands() {
       if (getControls().empty())
         return {operand_begin(), operand_end()};
-      return {getArgs().begin(), getArgs().end()};
+      return {getActuals().begin(), getActuals().end()};
     }
 
     bool applyToVariant() {

lib/Optimizer/CodeGen/TranslateToOpenQASM.cpp

@@ -176,7 +176,7 @@ static LogicalResult emitOperation(Emitter &emitter, quake::ApplyOp op) {
   // Separate classical and quantum arguments.
   SmallVector<Value> parameters;
   SmallVector<Value> targets;
-  for (auto arg : op.getArgs()) {
+  for (auto arg : op.getActuals()) {
     if (isa<quake::RefType, quake::VeqType>(arg.getType()))
       targets.push_back(arg);
     else

lib/Optimizer/Dialect/Quake/CanonicalPatterns.inc

@@ -316,25 +316,45 @@ struct ConcatSizePattern : public OpRewritePattern<quake::ConcatOp> {
     if (concat.getType().hasSpecifiedSize())
       return failure();
 
+    auto *ctx = rewriter.getContext();
+    auto loc = concat.getLoc();
+    auto refTy = quake::RefType::get(ctx);
     // Walk the arguments and sum them, if possible.
     std::size_t sum = 0;
+    SmallVector<Value> targets;
     for (auto opnd : concat.getTargets()) {
       if (auto veqTy = dyn_cast<quake::VeqType>(opnd.getType())) {
         if (!veqTy.hasSpecifiedSize())
           return failure();
         sum += veqTy.getSize();
+        targets.push_back(opnd);
         continue;
       }
-      assert(isa<quake::RefType>(opnd.getType()));
-      sum++;
+      if (auto stqTy = dyn_cast<quake::StruqType>(opnd.getType())) {
+        if (!stqTy.hasSpecifiedSize())
+          return failure();
+        sum += quake::getAllocationSize(stqTy);
+        auto arity = stqTy.getArity();
+        if (*arity) {
+          // Get each member for IR legalization.
+          for (auto [i, memTy] : llvm::enumerate(stqTy.getMembers())) {
+            auto mem = rewriter.create<quake::GetMemberOp>(loc, memTy, opnd, i);
+            targets.push_back(mem);
+          }
+        }
+        continue;
+      }
+      if (opnd.getType() == refTy) {
+        sum++;
+        targets.push_back(opnd);
+        continue;
+      }
+      return failure();
     }
 
     // Leans into the relax_size canonicalization pattern.
-    auto *ctx = rewriter.getContext();
-    auto loc = concat.getLoc();
     auto newTy = quake::VeqType::get(ctx, sum);
-    Value newOp =
-        rewriter.create<quake::ConcatOp>(loc, newTy, concat.getTargets());
+    Value newOp = rewriter.create<quake::ConcatOp>(loc, newTy, targets);
     auto noSizeTy = quake::VeqType::getUnsized(ctx);
     rewriter.replaceOpWithNewOp<quake::RelaxSizeOp>(concat, noSizeTy, newOp);
     return success();

lib/Optimizer/Dialect/Quake/QuakeOps.cpp

@@ -201,6 +201,55 @@ quake::InitializeStateOp quake::AllocaOp::getInitializedState() {
 // Apply
 //===----------------------------------------------------------------------===//
 
+LogicalResult quake::ApplyOp::verify() {
+  FunctionType asSig;
+  if (auto callee = getCallee()) {
+    auto fn =
+        SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(*this, *callee);
+    if (!fn)
+      return emitOpError("callee must be declared");
+    asSig = fn.getFunctionType();
+  } else {
+    Value callable = getIndirectCallee().front();
+    asSig = cast<cudaq::cc::CallableType>(callable.getType()).getSignature();
+  }
+
+  // Arity of callee's signature must be equal to number of arguments provided.
+  bool callingCallable = false;
+  if (getActuals().size() == asSig.getInputs().size() + 1) {
+    callingCallable = true;
+    if (!isa<cudaq::cc::CallableType>(getActuals().front().getType()))
+      return emitOpError("hidden argument must be callable");
+  } else if (getActuals().size() != asSig.getInputs().size()) {
+    return emitOpError("number of arguments must be consistent");
+  }
+
+  // Quantum reference type values are allowed to implicitly coerce to a relaxed
+  // veq type when they appear as arguments to a `quake.apply` op. Specifically,
+  // lowering the apply op is required to add a `quake.concat` op to manifest
+  // the type conversion.
+  auto isRelaxedVeq = [](Type ty1, Type ty2) {
+    if (auto veq2 = dyn_cast<quake::VeqType>(ty2))
+      return quake::isQuantumReferenceType(ty1) && !veq2.hasSpecifiedSize();
+    return false;
+  };
+
+  SmallVector<Type> actualTypes{getActuals().getTypes().begin() +
+                                    (callingCallable ? 1 : 0),
+                                getActuals().getTypes().end()};
+  // The args are the formal arguments and they must match.
+  for (auto [ty1, ty2] : llvm::zip(actualTypes, asSig.getInputs()))
+    if (ty1 != ty2 && !isRelaxedVeq(ty1, ty2))
+      return emitOpError("argument types must match");
+
+  // The results are the formal results and they must match.
+  for (auto [ty1, ty2] : llvm::zip(getResultTypes(), asSig.getResults()))
+    if (ty1 != ty2 && !isRelaxedVeq(ty1, ty2))
+      return emitOpError("result types must match");
+
+  return success();
+}
+
 void quake::ApplyOp::print(OpAsmPrinter &p) {
   if (getIsAdj())
     p << "<adj>";
@@ -213,7 +262,7 @@ void quake::ApplyOp::print(OpAsmPrinter &p) {
   p << ' ';
   if (!getControls().empty())
     p << '[' << getControls() << "] ";
-  p << getArgs() << " : ";
+  p << getActuals() << " : ";
   SmallVector<Type> operandTys{(*this)->getOperandTypes().begin(),
                                (*this)->getOperandTypes().end()};
   p.printFunctionalType(ArrayRef<Type>{operandTys}.drop_front(isDirect ? 0 : 1),
@@ -229,13 +278,14 @@ ParseResult quake::ApplyOp::parse(OpAsmParser &parser, OperationState &result) {
       return failure();
     result.addAttribute("is_adj", parser.getBuilder().getUnitAttr());
   }
+  OpAsmParser::UnresolvedOperand calleeOpnd;
   SmallVector<OpAsmParser::UnresolvedOperand> calleeOperand;
-  if (parser.parseOperandList(calleeOperand))
-    return failure();
-  bool isDirect = calleeOperand.empty();
-  if (calleeOperand.size() > 1)
-    return failure();
-  if (isDirect) {
+  bool isDirect;
+  if (parser.parseOptionalOperand(calleeOpnd).has_value()) {
+    isDirect = false;
+    calleeOperand.push_back(calleeOpnd);
+  } else {
+    isDirect = true;
     NamedAttrList attrs;
     SymbolRefAttr funcAttr;
     if (parser.parseCustomAttributeWithFallback(

lib/Optimizer/Transforms/ApplyOpSpecialization.cpp

@@ -105,15 +105,15 @@ struct ApplyOpAnalysis {
   void performAnalysis(Operation *op) {
     op->walk([&](quake::ApplyOp apply) {
       if (constProp) {
-        // If some of the arguments in getArgs() are constants, then materialize
-        // those constants in a clone of the variant. The specialized variant
-        // will then be able to perform better constant propagation even if not
-        // inlined.
+        // If some of the arguments in getActuals() are constants, then
+        // materialize those constants in a clone of the variant. The
+        // specialized variant will then be able to perform better constant
+        // propagation even if not inlined.
         auto calleeName = apply.getCallee()->getRootReference().str();
         if (func::FuncOp genericFunc =
                 module.lookupSymbol<func::FuncOp>(calleeName)) {
           SmallVector<Value> newArgs;
-          newArgs.append(apply.getArgs().begin(), apply.getArgs().end());
+          newArgs.append(apply.getActuals().begin(), apply.getActuals().end());
           IRMapping mapper;
           SmallVector<Value> preservedArgs;
           SmallVector<Type> inputTys;
@@ -314,8 +314,12 @@ struct ApplyOpPattern : public OpRewritePattern<quake::ApplyOp> {
   LogicalResult matchAndRewrite(quake::ApplyOp apply,
                                 PatternRewriter &rewriter) const override {
     std::string calleeOrigName;
-    if (apply.getCallee()) {
-      calleeOrigName = apply.getCallee()->getRootReference().str();
+    FunctionType calleeSignature;
+    if (auto callee = apply.getCallee()) {
+      calleeOrigName = callee->getRootReference().str();
+      auto fn =
+          SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(apply, *callee);
+      calleeSignature = fn.getFunctionType();
     } else {
       // Check if the first argument is a func.ConstantOp.
       auto calleeVals = apply.getIndirectCallee();
@@ -326,27 +330,31 @@ struct ApplyOpPattern : public OpRewritePattern<quake::ApplyOp> {
       if (!fc)
         return failure();
       calleeOrigName = fc.getValue().str();
+      calleeSignature = dyn_cast<FunctionType>(fc.getResult().getType());
     }
     auto calleeName = getVariantFunctionName(apply, calleeOrigName);
     auto *ctx = apply.getContext();
-    auto consTy = quake::VeqType::getUnsized(ctx);
+    auto unsizedVeqTy = quake::VeqType::getUnsized(ctx);
     SmallVector<Value> newArgs;
     if (!apply.getControls().empty()) {
-      auto consOp = rewriter.create<quake::ConcatOp>(apply.getLoc(), consTy,
-                                                     apply.getControls());
+      auto consOp = rewriter.create<quake::ConcatOp>(
+          apply.getLoc(), unsizedVeqTy, apply.getControls());
       newArgs.push_back(consOp);
     }
-    if (constProp) {
-      for (auto v : apply.getArgs()) {
-        if (auto c = v.getDefiningOp<arith::ConstantOp>())
-          continue;
-        newArgs.emplace_back(v);
-      }
-    } else {
-      newArgs.append(apply.getArgs().begin(), apply.getArgs().end());
+    for (auto [v, toTy] :
+         llvm::zip(apply.getActuals(), calleeSignature.getInputs())) {
+      if (constProp && v.getDefiningOp<arith::ConstantOp>())
+        continue;
+      Value arg = v;
+      if (arg.getType() != toTy)
+        arg =
+            rewriter.create<quake::ConcatOp>(apply.getLoc(), unsizedVeqTy, arg);
+      newArgs.emplace_back(arg);
     }
-    rewriter.replaceOpWithNewOp<func::CallOp>(apply, apply.getResultTypes(),
-                                              calleeName, newArgs);
+    LLVM_DEBUG(llvm::dbgs() << "replacing: " << apply << '\n');
+    [[maybe_unused]] auto result = rewriter.replaceOpWithNewOp<func::CallOp>(
+        apply, apply.getResultTypes(), calleeName, newArgs);
+    LLVM_DEBUG(llvm::dbgs() << "with " << result << '\n');
     return success();
   }
 
@@ -363,16 +371,18 @@ struct FoldCallable : public OpRewritePattern<quake::ApplyOp> {
       return failure();
 
     Value ind = apply.getIndirectCallee()[0];
-    if (auto callee = ind.getDefiningOp<cudaq::cc::InstantiateCallableOp>()) {
-      auto sym = callee.getCallee();
-      SmallVector<Value> newArguments = {ind};
-      newArguments.append(apply.getArgs().begin(), apply.getArgs().end());
-      rewriter.replaceOpWithNewOp<quake::ApplyOp>(
-          apply, apply.getResultTypes(), sym, apply.getIsAdj(),
-          apply.getControls(), newArguments);
-      return success();
-    }
-    return failure();
+    auto callee = ind.getDefiningOp<cudaq::cc::InstantiateCallableOp>();
+    if (!callee)
+      return failure();
+    auto sym = callee.getCallee();
+    SmallVector<Value> newArguments = {ind};
+    newArguments.append(apply.getActuals().begin(), apply.getActuals().end());
+    LLVM_DEBUG(llvm::dbgs() << "replacing " << apply << '\n');
+    [[maybe_unused]] auto result = rewriter.replaceOpWithNewOp<quake::ApplyOp>(
+        apply, apply.getResultTypes(), sym, apply.getIsAdj(),
+        apply.getControls(), newArguments);
+    LLVM_DEBUG(llvm::dbgs() << "as " << result << '\n');
+    return success();
   }
 };
 
@@ -529,7 +539,7 @@ class ApplySpecializationPass
                            apply.getControls().end());
         auto newApply = builder.create<quake::ApplyOp>(
             apply.getLoc(), apply.getResultTypes(), apply.getCalleeAttr(),
-            apply.getIsAdjAttr(), newControls, apply.getArgs());
+            apply.getIsAdjAttr(), newControls, apply.getActuals());
         apply->replaceAllUsesWith(newApply.getResults());
         apply->erase();
       } else if (isQuantumKernelCall(op)) {

lib/Optimizer/Transforms/PySynthCallableBlockArgs.cpp

@@ -118,7 +118,7 @@ class UpdateQuakeApplyOp : public OpConversionPattern<quake::ApplyOp> {
 
       rewriter.replaceOpWithNewOp<quake::ApplyOp>(
           op, TypeRange{}, FlatSymbolRefAttr::get(ctx, replacement.getName()),
-          adaptor.getIsAdj(), adaptor.getControls(), adaptor.getArgs());
+          adaptor.getIsAdj(), adaptor.getControls(), adaptor.getActuals());
       return success();
     }
     return failure();

test/Transforms/apply-0.qke

@@ -32,3 +32,57 @@ func.func @do_apply(%arg : !quake.ref, %brg : !quake.ref) {
 // CHECK:           quake.t<adj> [%[[VAL_0]]] %[[VAL_1]] : (!quake.veq<?>, !quake.ref) -> ()
 // CHECK:           return
 // CHECK:         }
+
+// The following test the implicit coercion of apply operands to unsized veq
+// type. This can be done uniformly by injecting a concat operation.
+
+func.func @utility_gate(%arg : !quake.veq<?>) {
+  quake.y %arg : (!quake.veq<?>) -> ()
+  quake.t %arg : (!quake.veq<?>) -> ()
+  return
+}
+
+func.func @apply_yski() {
+  %0 = quake.alloca !quake.ref
+  quake.apply @utility_gate %0 : (!quake.ref) -> ()
+  quake.dealloc %0 : !quake.ref
+  return
+}
+
+// CHECK-LABEL:   func.func @apply_yski() {
+// CHECK:           %[[VAL_0:.*]] = quake.alloca !quake.ref
+// CHECK:           %[[VAL_1:.*]] = quake.concat %[[VAL_0]] : (!quake.ref) -> !quake.veq<?>
+// CHECK:           call @utility_gate(%[[VAL_1]]) : (!quake.veq<?>) -> ()
+// CHECK:           quake.dealloc %[[VAL_0]] : !quake.ref
+// CHECK:           return
+// CHECK:         }
+
+func.func @apply_kaksi() {
+  %0 = quake.alloca !quake.veq<4>
+  quake.apply @utility_gate %0 : (!quake.veq<4>) -> ()
+  quake.dealloc %0 : !quake.veq<4>
+  return
+}
+
+// CHECK-LABEL:   func.func @apply_kaksi() {
+// CHECK:           %[[VAL_0:.*]] = quake.alloca !quake.veq<4>
+// CHECK:           %[[VAL_1:.*]] = quake.concat %[[VAL_0]] : (!quake.veq<4>) -> !quake.veq<?>
+// CHECK:           call @utility_gate(%[[VAL_1]]) : (!quake.veq<?>) -> ()
+// CHECK:           quake.dealloc %[[VAL_0]] : !quake.veq<4>
+// CHECK:           return
+// CHECK:         }
+
+func.func @apply_kolme() {
+  %0 = quake.alloca !quake.struq<!quake.ref, !quake.ref>
+  quake.apply @utility_gate %0 : (!quake.struq<!quake.ref, !quake.ref>) -> ()
+  quake.dealloc %0 : !quake.struq<!quake.ref, !quake.ref>
+  return
+}
+
+// CHECK-LABEL:   func.func @apply_kolme() {
+// CHECK:           %[[VAL_0:.*]] = quake.alloca !quake.struq<!quake.ref, !quake.ref>
+// CHECK:           %[[VAL_1:.*]] = quake.concat %[[VAL_0]] : (!quake.struq<!quake.ref, !quake.ref>) -> !quake.veq<?>
+// CHECK:           call @utility_gate(%[[VAL_1]]) : (!quake.veq<?>) -> ()
+// CHECK:           quake.dealloc %[[VAL_0]] : !quake.struq<!quake.ref, !quake.ref>
+// CHECK:           return
+// CHECK:         }

test/Transforms/apply_ctrl_veq_type.qke

@@ -6,14 +6,14 @@
 // the terms of the Apache License 2.0 which accompanies this distribution.   //
 // ========================================================================== //
 
-// Regression test for a bug where veq type specialization caused type mismatches
-// when creating control variants of functions that have inner calls to other
-// kernels expecting dynamic !quake.veq<?> types.
+// Regression test for a bug where veq type specialization caused type
+// mismatches when creating control variants of functions that have inner calls
+// to other kernels expecting dynamic !quake.veq<?> types.
 //
 // The bug scenario:
 // 1. main_kernel allocates !quake.veq<2> and relaxes it to !quake.veq<?>
 // 2. A controlled apply to thunk_kernel with a constant argument triggers
-//    constant propagation which would un-relax the veq type back to !quake.veq<2>
+// constant propagation which would un-relax the veq type back to !quake.veq<2>
 // 3. The thunk calls inner_kernel which expects !quake.veq<?>
 // 4. When the .ctrl variant is created, the inner call has a type mismatch
 //