Relax shape refinement implements

Development unity
#1

In Relax design notes, one of the key constructs for symbolic shape computations is match_shape. According to the design notes and early demos, match_shape was originally defined as a Relax expression R.match_shape. (I couldn’t find the right version of TVM to run the demo.)

Upon inspecting the latest TVM release candidate that support Relax source code, I noticed that match_shape is now defined as a Relax VM runtime function, R.call_packed("vm.builtin.match_shape"). This function is automatically matched through the transformation applied by the relax.transform.VMShapeLower pass. A relevant test case demonstrates this process.

Can someone please help to explain why this transition from an expression to a VM runtime function took place?

#2

Just to compare, older version Relax that has MatchShape as expression:

#3

we need to have low-level runtime support for such expressions like match shape, as well as symbolic shape computations, and VMShapeLower transform things into set of PackedFunc over the heap values, as well as runtume functions that execute these.

Note that we still have match_cast as the construct(which can be used for structure info refinement). The builtin runtime match_shape function is a mechanism to implement runtime support for the language construct in vm

This is one possible way to implement this feature, with a vm that is primarily centered around the PackedFunc invocation and works well for us so far. The testcase also likely explains how the lowering should take in place

#4

Thanks for reply! Indeed R.MatchCast is available for structure info refinement.