Merge pull request #34 from saulshanabrook/simplify-representation

Unified Representation and NumPy Overloading

Author: sdiehl

example_rewrite.py

@@ -5,21 +5,21 @@
 llvmlite.opaque_pointers_enabled = True
 
 from mlir_egglog import kernel
-from mlir_egglog.term_ir import Term, Sin, Cos, Mul, Add
+from mlir_egglog.term_ir import Term, sin, cos
 from egglog import rewrite, ruleset
 from mlir_egglog.optimization_rules import basic_math
 
 
 # A rewrite rule
 @ruleset
 def trig_double_angle(a: Term):
-    sin_a = Sin(a)
-    cos_a = Cos(a)
-    mul1 = Mul(sin_a, cos_a)
-    mul2 = Mul(cos_a, sin_a)
+    sin_a = sin(a)
+    cos_a = cos(a)
+    mul1 = sin_a * cos_a
+    mul2 = cos_a * sin_a
 
     # sin(a)*cos(a) + cos(a)*sin(a) -> 2 * sin(a)*cos(a)
-    yield rewrite(Add(mul1, mul2)).to(Mul(Term.lit_f32(2.0), mul1))
+    yield rewrite(mul1 + mul2).to(Term.lit_f32(2.0) * mul1)
 
 
 # Apply the rewrites

pyproject.toml

@@ -3,7 +3,7 @@ name = "mlir-egglog"
 version = "0.1.0"
 requires-python = ">=3.12"
 dependencies = [
-    "egglog>=8.0.1",
+    "egglog>=11.0.0",
     "llvmlite>=0.44.0",
     "numpy>=2.2.5",
     "pyyaml>=6.0.2",
@@ -12,7 +12,7 @@ dependencies = [
 [dependency-groups]
 dev = [
     "black>=25.1.0",
-    "mypy>=1.15.0",
+    "mypy>=1.17.1",
     "ruff>=0.9.10",
     "pytest>=8.3.5",
 ]
@@ -34,3 +34,7 @@ disallow_subclassing_any = "false"
 disallow_untyped_decorators = "false"
 disallow_any_generics = "true"
 follow_imports = "silent"
+
+[tool.uv.sources]
+# until https://github.com/python/mypy/pull/19600 is merged
+mypy = { git = "https://github.com/saulshanabrook/mypy", branch = "issue-19599" }

src/mlir_egglog/__init__.py

@@ -9,13 +9,7 @@
 # Version of the mlir-egglog package
 __version__ = "0.1.0"
 
-from mlir_egglog.expr_model import (  # noqa: E402
-    BinaryOp,
-    FloatLiteral,
-    IntLiteral,
-    Symbol,
-)
-from mlir_egglog.term_ir import Term, as_egraph  # noqa: E402
+from mlir_egglog.term_ir import Term  # noqa: E402
 from mlir_egglog.dispatcher import kernel  # noqa: E402
 
 __all__ = [
@@ -24,6 +18,5 @@
     "IntLiteral",
     "Symbol",
     "Term",
-    "as_egraph",
     "kernel",
 ]

src/mlir_egglog/builtin_functions.py

@@ -24,7 +24,7 @@ def __init__(
         py_func: types.FunctionType,
         rewrites: tuple[RewriteOrRule | Ruleset, ...] | None = None,
     ):
-        self.py_func = py_func
+        self.py_func = py_func  # type: ignore[assignment]
         self._compiled_func = None
         self._compiler = None
         self.rewrites = rewrites

src/mlir_egglog/dispatcher.py

@@ -4,21 +4,19 @@
 from egglog import EGraph, RewriteOrRule, Ruleset
 from egglog.egraph import UnstableCombinedRuleset
 
-from mlir_egglog.term_ir import Term, as_egraph
+from mlir_egglog.term_ir import Term
 from mlir_egglog.python_to_ir import interpret
-from mlir_egglog import builtin_functions as ns
-from mlir_egglog.expr_model import Expr
-from mlir_egglog.ir_to_mlir import convert_term_to_mlir
+from mlir_egglog.mlir_gen import MLIRGen
 
 # Rewrite rules
 from mlir_egglog.optimization_rules import basic_math, trig_simplify
 
 OPTS: tuple[Ruleset | RewriteOrRule, ...] = (basic_math, trig_simplify)
 
 
-def extract(ast: Expr, rules: tuple[RewriteOrRule | Ruleset, ...], debug=False) -> Term:
-    root = as_egraph(ast)
-
+def extract(
+    root: Term, rules: tuple[RewriteOrRule | Ruleset, ...], debug=False
+) -> Term:
     egraph = EGraph()
     egraph.let("root", root)
 
@@ -47,10 +45,21 @@ def compile(
     fn: FunctionType, rewrites: tuple[RewriteOrRule | Ruleset, ...] = OPTS, debug=True
 ) -> str:
     # Convert np functions according to the namespace map
-    exprtree = interpret(fn, {"np": ns})
+    exprtree = interpret(fn)
     extracted = extract(exprtree, rewrites, debug)
 
     # Get the argument spec
     argspec = inspect.signature(fn)
     params = ",".join(map(str, argspec.parameters))
     return convert_term_to_mlir(extracted, params)
+
+
+def convert_term_to_mlir(tree: Term, argspec: str) -> str:
+    """
+    Convert a term to MLIR.
+    """
+
+    argnames = map(lambda x: x.strip(), argspec.split(","))
+    argmap = {k: f"%arg_{k}" for k in argnames}
+    source = MLIRGen(tree, argmap).generate()
+    return source