[TRANSFORM] Fix virtual device annotation issue with BYOC subgraphs

[srkreddy1238]

Heterogeneous module partitioned by BYOC has functions nodes without any VirtualDevice definition (having FullyUnconstrained device). Ignoring the device here causes expr_virtual_devices_ being empty when PopVirtualDevice is called assuming above PushVirtualDevice is succeeded. PushVirtualDevice and PopVirtualDevice occurs as pairs across function body, hence it's better to insert the The Virtual Device for Uncontrained and Pop it subsequently.

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• No users to tag found in teams: transform _{See #10317 for details}
• Built docs for commit e24e5ea can be found here.

_{Generated by tvm-bot}

[csullivan]

Hi @srkreddy1238, I didn't quite understand the conclusion you make,

PushVirtualDevice and PopVirtualDevice occurs as pairs across function body, hence it's better to insert the The Virtual Device for Uncontrained and Pop it subsequently

Naively I can understand why we skip pushing an unconstrained virtual device: If something is fully unconstrained then it's virtual device shouldn't influence the function body visitation. What issue is this causing in BYOC?

Is there a unit test you can consider contributing that can help illustrate the issue?

[srkreddy1238]

PushVirtualDevice and PopVirtualDevice maintains the current VirtualDevice information on stack expr_virtual_devices_ . These calls are invoked across a block parsing (FunctionNodes, CallNode ..etc). The stack (expr_virtual_devices_) grows as we nest into the functions.

def @tvmgen_default_clml_main_17(%clml_17_i0: Tensor[(1, 2048), float32] /* ty=Tensor[(1, 2048), float32] */, Inline=1, Compiler="clml", global_symbol="tvmgen_default_clml_main_17", Primitive=1) -> Tensor[(1, 1000), float32] {
  %245 = fn (%FunctionVar_0_01: Tensor[(1, 2048), float32] /* ty=Tensor[(1, 2048), float32] */, PartitionedFromPattern="nn.dense_nn.bias_add_", Composite="clml.dense") -> Tensor[(1, 1000), float32] {
    %244 = nn.dense(%FunctionVar_0_01, meta[relay.Constant][257] /* ty=Tensor[(1000, 2048), float32] */, units=1000) /* ty=Tensor[(1, 1000), float32] */;
    add(%244, meta[relay.Constant][258] /* ty=Tensor[(1, 1000), float32] */) /* ty=Tensor[(1, 1000), float32] */
  } /* ty=fn (Tensor[(1, 2048), float32]) -> Tensor[(1, 1000), float32] */;
  %245(%clml_17_i0) /* ty=Tensor[(1, 1000), float32] */
}

In the above IR the composite function %245 = fn (%FunctionVar_0_01 will not have a VirtualDevice definition ( VirtualDevice(?)) and the PushVirtualDevice fails. Whereas at end of this function parsing the PopVirtualDevice removes the VirtualDevice corresponding to the outer scope function. From here all annotations fail due to empty stack.

I figured it out as part of large network. I am working on to facilitate a test case for the same.

[srkreddy1238]

@mbs-octoml

[csullivan]

Thanks @srkreddy1238, LGTM

[srkreddy1238]

@junrushao can you take a look on this ?

[csullivan]

cc @echuraev @elvin-n does this change look good to you and pass internal texture tests?

[echuraev]

LGTM. Thanks

[elvin-n]

cc @echuraev @elvin-n does this change look good to you and pass internal texture

Our tests runs in CI and if we see green CI it means this change does not affect them and whole flow. I cannot answer if this change affects origin Mark's idea and if it does not brake anything in other places. Again, assuming that we have many tests for VirtualDevices and they passed, this change should be safe.

[srkreddy1238]

@junrushao can you conclude this PR ? There are some follow up on this.

[junrushao]

Sure! It’s been approved for a while already. Not sure why it’s not merged yet