Merge branch 'main' into rename-do-something-fixture

Author: khalatepradnya

python/runtime/cudaq/platform/py_alt_launch_kernel.cpp

@@ -12,6 +12,7 @@
 #include "common/ArgumentConversion.h"
 #include "common/ArgumentWrapper.h"
 #include "common/Environment.h"
+#include "common/LayoutInfo.h"
 #include "cudaq/Optimizer/Builder/Marshal.h"
 #include "cudaq/Optimizer/Builder/Runtime.h"
 #include "cudaq/Optimizer/CAPI/Dialects.h"
@@ -169,31 +170,6 @@ py::args cudaq::simplifiedValidateInputArguments(py::args &args) {
   return processed;
 }
 
-std::pair<std::size_t, std::vector<std::size_t>>
-cudaq::getTargetLayout(mlir::ModuleOp mod, cudaq::cc::StructType structTy) {
-  mlir::StringRef dataLayoutSpec = "";
-  if (auto attr = mod->getAttr(cudaq::opt::factory::targetDataLayoutAttrName))
-    dataLayoutSpec = mlir::cast<mlir::StringAttr>(attr);
-  else
-    throw std::runtime_error("No data layout attribute is set on the module.");
-
-  auto dataLayout = llvm::DataLayout(dataLayoutSpec);
-  // Convert bufferTy to llvm.
-  llvm::LLVMContext context;
-  mlir::LLVMTypeConverter converter(structTy.getContext());
-  cudaq::opt::initializeTypeConversions(converter);
-  auto llvmDialectTy = converter.convertType(structTy);
-  mlir::LLVM::TypeToLLVMIRTranslator translator(context);
-  auto *llvmStructTy =
-      mlir::cast<llvm::StructType>(translator.translateType(llvmDialectTy));
-  auto *layout = dataLayout.getStructLayout(llvmStructTy);
-  auto strSize = layout->getSizeInBytes();
-  std::vector<std::size_t> fieldOffsets;
-  for (std::size_t i = 0, I = structTy.getMembers().size(); i != I; ++i)
-    fieldOffsets.emplace_back(layout->getElementOffset(i));
-  return {strSize, fieldOffsets};
-}
-
 void cudaq::handleStructMemberVariable(void *data, std::size_t offset,
                                        mlir::Type memberType,
                                        py::object value) {
@@ -626,75 +602,11 @@ cudaq::OpaqueArguments *cudaq::toOpaqueArgs(py::args &args, MlirModule mod,
 static void appendTheResultValue(ModuleOp module, const std::string &name,
                                  cudaq::OpaqueArguments &runtimeArgs,
                                  Type returnType) {
-  TypeSwitch<Type, void>(returnType)
-      .Case([&](IntegerType type) {
-        if (type.getIntOrFloatBitWidth() == 1) {
-          bool *ourAllocatedArg = new bool();
-          *ourAllocatedArg = 0;
-          runtimeArgs.emplace_back(ourAllocatedArg, [](void *ptr) {
-            delete static_cast<bool *>(ptr);
-          });
-          return;
-        }
-
-        long *ourAllocatedArg = new long();
-        *ourAllocatedArg = 0;
-        runtimeArgs.emplace_back(ourAllocatedArg, [](void *ptr) {
-          delete static_cast<long *>(ptr);
-        });
-      })
-      .Case([&](ComplexType type) {
-        Py_complex *ourAllocatedArg = new Py_complex();
-        ourAllocatedArg->real = 0.0;
-        ourAllocatedArg->imag = 0.0;
-        runtimeArgs.emplace_back(ourAllocatedArg, [](void *ptr) {
-          delete static_cast<Py_complex *>(ptr);
-        });
-      })
-      .Case([&](Float64Type type) {
-        double *ourAllocatedArg = new double();
-        *ourAllocatedArg = 0.;
-        runtimeArgs.emplace_back(ourAllocatedArg, [](void *ptr) {
-          delete static_cast<double *>(ptr);
-        });
-      })
-      .Case([&](Float32Type type) {
-        float *ourAllocatedArg = new float();
-        *ourAllocatedArg = 0.;
-        runtimeArgs.emplace_back(ourAllocatedArg, [](void *ptr) {
-          delete static_cast<float *>(ptr);
-        });
-      })
-      .Case([&](cudaq::cc::StdvecType ty) {
-        // Vector is a span: `{ data, length }`.
-        struct vec {
-          char *data;
-          std::size_t length;
-        };
-        vec *ourAllocatedArg = new vec{nullptr, 0};
-        runtimeArgs.emplace_back(
-            ourAllocatedArg, [](void *ptr) { delete static_cast<vec *>(ptr); });
-      })
-      .Case([&](cudaq::cc::StructType ty) {
-        auto [size, offsets] = cudaq::getTargetLayout(module, ty);
-        auto ourAllocatedArg = std::malloc(size);
-        runtimeArgs.emplace_back(ourAllocatedArg,
-                                 [](void *ptr) { std::free(ptr); });
-      })
-      .Case([&](cudaq::cc::CallableType ty) {
-        // Callables may not be returned from entry-point kernels. Append a
-        // dummy value as a placeholder.
-        runtimeArgs.emplace_back(nullptr, [](void *) {});
-      })
-      .Default([](Type ty) {
-        std::string msg;
-        {
-          llvm::raw_string_ostream os(msg);
-          ty.print(os);
-        }
-        throw std::runtime_error("Unsupported CUDA-Q kernel return type - " +
-                                 msg + ".\n");
-      });
+  auto [bufferSize, offsets] = cudaq::getResultBufferLayout(module, returnType);
+  if (bufferSize == 0)
+    return;
+  auto *buf = std::calloc(1, bufferSize);
+  runtimeArgs.emplace_back(buf, [](void *ptr) { std::free(ptr); });
 }
 
 // Launching the module \p mod will modify its content, such as by argument

python/utils/OpaqueArguments.h

@@ -111,10 +111,6 @@ void valueArgument(OpaqueArguments &argData, T *arg) {
 
 std::string mlirTypeToString(mlir::Type ty);
 
-/// @brief Return the size and member variable offsets for the input struct.
-std::pair<std::size_t, std::vector<std::size_t>>
-getTargetLayout(mlir::ModuleOp mod, cudaq::cc::StructType structTy);
-
 /// For the current struct member variable type, insert the value into the
 /// dynamically constructed struct.
 void handleStructMemberVariable(void *data, std::size_t offset,

runtime/common/LayoutInfo.cpp

@@ -12,6 +12,7 @@
 #include "cudaq/Optimizer/Builder/Runtime.h"
 #include "cudaq/Optimizer/CodeGen/QIROpaqueStructTypes.h"
 #include "cudaq/runtime/logger/logger.h"
+#include "llvm/ADT/TypeSwitch.h"
 #include "llvm/IR/DataLayout.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Types.h"
@@ -100,4 +101,65 @@ LayoutInfoType getLayoutInfo(const std::string &name, void *opt_module) {
     return extractLayout(name, quakeCode);
   return {};
 }
+
+LayoutInfoType getTargetLayout(ModuleOp mod, cc::StructType structTy) {
+  StringRef dataLayoutSpec = "";
+  if (auto attr = mod->getAttr(opt::factory::targetDataLayoutAttrName))
+    dataLayoutSpec = cast<StringAttr>(attr);
+  else
+    throw std::runtime_error("No data layout attribute is set on the module.");
+
+  auto dataLayout = llvm::DataLayout(dataLayoutSpec);
+  llvm::LLVMContext context;
+  LLVMTypeConverter converter(structTy.getContext());
+  opt::initializeTypeConversions(converter);
+  auto llvmDialectTy = converter.convertType(structTy);
+  LLVM::TypeToLLVMIRTranslator translator(context);
+  auto *llvmStructTy =
+      cast<llvm::StructType>(translator.translateType(llvmDialectTy));
+  auto *layout = dataLayout.getStructLayout(llvmStructTy);
+  auto strSize = layout->getSizeInBytes();
+  std::vector<std::size_t> fieldOffsets;
+  for (std::size_t i = 0, I = structTy.getMembers().size(); i != I; ++i)
+    fieldOffsets.emplace_back(layout->getElementOffset(i));
+  return {strSize, fieldOffsets};
+}
+
+LayoutInfoType getResultBufferLayout(ModuleOp mod, Type resultTy) {
+  std::size_t bufferSize = 0;
+  std::vector<std::size_t> fieldOffsets;
+
+  TypeSwitch<Type, void>(resultTy)
+      .Case([&](IntegerType ty) {
+        bufferSize =
+            ty.getIntOrFloatBitWidth() == 1 ? sizeof(bool) : sizeof(long);
+      })
+      .Case([&](ComplexType) { bufferSize = 2 * sizeof(double); })
+      .Case([&](Float64Type) { bufferSize = sizeof(double); })
+      .Case([&](Float32Type) { bufferSize = sizeof(float); })
+      .Case([&](cc::StdvecType) {
+        struct vec {
+          void *data;
+          std::size_t length;
+        };
+        bufferSize = sizeof(vec);
+      })
+      .Case([&](cc::StructType ty) {
+        auto [size, offsets] = getTargetLayout(mod, ty);
+        bufferSize = size;
+        fieldOffsets = std::move(offsets);
+      })
+      .Case([&](cc::CallableType) {})
+      .Default([](Type ty) {
+        std::string msg;
+        {
+          llvm::raw_string_ostream os(msg);
+          ty.print(os);
+        }
+        throw std::runtime_error("Unsupported CUDA-Q kernel return type - " +
+                                 msg + ".\n");
+      });
+
+  return {bufferSize, std::move(fieldOffsets)};
+}
 } // namespace cudaq

runtime/common/LayoutInfo.h

@@ -10,10 +10,28 @@
 #include <string>
 #include <vector>
 
+namespace mlir {
+class ModuleOp;
+class Type;
+} // namespace mlir
+
 namespace cudaq {
+namespace cc {
+class StructType;
+} // namespace cc
 
 using LayoutInfoType = std::pair<std::size_t, std::vector<std::size_t>>;
 
 LayoutInfoType getLayoutInfo(const std::string &name,
                              void *opt_module = nullptr);
+
+/// @brief Compute struct size and field offsets using the module's data layout.
+LayoutInfoType getTargetLayout(mlir::ModuleOp mod, cc::StructType structTy);
+
+/// @brief Compute the host-side buffer size (and struct field offsets, if
+/// applicable) for the given MLIR kernel return type. Returns {0, {}} for
+/// types that do not require a result buffer (e.g. `CallableType`). Throws on
+/// unsupported types.
+LayoutInfoType getResultBufferLayout(mlir::ModuleOp mod, mlir::Type resultTy);
+
 } // namespace cudaq