Skip to content

[Relay] Fixed VNNI batch matmul op strategy for meta schedule compatibility#10637

Merged
masahi merged 2 commits intoapache:mainfrom
masahi:vnni-bmm-metasch
Mar 16, 2022
Merged

[Relay] Fixed VNNI batch matmul op strategy for meta schedule compatibility#10637
masahi merged 2 commits intoapache:mainfrom
masahi:vnni-bmm-metasch

Conversation

@masahi
Copy link
Member

@masahi masahi commented Mar 16, 2022

Currently, meta schedule task extraction for VNNI batch_matmul fails because

  • When relay.backend.use_meta_schedule is True, is_auto_scheduler_enabled() is also True
  • When is_auto_scheduler_enabled() is True, this branch
    if is_dynamic(out_type) or is_auto_scheduler_enabled():
    strategy.add_implementation(
    wrap_compute_batch_matmul(
    topi.nn.batch_matmul, need_auto_scheduler_layout=True, need_out_dtype=True
    ),
    wrap_topi_schedule(topi.generic.nn.schedule_batch_matmul),
    name="batch_matmul.generic",
    plevel=10,
    )
    in op strategy is always taken, while I want to choose the x86-specific implementation of batch_matmul, topi.x86.batch_matmul in the else branch.

This PR fixes this problem by doing VNNI dispatch at the op strategy level before we enter is_auto_scheduler_enabled() == True branch, rather than inside topi.x86.batch_matmul.

cc @junrushao1994 This is my workaround for the issue we discussed last week.

@github-actions github-actions bot requested a review from junrushao March 16, 2022 05:43
Copy link
Contributor

@csullivan csullivan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@masahi masahi merged commit b3af0ef into apache:main Mar 16, 2022
pfk-beta pushed a commit to pfk-beta/tvm that referenced this pull request Apr 11, 2022
…bility (apache#10637)

* Fixed VNNI batch matmul op strategy for meta schedule compatibility

* update bert task extraction test
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants