Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[MetaSchedule] Fix a multilevel tiling error on dynamic relax workload #17465

Merged

Conversation

wrongtest-intellif
Copy link
Contributor

Below script would either segfault or report primitive error in MetaScheduleTuneIRMod with dynamic dim.

import tvm
from tvm import relax
from tvm.relax.frontend import nn
from tvm.relax.transform import *

class NNModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc2 = nn.Linear(784, 128)

    def forward(self, x):
        x = self.fc2(x)
        return x

mod, params = NNModule().export_tvm({"forward": {"x": nn.spec.Tensor(("n", 784), "float32")}})
mod.show()

# with tvm.target.Target("cuda -max_threads_per_block=4096 -max_shared_memory_per_block=8192 -arch=sm_70"):
with tvm.target.Target("llvm -num-cores 32"):
    mod = LegalizeOps()(mod)
    mod = AnnotateTIROpPattern()(mod)
    mod = FoldConstant()(mod)
    mod = FuseOps()(mod)
    mod = FuseTIR()(mod)
    mod = MetaScheduleTuneIRMod({}, "myworkingdir", 3)(mod)
    application_pass = MetaScheduleApplyDatabase( "myworkingdir")
    mod = application_pass(mod)
    
print(mod)
m = relax.build(mod, target="llvm", pipeline="default_build")

The change fix two related issues:

  1. GetLoopIntExtent() would return nullptr on non-const extent, thus the return value should check the nullity.

  2. tir::SamplePerfectTile() return -1 for dynamic factor part. While the Split primitive require None rv to denotes the dynamic dim inference. Thus it should handle this semantic gap in schedule interfaces.

@wrongtest-intellif wrongtest-intellif force-pushed the fix_ms_error_on_dynamic_workload branch from 06f5461 to 071f88e Compare October 15, 2024 11:05
@wrongtest-intellif wrongtest-intellif force-pushed the fix_ms_error_on_dynamic_workload branch from 071f88e to 8feed4d Compare October 15, 2024 11:16
@Hzfengsy Hzfengsy merged commit 58a43c8 into apache:main Oct 16, 2024
19 checks passed
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