[AOT] Calculate used memory at the callsite of primitive functions

[lhutton1]

Introduces a new pass in the AOT executor called "AnnotateUsedMemory" which applies liveness analysis to the callsite of each primitive function in order to calculate the total size of the live tensors at this point of execution. The result is provided as a function annotation called "used_memory", which can be consumed by later stages of the compiler (e.g. external codegens) to provide more information about the current memory consumption. This can be useful for some optimizations.

Note: this PR is dependent on #11091 so also shows the contents of that PR.

cc @Mousius @NicolaLancellotti @ekalda @manupa-arm

[lhutton1]

also cc @mbs-octoml @areusch

[lhutton1]

No functional changes here, simply moving manifest_lifetimes.cc ../ (outside scope of vm) and splitting into .cc/.h

[manupak]

Thanks @lhutton1 !

I did a take a first look. It looks broady good.
Few suggestions for more test cases and a question about using the the stack.

[areusch]

hi @lhutton1 @manupa-arm , just had a couple questions on this one

[lhutton1]

Thanks for the reviews @manupa-arm, @areusch, @altanh - hoping to have a revised version ready soon!

[areusch]

thanks @lhutton1 for the replies! ping us when this is ready for review again.

[lhutton1]

Apologies for the delay, this is ready for another look!

[areusch]

ok thanks @lhutton1 ! i'll defer to @altanh on this one

[altanh]

mostly LGTM with some small questions/nits!

[altanh]

Thanks for the changes! LGTM

[manupak]

LGTM!

[manupak]

@lhutton1 since it has been 18 days, should we re-run a round of CI -- just to be sure :)

[manupak]

Thanks @lhutton1 @altanh @areusch ! This is merged now

[zhaoyang-star]

Hi @lhutton1 , thanks for your contributition.
After running FuseOps pass, I want to get the memory usage per op or per primitive func by AnnotateUsedMemory pass for furture optimization. I get a resnet18 ir model, then put it as the input IRModule of AnnotateUsedMemory pass. The output IRModule has no used_memory attr. Test code as follow:

import pytest
from collections import OrderedDict
import numpy as np
import tvm
from tvm import relay
from tvm.relay import testing


def AnnotateUsedMemory():
    return relay.transform._ffi_api.AnnotateUsedMemory()


def _get_data(in_data_shapes, dtype="float32"):
    in_data = OrderedDict()
    for name, shape in in_data_shapes.items():
        in_data[name] = np.random.uniform(size=shape).astype(dtype)
    return in_data


def _run_relay(mod, params, in_data, pass_enabled):
    target = "llvm"
    dev = tvm.device("llvm", 0)
    in_data = [tvm.nd.array(value) for value in in_data.values()]

    if pass_enabled:
        mod = relay.transform.InferType()(mod)
        mod = relay.transform.ToANormalForm()(mod)
        mod = relay.transform.InferType()(mod)
        mod = AnnotateUsedMemory()(mod)
        # create primitive functions
        mod = relay.transform.FuseOps()(mod)

    print(f'\nmod when AnnotateUsedMemory is {pass_enabled}:\n {mod}')

    out_data = relay.create_executor(
        "graph", mod, device=dev, target=target).evaluate()(*in_data, **params)
    return out_data.numpy()


def _verify_results(mod, params, in_data, rtol=1e-5, atol=1e-5):
    before = _run_relay(mod, params, in_data, False)
    after = _run_relay(mod, params, in_data, True)
    np.testing.assert_allclose(before, after, rtol, atol)


def test_resnet():
    num_class = 1000
    in_data_shapes = OrderedDict({"data": (1, 3, 224, 224)})
    in_data = _get_data(in_data_shapes, dtype="float32")
    for n in [18]:  # 18, 34, 50, 101
        mod, params = tvm.relay.testing.resnet.get_workload(
            batch_size=1, num_classes=num_class, num_layers=n)
        _verify_results(mod, params, in_data)


if __name__ == "__main__":
    pytest.main([__file__])

I am not familar with AnnotateUsedMemory pass. Could memory usage per op or per primitive func be gotten by your pass? If not, how to get it based on your pass? Thanks in advance ^_^

[lhutton1]

Hi @zhaoyang-star, thanks for taking a look, its great to see this pass being used elsewhere. The pass currently expects the input to be a module of primitive functions so I would suggest running AnnotateUsedMemory after FuseOps similar to:

mod = relay.transform.InferType()(mod)
mod = relay.transform.FuseOps()(mod)
mod = relay.transform.InferType()(mod)
mod = relay.transform.ToANormalForm()(mod)
mod = relay.transform.InferType()(mod)
mod = AnnotateUsedMemory()(mod)

I did try running your example locally with the above change and this produced the relevant used_memory annotations. However, it looks like there is an issue while building the module after having run the AnnotateUsedMemory pass. Without digging too much into it I would suspect it's because this pass wasn't considered for the graph executor; only for the AOT executor. I believe changes similar to #11091 would be needed in the graph executor to support A-normal form. Hope this helps :)

[zhaoyang-star]

Hi @zhaoyang-star, thanks for taking a look, its great to see this pass being used elsewhere. The pass currently expects the input to be a module of primitive functions so I would suggest running AnnotateUsedMemory after FuseOps similar to:

mod = relay.transform.InferType()(mod)
mod = relay.transform.FuseOps()(mod)
mod = relay.transform.InferType()(mod)
mod = relay.transform.ToANormalForm()(mod)
mod = relay.transform.InferType()(mod)
mod = AnnotateUsedMemory()(mod)

I did try running your example locally with the above change and this produced the relevant used_memory annotations. However, it looks like there is an issue while building the module after having run the AnnotateUsedMemory pass. Without digging too much into it I would suspect it's because this pass wasn't considered for the graph executor; only for the AOT executor. I believe changes similar to #11091 would be needed in the graph executor to support A-normal form. Hope this helps :)

@lhutton1, I want to confirm: Did you reproduce the issue( no used_memory attr in the output log) using my script above? If you ran all right, could you please share your script? There is only one io_used_memory attr and no used_memory attr found after running my script.

If I placed the FuseOps before AnnotateUsedMemory just as you showed, there is a error Check failed: (tensor_type) is false:. You have mentioned maybe we should support ANF in graph executor to solve the error.

[lhutton1]

Hi @zhaoyang-star, yes I was able to reproduce the issue with your script. The script I have would be the same as yours just with the a different pass order as mentioned above. Placing FuseOps before AnnotateUsedMemory seems like the correct thing to do here; if you print out the module (mod) after the AnnotateUsedMemory pass you should be able to see the used_memory annotations. The Check failed: (tensor_type) is false: error comes later in the compilation so it seems as though some later optimization passes cannot deal with ANF yet.