[Compatible] Fix apache/tvm#12382

Author: comaniac

src/op/dialect/cutlass/cutlass_fusion.cc

@@ -78,14 +78,14 @@ OpEnv* Tune(const op::CallValues& call, OpEnv* op_env) {
   } else {
     std::vector<std::shared_ptr<TunableConfig>> tunable = env->ListTunableConfigs();
     const int number = 10, repeat = 1, min_repeat_ms = 0, cooldown_interval_ms = 0,
-              repeats_to_cooldow = 1;
+              repeats_to_cooldown = 1, limit_zero_time_iterations = 100;
     double min_time = std::numeric_limits<double>::max();
     for (auto& config : tunable) {
       env->SetTunableConfig(config);
       env->Init(call);
-      Array<FloatValue> result =
-          TimeEvaluator(TypedPackedFunc<void()>([&]() { env->Execute(call); }), call->device,
-                        number, repeat, min_repeat_ms, cooldown_interval_ms, repeats_to_cooldow)();
+      Array<FloatValue> result = TimeEvaluator(
+          TypedPackedFunc<void()>([&]() { env->Execute(call); }), call->device, number, repeat,
+          min_repeat_ms, limit_zero_time_iterations, cooldown_interval_ms, repeats_to_cooldown)();
       CHECK_EQ(result.size(), 1U);
       if (result[0]->value < min_time) {
         min_time = result[0]->value;

src/op/dialect/cutlass/timer.cc

@@ -22,13 +22,15 @@ tvm::runtime::Module MakeCutlassModule(PackedFunc pf) {
 }
 
 PackedFunc TimeEvaluator(PackedFunc pf, Device dev, int number, int repeat, int min_repeat_ms,
-                         int cooldown_interval_ms, int repeats_to_cooldown) {
+                         int limit_zero_time_iterations, int cooldown_interval_ms,
+                         int repeats_to_cooldown) {
   tvm::Device tvm_dev = dev;
   auto wrapper = [=](TVMArgs args, TVMRetValue* rv) mutable {
     const static PackedFunc rpv_eval = registry::GetPackedFunc("runtime.RPCTimeEvaluator");
     PackedFunc timer =
         rpv_eval(MakeCutlassModule(pf), "main", (int)tvm_dev.device_type, (int)tvm_dev.device_id,
-                 number, repeat, min_repeat_ms, cooldown_interval_ms, repeats_to_cooldown, "");
+                 number, repeat, min_repeat_ms, limit_zero_time_iterations, cooldown_interval_ms,
+                 repeats_to_cooldown, "");
     TVMRetValue timer_rv;
     timer.CallPacked(args, &timer_rv);
     const double* speed = reinterpret_cast<const double*>(timer_rv.operator std::string().data());

src/op/dialect/cutlass/timer.h

@@ -62,6 +62,9 @@ tvm::runtime::Module MakeCutlassModule(registry::PackedFunc pf);
  *                       minimum duration requirement of one `repeat`.
  *                       i.e., When the run time of one `repeat` falls below this time, the `number`
  *                       parameter will be automatically increased.
+ * \param limit_zero_time_iterations The maximum number of repeats when
+ *                                   measured time is equal to 0. It helps to avoid hanging
+ *                                   during measurements.
  * \param cooldown_interval_ms The cooldown interval in milliseconds between the number of
  *                             repeats defined by `repeats_to_cooldown`.
  * \param repeats_to_cooldown The number of repeats before the cooldown is activated.
@@ -70,6 +73,7 @@ tvm::runtime::Module MakeCutlassModule(registry::PackedFunc pf);
  */
 registry::PackedFunc TimeEvaluator(registry::PackedFunc pf, Device dev, int number = 10,
                                    int repeat = 1, int min_repeat_ms = 0,
+                                   int limit_zero_time_iterations = 100,
                                    int cooldown_interval_ms = 0, int repeats_to_cooldown = 1);
 
 }  // namespace cutlass