[Relay][Memory][VM]

[jroesch]

Explicitly manifest memory and tensor allocations in Relay. This is to enable future optimizations and transformations.

v0.0.3
fn (%x: Tensor[(10,), float32], %y: Tensor[(10,), float32]) -> Tensor[(10,), float32] {
%1 = fn (%p0: Tensor[(10,), float32], %p1: Tensor[(10,), float32], Primitive=1) -> Tensor[(10,), float32] {
%0 = add(%p0, %p0) /* ty=Tensor[(10,), float32] */;
subtract(%0, %p1) /* ty=Tensor[(10,), float32] */
};
let %x1: Tensor[(10,), float32] = %1(%x, %y) /* ty=Tensor[(10,), float32] */;
%x1
}

v0.0.3
fn (%x: Tensor[(10,), float32], %y: Tensor[(10,), float32]) -> Tensor[(10,), float32] {
%0 = annotation.alloc_storage(320);
let %x1: Tensor[(10,), float32] = let %output: Tensor[(10,), float32] = annotation.alloc_tensor(%0, meta[relay.Constant][0], meta[relay.attrs.AllocTensorAttrs][0]);
%2 = fn (%p0: Tensor[(10,), float32], %p1: Tensor[(10,), float32], Primitive=1) -> Tensor[(10,), float32] {
%1 = add(%p0, %p0) /* ty=Tensor[(10,), float32] */;
subtract(%1, %p1) /* ty=Tensor[(10,), float32] */
};
%3 = (%x, %y);
%4 = (%output,);
let %v = annotation.invoke_tvm_op(%2, %3, %4);
%output;
%x1
}

I have an initial version of the PR nearly ready for review. The general goal of this PR is to perform all calling convention regularization which has historically been done in the runtime in the . IR instead, so we can further optimize. The final remaining piece is the transform of dynamic shape handling code into the IR as well.

[ajtulloch]

@jroesch what kind of optimizations will this enable?

[wweic]

@ajtulloch This PR models memory allocation/deallocation explicitly in relay expression, it will open the opportunity to write passes to analyze memory operation and try allocation coalescing etc. It will also simplify vm bytecode generation.

[icemelon]

Any updates? @jroesch

[zhiics]

@jroesch I just started the review and had a quick glance on it. I will do a more careful review by this week. Jared, could you please provide a little more description for this PR, i.e. the mechanism of the memory optimization and why we introduce these annotations. That would help the review as well. Thank you.

[zhiics]

It looks good. I have a few questions:

• How is live memory released/killed currently?
• Can we reuse the allocated memory as the current memory planning pass does? Or we just get memory from the pool directly?

[zhiics]

I am half way done. I will spend more time tomorrow to look at the design.

[icemelon]

half-way through. will continue after the flight

[zhiics]

@jroesch I checked out the PR and debugged it. The compilation error is caused by converting tensors from ndarray to raw shape. On i386, I believe it is int32 but x86 it int64.

[tqchen]

some high level comments

[tqchen]

The mutational def was a bit dangerous, can we explicitly use the define_ref_methods and add the mutationable ref methods directly? This will also reduce one macro

[jroesch]

You can not do that as the definitions are incompatible overloads

[tqchen]

they are compatible, because operator->() and operator->()const are two different functions

[jroesch]

I got an error about return type overload not working when I tried to do it.

[wweic]

A few minor comments.

[zhiics]

LGTM