Support using complex numbers and lists in kernel functions

docs/sphinx/examples/cpp/other/builder/builder.cpp

@@ -17,6 +17,14 @@
 #include <cudaq/gradients.h>
 #include <cudaq/optimizers.h>
 
+static bool results_are_close(const cudaq::sample_result &f1,
+                              const cudaq::sample_result &f2) {
+  // Stub for a fancy comparison.
+  f1.dump();
+  f2.dump();
+  return true;
+}
+
 // This example demonstrates various uses for the `cudaq::builder`
 // type. This type enables one to dynamically construct callable
 // CUDA Quantum kernels via just-in-time compilation. The typical workflow
@@ -114,6 +122,37 @@ int main() {
     counts.dump();
   }
 
+  {
+    // In a simulated environment, it is sometimes useful to be able to
+    // specify an initial state vector. The initial state vector is 2 to the
+    // power `n` where `n` is the number of qubits.
+
+    // In this example, we create a kernel template `sim_kernel` that captures
+    // the variable `init_state` by reference.
+    auto sim_builder = cudaq::make_kernel();
+    std::vector<cudaq::simulation_scalar> init_state;
+    auto q = sim_builder.qalloc(init_state);
+    // Build the quantum circuit template here.
+    sim_builder.mz(q);
+
+    // Now we are ready to instantiate the kernel and invoke it. So we can set
+    // the `init_state` to a vector with 2 complex values (1 qubit) and
+    // get the results.
+    init_state = {{0.0, 1.0}, {1.0, 0.0}};
+    auto counts0 = cudaq::sample(sim_builder);
+
+    // Now suppose we have a different initial state with 4 complex values (2
+    // qubits). Let's rerun the kernel with the new `init_state`.
+    init_state = {{1.0, 0.0}, {0.0, 1.0}, {0.0, 1.0}, {1.0, 0.0}};
+    auto counts1 = cudaq::sample(sim_builder);
+
+    // Finally in this wholly contrived example, we test the results to make
+    // sure they are "close".
+    if (results_are_close(counts0, counts1)) {
+      printf("The two initial states generated results that are \"close\".\n");
+    }
+  }
+
   {
 
     // Let's do a final sampling task. Let's

include/cudaq/Frontend/nvqpp/ASTBridge.h

@@ -249,13 +249,8 @@ class QuakeBridgeVisitor
   bool TraverseConditionalOperator(clang::ConditionalOperator *x,
                                    DataRecursionQueue *q = nullptr);
   bool VisitReturnStmt(clang::ReturnStmt *x);
-  bool VisitCXXFunctionalCastExpr(clang::CXXFunctionalCastExpr *x) {
-    return true;
-  }
   bool TraverseInitListExpr(clang::InitListExpr *x,
                             DataRecursionQueue *q = nullptr);
-  bool TraverseCXXTemporaryObjectExpr(clang::CXXTemporaryObjectExpr *x,
-                                      DataRecursionQueue *q = nullptr);
 
   // These misc. statements are not (yet) handled by lowering.
   bool TraverseAsmStmt(clang::AsmStmt *x, DataRecursionQueue *q = nullptr);
@@ -287,16 +282,58 @@ class QuakeBridgeVisitor
   bool TraverseCXXConstructExpr(clang::CXXConstructExpr *x,
                                 DataRecursionQueue *q = nullptr);
   bool VisitCXXConstructExpr(clang::CXXConstructExpr *x);
-  bool VisitCXXTemporaryObjectExpr(clang::CXXTemporaryObjectExpr *x);
   bool VisitCXXOperatorCallExpr(clang::CXXOperatorCallExpr *x);
   bool WalkUpFromCXXOperatorCallExpr(clang::CXXOperatorCallExpr *x);
   bool TraverseDeclRefExpr(clang::DeclRefExpr *x,
                            DataRecursionQueue *q = nullptr);
   bool VisitDeclRefExpr(clang::DeclRefExpr *x);
   bool VisitFloatingLiteral(clang::FloatingLiteral *x);
+
+  // Cast operations.
+  bool TraverseCastExpr(clang::CastExpr *x, DataRecursionQueue *q = nullptr);
+  bool VisitCastExpr(clang::CastExpr *x);
+
   bool TraverseImplicitCastExpr(clang::ImplicitCastExpr *x,
-                                DataRecursionQueue *q = nullptr);
-  bool VisitImplicitCastExpr(clang::ImplicitCastExpr *x);
+                                DataRecursionQueue *q = nullptr) {
+    return TraverseCastExpr(x, q);
+  }
+  bool TraverseExplicitCastExpr(clang::ExplicitCastExpr *x,
+                                DataRecursionQueue *q = nullptr) {
+    return TraverseCastExpr(x, q);
+  }
+  bool TraverseCStyleCastExpr(clang::CStyleCastExpr *x,
+                              DataRecursionQueue *q = nullptr) {
+    return TraverseExplicitCastExpr(x, q);
+  }
+  bool TraverseCXXFunctionalCastExpr(clang::CXXFunctionalCastExpr *x,
+                                     DataRecursionQueue *q = nullptr) {
+    return TraverseExplicitCastExpr(x, q);
+  }
+  bool TraverseCXXAddrspaceCastExpr(clang::CXXAddrspaceCastExpr *x,
+                                    DataRecursionQueue *q = nullptr) {
+    return TraverseExplicitCastExpr(x, q);
+  }
+  bool TraverseCXXConstCastExpr(clang::CXXConstCastExpr *x,
+                                DataRecursionQueue *q = nullptr) {
+    return TraverseExplicitCastExpr(x, q);
+  }
+  bool TraverseCXXDynamicCastExpr(clang::CXXDynamicCastExpr *x,
+                                  DataRecursionQueue *q = nullptr) {
+    return TraverseExplicitCastExpr(x, q);
+  }
+  bool TraverseCXXReinterpretCastExpr(clang::CXXReinterpretCastExpr *x,
+                                      DataRecursionQueue *q = nullptr) {
+    return TraverseExplicitCastExpr(x, q);
+  }
+  bool TraverseCXXStaticCastExpr(clang::CXXStaticCastExpr *x,
+                                 DataRecursionQueue *q = nullptr) {
+    return TraverseExplicitCastExpr(x, q);
+  }
+  bool TraverseBuiltinBitCastExpr(clang::BuiltinBitCastExpr *x,
+                                  DataRecursionQueue *q = nullptr) {
+    return TraverseExplicitCastExpr(x, q);
+  }
+
   bool VisitInitListExpr(clang::InitListExpr *x);
   bool VisitIntegerLiteral(clang::IntegerLiteral *x);
   bool VisitCXXBoolLiteralExpr(clang::CXXBoolLiteralExpr *x);
