Strategy dialect for precise specification of complex rewrite rules

[apaszke]

@ntv recently brought to my mind that together with @ftynse (and now @Mogball) they started working on a dialect for driving pattern (and pass) applications. This is meant to be an improvement over the standard greedy driver provided by MLIR which does seem to be very constrained. I think this is all a super cool domain, and @ntv prompted me to dump some of my thoughts on this in an issue, so here we go!

Note that most of it is heavily inspired by the functional pearl paper that describes ELEVATE, only translated to MLIR. Finally, I’m far from being an MLIR expert, so parts of this proposal might be based in some significant misconceptions. Please point those out whenever you see them!

The strategy dialect

This is an outline of the operations I'd suggest to include in the new dialect:

Types:

• !strategy.strategy: Roughly represents a function with signature (!pdl.operation) -> (). Note that an application might fail which is represented as a side effect and not as a return value.

Operations:

• strategy.try: has a single-block region and one boolean result. If the execution of the region fails, it doesn't propagate the failure upwards but returns false. If the execution of the region succeeds returns true.
• strategy.apply: applies an external strategy to an operation (its single operand). The external strategies are identified by symbols that will be resolved to real passes and rewrite patterns during interpretation (e.g. @canonicalize could be backed by the canonicalization pass).
• strategy.reify: turns a single-block MLIR region that takes an operation as an argument into a !strategy.strategy value. Basically a lambda expression.
• strategy.apply_dynamic: applies a dynamic strategy to an operation (i.e. it takes two operands). If the dynamic strategy fails, the execution of this operation fails.

Many other useful types and operations can be inherited from PDL (!pdl.operation, !pdl.attribute, etc.).

The programs expressed in this dialect are meant to be interpreted sequentially, as usual in MLIR, and encode a sequence of rewrite applications, potentially orchestrated by control flow (e.g. using the scf dialect) and recursion.

Example program:

// Helper function for ignoring failure of strategy applications
func @try_apply(%s : !strategy.strategy, %op: !pdl.operation) {
  %success = strategy.try {
    strategy.apply_dynamic %s(%op)
  }
}

// Strategy transformer
func @or(%s1 : !strategy.strategy, %s2 : !strategy.strategy) -> !strategy.strategy {
  %ans = strategy.reify {
    ^bb0(%op: !pdl.operation):
      %success = strategy.try {
        strategy.apply_dynamic %s1(%op)  // pattern1
      }
      %failure = not %success
      scf.if %op {
        strategy.apply_dynamic %s2(%op)  // pattern2
      }
  }
  return %ans
}

func @my_strategy(%op: !pdl.operation) {
  // If canonicalization fails, we fail the whole strategy, and
  // none of the operations below this one will be applied.
  strategy.apply @canonicalize(%op)
  // Reify the MLIR blocks as !strategy.strategy first-class values
  %tile20 = strategy.reify {
    ^bb0(%op: !pdl.operation):
      %tile_size = arith.constant 20 : i64
      strategy.apply @linalg-tile-strategy(%op, %tile_size)
  }
  %tile5 = strategy.reify {
    ^bb0(%op: !pdl.operation):
      %tile_size = arith.constant 5 : i64
      strategy.apply @linalg-tile-strategy(%op, %tile_size)
  }
  // Try tiling with tile size of 20 or a tile size of 5 (and fail if both fail)
  %my_tiles_fallible = std.call @or(%tile20, %tile5)
  // Use the @try_apply helper to apply our strategy
  std.call @try_apply(%my_tiles_fallible, %op)
  return
}

// Functions can be recursive, so that e.g. the eager exhaustive application of
// a pattern is expressible as a first-class program. It might be ok
// to start without recursion and treat such patterns as builtins.
func @apply_eagerly(%op: !pdl.operation, %s: !strategy.strategy) {
  // Try to apply to this op (if possible)
  std.call @try_apply(%s, %op)
  // Try to apply to all children operations too
  for (%region in %op.regions) {
    for (%block in %region.blocks) {
      for (%subop in %block.operations) {
        std.call @apply_eagerly(%subop, %s)
      }
    }
  }
}

How to use it?

Initially it would make sense to start with a simple interpreter, but I imagine that eventually we could go much further. For example, we could try partially evaluating all of the ops from the strategy dialect and lower them all to a PDL program (assuming PDL is expressive enough).

What are patterns?

