Peak Memory Consumption doesnt change after applying Multi-Stream Execution

[ikaroso]

HI, I want to see the change of Peak Memory Consumption after applying multi-stream execution. I try ASAP and WaveFront, but it seems that the Peak Memory Consumption doesn't change at all.

And this is my code.

@yaoyaoding Hello, would you please help me solve this problem? THX :)

from raf._ffi.pass_ import ToGraphNormalForm, WavefrontStreamSchedule
import tvm
from raf._core.vm import VMCompiler
from raf._ffi.pass_ import EstimateMemory, InferType
from raf._core.device import Device



device="cuda"
model, _ = inception.get_model()

# Generate a dummy input.
in_out , _  = inception.get_input(batch_size=32, device="cuda")
m_in=in_out[0]
m_out=in_out[1]
mod = model._internal(m_in,m_out).mod
'''
with raf.ir.PassContext(opt_level=2, config={"raf.stream_schedule.policy": "wavefront"}):
    mod = RAFSequential([ToGraphNormalForm(), WavefrontStreamSchedule()])(mod)
'''
mod = RAFSequential([ToGraphNormalForm(), ASAPStreamSchedule()])(mod)


compiler = VMCompiler()
with tvm.transform.PassContext(opt_level=3):
    mod, _ = compiler.optimize(mod, device)
mod = InferType()(mod)
trace = [(name, mem.value) for name, mem in EstimateMemory(mod, Device(device), True)]
print(max(trace, key=lambda x: x[1])[1])

[yaoyaoding]

Hi @ikaroso,

The multi-stream execution accelerates the model if it contains multiple branches of small kernels, not aiming to reduce memory consumption.

The rematerialization pass is used to reduce the memory consumption and you can have a try.

[ikaroso]

Hello @yaoyaoding.
Thx for replying. I didn't describe the problem clearly. Well the actual problem is that multi-stream is supposed to increase the peak memory footprint (if I understand correctly). But result shows that the peak of the model doesn't change after applying multi-stream.

[yaoyaoding]

Yes, you are right that the multi-stream execution might increase memory consumption.

Example 1 (not increase peak memory)

a ----> b ---> d
  \          /
   ---> c --> 

We assume all operators (a-d) have a single output tensor with the same size (1 unit). Then both mutl-stream schedule Sequential(a, Parallel(b, c), d) and Sequential(a, b, c, d) would consume 3 units of peak memory.

Example 2 (increase peak memory):

a ----> b ---> d ---> f
  \                 /
   ---> c ---> e -->

Schedule Sequential(a, Parallel(b, c), Parallel(d, e), f) would consume 4 units of peak memory while schedule Sequential(a, b, c, d, e, f) would consume 3 units of peak memory.

In a more general case, we can write the memory consumption during model execution as

size(live tensors) + size(output tensors of operators that will be executed in next step)

The multi-stream execution might increase the second part of the above formula.

Above is a simple analysis of ideal implementation. There are two potential cases of what you observed for Inception:

1. The memory bottleneck of Incpetion does not occur in multi-stream execution (second part of the formula) but in the first part.
2. The implementation cases both schedules have the same peak memory.

If you want to further investigate, my suggestion is to just use Example 2 with sequential schedule and ASAP schedule to compare memory consumption. You can find the order of passes here starting at L838, the memory plan pass here, the memory estimator here. Print out the IRModule before and after memory plan might also be helpful. Let me know if anything confuses you or you find any mistake in above analysis.

[ikaroso]

@yaoyaoding very clear explanation and i think your analysis is correct. I ignored the "might" part and your analysis helps me understand deeply.

Before, I was doubting whether memory_estimator in raf can respond to the multi-stream execution. Now I guess the answer is "true", and I will use the experiment suggested by you to see the result.

Thx very much and I will post the result in here after the experiment.

[ikaroso]

@yaoyaoding i use example 2 to try ASAP and WaveFront, while the peak is still the same after multi-stream. Is there something wrong with my code?

"""
        x
        │
        ▼
      atan0
        │
     │      │
     ▼      ▼
    atan1  atan2
     │      │
     ▼      ▼
    atan3  atan4
     │      │
     │      │
     │      │
     │      │ 
     ▼      ▼
    concatenate
"""

