[Remat][Bugfix] Enhance the rematerialization pass to handle nested tuples

[zhouyuan1119]

Description

This PR addresses issue #8 and presents a small fix for the rematerialization pass to handle nested tuples. Specifically, it handles cases like the following:

...
let a = (x0, x1)
...
let b = some_call()
...
let c = (a, b)

Without this patch the remat pass would error out when processing c because of a simplified assumption we made in the first version of the pass. Notice that this fix does not address the case where a call node directly produces a nested tuple (which does not exist in current RAF as far as I know).

ResNet50 with AMP can now be compiled with rematerialization turned on, and the benchmarked latency is ~35ms per batch when the batch size is 16.

I will add a test for nested tuples later. Please review the changes first.

Checklist

• PR's title starts with a category (e.g. [BUGFIX], [MODEL], [TUTORIAL], [FEATURE], [DOC], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

cc @awslabs/raf-reviewer @comaniac

[comaniac]

The change LGTM.

[comaniac]

LGTM. Please remove the LOG

[comaniac]

Thanks @zhouyuan1119