dialects: (stim) initialise dialect

tests/dialects/stim/test_stim_printer_parser.py

@@ -0,0 +1,29 @@
+from io import StringIO
+
+import pytest
+
+from xdsl.dialects import stim
+from xdsl.dialects.stim.stim_printer_parser import StimPrinter
+from xdsl.dialects.test import TestOp
+from xdsl.ir import Block, Region
+
+
+def test_empty_circuit():
+    empty_block = Block()
+    empty_region = Region(empty_block)
+    module = stim.StimCircuitOp(empty_region)
+
+    assert module.stim() == ""
+
+
+def test_stim_circuit_ops_stim_printable():
+    op = TestOp()
+    block = Block([op])
+    region = Region(block)
+    module = stim.StimCircuitOp(region)
+
+    with pytest.raises(ValueError, match="Cannot print in stim format:"):
+        res_io = StringIO()
+        printer = StimPrinter(stream=res_io)
+
+        module.print_stim(printer)

tests/filecheck/dialects/stim/stim_ops.mlir

@@ -0,0 +1,20 @@
+// RUN: XDSL_ROUNDTRIP
+// RUN: XDSL_GENERIC_ROUNDTRIP
+
+// CHECK:       builtin.module {
+// CHECK-GENERIC:       "builtin.module"() ({
+
+stim.circuit {}
+// CHECK-NEXT:    stim.circuit {
+// CHECK-NEXT: }
+// CHECK-GENERIC-NEXT:    "stim.circuit"() ({
+// CHECK-GENERIC-NEXT:   }) : () -> ()
+
+stim.circuit attributes {"hello" = "world"} {}
+// CHECK-NEXT:    stim.circuit attributes {"hello" = "world"} {
+// CHECK-NEXT: }
+// CHECK-GENERIC-NEXT:    "stim.circuit"() ({
+// CHECK-GENERIC-NEXT:   }) {"hello" = "world"} : () -> ()
+
+// CHECK-NEXT:  }
+// CHECK-GENERIC-NEXT:  }) : () -> ()

xdsl/dialects/__init__.py

@@ -258,6 +258,11 @@ def get_stencil():
 
         return Stencil
 
+    def get_stim():
+        from xdsl.dialects.stim import Stim
+
+        return Stim
+
     def get_stream():
         from xdsl.dialects.stream import Stream
 
@@ -349,6 +354,7 @@ def get_transform():
         "snitch_stream": get_snitch_stream,
         "stablehlo": get_stablehlo,
         "stencil": get_stencil,
+        "stim": get_stim,
         "stream": get_stream,
         "symref": get_symref,
         "tensor": get_tensor,

xdsl/dialects/stim/__init__.py

@@ -0,0 +1,10 @@
+from xdsl.ir import Dialect
+
+from .ops import StimCircuitOp
+
+Stim = Dialect(
+    "stim",
+    [
+        StimCircuitOp,
+    ],
+)

xdsl/dialects/stim/ops.py

@@ -0,0 +1,50 @@
+from abc import ABC
+from io import StringIO
+
+from xdsl.ir import Region
+from xdsl.irdl import (
+    IRDLOperation,
+    PyRDLOpDefinitionError,
+    irdl_op_definition,
+    region_def,
+)
+
+from .stim_printer_parser import StimPrintable, StimPrinter
+
+
+class StimOp(IRDLOperation, ABC):
+    def print_stim(self, printer: StimPrinter) -> None:
+        raise (PyRDLOpDefinitionError("print_stim not implemented!"))
+
+
+@irdl_op_definition
+class StimCircuitOp(StimOp, IRDLOperation):
+    """
+    Base operation containing a stim program
+    """
+
+    name = "stim.circuit"
+
+    body = region_def("single_block")
+
+    assembly_format = "attr-dict-with-keyword $body"
+
+    def __init__(self, body: Region):
+        super().__init__(regions=[body])
+
+    def verify(self, verify_nested_ops: bool = True) -> None:
+        return
+
+    def print_stim(self, printer: StimPrinter):
+        for op in self.body.block.ops:
+            if not isinstance(op, StimPrintable):
+                raise ValueError(f"Cannot print in stim format: {op}")
+            op.print_stim(printer)
+        printer.print_string("")
+
+    def stim(self) -> str:
+        io = StringIO()
+        printer = StimPrinter(io)
+        self.print_stim(printer)
+        res = io.getvalue()
+        return res

xdsl/dialects/stim/stim_printer_parser.py

@@ -0,0 +1,26 @@
+import abc
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass(eq=False, repr=False)
+class StimPrinter:
+    stream: Any | None = field(default=None)
+
+    def print_string(self, text: str) -> None:
+        print(text, end="", file=self.stream)
+
+    @contextmanager
+    def in_braces(self):
+        self.print_string("{")
+        try:
+            yield
+        finally:
+            self.print_string("}")
+
+
+class StimPrintable(abc.ABC):
+    @abc.abstractmethod
+    def print_stim(self, printer: StimPrinter) -> None:
+        raise NotImplementedError()