[AOT] Enable A-Normal Form in the AOT executor

[lhutton1]

The sequence of calls produced by the AOT executor codegen is arbitrary, especially in the presence of 'branchy' networks. This makes it difficult to analyze memory usage for each call. By running the ToANormalForm pass to insert a series of let bindings before the lowering and codegen stages, we can establish an ordering for the evaluation of the external calls, thus allowing reliable analysis of memory usage.

Additionally contains a couple of error message improvements and comment corrections.

cc @manupa-arm @Mousius @areusch @mbs-octoml @ashutosh-arm @ekalda @NicolaLancellotti

[mbs-octoml]

In an ideal world we'd convert to ANF early on. In addition to making eval order explicit, it would also mean we could retire MixedModeVisitor since the deep recursion will be in lets which are trivial to deal with. However we can't do that without interrupting all uses of DFPatterns since the matcher does not 'see through' vars. You thankfully stay well out of that mess, so I'd support this.

That said the visit order is deterministic, right? Or pehaps we're accidentally depending on arbitrary undordered_set iteration order somewhere in the visitor machinery?

[lhutton1]

Thanks for taking a look @mbs-octoml, yeah our plan was to put our foot-in-the-door at this stage as it doesn't require too much modification and then enable ANF at an earlier stage in the future if need be.

I believe your thinking that the visit order is deterministic is correct, so then enabling ANF would just make the evaluation order explicit, like you said. Making the lowering pipeline to adhere to an explicit ordering at this stage can be helpful for passes which rely on an explicit ordering such as liveness analysis, rather than solely relying on the implementation of the visitor to provide such guarantees. @manupa-arm may be able to comment further here :)

[mbs-octoml]

LGTM but please do see what you can do with my suggestion in the storage allocator.

[MarisaKirisame]

Not much comment about the code, as I didnt know this part of the codebase well, but -
as the author of the ToANF pass very happy to see it replacing the graph mode :)
making recursive tuple possible is also great - it is a feature needed for reverse mode ad which is needed for training.

What do you guys want to do with analyzing memory for example? What branchy network do you guys have in mind? I was working on saving memory for deep learning and think it is very relevant to Dynamic Tensor Rematerialization - it can work even for full inference, but most 'straight-forward' network dont really have saving, only on network with a more complex structure.

[manupak]

@MarisaKirisame, as @lhutton1 mentioned before, we are trying to bring explicit ordering to middle-end of the stack for which we can rely on to make certain decisions in the backend passes. I suppose we dont have any special model in mind, but generally looking out for ones that would be generated from NAS's and common models such as Inception models. In future, I suppose it opens up possibilities to explore different variations of sequentialization given the topology of model.

[manupak]

Thanks @lhutton1 ! looking good.

Just a question/suggestion to see if can get rid of custom let-chain unrolling..

[manupak]

Won't we achieve the same effect by using something from here ? :

tvm/src/relay/transforms/device_aware_visitors.h

Lines 247 to 258 in f0854ae

	/*!
	* \brief Visit a let expression after the let-bound value and body have been visited.
	* Default implementation is a no-op.
	*/
	virtual void PostVisitLet_(const LetNode* let_node) { /* no-op */
	}
	/*!
	* \brief Visit the first let in a chain of let expressions after it has been visited.
	* Default implementation is a no-op.
	*/
	virtual void PostVisitLetBlock_(const LetNode* let_node) {}

[lhutton1]

hmm I can't see a way to make use of these here, perhaps I've missed something?

[manupak]

I was thinking instead of :

void PreVisitLetBinding_(const Var& var, const Expr& value) final {
    LOG(FATAL) << "let is not supported.";
    VisitExpr(value);
    StorageInfo si = GetStorage(value);
    storage_device_map_[var] = si;
  }

if we do PostVisitLet, wont we have the same effect (i.e. the body is visited)?

[mbs-octoml]

No objection from me if you want to just do the loop directly as I find that much clearer. In my travels I found switching to pre- and post-let-binding visitor overloads made the most sense when you find you need to accumulate intermediate rewritten bindings on a local stack in order to reconstruct the final let chain. ExpandANormalForm also achieves the same effect all be it by relying on the implicit memoization to capture the intermediate results (which hurts my head every time I come across it!)

[lhutton1]

@manupa-arm and I discussed this offline and it turns out we can just avoid let's being used in GetStorage

[manupak]

LGTM

[manupak]

Thanks @lhutton1 @mbs-octoml @ashutosh-arm !

This is merged now!