From 80b98681e94ffd15aee17933202734a1726b1595 Mon Sep 17 00:00:00 2001
From: Ben Howe <bhowe@nvidia.com>
Date: Thu, 8 Aug 2024 22:08:03 +0000
Subject: [PATCH 1/2] Update WireSetsToProfileQIR pass: remove WireSets and fix
 LLVM lowering problem

---
 include/cudaq/Optimizer/CodeGen/Passes.td     | 10 ++++
 .../CodeGen/WireSetsToProfileQIR.cpp          | 47 +++++++++++++++++--
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/include/cudaq/Optimizer/CodeGen/Passes.td b/include/cudaq/Optimizer/CodeGen/Passes.td
index d1232c24cd..2120ea0836 100644
--- a/include/cudaq/Optimizer/CodeGen/Passes.td
+++ b/include/cudaq/Optimizer/CodeGen/Passes.td
@@ -192,6 +192,16 @@ def WireSetToProfileQIR : Pass<"wireset-to-profile-qir", "mlir::func::FuncOp"> {
   ];
 }
 
+def WireSetToProfileQIRPost :
+    Pass<"wireset-to-profile-qir-post", "mlir::ModuleOp"> {
+  let summary = "Post processing for lowering wire sets to a profile of QIR";
+  let description = [{
+    This pass should be run immediately after wireset-to-profile-qir.
+  }];
+
+  let dependentDialects = ["cudaq::cc::CCDialect", "mlir::func::FuncDialect"];
+}
+
 def WireSetToProfileQIRPrep :
     Pass<"wireset-to-profile-qir-prep", "mlir::ModuleOp"> {
   let summary = "Prepare for lowering wire sets to a profile of QIR";
diff --git a/lib/Optimizer/CodeGen/WireSetsToProfileQIR.cpp b/lib/Optimizer/CodeGen/WireSetsToProfileQIR.cpp
index 7a96c0761d..d4a9713235 100644
--- a/lib/Optimizer/CodeGen/WireSetsToProfileQIR.cpp
+++ b/lib/Optimizer/CodeGen/WireSetsToProfileQIR.cpp
@@ -46,6 +46,7 @@
 
 namespace cudaq::opt {
 #define GEN_PASS_DEF_WIRESETTOPROFILEQIR
+#define GEN_PASS_DEF_WIRESETTOPROFILEQIRPOST
 #define GEN_PASS_DEF_WIRESETTOPROFILEQIRPREP
 #include "cudaq/Optimizer/CodeGen/Passes.h.inc"
 } // namespace cudaq::opt
@@ -145,6 +146,17 @@ struct ReturnWireRewrite : OpConversionPattern<quake::ReturnWireOp> {
   }
 };
 