import tvm
from raf._core.vm import VMCompiler
from raf._ffi.pass_ import EstimateMemory, InferType
from raf._core.device import Device
import raf
from raf.testing import randn
from raf.ir.pass_manager import RAFSequential
from raf._ffi.pass_ import ToGraphNormalForm, ASAPStreamSchedule,WavefrontStreamSchedule

device="cuda"
class Model(raf.Model):

    def build(self):
        pass

    @raf.model.trace
    def forward(self, x):
        p_0 = raf.atan(x)  # atan 0

        p_1 = raf.atan(p_0)  # atan 1
        p_1 = raf.atan(p_1)  # atan 3

        p_2 = raf.atan(p_0)  # atan 2
        p_2 = raf.atan(p_2)  # atan 4
        return raf.concatenate([p_1, p_2])



model = Model()
input_shape = [2, 2]
x, _ = randn(input_shape)
mod = model._internal(x).mod

compiler = VMCompiler()
with tvm.transform.PassContext(opt_level=3):
    mod, _ = compiler.optimize(mod, device)
mod = InferType()(mod)
trace = [(name, mem.value) for name, mem in EstimateMemory(mod, Device(device), True)]
print(max(trace, key=lambda x: x[1])[1])


mod = model._internal(x).mod
#print(mod)
#with raf.ir.PassContext(opt_level=2, config={"raf.stream_schedule.policy": "wavefront"}):
    #mod = RAFSequential([ToGraphNormalForm(), WavefrontStreamSchedule()])(mod)
mod = RAFSequential([ToGraphNormalForm(), ASAPStreamSchedule()])(mod)
#print(mod)

compiler = VMCompiler()
with tvm.transform.PassContext(opt_level=3):
    mod, _ = compiler.optimize(mod, device)
mod = InferType()(mod)
trace = [(name, mem.value) for name, mem in EstimateMemory(mod, Device(device), True)]
print(max(trace, key=lambda x: x[1])[1])

[yaoyaoding]

Try using 'add' instead of 'concatenate' as the output of 'concatenate' has double memory consumption. Thus, it is expected to have 4 units of peak memory for both sequential and ASAP schedules (occur when you run concatentate).

[ikaroso]

@yaoyaoding well i change 'concatenate' to 'add', the total peak is reduced compared to using 'concatenate'. But multi-stream still has the same peak as sequential.

concatenate

seq: 7.62939453125e-05
multi: 7.62939453125e-05

add:

seq: 3.0517578125e-05
multi: 3.0517578125e-05

i print the two IR moudle

def @main(%x: Tensor[(2, 2), float32] /* ty=Tensor[(2, 2), float32] */) {
  let %a1 = raf.op.atan(%x);
  let %a2 = raf.op.atan(%a1);
  let %a3 = raf.op.atan(%a2);
  let %a4 = raf.op.atan(%a1);
  let %a5 = raf.op.atan(%a4);
  let %a6 = raf.op.add(%a3, %a5, -114514i64 /* ty=() */, -114514i64 /* ty=() */);
  %a6
}

def @main(%x: Tensor[(2, 2), float32] /* ty=Tensor[(2, 2), float32] */) {
  let %x_0 = raf.op.set_stream(-114514i64, -114514i64);
  let %x_1 = raf.op.atan(%x);
  let %x_2 = raf.op.add_event(-114514i64);
  let %x_3 = raf.op.atan(%x_1);
  let %x_4 = raf.op.atan(%x_3);
  let %x_5 = raf.op.set_stream(-114514i64, -114514i64);
  let %x_6 = raf.op.wait_event(-114514i64);
  let %x_7 = raf.op.atan(%x_1);
  let %x_8 = raf.op.atan(%x_7);
  let %x_9 = raf.op.add_event(-114514i64);
  let %x_10 = raf.op.set_stream(-114514i64, -114514i64);
  let %x_11 = raf.op.wait_event(-114514i64);
  let %x_12 = raf.op.add(%x_4, %x_8, -114514i64 /* ty=() */, -114514i64 /* ty=() */);
  %x_12
}

and my code is like

import tvm
from raf._core.vm import VMCompiler
from raf._ffi.pass_ import EstimateMemory, InferType
from raf._core.device import Device
import raf
from raf.testing import randn
from raf.ir.pass_manager import RAFSequential
from raf._ffi.pass_ import ToGraphNormalForm, ASAPStreamSchedule,WavefrontStreamSchedule


