tvm
tvm copied to clipboard
[Bug] Per store feature extractor of meta schedule cannot handle TIR with storing constant tensor
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