[core] Refactor the lift array alloc pass into two proper phases.

include/cudaq/Optimizer/Builder/Intrinsics.h

@@ -94,41 +94,48 @@ class IRBuilder : public mlir::OpBuilder {
     return genCStringLiteral(loc, module, buffer);
   }
 
+  cc::GlobalOp genVectorOfConstants(mlir::Location loc, mlir::ModuleOp module,
+                                    llvm::StringRef name,
+                                    mlir::DenseElementsAttr values,
+                                    mlir::Type elementType);
+  cc::GlobalOp genVectorOfConstants(
+      mlir::Location loc, mlir::ModuleOp module, llvm::StringRef name,
+      const llvm::SmallVectorImpl<std::complex<double>> &values);
+  cc::GlobalOp genVectorOfConstants(
+      mlir::Location loc, mlir::ModuleOp module, llvm::StringRef name,
+      const llvm::SmallVectorImpl<std::complex<float>> &values);
+
   cc::GlobalOp
   genVectorOfConstants(mlir::Location loc, mlir::ModuleOp module,
-                       mlir::StringRef name,
-                       const std::vector<std::complex<double>> &values);
-  cc::GlobalOp
-  genVectorOfConstants(mlir::Location loc, mlir::ModuleOp module,
-                       mlir::StringRef name,
-                       const std::vector<std::complex<float>> &values);
-
+                       llvm::StringRef name,
+                       const llvm::SmallVectorImpl<double> &values);
   cc::GlobalOp genVectorOfConstants(mlir::Location loc, mlir::ModuleOp module,
-                                    mlir::StringRef name,
-                                    const std::vector<double> &values);
-  cc::GlobalOp genVectorOfConstants(mlir::Location loc, mlir::ModuleOp module,
-                                    mlir::StringRef name,
-                                    const std::vector<float> &values);
+                                    llvm::StringRef name,
+                                    const llvm::SmallVectorImpl<float> &values);
 
-  cc::GlobalOp genVectorOfConstants(mlir::Location loc, mlir::ModuleOp module,
-                                    mlir::StringRef name,
-                                    const std::vector<std::int64_t> &values);
+  cc::GlobalOp
+  genVectorOfConstants(mlir::Location loc, mlir::ModuleOp module,
+                       llvm::StringRef name,
+                       const llvm::SmallVectorImpl<std::int64_t> &values);
 
-  cc::GlobalOp genVectorOfConstants(mlir::Location loc, mlir::ModuleOp module,
-                                    mlir::StringRef name,
-                                    const std::vector<std::int32_t> &values);
+  cc::GlobalOp
+  genVectorOfConstants(mlir::Location loc, mlir::ModuleOp module,
+                       llvm::StringRef name,
+                       const llvm::SmallVectorImpl<std::int32_t> &values);
 
-  cc::GlobalOp genVectorOfConstants(mlir::Location loc, mlir::ModuleOp module,
-                                    mlir::StringRef name,
-                                    const std::vector<std::int16_t> &values);
+  cc::GlobalOp
+  genVectorOfConstants(mlir::Location loc, mlir::ModuleOp module,
+                       llvm::StringRef name,
+                       const llvm::SmallVectorImpl<std::int16_t> &values);
 
-  cc::GlobalOp genVectorOfConstants(mlir::Location loc, mlir::ModuleOp module,
-                                    mlir::StringRef name,
-                                    const std::vector<std::int8_t> &values);
+  cc::GlobalOp
+  genVectorOfConstants(mlir::Location loc, mlir::ModuleOp module,
+                       llvm::StringRef name,
+                       const llvm::SmallVectorImpl<std::int8_t> &values);
 
   cc::GlobalOp genVectorOfConstants(mlir::Location loc, mlir::ModuleOp module,
-                                    mlir::StringRef name,
-                                    const std::vector<bool> &values);
+                                    llvm::StringRef name,
+                                    const llvm::SmallVectorImpl<bool> &values);
 
   /// Load an intrinsic into \p module. The intrinsic to load has name \p name.
   /// This will automatically load any intrinsics that \p name depends upon.

include/cudaq/Optimizer/Transforms/Passes.td

@@ -160,7 +160,7 @@ def CombineQuantumAllocations :
   let dependentDialects = ["cudaq::cc::CCDialect", "quake::QuakeDialect"];
 }
 
