tvm icon indicating copy to clipboard operation
tvm copied to clipboard

[Bug] Per store feature extractor of meta schedule cannot handle TIR with storing constant tensor

Open shingjan opened this issue 3 years ago • 0 comments

As mentioned in the title, the per store feature extractor of meta schedule seems not be able to handle TIR with one simply block of storing constant tensor.

Actual behavior

Per store feature extractor cannot extract any feature out of this TIR.

Environment

Latest TVM main

Steps to reproduce

import tvm
from tvm import tir
from tvm.meta_schedule.tune import TuneConfig
from tvm.script import tir as T
from tvm.meta_schedule.postproc import VerifyGPUCode
from tvm.tir.schedule import Trace, Schedule
from tvm import meta_schedule as ms
from tvm.meta_schedule import TuneContext, MeasureCandidate
from tvm.meta_schedule.arg_info import TensorInfo
from tvm.target import Target
from typing import Callable, List

# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
# fmt: off
@tvm.script.ir_module
class Module:
    @T.prim_func
    def fused_zeros(var_T_full: T.handle) -> None:
        # function attr dict
        T.func_attr({"tir.noalias": True, "global_symbol": "main"})
        T_full = T.match_buffer(var_T_full, [1, 12, 4096], dtype="int64", align=128)
        # body
        # with T.block("root")
        for i0, i1, i2 in T.grid(1, 12, 4096):
            with T.block("T_full"):
                ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                T.reads()
                T.writes(T_full[ax0, ax1, ax2])
                T_full[ax0, ax1, ax2] = T.int64(0)
# fmt: on

if __name__ == """__main__""":
    extractor = ms.feature_extractor.PerStoreFeature()

    def _make_context(target) -> ms.TuneContext:
        return ms.TuneContext(
            target=target,
            num_threads=1,
        )

    def _make_candidate(f_sch: Callable[[], tir.Schedule]) -> ms.MeasureCandidate:
        return ms.MeasureCandidate(sch=f_sch(), args_info=[])

    def _create_schedule():
        func = Module["fused_zeros"]
        sch = tir.Schedule(func, debug_mask="all")
        return sch

    (feature,) = extractor.extract_from(
        _make_context(Target("nvidia/geforce-rtx-3070")),
        candidates=[_make_candidate(_create_schedule)],
    )
    feature = feature.numpy()

cc: @zxybazh @junrushao @vinx13

shingjan avatar Sep 14 '22 21:09 shingjan