Migrating Target Attributes to IRModule

[Mousius]

No description provided.

[leandron]

cc @jroesch @areusch

[junrushao]

Overall it looks good to me to migrate some attributes from Target to IRModule, which is more relevant. CC @areusch @zxybazh would love to hear from you guys too

[areusch]

@jroesch @mbs-octoml @tkonolige would be good to get your thoughts.

overall I agree we need to move these attributes out of Target. I'm not sure they should go on the IRModule--rather my initial thoughts on this was that we should just create a separate RuntimeConfig option and pass that to tvm.relay.build. I'm open to being swayed. I do think the attributes which should be moved are any attribute which doesn't influence autotvm. It looks like you've identified some of the right ones in your examples.

[tkonolige]

I think this is a good and necessary change. Thanks!

I agree with Andrew that it would be better to keep these parameters as a separate object. I think this makes a more consistent interface. We are all ready passing targets separately from the module, so it makes more logical sense to keep "RuntimeConfig" (or whatever you want to call it) separate. On the other hand, I'm also not opposed to moving target and "RuntimeConfig" into the IRModule. Passing around all these separate but related objects can be a pain and I think we are moving towards bundling them in the IRModule.

[mbs-octoml]

+1 overall. We've been thinking along the same lines for replacing the explicit TargetMap (and it's convention for designating the 'default' target) and the 'host target' + 'device target' convention with a single datastructure describing all the targets which are significant for compilation and which are to be used as defaults for the different situations (for shape computation, for residual loops after offloading to prims, for generic relay computation). I'm trying to write that up as an RFC and will ping y'all asap.

I'm 80% on the just-put-it-in-the-IRModule camp since that is the one datastructure we are already threading all the way through. But we also make very heavy use of implicit global stacks (eg pass contexts), as well as pass-everywhere arguments (eg existing targets, target pairs and target maps). Perhaps that's the only contentious thing we need to resolve here?

[areusch]

@mbs-octoml my only worry about putting things in an IRModule is that it's not a very user-friendly thing to do, and these are all user-facing options. for that reason i think the separate object makes sense and we can migrate the config into IRModule if we wish to place it there internally within the compiler.

[Mousius]

Thanks for the feedback @junrushao1994, @areusch, @tkonolige and @mbs-octoml, this is one of those design decisions where I was 50/50 on whether to keep relay.build minimal or to leverage it's existing capability as a factory function. I've amended the design to the alternative which hides most of the IRModule details behind relay.build which I agree is a better interface design:

tvm.relay.build(
    ir_module,
    target,
    target_host=target,
    executor=Executor({
      "kind": "aot",
      "unpacked-api": False
    }),
    runtime=Runtime({
      "kind": "c",
      "link-params": False
    }),
    params=params,
    mod_name="the_ultimate_cat_spotter",
)

Hopefully this addresses everyone's concerns 😸

[mbs-octoml]

[sorry, finger flub, sent too early]
@areusch Got it. Yes, the 'configuration' needs to be fully under the user's control.

I think I'm suggesting we need to collect the executor, runtime and various targets under a single CompilerConfig structure and bind that into the IRModule. That means it's just a mod.GetAttr<CompilerConfig> away.

Then I think I'll slightly pivot my target/device RFC into a 'handling multiple targets' RFC, and use the ComplierConfig as the carrier for what would replace the existing TargetMap and various defaulting conventions.

Am I making sense?

[Mousius]

Hi @mbs-octoml, apologies, I'm not sure how grouping them together on the IRModule is better than having them all as separate attributes? I had a similar question about TargetDevice vs kTarget + kDevice, it seems like we have a generic attribute system which goes through the function registry and this will make it harder to do that?

[mbs-octoml]

@Mousius, just a gut feeling that we should try to group statically known and related things together as an ObjectRef/Node pair rather than rely on string attribute names and dynamic type checking. After fiddling a bit with multi-target handling I think I'll want a few extra fields for capturing the default targets for primitives and residual relay/shape computation, but that's just two more fields which similarly can go into the mix of attributes. We don't have to sweat it now, do what feels right to you.

[Mousius]

@mbs-octoml I think we'll have more detail in each attribute as part of the various configurations, as you'll have:

SetAttr<kRuntime>(Runtime(...));
SetAttr<kTarget>(Target(...));

And as such I think I'd prefer to leave it as-is.

@junrushao1994 / @areusch / @tkonolige are there further concerns before we merge this?

[electriclilies]

Overall this seems like a good change!
One question I have is what the relationship between PassContext and the IRModule attributes will be in the future. I think that ideally we'd have all the information about transformations that will be done to the IRModule in one place.

Right now, the PassContext has some module-level configuration options, but it also has options that are extracted from the PassContext within passes and modifies the behavior of those passes. I think that we should at least move the module-level configuration options into the IRModule attributes, but I'm not sure about what we should do with the other options.

Another question is whether we want to distinguish between attributes / config options that are pass-invariant and attributes that passes can change. This would make it easier to make sure that the pass-invariant ones are being propagated properly.
(We don't need to resolve these questions before the PR goes in, just mentioning them while I'm thinking about it!)

[mbs-octoml]

LGTM

[leandron]

LGTM

[Mousius]

Hi @electriclilies,

Following up on your questions 😸 Overall I think we're thinking the same things in terms of direction of travel!

Overall this seems like a good change! One question I have is what the relationship between PassContext and the IRModule attributes will be in the future. I think that ideally we'd have all the information about transformations that will be done to the IRModule in one place.

I completely agree, I think the direction of travel should be towards removing it or integrating it with the IRModule so you have a representation of the compilation at any given time.

Right now, the PassContext has some module-level configuration options, but it also has options that are extracted from the PassContext within passes and modifies the behavior of those passes. I think that we should at least move the module-level configuration options into the IRModule attributes, but I'm not sure about what we should do with the other options.

This would be an RFC of interest to me, moving them into attributes and only passing the IRModule seems sensible - the Pass level configuration we should probably check to see if they could be more generic options or whether we can recompute the transforms efficiently enough they don't need to directly mutate the IRModule?

Another question is whether we want to distinguish between attributes / config options that are pass-invariant and attributes that passes can change. This would make it easier to make sure that the pass-invariant ones are being propagated properly. (We don't need to resolve these questions before the PR goes in, just mentioning them while I'm thinking about it!)

Agree, we should probably try to freeze attributes which aren't meant to change (hopefully most of them) and that would hopefully make things simpler for everyone. Unsure if TVM has support for this in attrs right now?

[areusch]

[sorry, finger flub, sent too early] @areusch Got it. Yes, the 'configuration' needs to be fully under the user's control.

I think I'm suggesting we need to collect the executor, runtime and various targets under a single CompilerConfig structure and bind that into the IRModule. That means it's just a mod.GetAttr<CompilerConfig> away.

Then I think I'll slightly pivot my target/device RFC into a 'handling multiple targets' RFC, and use the ComplierConfig as the carrier for what would replace the existing TargetMap and various defaulting conventions.

Am I making sense?

@mbs-octoml yeah this does make sense to me. sorry for the delay in reply here.