🎨 Rework unitary matrix proposal

Author: taminob

mlir/include/mlir/Dialect/QCO/IR/QCODialect.h

@@ -21,8 +21,6 @@
 #pragma clang diagnostic pop
 #endif
 
-#include "mlir/Dialect/Utils/MatrixUtils.h"
-
 #include <Eigen/Core>
 #include <mlir/Bytecode/BytecodeOpInterface.h>
 #include <mlir/Dialect/Arith/IR/Arith.h>
@@ -79,9 +77,7 @@ tryGetParameterAsDouble(UnitaryOpInterface op, size_t i);
  *                          not have a matrix definition, set this value to
  *                          nullptr.
  */
-template <size_t T, size_t P, typename UnitaryMatrixType,
-          UnitaryMatrixType (*MatrixDefinition)(UnitaryOpInterface)>
-class TargetAndParameterArityTrait {
+template <size_t T, size_t P> class TargetAndParameterArityTrait {
 public:
   template <typename ConcreteType>
   class Impl : public OpTrait::TraitBase<ConcreteType, Impl> {
@@ -128,13 +124,6 @@ class TargetAndParameterArityTrait {
       return this->getOperation()->getOperand(T + i);
     }
 
-    [[nodiscard]] UnitaryMatrixType getUnitaryMatrixDefinition() const
-      requires(MatrixDefinition != nullptr)
-    {
-      const auto* op = this->getConstOperation();
-      return MatrixDefinition(llvm::dyn_cast<UnitaryOpInterface>(op));
-    }
-
     Value getInputForOutput(Value output) {
       const auto& op = this->getOperation();
       for (size_t i = 0; i < T; ++i) {
@@ -168,29 +157,6 @@ class TargetAndParameterArityTrait {
 
 } // namespace mlir::qco
 
-//===----------------------------------------------------------------------===//
-// Operations Helpers
-//===----------------------------------------------------------------------===//
-
-namespace mlir::qco {
-
-[[nodiscard]] inline std::optional<double>
-tryGetParameterAsDouble(UnitaryOpInterface op, size_t i) {
-  using DummyArityType =
-      TargetAndParameterArityTrait<0, 0, Eigen::MatrixXcd, nullptr>;
-  const auto param = op.getParameter(i);
-  const auto floatAttr =
-      DummyArityType::Impl<arith::ConstantOp>::getStaticParameter(param);
-  if (!floatAttr) {
-    return std::nullopt;
-  }
-  return floatAttr.getValueAsDouble();
-}
-
-UnitaryOpInterface getControlledOp(UnitaryOpInterface op);
-
-} // namespace mlir::qco
-
 //===----------------------------------------------------------------------===//
 // Operations
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/QCO/IR/QCOInterfaces.td

@@ -109,25 +109,29 @@ def UnitaryOpInterface : OpInterface<"UnitaryOpInterface"> {
             "StringRef", "getBaseSymbol", (ins)
         >,
 
-        // Unitary matrix methods
+        // Internal helper methods, should not be used by user
         InterfaceMethod<
             "Returns the unitary matrix definition of the operation.",
-            "Eigen::MatrixXcd", "getUnitaryMatrix", (ins),
             [{
-              if constexpr (requires { $_op.getUnitaryMatrixDefinition(); }) {
-                return $_op.getUnitaryMatrixDefinition();
-              } else {
-                llvm::reportFatalUsageError("Operation '" + $_op.getBaseSymbol() + "' has no unitary matrix definition!");
-              }
-            }]
-        >,
-        InterfaceMethod<
-            "Returns the unitary matrix definition of the operation.",
-            "std::variant<Eigen::Matrix<std::complex<double>, 1, 1>, Eigen::Matrix2cd, Eigen::Matrix4cd, Eigen::MatrixXcd>",
-            "getFastUnitaryMatrixVariant", (ins),
+              std::variant<
+                std::nullopt_t,
+                Eigen::Matrix<std::complex<double>, 1, 1>,
+                Eigen::Matrix2cd,
+                Eigen::Matrix4cd,
+                Eigen::MatrixXcd
+              >
+            }],
+            "internalGetUnitaryVariant", (ins),
             [{
-              if constexpr (requires { $_op.getUnitaryMatrixDefinition(); }) {
-                return $_op.getUnitaryMatrixDefinition();
+              if constexpr (requires { $_op.getUnitaryMatrix().has_value(); }) {
+                // is matrix in std::optional
+                if (auto&& matrix = $_op.getUnitaryMatrix()) {
+                  return *matrix;
+                }
+                return std::nullopt;
+              } else if constexpr (requires { $_op.getUnitaryMatrix(); }) {
+                // is static matrix
+                return $_op.getUnitaryMatrix();
               } else {
                 llvm::reportFatalUsageError("Operation '" + $_op.getBaseSymbol() + "' has no unitary matrix definition!");
               }
@@ -137,11 +141,18 @@ def UnitaryOpInterface : OpInterface<"UnitaryOpInterface"> {
 
     let extraClassDeclaration = [{
       template<typename MatrixType>
-      MatrixType getFastUnitaryMatrix() const {
-        // removing const since it's safe to say UnitaryOpInterface::getFastUnitaryMatrixVariant()
-        // and OpType::getUnitaryMatrixDefinition() will not modify the state of this
-        auto&& matrix = const_cast<UnitaryOpInterface*>(this)->getFastUnitaryMatrixVariant();
-        return std::get<MatrixType>(matrix);
+      std::optional<MatrixType> getUnitaryMatrix() {
+        auto&& matrix = this->internalGetUnitaryVariant();
+
+        if constexpr (std::is_same_v<std::remove_cvref_t<MatrixType>, Eigen::MatrixXcd>) {
+          // cast to Eigen::MatrixXcd via visitor pattern
+          auto&& visitor = [](auto&& m) -> std::optional<MatrixType> { return m; };
+          return std::visit(visitor, matrix);
+        } else {
+          // if an explicit size has been requested, use checked access to variant
+          // to return the correct type (throws on wrong type)
+          return std::get<MatrixType>(matrix);
+        }
       }
     }];
 }