Decomp/pbc decomposable interface

mlir/include/PBC/IR/PBCOpInterfaces.h

@@ -14,10 +14,23 @@
 
 #pragma once
 
+#include "llvm/ADT/StringMap.h" // for DecomposableGate interface
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/Types.h"
+#include "mlir/Support/LLVM.h"
 
 //===----------------------------------------------------------------------===//
 // PBC interface declarations.
 //===----------------------------------------------------------------------===//
 
+namespace catalyst {
+namespace pbc {
+
+std::string defaultGetGraphOpId(mlir::Operation *op);
+
+}
+} // namespace catalyst
+
 #include "PBC/IR/PBCOpInterfaces.h.inc"

mlir/include/PBC/IR/PBCOpInterfaces.td

@@ -56,4 +56,48 @@ def PBCOpInterface : OpInterface<"PBCOpInterface"> {
 
 }
 
+def DecomposableGate : OpInterface<"DecomposableGate"> {
+    let description = [{
+        This interface provides a generic way to query decomposition data from a quantum operation.
+        An operation must implement this interface in order to compatible with the `graph-decomposition` pass.
+        Note that implementing this interface does not guarantee that an operation will decomposed successfully, only that it will be registered with the graph.
+    }];
+
+    let cppNamespace = "::catalyst::pbc";
+
+    let methods = [
+        InterfaceMethod<
+            "Return the name of the corresponding frontend operator (i.e. the rule supplier).", "std::string", "getOperatorName"
+        >,
+        InterfaceMethod<
+            "Return a map of frontend parameter name to shape and dtype of dynamic data.",
+            "llvm::StringMap<llvm::SmallVector<mlir::Type>>", "getDynamicShape"
+        >,
+        InterfaceMethod<
+            "Return a map of frontend parameter name to number of wires (excluding control wires).",
+            "llvm::StringMap<size_t>", "getWireLens"
+        >,
+        InterfaceMethod<
+            "Return a map of frontend parameter name to static data values.",
+            "mlir::DictionaryAttr", "getStaticData"
+        >,
+        InterfaceMethod<
+            "Return a string representation of any additional data an operator may need to provide"
+            " for uniqueness (ex. UID), or emptystring if not applicable.",
+            "std::string", "getExtraData", (ins), [{}], [{ return ""; }]
+        >,
+        InterfaceMethod<
+            "Return the operation's ID for use with the graph. "
+            "This should be computed by the following convention:\n"
+            "```\n"
+            "name + {map params to list of types} + {map wire args to number of wires} + {static data} + [extra data]\n"
+            "```\n"
+            "Note that extra data is currently specific to `OperatorOp`s, and should only appear when present."
+            "Each of these maps should be sorted lexicographically by key."
+            "In general, this should be computed by the default here and not overridden.",
+            "std::string", "getGraphOpId", (ins), [{}], [{ return ::catalyst::pbc::defaultGetGraphOpId($_op.getOperation()); }]
+        >
+    ];
+}
+
 #endif // PBCOP_INTERFACES

mlir/include/PBC/IR/PBCOps.td

