pattern rewrites

import%20marimo%0A%0A__generated_with%20%3D%20%220.23.15%22%0Aapp%20%3D%20marimo.App()%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%0A%20%20%20%20return%20(mo%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Aasync%20def%20_()%3A%0A%20%20%20%20%23%20Get%20the%20current%20notebook%20URL%2C%20drop%20the%20'blob'%20URL%20components%20that%20seem%20to%20be%20added%2C%0A%20%20%20%20%23%20and%20add%20the%20buildnumber%20that%20a%20makethedocs%20PR%20build%20seems%20to%20add.%20This%20allows%20to%20load%0A%20%20%20%20%23%20the%20wheel%20both%20locally%20and%20when%20deployed%20to%20makethedocs.%0A%20%20%20%20async%20def%20import_xdsl()%3A%0A%20%20%20%20%20%20%20%20import%20re%0A%20%20%20%20%20%20%20%20from%20urllib.parse%20import%20urlparse%0A%0A%20%20%20%20%20%20%20%20import%20micropip%0A%20%20%20%20%20%20%20%20from%20marimo%20import%20notebook_location%0A%0A%20%20%20%20%20%20%20%20url%20%3D%20str(notebook_location()).replace(%22blob%3A%22%2C%20%22%22)%0A%20%20%20%20%20%20%20%20print(f%22DEBUG%3A%20notebook%20url%20(full)%3A%20%7Burl%7D%22)%0A%0A%20%20%20%20%20%20%20%20url_parsed%20%3D%20urlparse(url)%0A%20%20%20%20%20%20%20%20scheme%20%3D%20url_parsed.scheme%0A%20%20%20%20%20%20%20%20netloc%20%3D%20url_parsed.netloc%0A%0A%20%20%20%20%20%20%20%20print(f%22DEBUG%3A%20notebook%20url%20(parsed)%3A%20%7Burl_parsed%7D%22)%0A%0A%20%20%20%20%20%20%20%20url%20%3D%20re.sub(%0A%20%20%20%20%20%20%20%20%20%20%20%20%22(%5B%5E%2F%5D)%2F(%5Ba-f0-9-%5D%2B-%5Ba-f0-9-%5D%2B-%5Ba-f0-9-%5D%2B-%5Ba-f0-9-%5D%2B)%22%2C%20%22%5C%5C1%2F%22%2C%20url%2C%20count%3D1%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20buildnumber%20%3D%20re.sub(%22.*--(%5B0-9%2B%5D%2B).*%22%2C%20%22%5C%5C1%22%2C%20url%2C%20count%3D1)%0A%0A%20%20%20%20%20%20%20%20new_url%20%3D%20scheme%20%2B%20%22%3A%2F%2F%22%20%2B%20netloc%0A%0A%20%20%20%20%20%20%20%20if%20buildnumber%20!%3D%20url%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20new_url%20%3D%20new_url%20%2B%20%22%2F%22%20%2B%20buildnumber%20%2B%20%22%2F%22%0A%0A%20%20%20%20%20%20%20%20print(f%22DEBUG%3A%20notebook%20url%20(trimmed)%3A%20%7Bnew_url%7D%22)%0A%0A%20%20%20%20%20%20%20%20await%20micropip.install(%22xdsl%20%40%20%22%20%2B%20new_url%20%2B%20%22%2Fxdsl-0.0.0-py3-none-any.whl%22)%0A%0A%20%20%20%20await%20import_xdsl()%0A%20%20%20%20from%20xdsl.utils%20import%20marimo%20as%20xmo%0A%0A%20%20%20%20return%20(xmo%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20from%20xdsl.parser%20import%20Parser%0A%20%20%20%20from%20xdsl.context%20import%20Context%0A%0A%20%20%20%20from%20xdsl.dialects.builtin%20import%20Builtin%0A%20%20%20%20from%20xdsl.dialects.arith%20import%20Arith%0A%20%20%20%20from%20xdsl.dialects.func%20import%20Func%0A%0A%20%20%20%20from%20xdsl.ir%20import%20Operation%0A%0A%20%20%20%20from%20xdsl.rewriter%20import%20Rewriter%0A%20%20%20%20from%20xdsl.pattern_rewriter%20import%20(%0A%20%20%20%20%20%20%20%20RewritePattern%2C%0A%20%20%20%20%20%20%20%20GreedyRewritePatternApplier%2C%0A%20%20%20%20%20%20%20%20PatternRewriteWalker%2C%0A%20%20%20%20)%0A%20%20%20%20from%20xdsl.dialects.arith%20import%20AddiOp%2C%20ConstantOp%2C%20MuliOp%0A%20%20%20%20from%20xdsl.dialects.builtin%20import%20ModuleOp%0A%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20AddiOp%2C%0A%20%20%20%20%20%20%20%20Arith%2C%0A%20%20%20%20%20%20%20%20Builtin%2C%0A%20%20%20%20%20%20%20%20ConstantOp%2C%0A%20%20%20%20%20%20%20%20Context%2C%0A%20%20%20%20%20%20%20%20Func%2C%0A%20%20%20%20%20%20%20%20GreedyRewritePatternApplier%2C%0A%20%20%20%20%20%20%20%20ModuleOp%2C%0A%20%20%20%20%20%20%20%20MuliOp%2C%0A%20%20%20%20%20%20%20%20Operation%2C%0A%20%20%20%20%20%20%20%20Parser%2C%0A%20%20%20%20%20%20%20%20PatternRewriteWalker%2C%0A%20%20%20%20%20%20%20%20RewritePattern%2C%0A%20%20%20%20%20%20%20%20Rewriter%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20Pattern%20Rewrites%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(Parser%2C%20ctx)%3A%0A%20%20%20%20_before_text%20%3D%20%22%22%22%5C%0A%20%20%20%20func.func%20%40my_func(%25x%3A%20index)%20-%3E%20index%20%7B%0A%20%20%20%20%20%20%25c0%20%3D%20arith.constant%200%20%3A%20index%0A%20%20%20%20%20%20%25c2%20%3D%20arith.constant%202%20%3A%20index%0A%20%20%20%20%20%20%25also_x%20%3D%20arith.addi%20%25c0%2C%20%25x%20%3A%20index%0A%20%20%20%20%20%20%25two_x%20%3D%20arith.muli%20%25also_x%2C%20%25c2%20%3A%20index%0A%20%20%20%20%20%20func.return%20%25two_x%20%3A%20index%0A%20%20%20%20%7D%0A%20%20%20%20%22%22%22%0A%20%20%20%20before_module%20%3D%20Parser(ctx%2C%20_before_text).parse_module()%0A%20%20%20%20None%0A%20%20%20%20return%20(before_module%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(Parser%2C%20ctx)%3A%0A%20%20%20%20_after_text%20%3D%20%22%22%22%5C%0A%20%20%20%20func.func%20%40my_func(%25x%3A%20index)%20-%3E%20index%20%7B%0A%20%20%20%20%20%20%25two_x%20%3D%20arith.addi%20%25x%2C%20%25x%20%3A%20index%0A%20%20%20%20%20%20func.return%20%25two_x%20%3A%20index%0A%20%20%20%20%7D%0A%20%20%20%20%22%22%22%0A%20%20%20%20after_module%20%3D%20Parser(ctx%2C%20_after_text).parse_module()%0A%20%20%20%20None%0A%20%20%20%20return%20(after_module%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(after_module%2C%20before_module%2C%20mo%2C%20xmo)%3A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20Most%20of%20the%20work%20in%20compilers%20is%20rewriting%20IR.%0A%20%20%20%20While%20some%20transformations%20have%20to%20take%20modules%20into%20account%2C%20some%20can%20be%20done%20locally%2C%20such%20as%20rewriting%20an%20addition%20with%20zero%20to%20just%20use%20the%20non-zero%20input.%0A%20%20%20%20These%20kinds%20of%20rewrites%20are%20called%20pattern%20rewrites%2C%20as%20they%20apply%20if%20some%20local%20pattern%20is%20matched.%0A%0A%20%20%20%20We%20can%20use%20xDSL's%20pattern%20rewriting%20API%20to%20rewrite%20this%20function%3A%0A%0A%20%20%20%20%7Bxmo.module_html(before_module)%7D%0A%0A%20%20%20%20Into%20this%20one%3A%0A%0A%20%20%20%20%7Bxmo.module_html(after_module)%7D%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Rationale%0A%0A%20%20%20%20A%20pattern%20rewrite%20is%20a%20compiler%20transformation%20that%20matches%20a%20DAG%20in%20the%20IR%2C%20and%20replace%20it%20with%20another%20DAG.%20For%20instance%2C%20simplifying%20%60x%20%2B%200%60%20to%20%60x%60%20is%20a%20common%20optimization%20that%20is%20represented%20as%20a%20pattern%20rewrite.%0A%0A%20%20%20%20As%20xDSL%20and%20MLIR%20allow%20the%20definition%20of%20higher-level%20dialects%2C%20pattern%20rewrites%20are%20very%20common%2C%20and%20can%20be%20used%20to%20both%20write%20high-level%20optimizations%2C%20and%20lowerings%20from%20a%20high-level%20dialect%20to%20a%20lower-level%20one.%20A%20general%20rationale%20for%20pattern%20rewrites%20in%20MLIR%20can%20be%20found%20%5Bhere%5D(https%3A%2F%2Fmlir.llvm.org%2Fdocs%2FRationale%2FRationaleGenericDAGRewriter%2F).%0A%0A%20%20%20%20Pattern%20rewrites%20are%20applied%20step%20by%20step%20on%20the%20IR.%20For%20instance%2C%20using%20the%20rewrite%20%60x%20%2B%200%20-%3E%20x%60%2C%20the%20IR%20%60x%20%2B%200%20%2B%200%60%20will%20be%20progressively%20transformed%20to%20%60x%20%2B%200%60%2C%20and%20then%20%60x%60%20using%20this%20single%20rewrite.%20Different%20application%20ordering%20and%20variations%20exist%20for%20rewrite%20patterns%2C%0A%20%20%20%20please%20see%20%60PatternRewriteWalker%60%20for%20more%20details.%0A%0A%20%20%20%20%23%23%20Defining%20a%20pattern%20rewrite%0A%0A%20%20%20%20Each%20Pattern%20rewrite%20is%20a%20class%20that%20inherits%20from%20%60PatternRewrite%60.%20It%20defines%20a%20single%20%60match_and_rewrite%60%20method%2C%20that%20is%20called%20to%20apply%20a%20pattern%20on%20an%20operation.%20The%20method%20will%20either%20return%20without%20any%20modification%20of%20the%20IR%2C%20or%20will%20modify%20the%20IR%20using%20the%20%60Rewriter%60.%20The%20most%20common%20operation%20to%20call%20on%20the%20%60Rewriter%60%20is%20%60replace_matched_op%60%2C%20which%20will%20replace%20the%20matched%20operation%20with%20a%20sequence%20of%20other%20operations%2C%20and%20replace%20the%20result%20values%20of%20the%20matched%20operation%20with%20new%20values.%0A%0A%20%20%20%20Here%20are%20two%20examples%20of%20pattern%20for%20%60x%20%2B%200%20-%3E%20x%60%20and%20%60x%20*%202%20-%3E%20x%20%2B%20x%60%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(AddiOp%2C%20ConstantOp%2C%20MuliOp%2C%20Operation%2C%20RewritePattern%2C%20Rewriter)%3A%0A%20%20%20%20%23%20The%20rewrite%20x%20%2B%200%20-%3E%20x%0A%20%20%20%20class%20AddZeroPattern(RewritePattern)%3A%0A%20%20%20%20%20%20%20%20def%20match_and_rewrite(self%2C%20op%3A%20Operation%2C%20rewriter%3A%20Rewriter)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Match%20an%20%60arith.addi%60%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20not%20isinstance(op%2C%20AddiOp)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Check%20if%20the%20right%20hand%20side%20is%20a%20constant%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20not%20isinstance(cst%20%3A%3D%20op.rhs.owner%2C%20ConstantOp)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Check%20if%20the%20constant%20is%200%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20cst.value.value.data%20!%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Replace%20%60x%20%2B%200%60%20with%20%60x%60%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20The%20first%20argument%20is%20the%20new%20operations%20to%20insert%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20The%20second%20argument%20is%20the%20replacement%20for%20the%20results%20of%20the%20matched%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20operation%20(op).%20As%20we%20replace%20%60x%20%2B%200%60%20with%20%60x%60%2C%20we%20replace%20the%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20%60arith.addi%60%20results%20with%20its%20first%20argument.%0A%20%20%20%20%20%20%20%20%20%20%20%20rewriter.replace(op%2C%20%5B%5D%2C%20new_results%3D%5Bop.lhs%5D)%0A%0A%20%20%20%20%23%20The%20rewrite%20x%20*%202%20-%3E%20x%20%2B%20x%0A%20%20%20%20class%20MulTwoPattern(RewritePattern)%3A%0A%20%20%20%20%20%20%20%20def%20match_and_rewrite(self%2C%20op%3A%20Operation%2C%20rewriter%3A%20Rewriter)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Match%20an%20%60arith.muli%60%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20not%20isinstance(op%2C%20MuliOp)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%0A%20%20%20%20%20%20%20%20%20%20%20%20x%20%3D%20op.lhs%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Check%20if%20the%20right%20hand%20side%20is%20a%20constant%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20not%20isinstance(cst%20%3A%3D%20op.rhs.owner%2C%20ConstantOp)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Check%20if%20the%20constant%20is%202%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20cst.value.value.data%20!%3D%202%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Replace%20%60x%20*%202%60%20with%20%60x%20%2B%20x%60%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20The%20results%20of%20the%20%60arith.muli%60%20operation%20are%20by%20default%20the%20results%20of%20the%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20last%20operation%20added%2C%20so%20here%20the%20%60arith.addi%60%0A%20%20%20%20%20%20%20%20%20%20%20%20add%20%3D%20AddiOp(x%2C%20x)%0A%20%20%20%20%20%20%20%20%20%20%20%20rewriter.replace(op%2C%20%5Badd%5D)%0A%0A%20%20%20%20return%20AddZeroPattern%2C%20MulTwoPattern%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Applying%20rewrite%20patterns%0A%0A%20%20%20%20There%20are%20two%20steps%20to%20apply%20rewrite%20patterns%20in%20xDSL.%0A%0A%20%20%20%20*%20First%2C%20combining%20multiple%20rewrite%20patterns%20in%20a%20single%20one%2C%20giving%20priorities%20between%20patterns.%0A%20%20%20%20*%20Then%2C%20walking%20the%20IR%20and%20applying%20the%20pattern%20recursively%20on%20all%20operations%20until%20no%20more%20pattern%20can%20be%20applied%20on%20any%20operation%3A%0A%0A%20%20%20%20Here%20is%20an%20example%20on%20how%20to%20apply%20rewrites%20on%20all%20operations%20in%20a%20module%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20AddZeroPattern%2C%0A%20%20%20%20GreedyRewritePatternApplier%2C%0A%20%20%20%20ModuleOp%2C%0A%20%20%20%20MulTwoPattern%2C%0A%20%20%20%20Operation%2C%0A%20%20%20%20PatternRewriteWalker%2C%0A)%3A%0A%20%20%20%20def%20apply_all_rewrites(module%3A%20ModuleOp%2C%20rewrites%3A%20list%5BOperation%5D)%3A%0A%20%20%20%20%20%20%20%20merged_pattern%20%3D%20GreedyRewritePatternApplier(AddZeroPattern()%2C%20MulTwoPattern())%0A%20%20%20%20%20%20%20%20walker%20%3D%20PatternRewriteWalker(merged_pattern)%0A%20%20%20%20%20%20%20%20walker.rewrite_module(module)%0A%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(Arith%2C%20Builtin%2C%20Context%2C%20Func)%3A%0A%20%20%20%20ctx%20%3D%20Context()%0A%20%20%20%20ctx.load_dialect(Builtin)%0A%20%20%20%20ctx.load_dialect(Arith)%0A%20%20%20%20ctx.load_dialect(Func)%0A%20%20%20%20return%20(ctx%2C)%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A