[GitHub] [tvm] shingjan opened a new issue, #12782: [Bug] Per store feature extractor of meta schedule cannot handle TIR with storing constant tensor

[GitBox <gi...@apache.org>]

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: commits-unsubscribe@tvm.apache.org.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org