[TE/TIR] Preserving loops of extent 1

Questions
1

How can I keep loops around that iterate over only a single element? I’ve tried disabling all the compilation passes in tvm.lower without success, which presumably implies it happens during CreateSchedule.

Thanks

2

The third argument in ScheduleOps is a boolean flag that indicates whether to keep trivial loops or not. There is no way to set it via tvm.lower or tvm.build, but if you want to convert the schedule to a statement by hand, you can do it there.

C++

Stmt ScheduleOps(Schedule sch, Map<IterVar, Range> dom_map_, bool debug_keep_trivial_loop);

Python:

TVM_REGISTER_GLOBAL("schedule.ScheduleOps").set_body([](TVMArgs args, TVMRetValue* ret) {
  if (args.size() == 2)
    *ret = ScheduleOps(args[0], args[1], false);
  else
    *ret = ScheduleOps(args[0], args[1], args[2]);
});
1 Like