Change target string to Target object in the TE compiler and interpreter

[electriclilies]

In this PR, I change target string to Target object in parts of the TE compiler and interpreter. This is follow-up work from #8802.

cc @junrushao1994 @Mousius @mbs-octoml @mikepapadim

[Mousius]

This is amazing @electriclilies, this is a much cleaner representation of the Targets, thanks for doing this 😸

[Mousius]

My suggestion got squashed because you're too fast 😸 I was suggesting moving the duplicated instantiation out:

Suggested change

	if (lowered_funcs.find(Target("ext_dev")) != lowered_funcs.end()) {
	lowered_funcs.Set(Target("ext_dev"), IRModule());
	Target ext_dev("ext_dev");
	if (lowered_funcs.find(ext_dev) != lowered_funcs.end()) {
	lowered_funcs.Set(ext_dev, IRModule());

[electriclilies]

Oh, my bad, I just didn't commit the file 😅

[Mousius]

Ahh, classic 😅

[junrushao]

🙏 Thank you for doing this! Really nice!

[Mousius]

Took a look at this locally, looks good to me 😸

[mbs-octoml]

LGTM

[Mousius]

Hmm, something weird is occurring here, unsure why those tests didn't run for me locally - but debugging test_any.py a bit I ran tvm::Dump(per_target_module_) and got (removed function body for brevity):

{
llvm -keys=cpu -link-params=0: IRModule({GlobalVar(intrp_fused_add): PrimFunc([placeholder, placeholder, T_add]) attrs={"from_legacy_te_schedule": (bool)1, "global_symbol": "intrp_fused_add", "tir.noalias": (bool)1} {}
}), 
llvm -keys=cpu -link-params=0: IRModule({GlobalVar(shape_func_add): PrimFunc([placeholder, placeholder, _broadcast_shape_func]) attrs={"from_legacy_te_schedule": (bool)1, "global_symbol": "shape_func_add", "tir.noalias": (bool)1} {}
})}

This should be gated in GetLoweredFunctions, with these lines (these are ran twice, once for normal cache and once for shape cache):

      if (!lowered_functions.count(target)) {
        lowered_functions.Set(target, IRModule(Map<GlobalVar, BaseFunc>({})));
      }

      lowered_functions[target]->Update(lowered_func->cached_func->funcs);

And only create a single IRModule per target, this indicates the Map is instead creating many keys even for the single Target it was passed in LowerTE (i.e. same Target object, I assume copied), unsure why this isn't breaking more.

This also means I was wrong about the Target("ext_dev") comparison, my apologies 😿. I believe this is due to ObjectHash used by Map only being able to do String or ObjectPtr comparisons?

@jroesch / @junrushao1994 any ideas? It would appear Map is behaving a little bit weirdly?

[mbs-octoml]

Are the hash and equality on Targets legit?
Also, possibly relevant, there's a bit of hackery to force shape functions to end up on the 'cpu' target which is hard coded, so it could be there's two 'llvm' targets in action that look the same but are distinct objects as far has hash/equality is concerned. I'd like to fix that by making device planning responsible for spelling out everything so there's no special cases left here.
Note ideally the code would be:
lowered_functions[target]->Add(...)
ie let the defaut ctor do it's thing if the key is not present.

[mbs-octoml]

(3oz of coffee later)
It will be the on-the-fly Target creation for the shape functions. By switching from String to Target we no longer implicitly identify distinct targets objects which have the same structure. To make this work we'll need to create all the Targets we need once and pass them around in the targets map. When lowering shape functions the target will need to be retrieved for the 'default cpu'. Again, eventually I'd like device planning to just spell this out, and for shape functions to have a well-defined place on which their appropriate Target annotation can be attached. As a stepping stone just passing in the 'default_cpu_target' or something would do?

[electriclilies]

@mbs-octoml @junrushao1994 @Mousius I implemented what Junru suggested-- just using std::unordered_map<Target, IRModule, TargetStrHash, TargetStrEqual> instead of tvm::Map<Target, IRModule> in a few places.
It looks like this will go green soon, can you take another look so we can get this merged? Thanks!

[junrushao]

It looks good to me. Only a nitpick

[junrushao]

Hey what about moving these methods temporarily to src/relay/backend/utils.h instead? Given these are only used in relay backend right now, I think it would be helpful to sort of prevent future developers to use them :-)

[electriclilies]

moved them!

[mbs-octoml]

nit: don't append a _ for local vars since the convention is it indicates a member var.

[electriclilies]

done!

[mbs-octoml]

LGTM. It will be a happy day when we nuke these maps.

[Mousius]

Minor nit, but otherwise excited to see this merged 😸

[Mousius]

Should probably remove this as it's not used in this file.

[jroesch]

We will send a follow up P that does this just for the sake of forward progress. Thanks!

[electriclilies]

Thanks all! I'll stick that change in a follow up PR @Mousius