(tvm) branch main updated: [MetaSchedule] Fix tune_tir crash with ScheduleError in RewriteParallelVectorizeUnroll (#18547)

[tlopex]

This is an automated email from the ASF dual-hosted git repository.

tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 7ed4d39a64 [MetaSchedule] Fix tune_tir crash with ScheduleError in 
RewriteParallelVectorizeUnroll (#18547)
7ed4d39a64 is described below

commit 7ed4d39a647b02f4e3d080a30a62f9dcc64b771a
Author: Dayuxiaoshui <158081477+dayuxiaos...@users.noreply.github.com>
AuthorDate: Thu Dec 25 20:31:37 2025 +0800

    [MetaSchedule] Fix tune_tir crash with ScheduleError in 
RewriteParallelVectorizeUnroll (#18547)
    
    This patch fixes issue #18423 where meta_schedule.tune_tir crashes
    during initial population sampling when RewriteParallelVectorizeUnroll
    postprocessor encounters blocks that violate compact dataflow
    requirements.
    
    The crash occurred when:
    - A block reads and writes to the same buffer
    - RewriteParallelVectorizeUnroll tries to parallelize/vectorize these
    loops
    - ScheduleError is thrown and propagates through parallel_for_dynamic
    
    Solution:
    - Added exception handling in RewriteParallelVectorizeUnroll::Apply to
    catch ScheduleError and skip problematic blocks instead of crashing
    - Added exception handling in ThreadedTraceApply::Apply to catch
    exceptions from postprocessors and treat them as normal failures
    
    This makes the tuning process more robust by gracefully handling
    schedule errors instead of crashing the entire tuning session.
---
 .../postproc/rewrite_parallel_vectorize_unroll.cc  | 39 ++++++++++++----------
 src/meta_schedule/utils.h                          | 14 +++++++-
 2 files changed, 35 insertions(+), 18 deletions(-)

diff --git a/src/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc 
b/src/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc
index d833af6142..5950ef742d 100644
--- a/src/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc
+++ b/src/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc
@@ -425,25 +425,30 @@ class RewriteParallelVectorizeUnrollNode : public 
PostprocNode {
         tir::ParsedAnnotation parsed = parsed_root;
         tir::AdjustParallelVectorize(sch, block_rv, loop_rvs, &parsed);
         const int loops_num = loop_rvs.size();
-        if (parsed.num_parallel_loops == loops_num && 
parsed.num_vectorize_loops == loops_num) {
-          // Fuse, split, vectorize and parallelize
-          tir::RewriteFuseSplitParallelVectorize(sch, &loop_rvs, 
parsed.max_vectorize_extent);
-        } else {
-          // Parallel
-          if (parsed.num_parallel_loops > 0) {
-            tir::RewriteParallel(sch, parsed.num_parallel_loops, &loop_rvs);
+        try {
+          if (parsed.num_parallel_loops == loops_num && 
parsed.num_vectorize_loops == loops_num) {
+            // Fuse, split, vectorize and parallelize
+            tir::RewriteFuseSplitParallelVectorize(sch, &loop_rvs, 
parsed.max_vectorize_extent);
+          } else {
+            // Parallel
+            if (parsed.num_parallel_loops > 0) {
+              tir::RewriteParallel(sch, parsed.num_parallel_loops, &loop_rvs);
+            }
+            // Vectorize
+            if (parsed.num_vectorize_loops > 0) {
+              tir::RewriteVectorize(sch, parsed.num_vectorize_loops, 
&loop_rvs);
+            }
           }
-          // Vectorize
-          if (parsed.num_vectorize_loops > 0) {
-            tir::RewriteVectorize(sch, parsed.num_vectorize_loops, &loop_rvs);
+          // AutoUnroll
+          if (parsed.unroll_explicit != -1 || parsed.unroll_implicit != -1) {
+            ICHECK(parsed.unroll_explicit == -1 || parsed.unroll_implicit == 
-1);
+            int unroll_explicit = parsed.unroll_explicit != -1;
+            int max_step = parsed.unroll_explicit + parsed.unroll_implicit + 1;
+            tir::RewriteUnroll(sch, unroll_explicit, max_step, block_rv, 
loop_rvs[0]);
           }
-        }
-        // AutoUnroll
-        if (parsed.unroll_explicit != -1 || parsed.unroll_implicit != -1) {
-          ICHECK(parsed.unroll_explicit == -1 || parsed.unroll_implicit == -1);
-          int unroll_explicit = parsed.unroll_explicit != -1;
-          int max_step = parsed.unroll_explicit + parsed.unroll_implicit + 1;
-          tir::RewriteUnroll(sch, unroll_explicit, max_step, block_rv, 
loop_rvs[0]);
+        } catch (const tir::ScheduleError& e) {
+          DLOG(WARNING) << "Failed to apply parallelization/vectorization: " 
<< e.what();
+          return false;
         }
       }
     }
diff --git a/src/meta_schedule/utils.h b/src/meta_schedule/utils.h
index ee94b1d2ab..96b18fb7e7 100644
--- a/src/meta_schedule/utils.h
+++ b/src/meta_schedule/utils.h
@@ -334,7 +334,19 @@ struct ThreadedTraceApply {
 
     for (int i = 0; i < n_; ++i) {
       Item& item = items_[i];
-      if (!item.postproc->Apply(sch)) {
+      bool success = true;
+      try {
+        if (!item.postproc->Apply(sch)) {
+          success = false;
+        }
+      } catch (const tir::ScheduleError& e) {
+        DLOG(WARNING) << "Postproc #" << i << " failed with ScheduleError: " 
<< e.what();
+        success = false;
+      } catch (const std::exception& e) {
+        DLOG(WARNING) << "Postproc #" << i << " failed with exception: " << 
e.what();
+        success = false;
+      }
+      if (!success) {
         item.fail_counter++;
         return std::nullopt;
       }