Revamp Hannk IR

[steven-johnson]

(I'm opening this as a draft because I'm not sure if the current design is acceptable; in particular, the implementation of OpMutator is ugly in ways that bother me a lot. I'm nevertheless tagging some people for 'review' because I'd like opinions on whether the ugliness is a dealbreaker.)

This PR attempts to refactor the Hannk IR (and transforms) to be a bit more principled and less error-prone. In particular:

• It removes (almost) all of the mutators from all of the Op hierarchy, with the exception of:
  • prepare() and execute(), which probably could be const but would require adding a handful of mutable annotations.
  • set_input(), which we still use in one place when updating consumed Tensors (this is a TODO)
• Ops are now considered 'immutable'; instead of mutating them in place, you should use an OpMutator, in which you construct a new, replacement Op when changes are needed.
• OpVisitor becomes all-const since it doesn't allow mutating anything.
• OpMutator is added; this is similar to Halide's IRMutator, with a couple of notable exceptions:
  • The op being mutated is passed in as a unique_ptr instead of as a raw ptr.
  • You can return nullptr, which simply means "delete this op from its parent group".

Using this, I rewrote all the transforms using OpMutator (and added a new flatten_groups pass); I also added a variant of the model-validator code from #6317. Everything in my test suite passes, and I haven't seen any regression in optimization (I'll post performance numbers as a followup comment).

The reason I'm hesitant to move forward with this is that the implementation of OpMutator is... ugly, due to the necessity of transferring unique_ptr ownership. I'm convinced that it's actually legal, safe C++, but I don't like it. (Nevertheless, I'd appreciate comments.)

[steven-johnson]

I'm updating this to "ready for review", even though I'm still not entirely happy with it, because after some reflection, I think it's still a net improvement.

[steven-johnson]

Monday Morning Afternoon Review Ping

[dsharletg]

I had hesitated to make ops immutable, because this design will require more memory allocations, which might be an issue on DSPs with weak scalar performance. But I also think we should deal with that later if it's an issue, rather than make the IR harder to work with now.

[dsharletg]

What's going on here? This looks suspicious... it looks like a getter that takes the op out of the group?

[steven-johnson]

Yep -- this is used by OpMutator::visit to mutate an OpGroup. It needs to take ownership of the unique_ptr so it can grant ownership to the new OpGroup. (The implementation with swap is probably overkill; I wanted to be sure I wasn't leaving the old OpGroup in an inconsistent state, but that's probably not essential.)

[dsharletg]

Should this be renamed then? take_op or something?

[steven-johnson]

Done

[dsharletg]

Comment why this takes an op as a parameter

[steven-johnson]

I'll comment here for starters (and move into code if that's what you want)...

The problem here is that we are trying to implement double-dispatch via double-virtual-method calls (in the same fashion that Halide's similar classes do); the gotcha is that we also want to transfer ownership of 'this' into the receiving method call (since it will ultimately either keep-or-replace it).

In Halide's IR, the IntrusivePtr approach allows us to transparently 'recover' a refcounted Expr from a specific pointer (e.g. const Call * implicitly converts to Expr and maintains the refcount correctly); with unique_ptr, there isn't any such way to do this, so we have to pass a redundant argument. It's not pretty, but the generated code is efficient.

For the record, I considered moving Op to use shared_ptr instead, which can sorta-kinda do the same sort of recovery via std::enable_shared_from_this, but it felt like it required too much obscure special knowledge to be friendly (e.g. explicit calls to the shared_from_this() method). I also considered moving Op to just use a cut-down variant of IntrusivePtr, but that would have been a pretty major chunk of code churn.

[dsharletg]

I guess more specifically, I just think there should be a brief comment here on how to actually use this function: what the arguments mean.

[steven-johnson]

Added some comments, but also overhauled the mutate()/accept()/visit() methods some to be cleaner and to make improper calls hard-to-impossible, PTAL

[dsharletg]

Is the reverse order a common use case? I am wondering if RemoveDeadOps should just override the default behavior rather than building it into the base class.

[steven-johnson]

No, it has exactly one use case currently. I initially did what you suggest here, but the implementation of OpMutator::visit ended up feeling ticklish enough to get right that I thought it more prudent to encapsulate it all in one spot, rather than encourage possibly-wrong code replication. (I don't feel strongly about this, though.)

[steven-johnson]

PTAL

[steven-johnson]

this design will require more memory allocations

Yeah, it's a fair concern. It may be that the Halide IRMutator pattern isn't the ideal one for what we want to do. But it does seem like a more tractable way to get IR that is reliably clean -- my hope here is that if we find the extra allocations to be problematic, we can reintroduce mutability on a selective basis.