Optimized Ion Dialect to minimize redundant code in IonOp (#1471)

**Context:** The current design of the IonOp redefines all the level attributes of a transition from scratch, which is quite inefficient when lowering to llvm IR. This can be optimized by labeling the pre-defined levels (upstate, downstate, estate) and just referencing them in the transition attribute. **Description of the Change:** - Added a string attribute `label` to Level. - Changed the levels of a transition from `LevelAttr` to `string` **Benefits:** Can reference the label of a level instead redefining the level attribute with all its parameters which results in reduced code size. **Possible Drawbacks:** **Related GitHub Issues:** --------- Co-authored-by: Haochen Wang <[email protected]>
PennyLaneAI · Jan 21, 2025 · 682f7b6 · 682f7b6
1 parent d33814c
commit 682f7b6
Show file tree

Hide file tree

Showing 8 changed files with 49 additions and 129 deletions.
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
@@ -30,6 +30,10 @@
 * Replace `ValueRange` with `ResultRange` and `Value` with `OpResult` to better align with the semantics of `**QubitResult()` functions like `getNonCtrlQubitResults()`. This change ensures clearer intent and usage. Improve the `matchAndRewrite` function by using `replaceAllUsesWith` instead of for loop.
   [(#1426)](https://github.com/PennyLaneAI/catalyst/pull/1426)
 
+* Improved ion dialect to reduce redundant code generated. Added a string attribute `label` to Level.
+  Also changed the levels of a transition from `LevelAttr` to `string`
+  [(#1471)](https://github.com/PennyLaneAI/catalyst/pull/1471)
+
 
 <h3>Documentation 📝</h3>
 
@@ -39,4 +43,4 @@ This release contains contributions from (in alphabetical order):
 
 Christina Lee
 Mehrdad Malekmohammadi
-Sengthai Heng
+Sengthai Heng
diff --git a/mlir/include/Ion/IR/IonOps.td b/mlir/include/Ion/IR/IonOps.td
@@ -50,6 +50,7 @@ def LevelAttr : Ion_Attr<"Level", "level"> {
   let summary = "A class to represent an atomic level.";
 
   let parameters = (ins
+    "mlir::StringAttr":$label,
     "mlir::IntegerAttr":$principal,
     "mlir::FloatAttr":$spin,
     "mlir::FloatAttr":$orbital,
@@ -63,6 +64,7 @@ def LevelAttr : Ion_Attr<"Level", "level"> {
 
   let builders = [
     AttrBuilderWithInferredContext<(ins     
+    "mlir::StringAttr":$label,
     "mlir::IntegerAttr":$principal,
     "mlir::FloatAttr":$spin,
     "mlir::FloatAttr":$orbital,
@@ -71,7 +73,7 @@ def LevelAttr : Ion_Attr<"Level", "level"> {
     "mlir::FloatAttr":$spin_orbital_nuclear,
     "mlir::FloatAttr":$spin_orbital_nuclear_magnetization,
     "mlir::FloatAttr":$energy), [{
-      return $_get(principal.getContext(), principal, spin, orbital, nuclear, spin_orbital, spin_orbital_nuclear, spin_orbital_nuclear_magnetization, energy);
+      return $_get(principal.getContext(), label, principal, spin, orbital, nuclear, spin_orbital, spin_orbital_nuclear, spin_orbital_nuclear_magnetization, energy);
     }]>
   ];
 
@@ -82,14 +84,14 @@ def TransitionAttr : Ion_Attr<"Transition", "transition"> {
   let summary = "A class to represent a atomic transition between two levels.";
 
   let parameters = (ins
-    "LevelAttr":$level_0,
-    "LevelAttr":$level_1,
+    "mlir::StringAttr":$level_0,
+    "mlir::StringAttr":$level_1,
     "mlir::FloatAttr":$einstein_a
   );
 
   let builders = [
-    AttrBuilderWithInferredContext<(ins "LevelAttr":$level_0,
-                                        "LevelAttr":$level_1,
+    AttrBuilderWithInferredContext<(ins "mlir::StringAttr":$level_0,
+                                        "mlir::StringAttr":$level_1,
                                         "mlir::FloatAttr":$einstein_a), [{
       return $_get(einstein_a.getContext(), level_0, level_1, einstein_a);
     }]>

diff --git a/mlir/include/Ion/Transforms/oqd_database_managers.hpp b/mlir/include/Ion/Transforms/oqd_database_managers.hpp
@@ -16,6 +16,7 @@
 
 #include <algorithm>
 #include <cassert>
+#include <set>
 
 #include <toml++/toml.hpp>
 
@@ -133,6 +134,7 @@ class OQDDatabaseManager {
         toml::node_view<toml::node> ionsToml = sourceTomlQubit["ions"];
 
         auto parseSingleLevel = [](auto level) {
+            std::string label = level["label"].as_string()->get();
             int64_t principal = level["principal"].as_integer()->get();
 
             std::vector<std::string> properties{"spin",
@@ -149,29 +151,22 @@ class OQDDatabaseManager {
                                return level[name].as_floating_point()->get();
                            });
 
-            return Level(principal, propertiesData[0], propertiesData[1], propertiesData[2],
+            return Level(label, principal, propertiesData[0], propertiesData[1], propertiesData[2],
                          propertiesData[3], propertiesData[4], propertiesData[5],
                          propertiesData[6]);
         };
 
-        auto parseSingleTransition = [](const auto &transition_entry,
-                                        const std::vector<Level> &allLevels) {
-            // FIXME: `allLevels` is hardcoded as {downstate, upstate, estate}
-            // Not super important, as the ion species is extremely unlikely to change, so
-            // hardcoding is fine
-
+        auto parseSingleTransition = [](const auto &transition_entry) {
             double einstein_a = transition_entry["einstein_a"].as_floating_point()->get();
             std::string level1 = transition_entry["level1"].as_string()->get();
             std::string level2 = transition_entry["level2"].as_string()->get();
 
-            std::map<std::string, int64_t> levelEncodings{
-                {"downstate", 0}, {"upstate", 1}, {"estate", 2}};
+            std::set<std::string> levelEncodings{"downstate", "upstate", "estate"};
             assert((levelEncodings.count(level1) & levelEncodings.count(level2)) &&
                    "Only \"downstate\", \"upstate\" and \"estate\" are allowed in the atom's "
                    "transition levels.");
 
-            return Transition(allLevels[levelEncodings[level1]], allLevels[levelEncodings[level2]],
-                              einstein_a);
+            return Transition(level1, level2, einstein_a);
         };
 
         for (auto &ion_it : *(ionsToml.as_table())) {
@@ -192,8 +187,7 @@ class OQDDatabaseManager {
             std::vector<Transition> transitions;
             auto *transitionsTable = data->at_path("transitions").as_table();
             for (auto &transition : *transitionsTable) {
-                transitions.push_back(
-                    parseSingleTransition(*(transition.second.as_table()), levels));
+                transitions.push_back(parseSingleTransition(*(transition.second.as_table())));
             }
 
             Ion ion(name, mass, charge, position, levels, transitions);

diff --git a/mlir/include/Ion/Transforms/oqd_database_types.hpp b/mlir/include/Ion/Transforms/oqd_database_types.hpp
@@ -63,13 +63,15 @@ struct PhononMode {
 struct Level {
     // This class represents an atomic level.
     // It contains the innate properties of the qubit.
+    std::string label;
     int64_t principal;
     double spin, orbital, nuclear, spin_orbital, spin_orbital_nuclear,
         spin_orbital_nuclear_magnetization, energy;
 
-    Level(int64_t _principal, double _spin, double _orbital, double _nuclear, double _spin_orbital,
-          double _spin_orbital_nuclear, double _spin_orbital_nuclear_magnetization, double _energy)
-        : principal(_principal), spin(_spin), orbital(_orbital), nuclear(_nuclear),
+    Level(std::string _label, int64_t _principal, double _spin, double _orbital, double _nuclear,
+          double _spin_orbital, double _spin_orbital_nuclear,
+          double _spin_orbital_nuclear_magnetization, double _energy)
+        : label(_label), principal(_principal), spin(_spin), orbital(_orbital), nuclear(_nuclear),
           spin_orbital(_spin_orbital), spin_orbital_nuclear(_spin_orbital_nuclear),
           spin_orbital_nuclear_magnetization(_spin_orbital_nuclear_magnetization), energy(_energy)
     {
@@ -79,10 +81,10 @@ struct Level {
 struct Transition {
     // This class represents a transition between two atomic levels.
     // It contains the innate properties of the qubit.
-    Level level_0, level_1;
+    std::string level_0, level_1;
     double einstein_a;
 
-    Transition(Level _level_0, Level _level_1, double _einstein_a)
+    Transition(std::string _level_0, std::string _level_1, double _einstein_a)
         : level_0(_level_0), level_1(_level_1), einstein_a(_einstein_a)
     {
     }

diff --git a/mlir/lib/Ion/Transforms/quantum_to_ion.cpp b/mlir/lib/Ion/Transforms/quantum_to_ion.cpp
@@ -45,18 +45,18 @@ struct QuantumToIonPass : impl::QuantumToIonPassBase<QuantumToIonPass> {
     LevelAttr getLevelAttr(MLIRContext *ctx, IRRewriter &builder, Level level)
     {
         return LevelAttr::get(
-            ctx, builder.getI64IntegerAttr(level.principal), builder.getF64FloatAttr(level.spin),
-            builder.getF64FloatAttr(level.orbital), builder.getF64FloatAttr(level.nuclear),
-            builder.getF64FloatAttr(level.spin_orbital),
+            ctx, builder.getStringAttr(level.label), builder.getI64IntegerAttr(level.principal),
+            builder.getF64FloatAttr(level.spin), builder.getF64FloatAttr(level.orbital),
+            builder.getF64FloatAttr(level.nuclear), builder.getF64FloatAttr(level.spin_orbital),
             builder.getF64FloatAttr(level.spin_orbital_nuclear),
             builder.getF64FloatAttr(level.spin_orbital_nuclear_magnetization),
             builder.getF64FloatAttr(level.energy));
     }
 
     TransitionAttr getTransitionAttr(MLIRContext *ctx, IRRewriter &builder, Transition transition)
     {
-        return TransitionAttr::get(ctx, getLevelAttr(ctx, builder, transition.level_0),
-                                   getLevelAttr(ctx, builder, transition.level_1),
+        return TransitionAttr::get(ctx, builder.getStringAttr(transition.level_0),
+                                   builder.getStringAttr(transition.level_1),
                                    builder.getF64FloatAttr(transition.einstein_a));
     }
 

diff --git a/mlir/test/Ion/Dialect.mlir b/mlir/test/Ion/Dialect.mlir
@@ -157,6 +157,7 @@ func.func @example_ion() -> !ion.ion {
         position=dense<[0, 1]>: tensor<2xi64>,
         levels=[
             #ion.level<
+                label="downstate",
                 principal=1,
                 spin=1.1,
                 orbital=2.2,
@@ -167,6 +168,7 @@ func.func @example_ion() -> !ion.ion {
                 energy=8.8
             >,
             #ion.level<
+                label="upstate",
                 principal=1,
                 spin=1.1,
                 orbital=2.2,
@@ -179,49 +181,13 @@ func.func @example_ion() -> !ion.ion {
         ],
         transitions=[
             #ion.transition<
-                level_0 = #ion.level<
-                    principal=1,
-                    spin=1.1,
-                    orbital=2.2,
-                    nuclear=3.3,
-                    spin_orbital=4.4,
-                    spin_orbital_nuclear=5.5,
-                    spin_orbital_nuclear_magnetization=6.6,
-                    energy=8.8
-                >,
-                level_1 = #ion.level<
-                    principal=1,
-                    spin=1.1,
-                    orbital=2.2,
-                    nuclear=3.3,
-                    spin_orbital=4.4,
-                    spin_orbital_nuclear=5.5,
-                    spin_orbital_nuclear_magnetization=6.6,
-                    energy=8.8
-                >,
+                level_0 = "downstate",
+                level_1 = "upstate",
                 einstein_a=10.10
             >,
             #ion.transition<
-                level_0 = #ion.level<
-                    principal=1,
-                    spin=1.1,
-                    orbital=2.2,
-                    nuclear=3.3,
-                    spin_orbital=4.4,
-                    spin_orbital_nuclear=5.5,
-                    spin_orbital_nuclear_magnetization=6.6,
-                    energy=8.8
-                >,
-                level_1 = #ion.level<
-                    principal=1,
-                    spin=1.1,
-                    orbital=2.2,
-                    nuclear=3.3,
-                    spin_orbital=4.4,
-                    spin_orbital_nuclear=5.5,
-                    spin_orbital_nuclear_magnetization=6.6,
-                    energy=8.8
-                >,
+                level_0 = "upstate",
+                level_1 = "downstate",
                 einstein_a=10.10
             >
         ]

diff --git a/mlir/test/Ion/QuantumToIon.mlir b/mlir/test/Ion/QuantumToIon.mlir
@@ -30,6 +30,7 @@ func.func @example_ion_two_qubit(%arg0: f64) -> !quantum.bit {
     // CHECK-SAME:    charge = 1.000000e+00
     // CHECK-SAME:    levels = [
     // CHECK-SAME:        #ion.level<
+    // CHECK-SAME:            label = "downstate",
     // CHECK-SAME:            principal = 6
     // CHECK-SAME:            spin = 4.000000e-01
     // CHECK-SAME:            orbital = 5.000000e-01
@@ -40,6 +41,7 @@ func.func @example_ion_two_qubit(%arg0: f64) -> !quantum.bit {
     // CHECK-SAME:            energy = 0.000000e+00
     // CHECK-SAME:        >,
     // CHECK-SAME:        #ion.level<
+    // CHECK-SAME:            label = "upstate",
     // CHECK-SAME:            principal = 6
     // CHECK-SAME:            spin = 1.400000e+00
     // CHECK-SAME:            orbital = 1.500000e+00
@@ -50,6 +52,7 @@ func.func @example_ion_two_qubit(%arg0: f64) -> !quantum.bit {
     // CHECK-SAME:            energy = 1.264300e+10
     // CHECK-SAME:        >,
     // CHECK-SAME:        #ion.level<
+    // CHECK-SAME:            label = "estate",
     // CHECK-SAME:            principal = 5
     // CHECK-SAME:            spin = 2.400000e+00
     // CHECK-SAME:            orbital = 2.500000e+00
@@ -65,72 +68,18 @@ func.func @example_ion_two_qubit(%arg0: f64) -> !quantum.bit {
     // CHECK-SAME:    position = dense<[1, 2, -1]> : vector<3xi64>
     // CHECK-SAME:    transitions = [
     // CHECK-SAME:        #ion.transition<
-    // CHECK-SAME:            level_0 = <
-    // CHECK-SAME:                principal = 6
-    // CHECK-SAME:                spin = 4.000000e-01
-    // CHECK-SAME:                orbital = 5.000000e-01
-    // CHECK-SAME:                nuclear = 6.000000e-01
-    // CHECK-SAME:                spin_orbital = 8.000000e-01
-    // CHECK-SAME:                spin_orbital_nuclear = 9.000000e-01
-    // CHECK-SAME:                spin_orbital_nuclear_magnetization = 1.000000e+00
-    // CHECK-SAME:                energy = 0.000000e+00
-    // CHECK-SAME:            >,
-    // CHECK-SAME:            level_1 = <
-    // CHECK-SAME:                principal = 5
-    // CHECK-SAME:                spin = 2.400000e+00
-    // CHECK-SAME:                orbital = 2.500000e+00
-    // CHECK-SAME:                nuclear = 2.600000e+00
-    // CHECK-SAME:                spin_orbital = 2.800000e+00
-    // CHECK-SAME:                spin_orbital_nuclear = 2.900000e+00
-    // CHECK-SAME:                spin_orbital_nuclear_magnetization = 3.000000e+00
-    // CHECK-SAME:                energy = 8.115200e+14
-    // CHECK-SAME:            >,
+    // CHECK-SAME:            level_0 = "downstate",
+    // CHECK-SAME:            level_1 = "estate",
     // CHECK-SAME:            einstein_a = 2.200000e+00
     // CHECK-SAME:        >,
     // CHECK-SAME:        #ion.transition<
-    // CHECK-SAME:            level_0 = <
-    // CHECK-SAME:                principal = 6
-    // CHECK-SAME:                spin = 4.000000e-01
-    // CHECK-SAME:                orbital = 5.000000e-01
-    // CHECK-SAME:                nuclear = 6.000000e-01
-    // CHECK-SAME:                spin_orbital = 8.000000e-01
-    // CHECK-SAME:                spin_orbital_nuclear = 9.000000e-01
-    // CHECK-SAME:                spin_orbital_nuclear_magnetization = 1.000000e+00
-    // CHECK-SAME:                energy = 0.000000e+00
-    // CHECK-SAME:            >,
-    // CHECK-SAME:            level_1 = <
-    // CHECK-SAME:                principal = 6
-    // CHECK-SAME:                spin = 1.400000e+00
-    // CHECK-SAME:                orbital = 1.500000e+00
-    // CHECK-SAME:                nuclear = 1.600000e+00
-    // CHECK-SAME:                spin_orbital = 1.800000e+00
-    // CHECK-SAME:                spin_orbital_nuclear = 1.900000e+00
-    // CHECK-SAME:                spin_orbital_nuclear_magnetization = 2.000000e+00
-    // CHECK-SAME:                energy = 1.264300e+10
-    // CHECK-SAME:            >,
+    // CHECK-SAME:            level_0 = "downstate",
+    // CHECK-SAME:            level_1 = "upstate",
     // CHECK-SAME:            einstein_a = 1.100000e+00
     // CHECK-SAME:        >,
     // CHECK-SAME:        #ion.transition<
-    // CHECK-SAME:            level_0 = <
-    // CHECK-SAME:                principal = 5
-    // CHECK-SAME:                spin = 2.400000e+00
-    // CHECK-SAME:                orbital = 2.500000e+00
-    // CHECK-SAME:                nuclear = 2.600000e+00
-    // CHECK-SAME:                spin_orbital = 2.800000e+00
-    // CHECK-SAME:                spin_orbital_nuclear = 2.900000e+00
-    // CHECK-SAME:                spin_orbital_nuclear_magnetization = 3.000000e+00
-    // CHECK-SAME:                energy = 8.115200e+14
-    // CHECK-SAME:            >,
-    // CHECK-SAME:            level_1 = <
-    // CHECK-SAME:                principal = 6
-    // CHECK-SAME:                spin = 1.400000e+00
-    // CHECK-SAME:                orbital = 1.500000e+00
-    // CHECK-SAME:                nuclear = 1.600000e+00
-    // CHECK-SAME:                spin_orbital = 1.800000e+00
-    // CHECK-SAME:                spin_orbital_nuclear = 1.900000e+00
-    // CHECK-SAME:                spin_orbital_nuclear_magnetization = 2.000000e+00
-    // CHECK-SAME:                energy = 1.264300e+10
-    // CHECK-SAME:            >,
+    // CHECK-SAME:            level_0 = "estate",
+    // CHECK-SAME:            level_1 = "upstate",
     // CHECK-SAME:            einstein_a = 3.300000e+00
     // CHECK-SAME:        >
     // CHECK-SAME:    ]

diff --git a/mlir/test/Ion/oqd_qubit_parameters.toml b/mlir/test/Ion/oqd_qubit_parameters.toml
@@ -26,6 +26,7 @@ mass = 171.0
 charge = +1.0
 position = [1, 2, -1]
 
+levels.downstate.label = "downstate"
 levels.downstate.principal = 6
 levels.downstate.spin = 0.4
 levels.downstate.orbital = 0.5
@@ -35,6 +36,7 @@ levels.downstate.spin_orbital_nuclear = 0.9
 levels.downstate.spin_orbital_nuclear_magnetization = 1.0
 levels.downstate.energy = 0.0
 
+levels.upstate.label = "upstate"
 levels.upstate.principal = 6
 levels.upstate.spin = 1.4
 levels.upstate.orbital = 1.5
@@ -44,6 +46,7 @@ levels.upstate.spin_orbital_nuclear = 1.9
 levels.upstate.spin_orbital_nuclear_magnetization = 2.0
 levels.upstate.energy = 12.643e9
 
+levels.estate.label = "estate"
 levels.estate.principal = 5
 levels.estate.spin = 2.4
 levels.estate.orbital = 2.5