Merge branch 'main' into use-zstd-layers

Author: mitchdz

.github/workflows/config/spelling_allowlist.txt

@@ -313,6 +313,7 @@ runtimes
 rvalue
 scalability
 scalable
+selectable
 sexualized
 shifter
 shifters

include/cudaq/Optimizer/Builder/Factory.h

@@ -89,19 +89,23 @@ mlir::Type genArgumentBufferType(mlir::Type ty);
 ///
 /// A kernel signature of
 /// ```c++
-/// i32_t operator() (i16_t, std::vector<double>, double);
+///   i32_t operator() (i16_t, std::vector<double>, double);
 /// ```
 /// will generate the LLVM struct
 /// ```llvm
-/// { i16, i64, double, i32 }
+///   { i16, i64, double, i32 }
 /// ```
 /// where the values of the vector argument are pass-by-value and appended to
 /// the end of the struct as a sequence of \i n double values.
 ///
 /// The leading `startingArgIdx + 1` parameters are omitted from the struct.
+///
+/// NB: It is DEEPLY INCORRECT to add a packed attribute to this data structure
+/// and pass it to other APIs, since there is absolutely, positively NO chance
+/// that foreign code will be able to decode this buffer correctly. To do so
+/// requires information that is erased by the NVQ++ compiler.
 cudaq::cc::StructType buildInvokeStructType(mlir::FunctionType funcTy,
-                                            std::size_t startingArgIdx = 0,
-                                            bool packed = false);
+                                            std::size_t startingArgIdx = 0);
 
 /// Return the LLVM-IR dialect type: `[length x i8]`.
 inline mlir::Type getStringType(mlir::MLIRContext *ctx, std::size_t length) {

lib/Optimizer/Builder/Factory.cpp

@@ -78,17 +78,17 @@ Type factory::genArgumentBufferType(Type ty) {
   return genBufferType</*isOutput=*/false>(ty);
 }
 
-cudaq::cc::StructType factory::buildInvokeStructType(FunctionType funcTy,
-                                                     std::size_t startingArgIdx,
-                                                     bool packed) {
+cudaq::cc::StructType
+factory::buildInvokeStructType(FunctionType funcTy,
+                               std::size_t startingArgIdx) {
   auto *ctx = funcTy.getContext();
   SmallVector<Type> eleTys;
   for (auto inTy : llvm::enumerate(funcTy.getInputs()))
     if (inTy.index() >= startingArgIdx)
       eleTys.push_back(genBufferType</*isOutput=*/false>(inTy.value()));
   for (auto outTy : funcTy.getResults())
     eleTys.push_back(genBufferType</*isOutput=*/true>(outTy));
-  return cudaq::cc::StructType::get(ctx, eleTys, packed);
+  return cudaq::cc::StructType::get(ctx, eleTys, /*packed=*/false);
 }
 
 Value factory::packIsArrayAndLengthArray(Location loc,

lib/Optimizer/CodeGen/VerifyQIRProfile.cpp

@@ -46,15 +46,15 @@ struct VerifyQIRProfilePass
     if (!func->hasAttr(cudaq::entryPointAttrName))
       return;
     auto *ctx = &getContext();
-    bool isBaseProfile = convertTo.getValue() == "qir-base";
+    const bool isBaseProfile = convertTo.getValue() == "qir-base";
     func.walk([&](Operation *op) {
       if (auto call = dyn_cast<LLVM::CallOp>(op)) {
         auto funcNameAttr = call.getCalleeAttr();
         if (!funcNameAttr)
           return WalkResult::advance();
         auto funcName = funcNameAttr.getValue();
-        if (!funcName.startswith("__quantum_") ||
-            funcName == cudaq::opt::QIRCustomOp) {
+        if (isBaseProfile && (!funcName.startswith("__quantum_") ||
+                              funcName.equals(cudaq::opt::QIRCustomOp))) {
           call.emitOpError("unexpected call in QIR base profile");
           passFailed = true;
           return WalkResult::advance();

runtime/common/BaseRemoteRESTQPU.h

@@ -34,6 +34,7 @@
 #include "cudaq/platform/qpu.h"
 #include "cudaq/platform/quantum_platform.h"
 #include "nvqpp_config.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/Bitcode/BitcodeReader.h"
 #include "llvm/Bitcode/BitcodeWriter.h"
 #include "llvm/IR/Module.h"
@@ -444,6 +445,87 @@ class BaseRemoteRESTQPU : public cudaq::QPU {
     auto func = m_module.lookupSymbol<mlir::func::FuncOp>(
         std::string(cudaq::runtime::cudaqGenPrefixName) + kernelName);
 
+    llvm::SmallVector<mlir::func::FuncOp> newFuncOpsWithDefinitions;
+    llvm::SmallSet<std::string, 4> deviceCallCallees;
+    // For every declaration without a definition, we need to try to find the
+    // function in the Quake registry and copy the functions into this module.
+    m_module.walk([&](mlir::func::FuncOp funcOp) {
+      if (!funcOp.isDeclaration()) {
+        // Skipping function because it already has a definition.
+        return mlir::WalkResult::advance();
+      }
+      // Definition doesn't exist, so we need to find it in the Quake registry.
+      mlir::StringRef fullFuncName = funcOp.getName();
+      mlir::StringRef kernelName = [fullFuncName]() {
+        mlir::StringRef retVal = fullFuncName;
+        // TODO - clean this up to not have to do this. Considering the module's
+        // map, or cudaq::details::getKernelName(). But make sure it works for
+        // standard C++ functions.
+
+        // Only get the portion before the first ".".
+        if (auto ix = fullFuncName.find("."); ix != mlir::StringRef::npos)
+          retVal = fullFuncName.substr(0, ix);
+        // Also strip out __nvqpp_mlirgen__function_ from the beginning of the
+        // function name.
+        if (retVal.starts_with(cudaq::runtime::cudaqGenPrefixName)) {
+          retVal = retVal.substr(cudaq::runtime::cudaqGenPrefixLength);
+        }
+        return retVal;
+      }();
+      std::string quakeCode =
+          kernelName.empty()
+              ? ""
+              : cudaq::get_quake_by_name(kernelName.str(),
+                                         /*throwException=*/false);
+      if (quakeCode.empty()) {
+        // Skipping function because it does not have a quake code.
+        return mlir::WalkResult::advance();
+      }
+      auto tmp_module =
+          parseSourceString<mlir::ModuleOp>(quakeCode, contextPtr);
+      auto tmpFuncOpWithDefinition =
+          tmp_module->lookupSymbol<mlir::func::FuncOp>(fullFuncName);
+      auto newNameAttr = mlir::StringAttr::get(m_module.getContext(),
+                                               fullFuncName.str() + ".stitch");
+      auto clonedFunc = tmpFuncOpWithDefinition.clone();
+      clonedFunc.setName(newNameAttr);
+      mlir::SymbolTable symTable(m_module);
+      symTable.insert(clonedFunc);
+      newFuncOpsWithDefinitions.push_back(clonedFunc);
+
+      if (failed(mlir::SymbolTable::replaceAllSymbolUses(
+              funcOp.getOperation(), newNameAttr, m_module.getOperation()))) {
+        throw std::runtime_error(
+            fmt::format("Failed to replace symbol uses for function {}",
+                        fullFuncName.str()));
+      }
+      return mlir::WalkResult::advance();
+    });
+
+    // For each one of the added functions, we need to traverse them to find
+    // device calls (in order to create declarations for them)
+    for (auto &funcOp : newFuncOpsWithDefinitions) {
+      mlir::OpBuilder builder(m_module);
+      builder.setInsertionPointToStart(m_module.getBody());
+      funcOp.walk([&](cudaq::cc::DeviceCallOp deviceCall) {
+        auto calleeName = deviceCall.getCallee();
+        // If the callee is already in the symbol table, nothing to do.
+        if (m_module.lookupSymbol<mlir::func::FuncOp>(calleeName))
+          return;
+
+        // Otherwise, we need to create a declaration for the callback function.
+        auto argTypes = deviceCall.getArgs().getTypes();
+        auto resTypes = deviceCall.getResultTypes();
+        auto funcType = builder.getFunctionType(argTypes, resTypes);
+
+        // Create a *declaration* (no body) for the callback function.
+        [[maybe_unused]] auto decl = builder.create<mlir::func::FuncOp>(
+            deviceCall.getLoc(), calleeName, funcType);
+        decl.setPrivate();
+        deviceCallCallees.insert(calleeName.str());
+      });
+    }
+
     // Create a new Module to clone the function into
     auto location = mlir::FileLineColLoc::get(&context, "<builder>", 1, 1);
     mlir::ImplicitLocOpBuilder builder(location, &context);
@@ -463,7 +545,7 @@ class BaseRemoteRESTQPU : public cudaq::QPU {
       // passes.
       if (auto lfunc = dyn_cast<mlir::func::FuncOp>(op)) {
         bool skip = lfunc.getName().ends_with(".thunk");
-        if (!skip)
+        if (!skip && !deviceCallCallees.contains(lfunc.getName().str()))
           for (auto &entry : mangledNameMap)
             if (lfunc.getName() ==
                 cast<mlir::StringAttr>(entry.getValue()).getValue()) {

runtime/common/RuntimeMLIR.cpp

@@ -7,6 +7,7 @@
  ******************************************************************************/
 
 #include "RuntimeMLIR.h"
+#include "ThunkInterface.h"
 #include "cudaq/Optimizer/Builder/Runtime.h"
 #include "cudaq/Optimizer/CodeGen/IQMJsonEmitter.h"
 #include "cudaq/Optimizer/CodeGen/OpenQASMEmitter.h"
@@ -16,6 +17,7 @@
 #include "cudaq/Optimizer/Dialect/Quake/QuakeDialect.h"
 #include "cudaq/Optimizer/InitAllDialects.h"
 #include "cudaq/Optimizer/InitAllPasses.h"
+#include "cudaq/Support/TargetConfig.h"
 #include "llvm/Bitcode/BitcodeWriter.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/MC/SubtargetFeature.h"
@@ -93,4 +95,22 @@ std::unique_ptr<MLIRContext> initializeMLIR() {
   return context;
 }
 
+std::optional<std::string> getEntryPointName(OwningOpRef<ModuleOp> &module) {
+  std::string name;
+  // FIXME: don't use a recursive `walk` to DFS for FuncOps, which appear as
+  // children, in a Module.
+  module->walk([&name](mlir::func::FuncOp op) {
+    if (op.getName().endswith(".thunk")) {
+      name = op.getName();
+      return mlir::WalkResult::interrupt();
+    }
+    return mlir::WalkResult::advance();
+  });
+
+  if (!name.empty())
+    return name;
+
+  return std::nullopt;
+}
+
 } // namespace cudaq