[MetaSchedule] MutateTileSize skip single-candidate SampleCategorical (…

…apache#14072)

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.

src/meta_schedule/mutator/mutate_tile_size.cc

@@ -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);

tests/python/unittest/test_meta_schedule_mutator_mutate_tile_size.py

@@ -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()