[Metaschedule] ApplyCustomRule only if present

[quic-sanirudh]

When a PrimFunc contains the schedule_rule attribute, all schedule rules are ignored except ApplyCustomRule(), regardless of whether ApplyCustomRule() itself is available or not.

This patch changes the code so that schedule_rule attribute is only considered when ApplyCustomRule() is one of the schedule_rules passed to post_order_apply and ignored otherwise in favor of the existing custom rules.

This was needed as some topi.nn computes contain the schedule_rule attribute, which causes them to be ignored by meta_schedule irrespective of whether a custom rule function is registered or not.

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @Hzfengsy, @elvin-n, @junrushao _{See #10317 for details}

_{Generated by tvm-bot}

[quic-sanirudh]

@junrushao @zxybazh Could you please take a look at this PR. I was writing a new schedule rule and realized that it wasn't being applied on some blocks just because they had the schedule_rule attribute by default in topi.

I wasn't sure of whether to check for the presence of just ApplyCustomRule(), or both ApplyCustomRule() and a registered schedule function for that attribute.

[zxybazh]

Thanks for contributing!

IIUC, you would like to apply non custom schedule rules to a block that has schedule_rule attribute. May I ask why is the attribute schedule_rule necessary if the custom schedule rule is not applied, i.e., can you remove the attribute instead?

On the other hand, if we set schedule_rule to None we can avoid certain block being used for schedule rule application, this PR would break the assumption here.

[quic-sanirudh]

Thanks for contributing!

IIUC, you would like to apply non custom schedule rules to a block that has schedule_rule attribute. May I ask why is the attribute schedule_rule necessary if the custom schedule rule is not applied, i.e., can you remove the attribute instead?

On the other hand, if we set schedule_rule to None we can avoid certain block being used for schedule rule application, this PR would break the assumption here.

Yes I'm trying to apply a non custom schedule rule to blocks that contain the schedule_rule attribute. Many generic topi implementations like conv2d and pooling ops contain the schedule_rule attirbute which is the compute picked up by most targets.

When tuning a relay IRModule with metaschedule, we would have to introduce an op strategy to point to a different compute for each of these ops, just to remove the attribute. The idea behind this PR was to try and avoid those changes to each op.

Also, there might be users new to metaschedule but familiar with generic TVM relay workflow and they might see that their ops are not tuned at all if their op strategy by default picks up this compute and they might not have any custom schedule functions registered. I ran into this as well, but luckily I was familiar with the meta schedule code enough to find the problem quickly.

As for setting schedule_rule to None, IIUC, it prevents from applying the custom rule and we still don't apply non custom schedule rules to those blocks as well since it has the schedule_rule attribute, which means the block doesn't get tuned at all. If that was the intention, we could filter out the blocks that have schedule_rule=None in the block collector above or we can check for that as well as part of the IsApplyCustomRule condition. Let me know if any of those changes might make sense. Since schedule_rule=None means no rules are applied, it could be independent of the ApplyCustomRule() and we could check for that in post_order_apply.cc instead.

I did not realize about schedule_rule=None condition earlier as there were no tests with that, so I could add a test for that as well to make it as an expected behavior if that's okay.

Alternative to this PR, we could remove those schedule_rule attributes from generic topi.nn ops and expect that the attribute should only be added to target specific op strategies. I'm not sure how that would affect existing targets now (potential performance degradation), as we don't know who relies on these attributes right now.

P.S. One last thing. I'm quite new to this entire code base, so please feel free to correct me anywhere and thanks for reviewing the change.

[zxybazh]

Hey thanks for sharing more details. And yes I agree a test for schedule rule as None would be helpful. From a design perspective, we would like to make sure custom schedule rules are applied without interference from other schedule rules. From user perspective, first this provides a default for most of the hardwares like x86 or cuda, and these schedule rules can be customized or replaced if you register a new schedule function into the tvm global registry in the same schedule rule name. Removing those topi default schedule rules wouldnot only make it harder for new users to get started but also impact many current tuning workflows.

IMHO, the use case to apply non-custom schedule rules to some topi nn function is actually more of a high-level usage because for the topi nn workloads that you pointed out, usual schedule rules would not be able to extract the best performance comparing to using the schedule rules we specified. On the other hand, these custom schedule rules are target specific, if you do have a good schedule rule set for certain workload, please feel free to send in new custom schedule rule for your hardware, it would benefit even more users :)

[quic-sanirudh]

@zxybazh Thanks for the explanation. I understand that this change breaks the assumption of schedule_rule and so it makes sense to close this PR without merging. I would like to ask a couple of last questions if that's okay?

1. Is there a way to modify the attributes returned by the generic topi compute?
   • I'm asking this because the compute is in general the same for most (if not all) targets. I couldn't find a way to just modify the attributes to remove schedule_rule attribute from the ComputeOp, which would force these targets to write the full compute function(or copy them from the generic implementation), just to remove the attribute.
2. Second question is related to adding these attributes to specific targets. If those custom schedule rules are beneficial for specific targets like x86 or cuda, wouldn't it be better if the attribute is added for those targets and keep the generic topi computes without any attributes.
   • To me, somehow it feels wrong to add attributes that benefit specific targets to the generic computes. Pleaser let me know your views on this.

Thanks again for the discussion, and I'll close this PR for now :)