[alexcgatea]

Adds native rules to kirin-rewrite: rules written in Rust, with no pattern, for work a #[rewrite_rule] pattern can't express, such as CSE, which has to remember the expressions it has seen. Covers the four open native-rule items of M2 in docs/design/rewrite-engine.md.

Stacked on alex/walk-results.

What's added

• Two traits
  • StatementRule<L> is handed each statement the driver visits.
  • BlockRule<L> is handed each block a visited statement owns: its Blocks, then every block of its CFGs. Graph bodies are not handed over.
  • Both return Result<WalkResult, RewriteError>, so they steer the sweep like pattern rules do.
• RuleSet methods: add_statement_rule(name, recipe) and add_block_rule(name, recipe). The name is used in errors.
• One ordered rule list: pattern rules and native rules share one list. At each visited statement, every rule is tried once, in the order it was added.

rules.add_pattern_rule(double_neg(stage))?;
rules.add_statement_rule("mul_to_shift", || MulToShift);
rules.add_block_rule("cse", Cse::default);

Decisions

• Fresh state. RuleSet stores a recipe (a closure) per native rule, not a live rule. A statement rule gets a fresh copy every sweep, and a block rule a fresh copy for every block.
  • The fixpoint check relies on this: "a sweep changed nothing" only means "the next one would change nothing too" if a sweep doesn't depend on state left over from the last one.
  • Saved state also can't refer to IR a previous sweep erased.
• Block rules run inside the repeated sweeps. A block-local CSE runs again each sweep with a fresh table, instead of once with one persistent table. Gap 4 in the design doc is updated to match.
• Blocks are listed when the block rule's turn comes, after the earlier rules at that statement ran, so a block rule sees their edits. A block is handed to the rule this sweep only if it is under its owner when the walk visits the owner. One created or moved under a statement the walk already passed waits for the next sweep, which that edit triggers.
• No edit checking. A block rule may edit its block's statements and redirect uses of its values anywhere, as CSE needs. It must not edit the owner's other blocks.
  • The driver doesn't check this.
  • It only skips a listed block that an earlier copy erased or merged away, instead of handing a dead block to the rule.
• Native rules and erasing the visited statement. Their edits can't be known before the run, so the driver doesn't check them up front the way it does for pattern rules. A native rule that erases the visited statement in an order that doesn't allow it fails the run with DriverError::Walk.
• Errors. A native rule's RewriteError becomes DriverError::Rewrite, carrying the rule's name and the visited statement.

API changes

• RuleSet::add is renamed add_pattern_rule, to sit next to the two new methods. Its only caller was RuleSet::build.
• RuleSet::rules() is now crate-private.

Known limitation

A block rule gets a fresh copy per block, so a CSE block rule starts each block with an empty table of the expressions seen so far. It therefore can't reuse a value computed in an enclosing block, for example merge a neg %x in an scf.if arm into an identical neg %x before the if. Reusing one needs no driver change: hoist speculatable ops out of the nested block first, so both are in the same block, or have the rule seed its table by reading the statements of the enclosing blocks that come before the nested block's owner.

Tests

• New test language ScfFunctionLanguage in kirin-test-languages, behind feature scf-function-language. kirin-rewrite gains dev-dependencies on kirin-scf and kirin-cf.

[planza-quera]

Thanks for the work!
I'm a bit unsure as to whether allowing writing rewrite rules in this way is ideal. At my previous work it was a bit hassle having essentially two codebases in which rewrite rules are written: native C++ mlir patterns, and DSL PDLL patterns.

I suppose in Kirin this is not so much of a problem since the "DSL" #[rewrite_rule] patterns can live alongside the native Rust rules. Since both kinds of rules can also be included in the same RuleSet and applied together in the same rewrite driver, my concerns mostly go away. (Sorry for the ramble)

Regarding implementing CSE (or even DCE) as a #[rewrite_rule], my idea was that this would be possible once we have a way of accessing analysis results from the rewrite rules.

[alexcgatea]

I agree @planza-quera, I wasn't initially sold on why we should also have rules written in Rust either. But I think the main benefit is that these Rust rules are able to store state so they are more powerful.

[zhenrongliew]

Regarding implementing CSE (or even DCE) as a #[rewrite_rule], my idea was that this would be possible once we have a way of accessing analysis results from the rewrite rules.

Opened a PR implementing the read-only view to access analysis results. There is currently some awkwardness/inconvenience since the analysis results keep borrowing the same pipeline the rewrite engine needs.