device="cuda"
class Model(raf.Model):

    def build(self):
        pass

    @raf.model.trace
    def forward(self, x):
        p_0 = raf.atan(x)  # atan 0

        p_1 = raf.atan(p_0)  # atan 1
        p_1 = raf.atan(p_1)  # atan 3

        p_2 = raf.atan(p_0)  # atan 2
        p_2 = raf.atan(p_2)  # atan 4
        return raf.add(p_1, p_2)


model = Model()
input_shape = [2, 2]
x, _ = randn(input_shape)
mod = model._internal(x).mod

compiler = VMCompiler()
with tvm.transform.PassContext(opt_level=3):
    mod, _ = compiler.optimize(mod, device)
mod = InferType()(mod)
trace = [(name, mem.value) for name, mem in EstimateMemory(mod, Device(device), True)]
print(max(trace, key=lambda x: x[1])[1])


mod = model._internal(x).mod
print(mod)
#with raf.ir.PassContext(opt_level=2, config={"raf.stream_schedule.policy": "wavefront"}):
    #mod = RAFSequential([ToGraphNormalForm(), WavefrontStreamSchedule()])(mod)
mod = RAFSequential([ToGraphNormalForm(), ASAPStreamSchedule()])(mod)
print(mod)

compiler = VMCompiler()
with tvm.transform.PassContext(opt_level=3):
    mod, _ = compiler.optimize(mod, device)
mod = InferType()(mod)
trace = [(name, mem.value) for name, mem in EstimateMemory(mod, Device(device), True)]
print(max(trace, key=lambda x: x[1])[1])

then I remove

with tvm.transform.PassContext(opt_level=3):
    mod, _ = compiler.optimize(mod, device)

and the peak of 'seq' and 'multi-stream' is still the same

[yaoyaoding]

Hi @ikaroso,

Thanks for your results. Could you print out the module after optimization (mod, _ = compiler.optimize(mod, device))? I am interested in the output of ManifestAlloc pass of both schedules. Based on the output of manifest alloc pass and the implementation of EstimateMemory analyzer, I believe we can find out what happened under the hood.

[ikaroso]

@yaoyaoding thx for your advice,

The result of 'seq' is :

def @main(%x: Tensor[(2, 2), float32] /* ty=Tensor[(2, 2), float32] */) -> Tensor[(2, 2), float32] {
  let %x_2 /* ty=fn (Tensor[(2, 2), float32], (), ()) -> Tensor[(2, 2), float32] */ = fn (%p0: Tensor[(2, 2), float32] /* ty=Tensor[(2, 2), float32] */, %p1: () /* ty=() */, %p2: () /* ty=() */, Primitive=1, Dialect="tvm") -> Tensor[(2, 2), float32] {
    %0 = raf.op.tvm.atan(%p0) /* ty=Tensor[(2, 2), float32] */;
    %1 = raf.op.tvm.atan(%0) /* ty=Tensor[(2, 2), float32] */;
    %2 = raf.op.tvm.atan(%0) /* ty=Tensor[(2, 2), float32] */;
    %3 = raf.op.tvm.atan(%1) /* ty=Tensor[(2, 2), float32] */;
    %4 = raf.op.tvm.atan(%2) /* ty=Tensor[(2, 2), float32] */;
    raf.op.tvm.add(%3, %4, %p1, %p2) /* ty=Tensor[(2, 2), float32] */
  } /* ty=fn (Tensor[(2, 2), float32], (), ()) -> Tensor[(2, 2), float32] */;
  let %x_3 /* ty=(Tensor[(2, 2), float32], (), ()) */ = (%x, -114514i64 /* ty=() */, -114514i64 /* ty=() */) /* ty=(Tensor[(2, 2), float32], (), ()) */;
  let %x_4 /* ty=float32 */ = raf.op.vm.alloc_storage(-114514i64, -114514i64 /* ty=int64 */, -114514i64 /* ty=int32 */, -114514i64 /* ty=int32 */, -114514i64 /* ty=int64 */, -114514i64);
  let %x_5 /* ty=Tensor[(2, 2), float32] */ = raf.op.vm.alloc_tensor(%x_4, -114514i64 /* ty=Tensor[(2, 2), int64] */, -114514i64 /* ty=int64 */, -114514i64 /* ty=Tensor[(2, 2), int64] */, -114514i64);
  let %x_6 /* ty=(Tensor[(2, 2), float32],) */ = (%x_5,) /* ty=(Tensor[(2, 2), float32],) */;
  let %x_7 /* ty=() */ = raf.op.vm.invoke_op(%x_2, %x_3, %x_6) /* ty=() */;
  let %x2 /* ty=Tensor[(2, 2), float32] */ = %x_5;
  %x2
}

