shingjan opened a new issue, #12782:
URL: https://github.com/apache/tvm/issues/12782
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
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