@@ -579,7 +616,6 @@ class QuakeBridgeVisitor
   llvm::DenseMap<clang::RecordType *, mlir::Type> records;
 
   // State Flags
-
   bool skipCompoundScope : 1 = false;
   bool isEntry : 1 = false;
   /// If there is a catastrophic error in the bridge (there is no rational way
@@ -589,6 +625,7 @@ class QuakeBridgeVisitor
   bool visitImplicitCode : 1 = false;
   bool inRecType : 1 = false;
   bool allowUnknownRecordType : 1 = false;
+  bool initializerIsGlobal : 1 = false;
 };
 } // namespace details
 

include/cudaq/Optimizer/Builder/Factory.h

@@ -159,6 +159,12 @@ inline mlir::Value createF64Constant(mlir::Location loc,
   return createFloatConstant(loc, builder, value, builder.getF64Type());
 }
 
+/// Return the integer value if \p v is an integer constant.
+std::optional<std::uint64_t> maybeValueOfIntConstant(mlir::Value v);
+
+/// Return the floating point value if \p v is a floating-point constant.
+std::optional<double> maybeValueOfFloatConstant(mlir::Value v);
+
 //===----------------------------------------------------------------------===//
 
 inline mlir::Block *addEntryBlock(mlir::LLVM::GlobalOp initVar) {

include/cudaq/Optimizer/CodeGen/CMakeLists.txt

@@ -6,11 +6,13 @@
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
 
+add_cudaq_dialect(CodeGen codegen)
+add_cudaq_dialect_doc(CodeGenDialect codegen)
+
 set(LLVM_TARGET_DEFINITIONS Passes.td)
 mlir_tablegen(Passes.h.inc -gen-pass-decls -name OptCodeGen)
 add_public_tablegen_target(OptCodeGenPassIncGen)
 
 set(LLVM_TARGET_DEFINITIONS Peephole.td)
 mlir_tablegen(Peephole.inc -gen-rewriters)
 add_public_tablegen_target(OptPeepholeIncGen)
-

include/cudaq/Optimizer/CodeGen/CodeGenDialect.td

@@ -0,0 +1,34 @@
+/********************************************************** -*- tablegen -*- ***
+ * 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.    *
+ ******************************************************************************/
+
+#ifndef CUDAQ_OPTIMIZER_CODEGEN_DIALECT
+#define CUDAQ_OPTIMIZER_CODEGEN_DIALECT
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+//===----------------------------------------------------------------------===//
+// Dialect definition.
+//===----------------------------------------------------------------------===//
+
+def CodeGenDialect : Dialect {
+  let name = "codegen";
+  let summary = "Code generation helpers";
+  let description = [{
+    Do not use this dialect outside of code generation.
+  }];
+
+  let cppNamespace = "cudaq::codegen";
+  let useDefaultTypePrinterParser = 1;
+  let useFoldAPI = kEmitFoldAdaptorFolder;
+
+  let extraClassDeclaration = [{
+    void registerTypes(); // register at least a bogo type.
+  }];
+}
+
+#endif

include/cudaq/Optimizer/CodeGen/CodeGenOps.td

@@ -0,0 +1,47 @@
+/********************************************************** -*- tablegen -*- ***
+ * 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.    *
+ ******************************************************************************/
+
+#ifndef CUDAQ_OPTIMIZER_CODEGEN_OPS
+#define CUDAQ_OPTIMIZER_CODEGEN_OPS
+
+include "cudaq/Optimizer/CodeGen/CodeGenDialect.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
+include "cudaq/Optimizer/Dialect/Common/Traits.td"
+include "cudaq/Optimizer/Dialect/CC/CCTypes.td"
+include "cudaq/Optimizer/Dialect/Quake/QuakeTypes.td"
+
+//===----------------------------------------------------------------------===//
+// The codegen quake dialect is a transitory set of operations used exclusively
+// during codegen. The ops defined here make the process of converting the Quake
+// code to a another target dialect (e.g. LLVM-IR) easier. They should not be
+// used outside of the codegen passes.
+//===----------------------------------------------------------------------===//
+
+class CGQOp<string mnemonic, list<Trait> traits = []>
+    : Op<CodeGenDialect, mnemonic, traits>;
+
+def cgq_RAIIOp : CGQOp<"qmem_raii", [MemoryEffects<[MemAlloc, MemWrite]>]> {
+  let summary = "Combine allocation and initialization of set of qubits.";
+  let description = [{
+    Used only in QIR code generation.
+  }];
+
+  let arguments = (ins
+    cc_PointerType:$initState,
+    TypeAttr:$allocType,
+    Optional<AnySignlessIntegerOrIndex>:$allocSize
+  );
+  let results = (outs VeqType);
+
+  let assemblyFormat = [{
+    $initState `(` $allocType ( `[` $allocSize^ `]` )? `)` `:`
+    functional-type(operands, results) attr-dict
+  }];
+}
+
+#endif

include/cudaq/Optimizer/CodeGen/CodeGenTypes.td

@@ -0,0 +1,34 @@
+/********************************************************** -*- tablegen -*- ***
+ * 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.    *
+ ******************************************************************************/
+
+#ifndef CUDAQ_OPTIMIZER_CODEGEN_TYPES
+#define CUDAQ_OPTIMIZER_CODEGEN_TYPES
+
+include "cudaq/Optimizer/CodeGen/CodeGenDialect.td"
+include "mlir/Interfaces/DataLayoutInterfaces.td"
+include "mlir/IR/AttrTypeBase.td"
+
+//===----------------------------------------------------------------------===//
+// BaseType
+//===----------------------------------------------------------------------===//
+
+class CodeGenType<string name, string typeMnemonic, list<Trait> traits = [],
+             string base = "mlir::Type">
+    : TypeDef<CodeGenDialect, name, traits, base> {
+  let mnemonic = typeMnemonic;
+}
+
+// There are no codegen dialect types, but the dialect tablegen generates type
+// boilerplate for the dialect anyway.
+
+def codegen_DoNotUseType : CodeGenType<"DoNotUse", "dummy"> {
+  let summary = "";
+  let description = [{ }];
+}
+
+#endif // CUDAQ_OPTIMIZER_CODEGEN_TYPES

include/cudaq/Optimizer/CodeGen/Passes.h

@@ -52,6 +52,9 @@ std::unique_ptr<mlir::Pass> createRemoveMeasurementsPass();
 /// Register target pipelines.
 void registerTargetPipelines();
 
+/// Register CodeGenDialect with the provided DialectRegistry.
+void registerCodeGenDialect(mlir::DialectRegistry &registry);
+
 // declarative passes
 #define GEN_PASS_DECL
 #define GEN_PASS_REGISTRATION

include/cudaq/Optimizer/CodeGen/Passes.td

@@ -19,7 +19,9 @@ def QuakeToQIR : Pass<"quake-to-qir", "mlir::ModuleOp"> {
     QIR qubits.
   }];
 