and the result of 'asap' is :

def @main(%x: Tensor[(2, 2), float32] /* ty=Tensor[(2, 2), float32] */) -> Tensor[(2, 2), float32] {
  let %x_2 /* ty=fn (Tensor[(2, 2), float32], (), ()) -> Tensor[(2, 2), float32] */ = fn (%p0: Tensor[(2, 2), float32] /* ty=Tensor[(2, 2), float32] */, %p1: () /* ty=() */, %p2: () /* ty=() */, Primitive=1, Dialect="tvm") -> Tensor[(2, 2), float32] {
    %0 = raf.op.tvm.atan(%p0) /* ty=Tensor[(2, 2), float32] */;
    %1 = raf.op.tvm.atan(%0) /* ty=Tensor[(2, 2), float32] */;
    %2 = raf.op.tvm.atan(%0) /* ty=Tensor[(2, 2), float32] */;
    %3 = raf.op.tvm.atan(%1) /* ty=Tensor[(2, 2), float32] */;
    %4 = raf.op.tvm.atan(%2) /* ty=Tensor[(2, 2), float32] */;
    raf.op.tvm.add(%3, %4, %p1, %p2) /* ty=Tensor[(2, 2), float32] */
  } /* ty=fn (Tensor[(2, 2), float32], (), ()) -> Tensor[(2, 2), float32] */;
  let %x_3 /* ty=(Tensor[(2, 2), float32], (), ()) */ = (%x, -114514i64 /* ty=() */, -114514i64 /* ty=() */) /* ty=(Tensor[(2, 2), float32], (), ()) */;
  let %x_4 /* ty=float32 */ = raf.op.vm.alloc_storage(-114514i64, -114514i64 /* ty=int64 */, -114514i64 /* ty=int32 */, -114514i64 /* ty=int32 */, -114514i64 /* ty=int64 */, -114514i64);
  let %x_5 /* ty=Tensor[(2, 2), float32] */ = raf.op.vm.alloc_tensor(%x_4, -114514i64 /* ty=Tensor[(2, 2), int64] */, -114514i64 /* ty=int64 */, -114514i64 /* ty=Tensor[(2, 2), int64] */, -114514i64);
  let %x_6 /* ty=(Tensor[(2, 2), float32],) */ = (%x_5,) /* ty=(Tensor[(2, 2), float32],) */;
  let %x_7 /* ty=() */ = raf.op.vm.invoke_op(%x_2, %x_3, %x_6) /* ty=() */;
  let %x2 /* ty=Tensor[(2, 2), float32] */ = %x_5;
  %x2
}

and my code is :

device="cuda"
class Model(raf.Model):

    def build(self):
        pass

    @raf.model.trace
    def forward(self, x):
        p_0 = raf.atan(x)  # atan 0

        p_1 = raf.atan(p_0)  # atan 1
        p_1 = raf.atan(p_1)  # atan 3

        p_2 = raf.atan(p_0)  # atan 2
        p_2 = raf.atan(p_2)  # atan 4
        return raf.add(p_1, p_2)


model = Model()
input_shape = [2, 2]
x, _ = randn(input_shape)
mod = model._internal(x).mod

compiler = VMCompiler()
with tvm.transform.PassContext(opt_level=3):
    mod, _ = compiler.optimize(mod, device)
print(mod)
mod = InferType()(mod)
trace = [(name, mem.value) for name, mem in EstimateMemory(mod, Device(device), True)]
print(max(trace, key=lambda x: x[1])[1])


mod = model._internal(x).mod
#print(mod)
#with raf.ir.PassContext(opt_level=2, config={"raf.stream_schedule.policy": "wavefront"}):
    #mod = RAFSequential([ToGraphNormalForm(), WavefrontStreamSchedule()])(mod)
