Failures with `-DMLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=1`

[christopherbate]

What happened?

Cannot pass check-stablehlo-(linalg|quick|python|ci) if MLIR is built with CMake option
-DMLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=1.

When -DMLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=1 is enabled, the pattern-based infrastructure (greedy rewrites and dialect convesion) enforce additional invariants and will crash if the invariants are violated.

• When using the greedy rewrite driver, after each pattern is applied the IR must successfully verify. Sometimes the application of multiple patterns can hide issues where invalid IR is being created.

  • In Stablehlo to Linalg conversion, invalid linalg.map ops are being created where the result type does not match the DPS init argument type. Without expensive checks enabled, this gets ignored and happens to be corrected after all patterns are finished running. But under different conditions, this may not always be true. To correct the issue, a tensor.cast should be inserted.

• In dynamism rewriter, a stablehlo.real_dynamic_slice can be rewritten to a stablehlo.slice that is invalid (because the input program may have been "valid" but "incorrect", e.g. if the size of the slice is a stablehlo.constant that does not equal the actual specified static result type. I guess the upstream rule is that such canonicalization patterns should return failure vs. creating IR that does not verify. C.f. tensor.extract_slice canonicalizer starting from here. Without expensive checks enabled, this will raise a diagnostic at the end of the pass, but with expensive checks enabled it will crash with an assertion. This seems to be an area where MLIR upstream is particularly opinionated about what rewrite patterns should do.

• Some patterns not notifying the rewriter when making modifications (result type updates in Stablehlo type refinement pass being main issue I noticed)