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] MutateTileSize skip single-candidate SampleCategorical #14072

Merged
merged 1 commit into from
Feb 22, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
[MetaSchedule] MutateTileSize skip single-candidate SampleCategorical
When handling SampleCategorical instructions, our MutateTileSize mutator
mutates it by asserting that "there should be at least one candidate
other than the current decision". However, this argument fails when a
SampleCategorical instruction has only one candidate, since there is no
other candidate besides the unique one. On such cases, the mutator will
lead to exception.

For such cases, the expected behavior is that the mutator ignores such
SampleCategorical instructions. This PR updates the mutator in this way
and provides a regression test.
  • Loading branch information
MasterJH5574 committed Feb 21, 2023
commit d613d621a19c47269f32d6e19df7bd649b337eb7
7 changes: 7 additions & 0 deletions src/meta_schedule/mutator/mutate_tile_size.cc
Original file line number Diff line number Diff line change
@@ -128,6 +128,13 @@ void FindSampleVectorize(const Trace& trace, std::vector<Instruction>* inst,
if (inst->kind.same_as(inst_sample_categorical)) {
ICHECK_EQ(inst->outputs.size(), 1);
if (annotated.count(inst->outputs[0].get())) {
ICHECK_EQ(inst->attrs.size(), 2);
std::vector<double> probs =
support::AsVector<FloatImm, double>(Downcast<Array<FloatImm>>(inst->attrs[1]));
if (probs.size() == 1) {
// Skip mutating the sampling instructions who have only single candidate.
continue;
}
const auto* d = TVM_TYPE_AS(decision, IntImmNode);
instructions.push_back(inst);
decisions.push_back(d->value);
Original file line number Diff line number Diff line change
@@ -95,5 +95,18 @@ def test_mutate_tile_size_matmul():
assert len(results) > 15


def test_mutate_sample_categorical_single_candidate():
mutator = _make_mutator(
target=Target("llvm --num-cores=16"),
)
sch = Schedule(matmul, debug_mask="all")
sch.sample_categorical(candidates=[1], probs=[1.0], decision=0)

# The mutator finds the SampleCategorical has only one candidate, and thus skips it.
trace = mutator.apply(sch.trace)
assert trace is None


if __name__ == "__main__":
test_mutate_tile_size_matmul()
test_mutate_sample_categorical_single_candidate()