mod = RAFSequential([ToGraphNormalForm(), ASAPStreamSchedule()])(mod)
#print(mod)

compiler = VMCompiler()
with tvm.transform.PassContext(opt_level=3):
    mod, _ = compiler.optimize(mod, device)
print(mod)
mod = InferType()(mod)
trace = [(name, mem.value) for name, mem in EstimateMemory(mod, Device(device), True)]
print(max(trace, key=lambda x: x[1])[1])

it seems that those two 'mod' are the same?

[yaoyaoding]

Hi @ikaroso,

Try replaceing

with tvm.transform.PassContext(opt_level=3):

with something like

with raf.ir.PassContext(opt_level=2, config={"raf.stream_schedule.policy": "asap"}):

to use multi-stream schedule. The compiler.optimize would run passes that remove the schedule operators (I guess in dead code elimination pass) and reschedule using the sequential schedule.

Usually, we directly call these pass like
mod = RAFSequential([ToGraphNormalForm(), ASAPStreamSchedule()])(mod) for only debug purpose.

[ikaroso]

Hello @yaoyaoding ,
Yeah, seems that it worked, but I am not sure whether it is correct.
I get the result below.

seq: 1.52587890625e-05
multi: 3.0517578125e-05

and my code is below. As can be seen, I remove compiler.optimize for mod1, and follow your advice to use asap in mod2. Is it correct to use multi-stream schedule through:

compiler = VMCompiler()
with raf.ir.PassContext(opt_level=2, config={"raf.stream_schedule.policy": "asap"}):
    mod2, _ = compiler.optimize(mod2, device)

If we assume the peak of multi is 4 unit, then the peak of seq should be 3 unit, which is 2.288e-05. While the result of seq is 1.52587890625e-05, which is 2 unit. I guess RAF uses a different peak calculation style?

model = Model()
input_shape = [2, 2]
x, _ = randn(input_shape)
mod1 = model._internal(x).mod
mod1 = InferType()(mod)
trace = [(name, mem.value) for name, mem in EstimateMemory(mod1, Device(device), True)]
peak_seq=max(trace, key=lambda x: x[1])[1]

model=Model()
mod2 = model._internal(x).mod
compiler = VMCompiler()
with raf.ir.PassContext(opt_level=2, config={"raf.stream_schedule.policy": "asap"}):
    mod2, _ = compiler.optimize(mod2, device)
mod2 = InferType()(mod)
trace = [(name, mem.value) for name, mem in EstimateMemory(mod2, Device(device), True)]
peak_multi=max(trace, key=lambda x: x[1])[1]

print("seq: {}\nmulti: {}".format(peak_seq,peak_multi))

And I have some extra questions.

1. Does compiler.optimize do something to increase the peak of a mod?
   Since the result of :

mod1 = model._internal(x).mod
mod1 = InferType()(mod)
trace = [(name, mem.value) for name, mem in EstimateMemory(mod1, Device(device), True)]
peak_seq=max(trace, key=lambda x: x[1])[1]

is 1.52587890625e-05, while the result of :

mod1 = model._internal(x).mod
compiler = VMCompiler()
with tvm.transform.PassContext(opt_level=3):
    mod1, _ = compiler.optimize(mod1, device)
print(mod1)
mod1 = InferType()(mod1)
trace = [(name, mem.value) for name, mem in EstimateMemory(mod1, Device(device), True)]
peak_seq=max(trace, key=lambda x: x[1])[1]

is 3.0517578125e-05, which is the same peak after applying ASAP.

2. You said that mod = RAFSequential([ToGraphNormalForm(), ASAPStreamSchedule()])(mod) is usually for debug use. And I guess it will not have a actual effect? Since the result of :

mod = model._internal(x).mod
mod = RAFSequential([ToGraphNormalForm(), ASAPStreamSchedule()])(mod)
mod = InferType()(mod)
trace = [(name, mem.value) for name, mem in EstimateMemory(mod, Device(device), True)]
print(max(trace, key=lambda x: x[1])[1])

is 1.52587890625e-05, which is same as a mod without any optimization.

[yaoyaoding]

Hi @ikaroso,

1. "Is it correct to use multi-stream schedule",
   yes.
