[Transform] RewriteDataflowReshape to op and VMBuiltinLower handling

[MasterJH5574]

As discussed in #407 (comment), we update the behavior of pass RewriteDataflowReshape.

In short, prior to this PR, the pass transforms calls of reshape PrimFunc in dataflow blocks to direct calls of runtime packed func “vm.builtin.reshape.” The consequence of this behavior is that the memory planning pass has to check the reshape op by string comparison of ExternFunc.global_symbol, which is not ideal.

Therefore, this PR changes the RewriteDataflowReshape’s behavior, transforming calls of reshape PrimFunc to our high-level reshape op “relax.reshape,” and let the VMBuiltinLower pass to lowers the op to calls of “vm.builtin.reshape.”

[masahi]

Unrelated, but the following line that assumes the callee of call_tir is GlobalVar is incorrect, since it can also be ExternFunc after RunCodegen pass in the BYOC flow.

https://github.com/tlc-pack/relax/blob/relax/src/relax/transform/rewrite_dataflow_reshape.cc#L71

Can you add a check call->args[0]->IsInstance<ExternFuncNode>() at L68?

[MasterJH5574]

Unrelated, but the following line that assumes the callee of call_tir is GlobalVar is incorrect, since it can also be ExternFunc after RunCodegen pass in the BYOC flow.

https://github.com/tlc-pack/relax/blob/relax/src/relax/transform/rewrite_dataflow_reshape.cc#L71

Can you add a check call->args[0]->IsInstance<ExternFuncNode>() at L68?

@masahi Thanks! Sorry it was my negligence. Will add in the next commit.

[MasterJH5574]

@masahi Updated :-)

[slyubomirsky]

We should probably document when in compilation this rewrite (and any other passes whose order is important) should happen, since there will be dependencies on it (like #407).

[MasterJH5574]

We should probably document when in compilation this rewrite (and any other passes whose order is important) should happen, since there will be dependencies on it (like #407).

@slyubomirsky Thanks! Added one note here

relax/include/tvm/relax/transform.h

Lines 113 to 115 in 3c79b61

	* \note The pass is applied at the first stage of Relax VM build, before
	* rewriting call_tir, as this pass requires dataflow information.
	*/

relax/python/tvm/relax/transform/transform.py

Lines 117 to 122 in 3c79b61

	Notes
	-----
	The pass is applied at the first stage of Relax VM build, before
	rewriting call_tir, as this pass requires dataflow information.
	"""
	return _ffi_api.RewriteDataflowReshape() # type: ignore

[MasterJH5574]

Unrelated, but the following line that assumes the callee of call_tir is GlobalVar is incorrect, since it can also be ExternFunc after RunCodegen pass in the BYOC flow.

https://github.com/tlc-pack/relax/blob/relax/src/relax/transform/rewrite_dataflow_reshape.cc#L71

Can you add a check call->args[0]->IsInstance<ExternFuncNode>() at L68?

@masahi Hi Masa. Sorry that I didn’t think too much on your comment before. One quick question: why we use call_tir to call an ExternFunc, while not directly use the ExternFunc as CallNode::op? call_tir is supposed to call TIR PrimFunc IMO. And if we want to call some ExternFunc, we can directly call that ExternFunc via Call(op=extern_func, args=my_args, ...).

[masahi]

I don't know, and I wondered too. cc @sunggg who wrote the code below.

relax/src/relax/transform/run_codegen.cc

Line 84 in 83adb87

static const Op& call_op = Op::Get("relax.call_tir");

[MasterJH5574]

@masahi I just got a quick answer.

Here we use call_tir calling ExternFunc for call_dps_packed. This means the interface of the ExternFunc is in DPS style, and thus when we use call_tir to call it, we write a normal call instead of a DPS call, and don’t need to allocate the result memory in ahead.

We may need a dedicated op for call_dps_packed (#430), as using call_tir both normal call_tir and call_dps_packed is too much confusing. Nevertheless, we can leave it as it is now. Sorry for the bothering.