+struct WireSetRewrite : OpConversionPattern<quake::WireSetOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(quake::WireSetOp wireSetOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    rewriter.eraseOp(wireSetOp);
+    return success();
+  }
+};
+
 struct MzRewrite : OpConversionPattern<quake::MzOp> {
   using Base = OpConversionPattern;
   explicit MzRewrite(TypeConverter &typeConverter, unsigned &counter,
@@ -195,13 +207,18 @@ struct DiscriminateRewrite : OpConversionPattern<quake::DiscriminateOp> {
     // NB: This is thread safe as it should never do an insertion, just a
     // lookup.
     auto nameObj = irb.genCStringLiteralAppendNul(loc, mod, iter->second);
+    auto arrI8Ty = mlir::LLVM::LLVMArrayType::get(rewriter.getI8Type(),
+                                                  iter->second.size() + 1);
+    auto ptrArrTy = cudaq::cc::PointerType::get(arrI8Ty);
+    Value nameVal = rewriter.create<cudaq::cc::AddressOfOp>(loc, ptrArrTy,
+                                                            nameObj.getName());
     auto cstrTy = cudaq::cc::PointerType::get(rewriter.getI8Type());
-    Value nameVal =
-        rewriter.create<cudaq::cc::AddressOfOp>(loc, cstrTy, nameObj.getName());
+    Value nameValCStr =
+        rewriter.create<cudaq::cc::CastOp>(loc, cstrTy, nameVal);
 
     rewriter.create<func::CallOp>(
         loc, std::nullopt, "__quantum__rt__result_record_output",
-        ValueRange{adaptor.getMeasurement(), nameVal});
+        ValueRange{adaptor.getMeasurement(), nameValCStr});
     if (isAdaptiveProfile) {
       std::string funcName = toQisBodyName(std::string("read_result"));
       rewriter.replaceOpWithNewOp<func::CallOp>(
@@ -371,6 +388,29 @@ struct WireSetToProfileQIRPrepPass
     LLVM_DEBUG(llvm::dbgs() << "Module after prep:\n"; op->dump());
   }
 };
+
+struct WireSetToProfileQIRPostPass
+    : public cudaq::opt::impl::WireSetToProfileQIRPostBase<
+          WireSetToProfileQIRPostPass> {
+  using WireSetToProfileQIRPostBase::WireSetToProfileQIRPostBase;
+
+  void runOnOperation() override {
+    ModuleOp op = getOperation();
+    auto *ctx = &getContext();
+    RewritePatternSet patterns(ctx);
+    QuakeTypeConverter quakeTypeConverter;
+    patterns.insert<WireSetRewrite>(quakeTypeConverter, ctx);
+    ConversionTarget target(*ctx);
+    target.addLegalDialect<arith::ArithDialect, cudaq::cc::CCDialect,
+                           func::FuncDialect, LLVM::LLVMDialect>();
+    target.addIllegalDialect<quake::QuakeDialect>();
+
+    LLVM_DEBUG(llvm::dbgs() << "Module before:\n"; op.dump());
+    if (failed(applyPartialConversion(op, target, std::move(patterns))))
+      signalPassFailure();
+    LLVM_DEBUG(llvm::dbgs() << "Module after:\n"; op.dump());
+  }
+};
 } // namespace
 
 void cudaq::opt::addWiresetToProfileQIRPipeline(OpPassManager &pm,
@@ -380,6 +420,7 @@ void cudaq::opt::addWiresetToProfileQIRPipeline(OpPassManager &pm,
   if (!profile.empty())
     wopt.convertTo = profile.str();
   pm.addNestedPass<func::FuncOp>(cudaq::opt::createWireSetToProfileQIR(wopt));
+  pm.addPass(cudaq::opt::createWireSetToProfileQIRPost());
 }
 
 // Pipeline option: let the user specify the profile name.

From ab08dfe310358446a6addc6290057060e64f4145 Mon Sep 17 00:00:00 2001
From: Ben Howe <bhowe@nvidia.com>
Date: Thu, 8 Aug 2024 22:24:16 +0000
Subject: [PATCH 2/2] Address WireSetsToProfileQIR.cpp PR comment

---
 lib/Optimizer/CodeGen/WireSetsToProfileQIR.cpp | 2 --
 1 file changed, 2 deletions(-)

diff --git a/lib/Optimizer/CodeGen/WireSetsToProfileQIR.cpp b/lib/Optimizer/CodeGen/WireSetsToProfileQIR.cpp
index d4a9713235..71a38af828 100644
--- a/lib/Optimizer/CodeGen/WireSetsToProfileQIR.cpp
+++ b/lib/Optimizer/CodeGen/WireSetsToProfileQIR.cpp
@@ -401,8 +401,6 @@ struct WireSetToProfileQIRPostPass
     QuakeTypeConverter quakeTypeConverter;
     patterns.insert<WireSetRewrite>(quakeTypeConverter, ctx);
     ConversionTarget target(*ctx);
-    target.addLegalDialect<arith::ArithDialect, cudaq::cc::CCDialect,
-                           func::FuncDialect, LLVM::LLVMDialect>();
     target.addIllegalDialect<quake::QuakeDialect>();
 
     LLVM_DEBUG(llvm::dbgs() << "Module before:\n"; op.dump());