2. "I guess RAF uses a different peak calculation style."
   Better to read the code of estimate_memory.cc. I briefly read the code and found current implementation is to trace the memory alloc op and free op and record the memory consumption.
3. "Does compiler.optimize do something to increase the peak of a mod?"
   I do not think so.
4. "And I guess it will not have an actual effect?"
   It is strange that you can get a non-zero memory number if you directly run the snippet of code as the two passes do not inject the alloc memory and free memory op to the mod.

Sorry I am busy on my own project recent days so I can not spare time to try and run the demo code (I need to setup and build raf for that). The number you get 1.52587890625e-05 is close to 16384 (1.5258789*10^-5 * 1024 * 1024 * 1024 Bytes) which is 2^14. When we allocate storage in raf, the memory pool will round it up to some size like 2^k. Thus it makes no sense to expect 3/4 or some other ratio in our ideal analysis. My suggestion is, print the mod after optimization (there should be memory allocation/free operator in it), and read the code of memory estimater.

Thanks.

[ikaroso]

Hello @yaoyaoding ,

Thanks for your reply and I really respect it. I will follow your advice to see what is going on in the example code.
Don't be sorry since you have helped me a lot and I really appreciate it.

To make sure that I understand correctly. I should compare the difference between:

with raf.ir.PassContext(opt_level=2, config={"raf.stream_schedule.policy": "asap"}):
    mod, _ = compiler.optimize(mod, device)

and

with tvm.transform.PassContext(opt_level=3):
    mod, _ = compiler.optimize(mod, device)

Thank you.

[yaoyaoding]

Hi @ikaroso,

Compare

with raf.ir.PassContext(opt_level=2, config={"raf.stream_schedule.policy": "asap"}):
    mod, _ = compiler.optimize(mod, device)
print(mod)

with

with raf.ir.PassContext(opt_level=2, config={"raf.stream_schedule.policy": "sequential"}):
    mod, _ = compiler.optimize(mod, device)
print(mod)

[ikaroso]

Hi @yaoyaoding ,
I compare asap and sequential and they have the same peak memory footprint 3.0517578125e-05 MBs.

From the two mod, I guess they do allocate the same size of memory (from alloc_storage and alloc_tensor), though I don't totally understand it.
SEQ:

def @main(%x: Tensor[(2, 2), float32] /* ty=Tensor[(2, 2), float32] */) -> Tensor[(2, 2), float32] {
  let %x_2 /* ty=fn (Tensor[(2, 2), float32], (), ()) -> Tensor[(2, 2), float32] */ = fn (%p0: Tensor[(2, 2), float32] /* ty=Tensor[(2, 2), float32] */, %p1: () /* ty=() */, %p2: () /* ty=() */, Primitive=1, Dialect="tvm") -> Tensor[(2, 2), float32] {
    %0 = raf.op.tvm.atan(%p0) /* ty=Tensor[(2, 2), float32] */;
    %1 = raf.op.tvm.atan(%0) /* ty=Tensor[(2, 2), float32] */;
    %2 = raf.op.tvm.atan(%0) /* ty=Tensor[(2, 2), float32] */;
    %3 = raf.op.tvm.atan(%1) /* ty=Tensor[(2, 2), float32] */;
    %4 = raf.op.tvm.atan(%2) /* ty=Tensor[(2, 2), float32] */;
    raf.op.tvm.add(%3, %4, %p1, %p2) /* ty=Tensor[(2, 2), float32] */
  } /* ty=fn (Tensor[(2, 2), float32], (), ()) -> Tensor[(2, 2), float32] */;
  let %x_3 /* ty=(Tensor[(2, 2), float32], (), ()) */ = (%x, -114514i64 /* ty=() */, -114514i64 /* ty=() */) /* ty=(Tensor[(2, 2), float32], (), ()) */;
  let %x_4 /* ty=float32 */ = raf.op.vm.alloc_storage(-114514i64, -114514i64 /* ty=int64 */, -114514i64 /* ty=int32 */, -114514i64 /* ty=int32 */, -114514i64 /* ty=int64 */, -114514i64);
  let %x_5 /* ty=Tensor[(2, 2), float32] */ = raf.op.vm.alloc_tensor(%x_4, -114514i64 /* ty=Tensor[(2, 2), int64] */, -114514i64 /* ty=int64 */, -114514i64 /* ty=Tensor[(2, 2), int64] */, -114514i64);
  let %x_6 /* ty=(Tensor[(2, 2), float32],) */ = (%x_5,) /* ty=(Tensor[(2, 2), float32],) */;
  let %x_7 /* ty=() */ = raf.op.vm.invoke_op(%x_2, %x_3, %x_6) /* ty=() */;
  let %x2 /* ty=Tensor[(2, 2), float32] */ = %x_5;
  %x2
}

