This tool proposes a dialect-agnostic way to apply equality saturation optimizations to LLVM's Multi-Level IR (MLIR). The equality saturation engine used is Egglog.
Clone LLVM with commit 7b074fc9362a4a6a5f172dd8936a22fb25eff96b.
git clone https://github.com/llvm/llvm-project.git llvm
git checkout 7b074fc9362a4a6a5f172dd8936a22fb25eff96bThe build LLVM core and MLIR in Release mode. This may take a while.
cd llvm
mkdir build-release
cmake -S llvm -B build-release -G Ninja \
-DLLVM_ENABLE_PROJECTS="mlir" \
-DCMAKE_BUILD_TYPE=Release \
-DLLVM_USE_LINKER=lld
cmake --build build-releaseIf you have issues building, follow the build guide.
Egglog is the equality saturation engine. This PR has a feature DialEgg needs.
git clone -b cost-action https://github.com/saulshanabrook/egg-smol.git egglog
cd egglog
git checkout b6e1c96ed7335366e90056ea0a24ef425dfbb8fbThen build it in release mode
cargo build --releaseWithin the root directory of this repo, build DialEgg:
mkdir build
cmake -S . -B build -DLLVM_DIR=[path to llvm cmake dir] -DMLIR_DIR=[path to mlir cmake dir]
cmake --build buildOptimize an example MLIR file in the test directory.
MLIR in test/classic/classic.mlir:
func.func @classic(%a: i32) -> i32 {
%c2 = arith.constant 2 : i32
%mul = arith.muli %a, %c2 : i32
%div = arith.divsi %mul, %c2 : i32
func.return %div : i32
}Egglog ops and rules in test/classic/classic.egg:
(include "src/base.egg")
;;;; arith dialect ;;;;
(function arith_constant (AttrPair Type) Op)
(function arith_muli (Op Op AttrPair Type) Op)
(function arith_divsi (Op Op Type) Op)
(function arith_shli (Op Op AttrPair Type) Op)
;; OPS HERE ;;
;; RULES HERE ;;
(ruleset rules)
(let c1 (arith_constant (NamedAttr "value" (IntegerAttr 1 (I32))) (I32))) ; 1
(let c2 (arith_constant (NamedAttr "value" (IntegerAttr 2 (I32))) (I32))) ; 2
(rewrite (arith_divsi ?x ?x (I32)) c1 :ruleset rules) ; x / x = 1
(rewrite (arith_muli ?x c1 ?a (I32)) ?x :ruleset rules) ; x * 1 = x
(rewrite ; x * 2 = x << 1
(arith_muli ?x c2 ?a (I32)) ; x * 2
(arith_shli ?x c1 ?a (I32)) ; x << 1
:ruleset rules
)
(rewrite ; (xy) / z = x (z / y)
(arith_divsi (arith_muli ?x ?y ?a ?t) ?z ?t)
(arith_muli ?x (arith_divsi ?y ?z ?t) ?a ?t)
:ruleset rules
)
(run-schedule (saturate rules))
;; EXTRACTS HERE ;;Run this and the output will be the optimized MLIR.
./build/egg-opt --eq-sat test/classic/classic.mlir --egg test/classic/classic.eggResult:
func.func @classic(%arg0: i32) -> i32 {
return %arg0 : i32
}- Better support for ops with variadic number of operands and attributes
- Better support for ops with optional operands and attributes