Add Pauli algebra dialect example to docs (#365)

There's still an issue with a reference to `Fold`, that mkdocs can't
resolve for some reason.

---------

Co-authored-by: Xiu-zhe (Roger) Luo <[email protected]>

Author: david-pl

docs/cookbook/paulilang/pauli.md

@@ -0,0 +1,10 @@
+document$.subscribe(({ body }) => {
+    renderMathInElement(body, {
+      delimiters: [
+        { left: "$$",  right: "$$",  display: true },
+        { left: "$",   right: "$",   display: false },
+        { left: "\\(", right: "\\)", display: false },
+        { left: "\\[", right: "\\]", display: true }
+      ],
+    })
+  })

docs/scripts/katex.js

@@ -6,6 +6,7 @@ This folder contains examples of how to use the Kirin library. Each example is a
 
 - `simple.py`: A simple example that demonstrates how to create a simple Kirin dialect group and its kernel.
 - `food`: A more sophisticated example but without any domain specifics. It demonstrates how to create a new Kirin dialect and combine it with existing dialects with custom analysis and rewrites.
+- `pauli`: An example that implements a dialect with rewrites that simplifies products of Pauli matrices.
 
 ## Examples outside this folder with more domain-specific contents
 

example/README.md

@@ -0,0 +1,3 @@
+from kirin import ir
+
+_dialect = ir.Dialect("pauli")

example/pauli/dialect.py

@@ -0,0 +1,13 @@
+from dialect import _dialect
+
+from kirin import ir
+from kirin.prelude import basic_no_opt
+
+
+@ir._dialectgroup(basic_no_opt.add(dialect=_dialect))
+def pauli(self):
+    def run_pass(mt):
+        # TODO
+        pass
+
+    return run_pass

example/pauli/group.py

@@ -0,0 +1,30 @@
+import numpy as np
+
+from kirin.interp import MethodTable, impl
+
+from .stmts import X, Y, Z, Id
+from .dialect import _dialect
+
+
+@_dialect.register
+class PauliMethods(MethodTable):
+    X_mat = np.array([[0, 1], [1, 0]])
+    Y_mat = np.array([[0, -1j], [1j, 0]])
+    Z_mat = np.array([[1, 0], [0, -1]])
+    Id_mat = np.array([[1, 0], [0, 1]])
+
+    @impl(X)  # (1)!
+    def x(self, interp, frame, stmt: X):
+        return (stmt.pre_factor * self.X_mat,)
+
+    @impl(Y)
+    def y(self, interp, frame, stmt: Y):
+        return (self.Y_mat * stmt.pre_factor,)
+
+    @impl(Z)
+    def z(self, interp, frame, stmt: Z):
+        return (self.Z_mat * stmt.pre_factor,)
+
+    @impl(Id)
+    def id(self, interp, frame, stmt: Id):
+        return (self.Id_mat * stmt.pre_factor,)

example/pauli/interp.py

@@ -0,0 +1,80 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.rewrite import abc, result
+from kirin.dialects import py
+
+from .stmts import X, Y, Z, Id, PauliOperator
+
+
+@dataclass
+class RewritePauliMult(abc.RewriteRule):
+    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+        if not isinstance(node, py.binop.Mult):
+            return result.RewriteResult()
+
+        if not isinstance(node.lhs.owner, PauliOperator) and not isinstance(
+            node.rhs.owner, PauliOperator
+        ):
+            return result.RewriteResult()
+
+        if isinstance(node.lhs.owner, py.Constant):
+            assert isinstance(node.rhs.owner, PauliOperator)  # make the linter happy
+            new_op = self.number_pauli_mult(node.lhs.owner, node.rhs.owner)
+            node.replace_by(new_op)
+            return result.RewriteResult(has_done_something=True)
+        elif isinstance(node.rhs.owner, py.Constant):
+            assert isinstance(node.lhs.owner, PauliOperator)  # make the linter happy
+            new_op = self.number_pauli_mult(node.rhs.owner, node.lhs.owner)
+            node.replace_by(new_op)
+            return result.RewriteResult(has_done_something=True)
+
+        if not isinstance(node.lhs.owner, PauliOperator) or not isinstance(
+            node.rhs.owner, PauliOperator
+        ):
+            return result.RewriteResult()
+
+        new_op = self.pauli_pauli_mult(node.lhs.owner, node.rhs.owner)
+        node.replace_by(new_op)
+        return result.RewriteResult(has_done_something=True)
+
+    @staticmethod
+    def number_pauli_mult(lhs: py.Constant, rhs: PauliOperator) -> PauliOperator:
+        num = lhs.value.unwrap() * rhs.pre_factor
+        return type(rhs)(pre_factor=num)
+
+    @staticmethod
+    def pauli_pauli_mult(lhs: PauliOperator, rhs: PauliOperator) -> PauliOperator:
+        num = rhs.pre_factor * lhs.pre_factor
+
+        if isinstance(lhs, type(rhs)):
+            return Id(pre_factor=num)
+
+        if isinstance(lhs, type(rhs)):
+            return Id(pre_factor=num)
+
+        if isinstance(lhs, Id):
+            return type(rhs)(pre_factor=num)
+
+        if isinstance(rhs, Id):
+            return type(lhs)(pre_factor=num)
+
+        if isinstance(lhs, X):
+            if isinstance(rhs, Y):
+                return Z(pre_factor=1j * num)
+            elif isinstance(rhs, Z):
+                return Y(pre_factor=-1j * num)
+
+        if isinstance(lhs, Y):
+            if isinstance(rhs, X):
+                return Z(pre_factor=-1j * num)
+            elif isinstance(rhs, Z):
+                return X(pre_factor=1j * num)
+
+        if isinstance(lhs, Z):
+            if isinstance(rhs, Y):
+                return X(pre_factor=-1j * num)
+            elif isinstance(rhs, X):
+                return Y(pre_factor=1j * num)
+
+        raise RuntimeError("How on earth did we end up here?")

example/pauli/rewrite.py

@@ -0,0 +1,11 @@
+from .group import pauli
+from .stmts import X, Y, Z
+
+
+@pauli
+def main():
+    ex = (X() + 2 * Y()) * Z()
+    return ex
+
+
+main.print()

example/pauli/script.py

@@ -0,0 +1,35 @@
+from numbers import Number
+
+import numpy as np
+
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+
+from .dialect import _dialect
+
+
+@statement
+class PauliOperator(ir.Statement):
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
+    pre_factor: Number = info.attribute(default=1)
+    result: ir.ResultValue = info.result(types.PyClass(np.matrix))
+
+
+@statement(dialect=_dialect)
+class X(PauliOperator):
+    pass
+
+
+@statement(dialect=_dialect)
+class Y(PauliOperator):
+    pass
+
+
+@statement(dialect=_dialect)
+class Z(PauliOperator):
+    pass
+
+
+@statement(dialect=_dialect)
+class Id(PauliOperator):
+    pass

example/pauli/stmts.py

@@ -35,6 +35,7 @@ nav:
         - Advanced Rewriting: cookbook/foodlang/cf_rewrite.md
         - Food Pricing Analysis Pass: cookbook/foodlang/analysis.md
         - Receipt Codegen: cookbook/foodlang/codegen.md
+      - Pauli Algebra: cookbook/paulilang/pauli.md
 
   - Blog:
     - blog/index.md
@@ -144,6 +145,7 @@ plugins:
 
 extra_css:
   - stylesheets/extra.css
+  - https://unpkg.com/katex@0/dist/katex.min.css
 
 markdown_extensions:
   - abbr
@@ -158,6 +160,8 @@ markdown_extensions:
   - pymdownx.tilde
   - pymdownx.tabbed:
       alternate_style: true
+  - pymdownx.arithmatex:
+      generic: true
 
 copyright: Copyright © 2024 Kirin contributors
 
@@ -169,3 +173,8 @@ extra:
       link: https://x.com/QueraComputing
     - icon: fontawesome/brands/linkedin
       link: https://www.linkedin.com/company/quera-computing-inc/
+
+extra_javascript:
+  - scripts/katex.js
+  - https://unpkg.com/katex@0/dist/katex.min.js
+  - https://unpkg.com/katex@0/dist/contrib/auto-render.min.js