ASAP:

def @main(%x: Tensor[(2, 2), float32] /* ty=Tensor[(2, 2), float32] */) -> Tensor[(2, 2), float32] {
  let %x_0 /* ty=() */ = raf.op.set_stream(-114514i64 /* ty=int64 */, -114514i64 /* ty=int64 */) /* ty=() */;
  let %x_4 /* ty=fn (Tensor[(2, 2), float32], (), ()) -> Tensor[(2, 2), float32] */ = fn (%p0: Tensor[(2, 2), float32] /* ty=Tensor[(2, 2), float32] */, %p1: () /* ty=() */, %p2: () /* ty=() */, Primitive=1, Dialect="tvm") -> Tensor[(2, 2), float32] {
    %0 = raf.op.tvm.atan(%p0) /* ty=Tensor[(2, 2), float32] */;
    %1 = raf.op.tvm.atan(%0) /* ty=Tensor[(2, 2), float32] */;
    %2 = raf.op.tvm.atan(%0) /* ty=Tensor[(2, 2), float32] */;
    %3 = raf.op.tvm.atan(%1) /* ty=Tensor[(2, 2), float32] */;
    %4 = raf.op.tvm.atan(%2) /* ty=Tensor[(2, 2), float32] */;
    raf.op.tvm.add(%3, %4, %p1, %p2) /* ty=Tensor[(2, 2), float32] */
  } /* ty=fn (Tensor[(2, 2), float32], (), ()) -> Tensor[(2, 2), float32] */;
  let %x_5 /* ty=(Tensor[(2, 2), float32], (), ()) */ = (%x, -114514i64 /* ty=() */, -114514i64 /* ty=() */) /* ty=(Tensor[(2, 2), float32], (), ()) */;
  let %x_6 /* ty=float32 */ = raf.op.vm.alloc_storage(-114514i64, -114514i64 /* ty=int64 */, -114514i64 /* ty=int32 */, -114514i64 /* ty=int32 */, -114514i64 /* ty=int64 */, -114514i64);
  let %x_7 /* ty=Tensor[(2, 2), float32] */ = raf.op.vm.alloc_tensor(%x_6, -114514i64 /* ty=Tensor[(2, 2), int64] */, -114514i64 /* ty=int64 */, -114514i64 /* ty=Tensor[(2, 2), int64] */, -114514i64);
  let %x_8 /* ty=(Tensor[(2, 2), float32],) */ = (%x_7,) /* ty=(Tensor[(2, 2), float32],) */;
  let %x_9 /* ty=() */ = raf.op.vm.invoke_op(%x_4, %x_5, %x_8) /* ty=() */;
  let %x_1 /* ty=Tensor[(2, 2), float32] */ = %x_7;
  %x_1
}

And the code is:

model = Model()
input_shape = [2, 2]
x, _ = randn(input_shape)
mod1 = model._internal(x).mod

compiler = VMCompiler()
with raf.ir.PassContext(opt_level=2, config={"raf.stream_schedule.policy": "sequential"}):
    mod1, _ = compiler.optimize(mod1, device)
print(mod1)
mod1 = InferType()(mod1)
trace = [(name, mem.value) for name, mem in EstimateMemory(mod1, Device(device), True)]
peak_seq=max(trace, key=lambda x: x[1])[1]

model=Model()
mod2 = model._internal(x).mod
compiler = VMCompiler()
with raf.ir.PassContext(opt_level=2, config={"raf.stream_schedule.policy": "asap"}):
    mod2, _ = compiler.optimize(mod2, device)
print(mod2)
mod2 = InferType()(mod2)
trace = [(name, mem.value) for name, mem in EstimateMemory(mod2, Device(device), True)]
peak_multi=max(trace, key=lambda x: x[1])[1]

print("seq: {}\nmulti: {}".format(peak_seq,peak_multi))