Patterns can be applied to a particular "root" operation and either fail to match, or mutate the IR according to their definition. Operation passes work similarly, except they always succeed (at least as far as I understand them). In the new dialect both passes and patterns defined in C++ would be treated as externally defined strategies, which are functions that apply to an operation handle and can throw as a side effect.

[nicolasvasilache]

Nice!

apply/apply_dynamic seem analogous to call/call_indirect, right?

I think I see recursion being useful in cases where we write super generic stuff that targets an abstract machine model.
I also have in the back of my mind some Darkroom / Terra use cases but eventually I think the tradeoff between nesting/recursion and more structured IR design favors IR design: the apparent simplicity really gets moved to complex combinations of small things.. maybe I am diverging a bit :)
Could you share some other use cases that you could solve well thanks to recursion?

My inclination is to start without recursion and see where this brings us, as we can likely start scaling in a more naive fashion on fixed HW models. Generalization can come later.

I have a somewhat similar thinking for failures: I wonder how much "can fail" is a property of application of a complex pattern to "too low-level IR objects". Experience with the current system seems to show that our very progressive lowering approach from high-level IR allows us to not worry about rollbacks. There seems to be a connection to more static vs more dynamic programs; i.e. can you hoist all your conditions upfront or not.

I imagine we'll encounter plenty of use cases for failures as we scale up but similarly, to get off the ground, I think we can do without most of the infra complexity.

How about a proto ? :)

[apaszke]

apply/apply_dynamic seem analogous to call/call_indirect, right?

Correct!

Could you share some other use cases that you could solve well thanks to recursion?

Recursion is mostly useful for expressing rewrites that need to traverse down arbitrary unknown operations, and we certainly could start without it, or only with some special forms (e.g. exhaustive greedy application like in the default driver).

I have a somewhat similar thinking for failures: I wonder how much "can fail" is a property of application of a complex pattern to "too low-level IR objects".

I think that you're right in that proper IR design can let you avoid thinking about failures in many cases, but I don't think it's always possible. Linalg might be a neat enough use-case and design such that you can get away without it, but you might run into trouble trying to scale it up to arbitrary MLIR transforms. And IMO things like loop-interchange should be also expressible as rewrites, even though this is out of scope for linalg.

How about a proto ? :)

Sounds good to me! :)

[Mogball]

!strategy.strategy: Roughly represents a function with signature (!pdl.operation) -> (). Note that an application might fail which is represented as a side effect and not as a return value.

What happens if the strategy mutates the operation, or replaces it entirely? E.g.

strategy.apply @canonicalize(%op) // %op is canonicalized away/replaced
strategy.apply @myStrategy(%op)  // %op is now invalid?

I will also propose that "strategies" can return IR handles, i.e.

%tiledOp = strategy.apply @tile(%op, %tileSize) : (!pdl.operation, i64) -> (!pdl.operation)
strategy.apply @canonicalize(%tiledOp) : (!pdl.operation) -> ()

[apaszke]

Extending strategies to have additional return values is definitely possible (in functional terms those would run in the Either monad, which roughly means that every strategy either succeeds with a value or fails with an error).

The question of what happens if an op is destroyed is a good one and it does tie to the difficulties around mutability. I think there are multiple ways we can potentially handling that. First, we could consider this to be a programmer error and rooting a strategy at a deleted operation could cause the overall strategy to fail. Alternatively, we could make it so that patterns cannot delete operations and they are expected to be later cleaned up by DCE if they end up being unnecessary.

[ftynse]

I would add a strategy.define operation that, similarly to function definitions, associates a region with a symbol name. This is a symbol-level counterpart to strategy.reify that produces a value. Relying on builtin.func for this purpose will quickly turn problematic.

In general, do you think we can survive with non-parametric strategies (i.e. ones that only take a !pdl.operation) or will there be some parametricity?

Plus the usual concerns about complexity of the rewrite engine and action-at-a-distance on the IR that make debugging hard should something go wrong. I suppose these cannot be answered without a prototype.

[Mogball]

In general, do you think we can survive with non-parametric strategies (i.e. ones that only take a !pdl.operation) or will there be some parametricity?

I think adding parameters to patterns is relatively straightforward (attributes on apply?).

Alternatively, we could make it so that patterns cannot delete operations and they are expected to be later cleaned up by DCE if they end up being unnecessary.

This sounds to me like too strict of a restriction. We can instead have the pattern explicitly return the operation handles that can be chained.