@@ -136,7 +136,8 @@ def FabricateOp : PBC_Op<"fabricate",
 }
 
 def PPRotationOp : PBC_Op<"ppr", [PBCOpInterface, AttrSizedOperandSegments,
-    DeclareOpInterfaceMethods<ResourceQuantumOpInterface, ["getResourceAdjointFlag", "getResourceDetailedName"]>]> {
+    DeclareOpInterfaceMethods<ResourceQuantumOpInterface, ["getResourceAdjointFlag", "getResourceDetailedName"]>,
+    DeclareOpInterfaceMethods<DecomposableGate>]> {
     let summary = "Pauli Product Rotation on qubits.";
 
     let description = [{
@@ -249,12 +250,22 @@ def PPRotationOp : PBC_Op<"ppr", [PBCOpInterface, AttrSizedOperandSegments,
         bool isClifford(){
             return hasPiOverTwoRotation() || hasPiOverFourRotation();  // π/2 and π/4 rotations are Clifford
         };
+        std::string getPauliWord()
+        {
+            std::string pauliWord;
+
+            for (mlir::Attribute pauliChar : getPauliProduct()) {
+                pauliWord += cast<mlir::StringAttr>(pauliChar).getValue();
+            }
+            return pauliWord;
+        }
     }];
     let extraClassDeclaration = extraBaseClassDeclaration;
 }
 
 def PPRotationArbitraryOp : PBC_Op<"ppr.arbitrary", [PBCOpInterface, AttrSizedOperandSegments,
-    DeclareOpInterfaceMethods<ResourceQuantumOpInterface, ["getResourceDetailedName"]>]> {
+    DeclareOpInterfaceMethods<ResourceQuantumOpInterface, ["getResourceDetailedName"]>,
+    DeclareOpInterfaceMethods<DecomposableGate>]> {
     let summary = "Pauli Product Rotation with arbitrary angle.";
     let description = [{
         The PPRotationArbitraryOp represents a Pauli product rotation operation with an arbitrary angle.
@@ -354,10 +365,24 @@ def PPRotationArbitraryOp : PBC_Op<"ppr.arbitrary", [PBCOpInterface, AttrSizedOp
 
     let hasVerifier = 1;
     let hasCanonicalizeMethod = 1;
+
+    code extraBaseClassDeclaration = [{
+        std::string getPauliWord()
+        {
+            std::string pauliWord;
+
+            for (mlir::Attribute pauliChar : getPauliProduct()) {
+                pauliWord += cast<mlir::StringAttr>(pauliChar).getValue();
+            }
+            return pauliWord;
+        }
+    }];
+    let extraClassDeclaration = extraBaseClassDeclaration;
 }
 
 def PPMeasurementOp : PBC_Op<"ppm", [PBCOpInterface,
-    DeclareOpInterfaceMethods<ResourceQuantumOpInterface, ["getResourceAdjointFlag", "getResourceDetailedName"]>]> {
+    DeclareOpInterfaceMethods<ResourceQuantumOpInterface, ["getResourceAdjointFlag", "getResourceDetailedName"]>,
+    DeclareOpInterfaceMethods<DecomposableGate>]> {
     let summary = "Pauli Product Measurement on qubits.";
 
     let description = [{
@@ -424,6 +449,18 @@ def PPMeasurementOp : PBC_Op<"ppm", [PBCOpInterface,
     }];
 
     let hasVerifier = 1;
+    code extraBaseClassDeclaration = [{
+        std::string getPauliWord()
+        {
+            std::string pauliWord;
+
+            for (mlir::Attribute pauliChar : getPauliProduct()) {
+                pauliWord += cast<mlir::StringAttr>(pauliChar).getValue();
+            }
+            return pauliWord;
+        }
+    }];
+    let extraClassDeclaration = extraBaseClassDeclaration;
 }
 
 def SelectPPMeasurementOp : PBC_Op<"select.ppm", [DeclareOpInterfaceMethods<ResourceQuantumOpInterface, ["getResourceDetailedName"]>]> {

mlir/lib/PBC/IR/PBCOpInterfaces.cpp

@@ -14,11 +14,160 @@
 
 #include "PBC/IR/PBCOpInterfaces.h"
 
+#include <cstddef>
+#include <cstdint>
+#include <string>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
+#include "mlir/IR/Types.h"
+#include "mlir/Support/LLVM.h"
+
 using namespace mlir;
 using namespace catalyst::pbc;
 
+#include "PBC/IR/PBCOpInterfaces.cpp.inc"
+//===----------------------------------------------------------------------===//
+// Helpers
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+void printAttr(mlir::Attribute attr, llvm::raw_string_ostream &ss) {
+    llvm::TypeSwitch<mlir::Attribute, void>(attr)
+        .Case<mlir::DictionaryAttr>([&](mlir::DictionaryAttr dict) {
+            ss << "{";
+            for (auto [i, entry] : llvm::enumerate(dict)) {
+                if (i > 0) {
+                    ss << ",";
+                }
+
+                ss << entry.getName().str() << ":";
+                printAttr(entry.getValue(), ss);
+            }
+            ss << "}";
+        })
+        .Case<mlir::ArrayAttr>([&](mlir::ArrayAttr arr) {
+            ss << "[";
+            for (auto [i, attr] : llvm::enumerate(arr)) {
+                if (i > 0) {
+                    ss << ",";
+                }
+                printAttr(attr, ss);
+            }
+            ss << "]";
+        })
+        .Case<mlir::StringAttr>([&](mlir::StringAttr attr) { ss << attr.str(); })
+        .Case<mlir::IntegerAttr>([&](mlir::IntegerAttr attr) { ss << attr.getInt(); })
+        .Case<mlir::FloatAttr>([&](mlir::FloatAttr attr) { ss << attr.getValueAsDouble(); })
+        .Default([&](mlir::Attribute attr) { attr.print(ss); });
+}
+
+void printShapedType(ArrayRef<int64_t> shape, int64_t dim, Type elementType,
+                     llvm::raw_string_ostream &ss) {
+    // Rank-0 tensors (e.g. tensor<f64>) have an empty shape; print the
+    // element type directly instead of indexing into the empty ArrayRef.
+    if (shape.empty()) {
+        ss << elementType;
+        return;
+    }
+
+    int64_t length = shape[dim];
+    auto printList = [&](auto printItem) {
+        ss << "[";
+        for (int64_t i = 0; i < length; i++) {
+            printItem();
+            if (i != length - 1) {
+                ss << ",";
+            }
+        }
+        ss << "]";
+    };
+
+    if (static_cast<int64_t>(shape.size()) == dim + 1) {
+        printList([&]() { ss << elementType; });
+    } else {
+        printList([&]() { printShapedType(shape, dim + 1, elementType, ss); });
+    }
+}
+
+void printType(mlir::Type type, llvm::raw_string_ostream &ss) {
+    llvm::TypeSwitch<mlir::Type, void>(type)
+        .Case<mlir::ShapedType>([&](mlir::ShapedType shapedType) {
+            printShapedType(shapedType.getShape(), 0, shapedType.getElementType(), ss);
+        })
+        .Default([&](mlir::Type other) { other.print(ss); });
+}
+
+template <typename T, typename PrintFunc>
+void printSortedMap(const llvm::StringMap<T> &map, llvm::raw_string_ostream &ss,
+                    PrintFunc printValue) {
+    llvm::SmallVector<llvm::StringRef> keys;
+    for (const llvm::StringRef key : map.keys()) {
+        keys.push_back(key);
+    }
+    llvm::sort(keys);
+
+    ss << "{";
+    for (auto [i, key] : llvm::enumerate(keys)) {
+        if (i > 0) {
+            ss << ",";
+        }
+        ss << key << ":";
+        printValue(map.lookup(key), ss);
+    }
+    ss << "}";
+}
+
+void printDynamicShape(const llvm::StringMap<llvm::SmallVector<mlir::Type>> &map,
+                       llvm::raw_string_ostream &ss) {
+    printSortedMap(map, ss, [](const auto &types, llvm::raw_string_ostream &stream) {
+        stream << "[";
+        for (auto [j, type] : llvm::enumerate(types)) {
+            if (j > 0) {
+                stream << ",";
+            }
+            printType(type, stream);
+        }
+        stream << "]";
+    });
+}
+
+void printWireLens(const llvm::StringMap<size_t> &map, llvm::raw_string_ostream &ss) {
+    printSortedMap(map, ss, [](size_t len, llvm::raw_string_ostream &stream) { stream << len; });
+}
+
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // PBC interface definitions.
 //===----------------------------------------------------------------------===//
 
-#include "PBC/IR/PBCOpInterfaces.cpp.inc"
+namespace catalyst {
+namespace pbc {
+
+std::string defaultGetGraphOpId(Operation *op) {
+    std::string out;
+    llvm::raw_string_ostream ss(out);
+
+    DecomposableGate gate = cast<DecomposableGate>(op);
+
+    ss << gate.getOperatorName();
+    printDynamicShape(gate.getDynamicShape(), ss);
+    printWireLens(gate.getWireLens(), ss);
+    printAttr(gate.getStaticData(), ss);
+    if (gate.getExtraData() != "") {
+        ss << '[' << gate.getExtraData() << ']';
+    }
+    ss.flush();
+
+    return out;
+}
+
+} // namespace pbc
+} // namespace catalyst

mlir/lib/PBC/IR/PBCOps.cpp

@@ -231,6 +231,58 @@ void LayerOp::print(OpAsmPrinter &p) {
                   /*printBlockTerminators=*/!getInitArgs().empty());
 }
 
+//===----------------------------------------------------------------------===//
+// PBC op interface methods.
+//===----------------------------------------------------------------------===//
+
+// PPRotationOp
+std::string PPRotationOp::getOperatorName() { return "PauliRot"; }
+llvm::StringMap<llvm::SmallVector<mlir::Type>> PPRotationOp::getDynamicShape() {
+    return {{"theta", {mlir::Float64Type::get(getContext())}}};
+}
+llvm::StringMap<size_t> PPRotationOp::getWireLens() { return {{"wires", getInQubits().size()}}; }
+mlir::DictionaryAttr PPRotationOp::getStaticData() {
+    mlir::MLIRContext *ctx = getContext();
+    mlir::NamedAttribute pauliWordEntry = mlir::NamedAttribute(
+        mlir::StringAttr::get(ctx, "pauli_word"), mlir::StringAttr::get(ctx, getPauliWord()));
+    return mlir::DictionaryAttr::get(ctx, {pauliWordEntry});
+}
+
+// PPRotationArbitraryOp
+std::string PPRotationArbitraryOp::getOperatorName() { return "PauliRot"; }
+llvm::StringMap<llvm::SmallVector<mlir::Type>> PPRotationArbitraryOp::getDynamicShape() {
+    return {{"theta", {mlir::Float64Type::get(getContext())}}};
+}
+llvm::StringMap<size_t> PPRotationArbitraryOp::getWireLens() {
+    return {{"wires", getInQubits().size()}};
+}
+mlir::DictionaryAttr PPRotationArbitraryOp::getStaticData() {
+    mlir::MLIRContext *ctx = getContext();
+    mlir::NamedAttribute pauliWordEntry = mlir::NamedAttribute(
+        mlir::StringAttr::get(ctx, "pauli_word"), mlir::StringAttr::get(ctx, getPauliWord()));
+    return mlir::DictionaryAttr::get(ctx, {pauliWordEntry});
+}
+
+// PPMeasurementOp
+std::string PPMeasurementOp::getOperatorName() { return "PauliMeasure"; }
+llvm::StringMap<llvm::SmallVector<mlir::Type>> PPMeasurementOp::getDynamicShape() { return {}; }
+llvm::StringMap<size_t> PPMeasurementOp::getWireLens() { return {{"wires", getInQubits().size()}}; }
+mlir::DictionaryAttr PPMeasurementOp::getStaticData() {
+    mlir::MLIRContext *ctx = getContext();
+    mlir::NamedAttribute pauliWordEntry = mlir::NamedAttribute(
+        mlir::StringAttr::get(ctx, "pauli_word"), mlir::StringAttr::get(ctx, getPauliWord()));
+
+    // meas_uid and postselect are not lowered into PPMeasurementOp, so we set them to None
+    mlir::StringAttr noneStr = mlir::StringAttr::get(ctx, "None");
+    mlir::NamedAttribute measUidEntry =
+        mlir::NamedAttribute(mlir::StringAttr::get(ctx, "meas_uid"), noneStr);
+
+    mlir::NamedAttribute postselectEntry =
+        mlir::NamedAttribute(mlir::StringAttr::get(ctx, "postselect"), noneStr);
+
+    return mlir::DictionaryAttr::get(ctx, {measUidEntry, pauliWordEntry, postselectEntry});
+}
+
 //===----------------------------------------------------------------------===//
 // Implement ResourceQuantumOpInterface methods.
 //===----------------------------------------------------------------------===//

mlir/unittests/CMakeLists.txt

@@ -11,6 +11,7 @@ endfunction()
 
 if (CATALYST_GTEST_AVAILABLE)
   add_subdirectory(Example)
+  add_subdirectory(PBC)
   add_subdirectory(QRef)
   add_subdirectory(Quantum)
   add_subdirectory(Utils)

mlir/unittests/PBC/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(Interfaces)

mlir/unittests/PBC/Interfaces/CMakeLists.txt

@@ -0,0 +1,13 @@
+add_catalyst_unittest(CatalystPBCInterfaceTests
+    TestDecomposableGateInterface.cpp
+)
+
+get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+
+target_link_libraries(CatalystPBCInterfaceTests PRIVATE
+    ${dialect_libs}
+    MLIRQuantum
+    MLIRQRef
+    MLIRPBC
+    MLIRMBQC
+)