-  let dependentDialects = ["mlir::LLVM::LLVMDialect"];
+  let dependentDialects = [
+    "cudaq::codegen::CodeGenDialect", "mlir::LLVM::LLVMDialect"
+  ];
   let constructor = "cudaq::opt::createConvertToQIRPass()";
 }
 

include/cudaq/Optimizer/CodeGen/QIRFunctionNames.h

@@ -43,6 +43,10 @@ constexpr static const char QIRArrayGetElementPtr1d[] =
     "__quantum__rt__array_get_element_ptr_1d";
 constexpr static const char QIRArrayQubitAllocateArray[] =
     "__quantum__rt__qubit_allocate_array";
+constexpr static const char QIRArrayQubitAllocateArrayWithStateFP64[] =
+    "__quantum__rt__qubit_allocate_array_with_state_fp64";
+constexpr static const char QIRArrayQubitAllocateArrayWithStateFP32[] =
+    "__quantum__rt__qubit_allocate_array_with_state_fp32";
 constexpr static const char QIRQubitAllocate[] =
     "__quantum__rt__qubit_allocate";
 constexpr static const char QIRArrayQubitReleaseArray[] =

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

@@ -794,6 +794,35 @@ def cc_AddressOfOp : CCOp<"address_of", [Pure,
   }];
 }
 
+def cc_GlobalOp : CCOp<"global", [IsolatedFromAbove, Symbol]> {
+  let summary = "Create a global constant or variable";
+  let description = [{
+    A GlobalOp is used to create a global variable or constant that can be
+    referenced by symbol using the AddressOfOp. The type of this op is always
+    implicitly a cc::PointerType.
+
+    For example, this op may be used to define arrays of doubles, which may in
+    turn be used as initial state vectors for quantum memory (VeqType).
+  }];
+
+  let arguments = (ins
+    TypeAttr:$global_type,
+    StrAttr:$sym_name,
+    OptionalAttr<AnyAttr>:$value,
+    UnitAttr:$constant,
+    UnitAttr:$external
+  );
+
+  let hasCustomAssemblyFormat = 1;
+
+  let extraClassDeclaration = [{
+    cudaq::cc::PointerType getType() {
+      auto globalTy = getGlobalType();
+      return cudaq::cc::PointerType::get(globalTy);
+    }
+  }];
+}
+
 def cc_ExtractValueOp : CCOp<"extract_value", [Pure]> {
   let summary = "Extract a value from an aggregate value.";
   let description = [{