{"id":26888829,"url":"https://github.com/sdiehl/mlir-egglog","last_synced_at":"2025-07-09T10:40:57.672Z","repository":{"id":282463841,"uuid":"948489037","full_name":"sdiehl/mlir-egglog","owner":"sdiehl","description":"A toy compiler for NumPy array expressions that uses e-graphs and MLIR","archived":false,"fork":false,"pushed_at":"2025-03-31T10:15:54.000Z","size":295,"stargazers_count":33,"open_issues_count":0,"forks_count":2,"subscribers_count":1,"default_branch":"main","last_synced_at":"2025-03-31T11:24:52.745Z","etag":null,"topics":["e-graphs","egglog","equality-saturation","mlir","term-rewriting"],"latest_commit_sha":null,"homepage":"https://www.stephendiehl.com/posts/mlir_egraphs/","language":"Python","has_issues":false,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/sdiehl.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2025-03-14T12:32:36.000Z","updated_at":"2025-03-31T10:15:51.000Z","dependencies_parsed_at":"2025-03-14T20:24:00.273Z","dependency_job_id":"c0801c13-44ec-40bf-bf8b-d53ed0297b78","html_url":"https://github.com/sdiehl/mlir-egglog","commit_stats":null,"previous_names":["sdiehl/mlir-egglog"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sdiehl%2Fmlir-egglog","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sdiehl%2Fmlir-egglog/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sdiehl%2Fmlir-egglog/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sdiehl%2Fmlir-egglog/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/sdiehl","download_url":"https://codeload.github.com/sdiehl/mlir-egglog/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":246531968,"owners_count":20792736,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["e-graphs","egglog","equality-saturation","mlir","term-rewriting"],"created_at":"2025-03-31T20:01:56.640Z","updated_at":"2025-07-09T10:40:57.649Z","avatar_url":"https://github.com/sdiehl.png","language":"Python","funding_links":[],"categories":["Python","Applications"],"sub_categories":["Program Optimization"],"readme":"\u003cp align=\"center\"\u003e\n    \u003cimg src=\".github/logo.png\" width=\"500px\" alt=\"mlir-egglog\"\u003e\n\u003c/p\u003e\n\n# MLIR Egglog\n\nA toy specializing compiler for NumPy expressions that uses MLIR as a target and can use equality saturation (e-graphs) to do term rewriting on the intermediate representation, enabling extremely precise and composable optimizations of mathematical expressions before lowering to MLIR.\n\nWe use the embedded Datalog DSL [`egglog`](https://github.com/egraphs-good/egglog) to express and compose rewrite rules in pure Python and the [`egg`](https://docs.rs/egg/latest/egg/) library to extract optimized syntax trees from the e-graph.\n\nThe whole project is just under 1500 lines of code, and is designed to be a simple and easy to understand example of how to integrate e-graphs into a compiler pipeline.\n\n## What is Egg and Equality Saturation?\n\nThink of an e-graph as this magical data structure that's like a super-powered hash table of program expressions. Instead of just storing one way to write a program, it efficiently stores ALL equivalent ways to write it.\n\nEquality saturation is the process of filling this e-graph with all possible equivalent programs by applying rewrite rules until we can't find any more rewrites (that's the \"saturation\" part). The cool part? We can explore tons of different optimizations simultaneously, rather than having to pick a specific sequence of transformations. The you can apply a cost function over the entire e-graph to find the best solution. \n\nTraditionally you'd have to muddle through with a fixed-point iteration system and tons of top-down/bottom-up rewrite rule contingent on application orders, but e-graphs make it much more efficient and declarative.\n\n## Installation\n\nOn MacOS, install LLVM 20 which includes MLIR:\n\n```shell\nbrew install llvm@20\n```\n\nOn Linux, install the dependencies (setup instructions [here](https://apt.llvm.org/)):\n\n```shell\nsudo apt-get install -y llvm-20 llvm-20-dev llvm-20-tools mlir-20-tools\n```\n\nThen to use the library:\n\n```shell\ngit clone https://github.com/sdiehl/mlir-egglog.git\ncd mlir-egglog\npoetry install\npoetry run python example.py\n```\n\n## Usage\n\n```python\nfrom mlir_egglog import kernel\n\n@kernel(\"float32(float32)\")\ndef fn(x : float) -\u003e float:\n    # sinh(x + y) = sinh(x) * cosh(y) + cosh(x) * sinh(y)\n    return np.sinh(x) * np.cosh(x) + np.cosh(x) * np.sinh(x)\n\nout = fn(np.array([1, 2, 3]))\nprint(out)\n```\n\n## Custom Rewrite Rules\n\nYou can create your own optimization rules using the `ruleset` decorator. Here's a complete example that optimizes away addition with zero:\n\n```python\nfrom mlir_egglog import kernel\nfrom mlir_egglog.term_ir import Term, Add\nfrom egglog import rewrite, ruleset, RewriteOrRule, i64, f64\nfrom typing import Generator\n\n@ruleset\ndef float_rules(x: Term, y: Term, z: Term, i: i64, f: f64):\n    yield rewrite(Add(x, Term.lit_f32(0.0))).to(x)\n    yield rewrite(Add(Term.lit_f32(0.0), x)).to(x)\n\n@kernel(\"float32(float32)\", rewrites=(basic_math, float_rules))\ndef custom_fn(x):\n    return x + 0.0  # This addition will be optimized away!\n\ntest_input = np.array([1.0, 2.0, 3.0], dtype=np.float32)\nresult = custom_fn(test_input)\nprint(result)\n```\n\nThe rewrite rules are applied during compilation, so there's no runtime overhead. The generated MLIR code will be as if you just wrote `return x`. You can combine multiple rulesets to build up more complex program optimizations.\n\n## Codebase\n\nHere's the recommended order to understand the codebase:\n\n**Foundation Layer** - Expression representation and manipulation\n\n1. [`memory_descriptors.py`](src/mlir_egglog/memory_descriptors.py) - Basic memory management utilities for handling NumPy arrays and MLIR memory references\n2. [`expr_model.py`](src/mlir_egglog/expr_model.py) - Core expression model defining the base classes for mathematical expressions\n3. [`builtin_functions.py`](src/mlir_egglog/builtin_functions.py) - Implementation of basic mathematical functions and operations\n4. [`term_ir.py`](src/mlir_egglog/term_ir.py) - Intermediate representation for the egraph system with cost models for operations\n\n**Transformation Layer** - Code transformation and lowering\n\n5. [`python_to_ir.py`](src/mlir_egglog/python_to_ir.py) - Converts Python functions to the internal IR representation\n6. [`ir_to_mlir.py`](src/mlir_egglog/ir_to_mlir.py) - Transforms internal IR to MLIR representation\n7. [`basic_simplify.py`](src/mlir_egglog/basic_simplify.py) - Basic mathematical simplification rules\n8. [`trig_simplify.py`](src/mlir_egglog/trig_simplify.py) - Trigonometric function simplification rules\n\n**Optimization Layer** - Optimization and compilation\n\n9. [`egglog_optimizer.py`](src/mlir_egglog/egglog_optimizer.py) - Core optimization engine using egg-rewrite rules\n10. [`mlir_backend.py`](src/mlir_egglog/mlir_backend.py) - MLIR compilation pipeline and optimization passes\n11. [`llvm_runtime.py`](src/mlir_egglog/llvm_runtime.py) - LLVM runtime initialization and management\n\n**Execution Layer** - Runtime execution\n\n12. [`jit_engine.py`](src/mlir_egglog/jit_engine.py) - JIT compilation engine for executing optimized code\n13. [`dispatcher.py`](src/mlir_egglog/dispatcher.py) - High-level interface for function compilation and execution\n\n## License\n\nThis project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsdiehl%2Fmlir-egglog","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fsdiehl%2Fmlir-egglog","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsdiehl%2Fmlir-egglog/lists"}