[Generic] Forward MS and AS rewrites for generic schedules

[AndrewZhaoLuo]

Hitting generic strategies for ARM targets, which don't properly forward layout rewrite info.

[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.

• No users to tag found in teams: generic _{See #10317 for details}

_{Generated by tvm-bot}

[AndrewZhaoLuo]

cc @tkonolige

[tkonolige]

I don't think this is correct. Specifying need_meta_schedule_layout should be set from the specific target strategy. Like here: https://github.com/apache/tvm/blob/main/python/tvm/relay/op/strategy/arm_cpu.py#L581-L582

[AndrewZhaoLuo]

I don't think this is correct. Specifying need_meta_schedule_layout should be set from the specific target strategy. Like here: https://github.com/apache/tvm/blob/main/python/tvm/relay/op/strategy/arm_cpu.py#L581-L582

Is the a concern of correctness or style? There is no target-specific strategy in this case (batch_matmul on arm) so it hits the generic strategies.

[tkonolige]

Currently auto/metaschedule layout transform is enabled on a target-by-target basis. You're changing the default to be enabled for all targets that do not specify. In the past we've had some problems with layout transform not working on all targets, so I don't think this is a safe default.

[AndrewZhaoLuo]

Hmm so my bug though has us hit a generic strategy, but metaschedule still rewrites the layout during tuning.

During compilation step we get an error since the rewriting is not forwarded properly. Is layout rewriting also enabled on a per-target basis? Otherwise all targets which hit generic strategy will get same problem.

[junrushao]

For context: tuning-based layout rewriting on Relay is quite hacky and buggy because Relay is not designed to support this feature (even if it's necessary to deliver better performance). Relax natively supports this optimization and will be more likely to work decently.

[AndrewZhaoLuo]

@tkonolige would you be fine if I made copies of the generic strategies for my target, and properly forward rewrite info there?

We still have the problem with generic strategies getting rewritten, but my problem will be solved.

[AndrewZhaoLuo]

Ah so chatting with @tkonolige and others a bit more, this problem is a bit more complicated and maybe exposes some past inconsistency with tvm.target.Target. Anyway, here are my findings.

1. If we use generic strategy we should not allow LayoutRewrite (this is different from AlterOpLayout) pass used in MS. This is because LayoutRewrite really only makes sense for CPU platforms where non-managed cache means tensor layout should match access patterns from the get go. It is also buggy potentially so turning it on for only CPU can help manage risk area for bugs.

2. I was constructing my target object in TVM incorrectly. arm_cpu key should imply cpu key. Doing tvm.target.Target("...") does not do this properly (it implies arm_cpu but not cpu). Evidently I have been doing this wrong all my life. We should be doing tvm.target.arm_cpu("..."). There was a cpu strategy that forwards information correctly

Based on this here are future work:

1. This PR should be closed as if we hit generic strategy, we are not in CPU and we should not be doing LayoutRewrite
2. We should probably look into making tvm.target.Target make it so arm_cpu also implies cpu. I've been doing things this way all my life and no doubt have accidentally dispatched to generic strategies on accident.
3. Use of target through codebase might be a bit inconsistent. In some parts the hardware target is used (e.g. cpu, arm_cpu, gpu). In other places the compilation target is used (e.g. llvm, cuda). This is fine as long as usage is consistent between the two or you really care about the hardware/compilation target.

[AndrewZhaoLuo]

#13775 <-- is probably the best solution to my problem for now.