-def ConstPropComplex : Pass<"const-prop-complex", "mlir::ModuleOp"> {
+def ConstPropComplex : Pass<"const-prop-complex", "mlir::func::FuncOp"> {
   let summary = "Create and propagate complex constants.";
   let description = [{
     Rewrite the complex.CreateOp to complex.ConstantOp when possible.
@@ -383,11 +383,11 @@ def GenerateKernelExecution : Pass<"kernel-execution", "mlir::ModuleOp"> {
 
 def GetConcreteMatrix : Pass<"get-concrete-matrix", "mlir::func::FuncOp"> {
   let summary =
-    "Replace the unitary matrix generator function with concrete matrix.";
+    "Replace the unitary matrix generator function with a constant matrix.";
   let description = [{
     Given a custom operation whose generator attribute is another function 
     within the module, such that if `LiftArrayAlloc` pass has run, there will
-    be a global constant within the module which holds the concrete matrix 
+    be a global constant within the module which holds the constant matrix 
     representation for the custom operation. This pass will find that global
     variable and update the custom operation to directly point to it. 
 
@@ -424,6 +424,22 @@ def GetConcreteMatrix : Pass<"get-concrete-matrix", "mlir::func::FuncOp"> {
   }];
 }
 
+// GlobalizeArrayValues must be a module pass because it may promoted array
+// constants from functions to global constants (changes their scope).
+def GlobalizeArrayValues : Pass<"globalize-array-values", "mlir::ModuleOp"> {
+  let summary = "Convert const_array ops to globals.";
+  let description = [{
+    Often a `const_array` op can be canonicalized into scalar constants that
+    are then constant propagated to their uses in the quake ops. When this
+    happens, the `const_array` may become unused and be eliminated.
+
+    However, there can also be cases where the `const_array` remains alive, such
+    as when it is used in a `state_init` op. In such cases, we may be able to go
+    ahead and replace the `const_array` with a global constant. This pass makes
+    such conversions.
+  }];
+}
+
 // LambdaLifting is a module pass because it may modify the ModuleOp and add
 // new FuncOps.
 def LambdaLifting : Pass<"lambda-lifting", "mlir::ModuleOp"> {
@@ -439,7 +455,7 @@ def LambdaLifting : Pass<"lambda-lifting", "mlir::ModuleOp"> {
   let constructor = "cudaq::opt::createLambdaLiftingPass()";
 }
 
-def LiftArrayAlloc : Pass<"lift-array-value", "mlir::ModuleOp"> {
+def LiftArrayAlloc : Pass<"lift-array-alloc", "mlir::func::FuncOp"> {
   let summary = "Convert constant arrays built on the stack to array values";
   let description = [{
     The bridge or other passes may generate inline code to build an array of
@@ -476,6 +492,10 @@ def LiftArrayAlloc : Pass<"lift-array-value", "mlir::ModuleOp"> {
     updated or escapes the function, it cannot be replaced by a value. If
     it is elements are accessed in a read-only way, it is a legal transform
     and will enable further constant folding in other passes.
+
+    See the globalize array values pass for converting `const_array` values
+    to global constants. Conversion to globals is intentionally deferred to
+    allow constant propagation to take place correctly.
   }];
 
   let dependentDialects = ["mlir::complex::ComplexDialect"];

lib/Optimizer/Builder/Intrinsics.cpp

@@ -450,16 +450,16 @@ LogicalResult IRBuilder::loadIntrinsic(ModuleOp module, StringRef intrinName) {
 }
 
 template <typename T>
-DenseElementsAttr createDenseElementsAttr(const std::vector<T> &values,
+DenseElementsAttr createDenseElementsAttr(const SmallVectorImpl<T> &values,
                                           Type eleTy) {
   auto newValues = ArrayRef<T>(values.data(), values.size());
   auto tensorTy = RankedTensorType::get(values.size(), eleTy);
   return DenseElementsAttr::get(tensorTy, newValues);
 }
 
-DenseElementsAttr createDenseElementsAttr(const std::vector<bool> &values,
-                                          Type eleTy) {
-  std::vector<std::byte> converted;
+static DenseElementsAttr
+createDenseElementsAttr(const SmallVectorImpl<bool> &values, Type eleTy) {
+  SmallVector<std::byte> converted;
   for (auto it = values.begin(); it != values.end(); it++) {
     bool value = *it;
     converted.push_back(std::byte(value));
@@ -470,83 +470,99 @@ DenseElementsAttr createDenseElementsAttr(const std::vector<bool> &values,
   return DenseElementsAttr::get(tensorTy, newValues);
 }
 
-template <typename A>
-cc::GlobalOp buildVectorOfConstantElements(Location loc, ModuleOp module,
-                                           StringRef name,
-                                           const std::vector<A> &values,
-                                           IRBuilder &builder, Type eleTy) {
+static cc::GlobalOp buildVectorOfConstantElements(Location loc, ModuleOp module,
+                                                  StringRef name,
+                                                  DenseElementsAttr &arrayAttr,
+                                                  IRBuilder &builder,
+                                                  Type eleTy) {
   if (auto glob = module.lookupSymbol<cc::GlobalOp>(name))
     return glob;
   auto *ctx = builder.getContext();
   OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPointToEnd(module.getBody());
-  auto globalTy = cc::ArrayType::get(ctx, eleTy, values.size());
-
-  auto arrayAttr = createDenseElementsAttr(values, eleTy);
+  auto globalTy = cc::ArrayType::get(ctx, eleTy, arrayAttr.size());
   return builder.create<cudaq::cc::GlobalOp>(loc, globalTy, name, arrayAttr,
                                              /*constant=*/true,
                                              /*external=*/false);
 }
 
+template <typename A>
+cc::GlobalOp buildVectorOfConstantElements(Location loc, ModuleOp module,
+                                           StringRef name,
+                                           const SmallVectorImpl<A> &values,
+                                           IRBuilder &builder, Type eleTy) {
+  auto arrayAttr = createDenseElementsAttr(values, eleTy);
+  return buildVectorOfConstantElements(loc, module, name, arrayAttr, builder,
+                                       eleTy);
+}
+
+cc::GlobalOp IRBuilder::genVectorOfConstants(Location loc, ModuleOp module,
+                                             StringRef name,
+                                             DenseElementsAttr values,
+                                             Type elementType) {
+  return buildVectorOfConstantElements(loc, module, name, values, *this,
+                                       elementType);
+}
+
 cc::GlobalOp IRBuilder::genVectorOfConstants(
     Location loc, ModuleOp module, StringRef name,
-    const std::vector<std::complex<double>> &values) {
+    const SmallVectorImpl<std::complex<double>> &values) {
   return buildVectorOfConstantElements(loc, module, name, values, *this,
                                        ComplexType::get(getF64Type()));
 }
 
 cc::GlobalOp IRBuilder::genVectorOfConstants(
     Location loc, ModuleOp module, StringRef name,
-    const std::vector<std::complex<float>> &values) {
+    const SmallVectorImpl<std::complex<float>> &values) {
   return buildVectorOfConstantElements(loc, module, name, values, *this,
                                        ComplexType::get(getF32Type()));
 }
 
 cc::GlobalOp
 IRBuilder::genVectorOfConstants(Location loc, ModuleOp module, StringRef name,
-                                const std::vector<double> &values) {
+                                const SmallVectorImpl<double> &values) {
   return buildVectorOfConstantElements(loc, module, name, values, *this,
                                        getF64Type());
 }
 
-cc::GlobalOp IRBuilder::genVectorOfConstants(Location loc, ModuleOp module,
-                                             StringRef name,
-                                             const std::vector<float> &values) {
+cc::GlobalOp
+IRBuilder::genVectorOfConstants(Location loc, ModuleOp module, StringRef name,
+                                const SmallVectorImpl<float> &values) {
   return buildVectorOfConstantElements(loc, module, name, values, *this,
                                        getF32Type());
 }
 
 cc::GlobalOp
 IRBuilder::genVectorOfConstants(Location loc, ModuleOp module, StringRef name,
-                                const std::vector<std::int64_t> &values) {
+                                const SmallVectorImpl<std::int64_t> &values) {
   return buildVectorOfConstantElements(loc, module, name, values, *this,
                                        getI64Type());
 }
 
 cc::GlobalOp
 IRBuilder::genVectorOfConstants(Location loc, ModuleOp module, StringRef name,
-                                const std::vector<std::int32_t> &values) {
+                                const SmallVectorImpl<std::int32_t> &values) {
   return buildVectorOfConstantElements(loc, module, name, values, *this,
                                        getI32Type());
 }
 
 cc::GlobalOp
 IRBuilder::genVectorOfConstants(Location loc, ModuleOp module, StringRef name,
-                                const std::vector<std::int16_t> &values) {
+                                const SmallVectorImpl<std::int16_t> &values) {
   return buildVectorOfConstantElements(loc, module, name, values, *this,
                                        getI16Type());
 }
 
 cc::GlobalOp
 IRBuilder::genVectorOfConstants(Location loc, ModuleOp module, StringRef name,
-                                const std::vector<std::int8_t> &values) {
+                                const SmallVectorImpl<std::int8_t> &values) {
   return buildVectorOfConstantElements(loc, module, name, values, *this,
                                        getI8Type());
 }
 
-cc::GlobalOp IRBuilder::genVectorOfConstants(Location loc, ModuleOp module,
-                                             StringRef name,
-                                             const std::vector<bool> &values) {
+cc::GlobalOp
+IRBuilder::genVectorOfConstants(Location loc, ModuleOp module, StringRef name,
+                                const SmallVectorImpl<bool> &values) {
   return buildVectorOfConstantElements(loc, module, name, values, *this,
                                        getI1Type());
 }

lib/Optimizer/CodeGen/Pipelines.cpp

@@ -47,7 +47,8 @@ void cudaq::opt::addPipelineTranslateToOpenQASM(PassManager &pm) {
   pm.addNestedPass<func::FuncOp>(createClassicalMemToReg());
   pm.addPass(createLoopUnroll());
   pm.addPass(createCanonicalizerPass());
-  pm.addPass(createLiftArrayAlloc());
+  pm.addNestedPass<func::FuncOp>(createLiftArrayAlloc());
+  pm.addPass(createGlobalizeArrayValues());
   pm.addPass(createStatePreparation());
 }
 

lib/Optimizer/Transforms/CMakeLists.txt

@@ -30,6 +30,7 @@ add_cudaq_library(OptTransforms
   GenKernelExecution.cpp
   GenDeviceCodeLoader.cpp
   GetConcreteMatrix.cpp
+  GlobalizeArrayValues.cpp
   LambdaLifting.cpp
   LiftArrayAlloc.cpp
   LinearCtrlRelations.cpp

lib/Optimizer/Transforms/ConstPropComplex.cpp

@@ -172,29 +172,22 @@ class ConstPropComplexPass
 
   void runOnOperation() override {
     auto *ctx = &getContext();
-    auto module = getOperation();
-    for (Operation &op : *module.getBody()) {
-      auto func = dyn_cast<func::FuncOp>(op);
-      if (!func)
-        continue;
-      DominanceInfo domInfo(func);
-      std::string funcName = func.getName().str();
-      RewritePatternSet patterns(ctx);
-      patterns
-          .insert<ComplexCreatePattern, FloatCastPattern, FloatExtendPattern,
-                  FloatTruncatePattern, ComplexRePattern, ComplexImPattern>(
-              ctx);
-
-      LLVM_DEBUG(llvm::dbgs()
-                 << "Before lifting constant array: " << func << '\n');
-
-      if (failed(applyPatternsAndFoldGreedily(func.getOperation(),
-                                              std::move(patterns))))
-        signalPassFailure();
-
-      LLVM_DEBUG(llvm::dbgs()
-                 << "After lifting constant array: " << func << '\n');
-    }
+    auto func = getOperation();
+    DominanceInfo domInfo(func);
+    RewritePatternSet patterns(ctx);
+    patterns.insert<ComplexCreatePattern, FloatCastPattern, FloatExtendPattern,
+                    FloatTruncatePattern, ComplexRePattern, ComplexImPattern>(
+        ctx);
+
+    LLVM_DEBUG(llvm::dbgs() << "Before constant propagation of complex values: "
+                            << func << '\n');
+
+    if (failed(applyPatternsAndFoldGreedily(func.getOperation(),
+                                            std::move(patterns))))
+      signalPassFailure();
+
+    LLVM_DEBUG(llvm::dbgs() << "After constant propagation of complex values: "
+                            << func << '\n');
   }
 };
 } // namespace