[Bug] Metaschedule nonexistent error message

[tkonolige]

Using this test script:

import tvm
import tempfile
import tvm.script.tir as T

@T.prim_func
def bad_message(
    B: T.Buffer[(64), "float32"],
    B_pack: T.Buffer[(8, 8), "float32"],
):
    with T.block("root"):
        for jo in range(8):
            for ji in range(8):
                B_pack[jo, ji] = B[jo * 8 + ji]


def test_bad_message():
    with tempfile.TemporaryDirectory() as tmp:
        database = tvm.meta_schedule.database.JSONDatabase(
            tmp + "/workload.json", tmp + "/record.json"
        )
        s = tvm.meta_schedule.tune_tir(
            bad_message,
            target="llvm -mcpu=znver3 -num-cores=12",
            config=tvm.meta_schedule.TuneConfig(
                max_trials_global=1000,
                max_trials_per_task=1000,
                num_trials_per_iter=64,
            ),
            work_dir=tmp,
            database=database,
            runner=tvm.meta_schedule.runner.LocalRunner(),
        )

if __name__ == "__main__":
    test_bad_message()

gives the following message:

2022-07-19 11:46:33.549 INFO Logging directory: /tmp/tmpss9867du/logs
2022-07-19 11:46:33.550 INFO Logging directory: /tmp/tmpss9867du/logs
2022-07-19 11:46:33.550 INFO Working directory: /tmp/tmpss9867du
2022-07-19 11:46:33.551 INFO LocalBuilder: max_workers = 12
2022-07-19 11:46:33.900 INFO Initializing Task #0: "main"
Traceback (most recent call last):
  File "test_bad_error_message.py", line 35, in <module>
    test_bad_message()
  File "test_bad_error_message.py", line 21, in test_bad_message
    s = tvm.meta_schedule.tune_tir(
  File "tvm/python/tvm/meta_schedule/tune.py", line 414, in tune_tir
    database = tune_extracted_tasks(
  File "tvm/python/tvm/meta_schedule/tune.py", line 350, in tune_extracted_tasks
    task_scheduler.tune()
  File "tvm/python/tvm/meta_schedule/task_scheduler/task_scheduler.py", line 72, in tune
    _ffi_api.TaskSchedulerTune(self)  # type: ignore # pylint: disable=no-member
  File "tvm/_ffi/_cython/./packed_func.pxi", line 331, in tvm._ffi._cy3.core.PackedFuncBase.__call__
  File "tvm/_ffi/_cython/./packed_func.pxi", line 262, in tvm._ffi._cy3.core.FuncCall
  File "tvm/_ffi/_cython/./packed_func.pxi", line 251, in tvm._ffi._cy3.core.FuncCall3
  File "tvm/_ffi/_cython/./base.pxi", line 181, in tvm._ffi._cy3.core.CHECK_CALL
tvm._ffi.base.TVMError: TVMError:

Its unclear what exactly the error is here.

@junrushao1994 @vinx13

[vinx13]

cc @zxybazh

[zxybazh]

Hi thanks for reporting this issue, I think you may correct the PrimFunc by removing the root block definition and variable binding as follows:

@T.prim_func
def bad_message(
    B: T.Buffer[(64), "float32"],
    B_pack: T.Buffer[(8, 8), "float32"],
):
    for jo in range(8):
        for ji in range(8):
            with T.block():
                vo, vi = T.axis.remap("SS", [jo, ji])
                B_pack[vo, vi] = B[vo * 8 + vi]

As for the error here, I did a bit dig in and found it was triggered when Autoinline was applied to the root block, which would throw a RootBlockError because the primitive does not operate on the root block. Would like to confirm with @vinx13 if this is designated behaviour? It can be reproduced by adding with T.block("root"): before the for loop in primfunc.

[zxybazh]

Corrected TIR would not cause the problem, please define the write buffer for root block as follows:

@tvm.script.ir_module
class Module:
    @T.prim_func
    def bad_message(
        B: T.Buffer[(64), "float32"],
        B_pack: T.Buffer[(8, 8), "float32"],
    ):
        with T.block("root"):
            T.reads()
            T.writes()
            for jo in range(8):
                for ji in range(8):
                    with T.block():
                        vo, vi = T.axis.remap("SS", [jo, ji])
                        B_pack[vo, vi] = B[vo * 8 + vi]

And still the problem lies in bad message for error reporting, would work on that later.