Support measurements of subveqs for braket

include/cudaq/Optimizer/Transforms/Passes.td

@@ -141,6 +141,32 @@ def CheckKernelCalls : Pass<"check-kernel-calls", "mlir::func::FuncOp"> {
   }];
 }
 
+def CombineMeasurements :
+    Pass<"combine-measurements", "mlir::func::FuncOp"> {
+  let summary = "Extends mesurements on subveqs adds output names";
+  let description = [{
+    Replace a pattern such as:
+    ```
+    func.func @kernel() attributes {"cudaq-entrypoint"} {
+      %1 = ... : !quake.veq<4>
+      %2 = quake.subveq %1, %c2, %c3 : (!quake.veq<4>, i32, i32) ->
+            !quake.veq<2>
+      %measOut = quake.mz %2 : (!quake.veq<2>) -> !cc.stdvec<!quake.measure>
+    }
+    ```
+    with:
+    ```
+    func.func @kernel() attributes {"cudaq-entrypoint", ["output_names",
+    "[[[0,[1,\22q0\22]],[1,[2,\22q1\22]]]]"]} {
+      %1 = ... : !quake.veq<4>
+      %measOut = quake.mz %1 : (!quake.veq<4>) -> !cc.stdvec<!quake.measure>
+    }
+    ```
+  }];
+  let dependentDialects = ["cudaq::cc::CCDialect", "quake::QuakeDialect"];
+}
+
+
 def CombineQuantumAllocations :
     Pass<"combine-quantum-alloc", "mlir::func::FuncOp"> {
   let summary = "Combines quake alloca operations.";

lib/Optimizer/CodeGen/Pipelines.cpp

@@ -45,20 +45,7 @@ void cudaq::opt::addPipelineTranslateToOpenQASM(PassManager &pm) {
   pm.addPass(createCanonicalizerPass());
   pm.addPass(createCSEPass());
   pm.addNestedPass<func::FuncOp>(createClassicalMemToReg());
-  pm.addPass(createLoopUnroll());
-  pm.addPass(createCanonicalizerPass());
-  pm.addNestedPass<func::FuncOp>(createLiftArrayAlloc());
-  pm.addPass(createGlobalizeArrayValues());
-  pm.addPass(createStatePreparation());
-  pm.addNestedPass<func::FuncOp>(createGetConcreteMatrix());
-  pm.addPass(createUnitarySynthesis());
   pm.addPass(createSymbolDCEPass());
-  pm.addPass(createCanonicalizerPass());
-  pm.addPass(createCSEPass());
-  pm.addNestedPass<func::FuncOp>(createMultiControlDecompositionPass());
-  pm.addPass(createDecompositionPass(
-      {.enabledPatterns = {"CCZToCX", "RxAdjToRx", "RyAdjToRy", "RzAdjToRz"}}));
-  pm.addPass(createCanonicalizerPass());
 }
 
 void cudaq::opt::addPipelineTranslateToIQMJson(PassManager &pm) {

lib/Optimizer/CodeGen/TranslateToOpenQASM.cpp

@@ -23,13 +23,12 @@ using namespace cudaq;
 //===----------------------------------------------------------------------===//
 // Helper functions
 //===----------------------------------------------------------------------===//
-
 /// Translates operation names into OpenQASM gate names
 static LogicalResult translateOperatorName(quake::OperatorInterface optor,
                                            StringRef &name) {
   StringRef qkeName = optor->getName().stripDialect();
   if (optor.getControls().size() == 0) {
-    name = StringSwitch<StringRef>(qkeName).Case("r1", "u1").Default(qkeName);
+    name = StringSwitch<StringRef>(qkeName).Default(qkeName);
   } else if (optor.getControls().size() == 1) {
     name = StringSwitch<StringRef>(qkeName)
                .Case("h", "ch")

lib/Optimizer/Transforms/CMakeLists.txt

@@ -17,6 +17,7 @@ add_cudaq_library(OptTransforms
   ApplyOpSpecialization.cpp
   ArgumentSynthesis.cpp
   BasisConversion.cpp
+  CombineMeasurements.cpp
   CombineQuantumAlloc.cpp
   ConstPropComplex.cpp
   Decomposition.cpp

lib/Optimizer/Transforms/CombineMeasurements.cpp

@@ -0,0 +1,255 @@
+/*******************************************************************************
+ * Copyright (c) 2022 - 2024 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 "PassDetails.h"
+#include "cudaq/Optimizer/CodeGen/QIRAttributeNames.h"
+#include "cudaq/Optimizer/Dialect/CC/CCOps.h"
+#include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h"
+#include "cudaq/Optimizer/Dialect/Quake/QuakeTypes.h"
+#include "cudaq/Optimizer/Transforms/Passes.h"
+#include "nlohmann/json.hpp"
+#include "llvm/Support/Debug.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/Passes.h"
+
+namespace cudaq::opt {
+#define GEN_PASS_DEF_COMBINEMEASUREMENTS
+#include "cudaq/Optimizer/Transforms/Passes.h.inc"
+} // namespace cudaq::opt
+
+#define DEBUG_TYPE "combine-measurements"
+
+using namespace mlir;
+
+namespace {
+
+// After combine-quantum-alloc, we have one top allocation per function.
+// The following type is used to store qubit mapping from result qubit
+// index to the actual qubit index and register name.
+// map[result] --> [qb,regName]
+// Note: register name is currently not used in `OpenQasm2` backends,
+// so we supply a bogus name.
+using OutputNamesType =
+    std::map<std::size_t, std::pair<std::size_t, std::string>>;
+
+struct Analysis {
+  Analysis() = default;
+  Analysis(const Analysis &) = delete;
+  Analysis(Analysis &&) = delete;
+  Analysis &operator=(const Analysis &) = delete;
+
+  mlir::DenseMap<mlir::Value, std::size_t> measurements;
+  OutputNamesType resultQubitVals;
+  quake::MzOp lastMeasurement;
+
+  bool empty() const { return measurements.empty(); }
+
+  LogicalResult analyze(func::FuncOp func) {
+    quake::AllocaOp qalloc;
+    std::size_t currentOffset = 0;
+
+    for (auto &block : func.getRegion()) {
+      for (auto &op : block) {
+        if (auto alloc = dyn_cast_or_null<quake::AllocaOp>(&op)) {
+          if (qalloc)
+            return op.emitError("Multiple qalloc statements found");
+
+          qalloc = alloc;
+        } else if (auto measure = dyn_cast_or_null<quake::MzOp>(&op)) {
+          if (!measure.use_empty()) {
+            measure.emitWarning("Measurements with uses are not supported");
+            return success();
+          }
+
+          auto veqOp = measure.getOperand(0);
+          auto ty = veqOp.getType();
+
+          std::size_t size = 0;
+          if (auto veqTy = dyn_cast<quake::RefType>(ty))
+            size = 1;
+          else if (auto veqTy = dyn_cast<quake::VeqType>(ty)) {
+            size = veqTy.getSize();
+            if (size == 0)
+              return op.emitError("Unknown measurement size");
+          }
+
+          measurements[measure.getMeasOut()] = currentOffset;
+          lastMeasurement = measure;
+          currentOffset += size;
+        }
+      }
+    }
+
+    return success();
+  }
+};
+
+class ExtendQubitMeasurePattern : public OpRewritePattern<quake::MzOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  explicit ExtendQubitMeasurePattern(MLIRContext *ctx, Analysis &analysis)
+      : OpRewritePattern(ctx), analysis(analysis) {}
+
+  // Replace a pattern such as:
+  // ```
+  // %0 = ...: !quake.veq<2>
+  // %1 = quake.extract_ref %0[0] : (!quake.veq<2>) -> !quake.ref
+  // %measOut = quake.mz %1 : (!quake.ref) -> !quake.measure
+  // ```
+  // with:
+  // ```
+  //   %1 = ... : !quake.veq<4>
+  //   %measOut = quake.mz %1 : (!quake.veq<4>) -> !cc.stdvec<!quake.measure>
+  // ```
+  // And collect output names information:  `"[[[0,[1,"q0"]],[1,[2,"q1"]]]]"`
+  LogicalResult matchAndRewrite(quake::MzOp measure,
+                                PatternRewriter &rewriter) const override {
+
+    auto veqOp = measure.getOperand(0);
+    if (auto extract = veqOp.getDefiningOp<quake::ExtractRefOp>()) {
+      auto veq = extract.getVeq();
+      std::size_t idx;
+
+      if (extract.hasConstantIndex())
+        idx = extract.getConstantIndex();
+      else if (auto cst =
+                   extract.getIndex().getDefiningOp<arith::ConstantIntOp>())
+        idx = static_cast<std::size_t>(cst.value());
+      else
+        return extract.emitError("Non-constant index in ExtractRef");
+
+      auto offset = idx + analysis.measurements[measure.getMeasOut()];
+      analysis.resultQubitVals[offset] =
+          std::make_pair(idx, std::to_string(idx));
+
+      auto resultType = cudaq::cc::StdvecType::get(measure.getType(0));
+      if (measure == analysis.lastMeasurement)
+        rewriter.replaceOpWithNewOp<quake::MzOp>(measure, TypeRange{resultType},
+                                                 ValueRange{veq},
+                                                 measure.getRegisterNameAttr());
+      else if (measure.use_empty())
+        rewriter.eraseOp(measure);
+    }
+
+    return failure();
+  }
+
+private:
+  Analysis &analysis;
+};
+
+class ExtendVeqMeasurePattern : public OpRewritePattern<quake::MzOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  explicit ExtendVeqMeasurePattern(MLIRContext *ctx, Analysis &analysis)
+      : OpRewritePattern(ctx), analysis(analysis) {}
+
+  // Replace a pattern such as:
+  // ```
+  //   %1 = ... : !quake.veq<4>
+  //   %2 = quake.subveq %1, %c1, %c2 : (!quake.veq<4>, i32, i32) ->
+  //        !quake.veq<2>
+  //   %measOut = quake.mz %2 : (!quake.veq<2>) -> !cc.stdvec<!quake.measure>
+  // ```
+  // with:
+  // ```
+  //   %1 = ... : !quake.veq<4>
+  //   %measOut = quake.mz %1 : (!quake.veq<4>) -> !cc.stdvec<!quake.measure>
+  // ```
+  // And collect output names information:  `"[[[0,[1,"q0"]],[1,[2,"q1"]]]]"`
+  LogicalResult matchAndRewrite(quake::MzOp measure,
+                                PatternRewriter &rewriter) const override {
+
+    auto veqOp = measure.getOperand(0);
+    if (auto subveq = veqOp.getDefiningOp<quake::SubVeqOp>()) {
+      Value lowOp = subveq.getLow();
+      Value highOp = subveq.getHigh();
+
+      auto constLow = lowOp.getDefiningOp<arith::ConstantIntOp>();
+      if (!constLow)
+        return subveq.emitError("Non-constant low index in subveq");
+      auto constHigh = highOp.getDefiningOp<arith::ConstantIntOp>();
+      if (!constHigh)
+        return subveq.emitError("Non-constant high index in subveq");
+
+      auto low = static_cast<std::size_t>(constLow.value());
+      auto high = static_cast<std::size_t>(constHigh.value());
+
+      for (std::size_t i = low; i <= high; i++) {
+        auto start = analysis.measurements[measure.getMeasOut()];
+        auto offset = i - low + start;
+        analysis.resultQubitVals[offset] = std::make_pair(i, std::to_string(i));
+      }
+      if (measure == analysis.lastMeasurement)
+        rewriter.replaceOpWithNewOp<quake::MzOp>(
+            measure, measure.getResultTypes(), ValueRange{subveq.getVeq()},
+            measure.getRegisterNameAttr());
+      else if (measure.use_empty())
+        rewriter.eraseOp(measure);
+
+      return success();
+    }
+
+    return failure();
+  }
+
+private:
+  Analysis &analysis;
+};
+
+class CombineMeasurementsPass
+    : public cudaq::opt::impl::CombineMeasurementsBase<
+          CombineMeasurementsPass> {
+public:
+  using CombineMeasurementsBase::CombineMeasurementsBase;
+
+  void runOnOperation() override {
+    auto *ctx = &getContext();
+    func::FuncOp func = getOperation();
+    OpBuilder builder(func);
+
+    LLVM_DEBUG(llvm::dbgs() << "Function before combining measurements:\n"
+                            << func << "\n\n");
+
+    // Analyze the function to find all qubit mappings.
+    Analysis analysis;
+    if (failed(analysis.analyze(func))) {
+      func.emitOpError("Combining measurements failed");
+      signalPassFailure();
+    }
+
+    if (analysis.empty())
+      return;
+
+    // Extend measurement into one last full measurement.
+    RewritePatternSet patterns(ctx);
+    patterns.insert<ExtendQubitMeasurePattern, ExtendVeqMeasurePattern>(
+        ctx, analysis);
+    if (failed(applyPatternsAndFoldGreedily(func.getOperation(),
+                                            std::move(patterns)))) {
+      func.emitOpError("Combining measurements failed");
+      signalPassFailure();
+    }
+
+    // Add output names mapping attribute.
+    if (!analysis.resultQubitVals.empty()) {
+      nlohmann::json resultQubitJSON{analysis.resultQubitVals};
+      func->setAttr(cudaq::opt::QIROutputNamesAttrName,
+                    builder.getStringAttr(resultQubitJSON.dump()));
+    }
+
+    LLVM_DEBUG(llvm::dbgs() << "Function after combining measurements:\n"
+                            << func << "\n\n");
+  }
+};
+} // namespace

lib/Optimizer/Transforms/DecompositionPatterns.cpp

@@ -623,6 +623,34 @@ struct SToR1 : public OpRewritePattern<quake::SOp> {
   }
 };
 
+// quake.s<adj> target
+// ─────────────────────────────────
+// (quake.z * quake.s) target
+struct SAdjToSZ : public OpRewritePattern<quake::SOp> {
+  using OpRewritePattern<quake::SOp>::OpRewritePattern;
+
+  void initialize() { setDebugName("SAdjToSZ"); }
+
+  LogicalResult matchAndRewrite(quake::SOp op,
+                                PatternRewriter &rewriter) const override {
+    if (!op.isAdj())
+      return failure();
+
+    // Op info
+    auto loc = op->getLoc();
+    auto parameters = op.getParameters();
+    SmallVector<Value> controls(op.getControls());
+    Value target = op.getTarget();
+
+    QuakeOperatorCreator qRewriter(rewriter);
+    qRewriter.create<quake::ZOp>(loc, parameters, controls, target);
+    qRewriter.create<quake::SOp>(loc, parameters, controls, target);
+    qRewriter.selectWiresAndReplaceUses(op, controls, target);
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // TOp decompositions
 //===----------------------------------------------------------------------===//
@@ -1557,6 +1585,7 @@ void cudaq::populateWithAllDecompositionPatterns(RewritePatternSet &patterns) {
     // SOp patterns
     SToPhasedRx,
     SToR1,
+    SAdjToSZ,
     // TOp patterns
     TToPhasedRx,
     TToR1,

python/cudaq/kernel/ast_bridge.py

@@ -4032,4 +4032,5 @@ def compile_to_mlir(astModule, metadata,
     if len(bridge.dependentCaptureVars):
         extraMetaData['dependent_captures'] = bridge.dependentCaptureVars
 
+    bridge.module.dump()
     return bridge.module, bridge.argTypes, extraMetaData