[Relax] Allow DeadCodeElimination within ApplyPassToFunction

[Lunderberg]

The tvm.ir.transform.ApplyPassToFunction allows a transform to be applied selectively to some portions of a IRModule, without applying to the entire IRModule. For example, to apply an optimization pass (e.g. relax.transform.ExpandMatmulOfSum) or an interface-altering pass (e.g. relax.transform.BundleModelParams) to specific functions. It does so by generating an intermediate IRModule containing only the functions specified, applying the transform to that intermediate, then merging the results.

When using ApplyPassToFunction to apply DeadCodeElimination, or a pipeline containing DeadCodeElimination, this intermediate IRModule may contain calls to GlobalVar instances that are not within the intermediate IRModule. Prior to this commit, this resulted in an error being thrown when collecting the call graph. This commit updates DeadCodeElimination to instead handle incomplete call-graph collection.

[slyubomirsky]

Hmmm. This approach seems fine for the specific case of DCE, but I wonder if the design of ApplyPassToFunction needs further thought, since the fact that it calls passes on a potentially malformed module could be an issue for other passes as well. Maybe the pass infrastructure/interfaces could try to accommodate the case of not processing the entire module? A shorter-term change might be ensuring that all module-level passes check whether a function is defined when looking up global vars (we could make it a helper function in the ExprMutator base, even, since most of the examples I found were using the ExprMutator's mod_ field for this purpose). Cursory search for other examples of directly calling Lookup on mod_: BindSymbolicVars, FuseTIR, RewriteCudaGraph, MergeCompositeFunctions, LambdaLift, DetectRecursion (analysis), with a few others that are related to training or autotuning.

We should be careful about introducing another possible pitfall to pass design (in this case having to accommodate a violated invariant).

[Lunderberg]

Agreed, it would be good to avoid having too many special cases. The intent was to allow optimization passes, or API-altering passes, to be part of an end-to-end pipeline while only impacting a few functions within the module. (e.g. Applying LazyGetInput to the result of LiftTransformParams, but not to any of the inference functions.)

If it requires more special cases in general, I think the ApplyPassToFunction could be improved by replacing all non-matching functions with relax::ExternFunc, rather than filtering them out altogether. That way, the inner transform could operate on a well-formed IRModule, but wouldn't be able to mutate other functions.

[slyubomirsky]

Might as well approve these changes since ApplyPassToFunction is already there and they address it. I don't think a huge redesign would necessarily be in order, but we should give some thought to the intended design for it.