updt

Author: LeiWang1999

src/op/parallel.cc

@@ -272,6 +272,8 @@ LayoutMap ParallelOpNode::InferLayout(const LayoutInferArgs &T,
   };
   // Step 1: try to infer loop's partition from a source fragment
   Buffer source_buffer, read_source_buffer;
+  Buffer replicated_write_buffer; // Backup: fully replicated write buffer
+
   for (const auto &[buffer, indices] : indice_map_) {
     if (T.layout_map.count(buffer)) {
       // skip reducers with rep=ALL
@@ -280,9 +282,26 @@ LayoutMap ParallelOpNode::InferLayout(const LayoutInferArgs &T,
         continue;
 
       auto frag = T.layout_map[buffer].as<Fragment>().value();
+      bool is_fully_replicated = buffer_is_completed_replicated(buffer);
+      bool is_reducer = reducer_info_map_.count(buffer->data);
 
       if (buffer_is_write_.count(buffer)) {
-        source_buffer = buffer;
+        // Allow fully replicated write buffers if they are reducers
+        // (reducers need replication for correctness)
+        if (!is_fully_replicated || is_reducer) {
+          // Prefer non-replicated write buffers, but also allow replicated
+          // reducers
+          if (!source_buffer.defined() ||
+              (!is_fully_replicated && source_buffer.defined())) {
+            source_buffer = buffer;
+          }
+        } else if (!replicated_write_buffer.defined()) {
+          // Keep fully replicated NON-reducer write buffer as backup
+          replicated_write_buffer = buffer;
+          DLOG(INFO) << "Found fully replicated non-reducer write buffer "
+                     << buffer << " as backup for loop layout inference"
+                     << '\n';
+        }
       } else {
         // Keep the buffer with largest number of indices
         // (which means the inference based on that buffer is more accurate)
@@ -308,6 +327,18 @@ LayoutMap ParallelOpNode::InferLayout(const LayoutInferArgs &T,
     Fragment src_layout = T.layout_map[buffer].as<Fragment>().value();
     DLOG(INFO) << "[compute_loop_layout_from_buffer] infer from buffer `"
                << buffer << "` of layout " << src_layout->DebugOutput() << '\n';
+
+    // Check if this buffer is a reducer
+    bool is_reducer = reducer_info_map_.count(buffer->data);
+
+    // Defensive check: warn if attempting to infer from fully replicated buffer
+    // But allow it for reducers (they need replication for correctness)
+    if (src_layout->IsCompletedReplicated() && !is_reducer) {
+      DLOG(WARNING)
+          << "Attempting to infer loop layout from fully replicated buffer "
+          << buffer << ", this may cause incorrect replication propagation";
+    }
+
     Fragment result;
     if (IsCommonAccessIndice(buffer)) {
       result = src_layout;
@@ -318,6 +349,34 @@ LayoutMap ParallelOpNode::InferLayout(const LayoutInferArgs &T,
       PrimExpr loop_var_to_thread =
           src_layout->ForwardThread(indice_map_[buffer], rep);
       loop_var_to_thread = analyzer_.Simplify(loop_var_to_thread);
+
+      // Check if loop_var_to_thread only depends on rep (not on loop_vars)
+      // This indicates a fully replicated buffer that shouldn't determine loop
+      // layout UNLESS it's a reducer (reducers need full replication)
+      bool uses_loop_var = false;
+      PostOrderVisit(loop_var_to_thread, [&](const ObjectRef &objref) {
+        if (auto var = objref.as<Var>()) {
+          for (const auto &loop_var : loop_vars_) {
+            if (var->same_as(loop_var->var)) {
+              uses_loop_var = true;
+              break;
+            }
+          }
+        }
+      });
+
+      if (!uses_loop_var && !is_reducer) {
+        // loop_var_to_thread only depends on rep, not on loop variables
+        // And it's not a reducer, so this is likely an index offset case
+        DLOG(WARNING) << "Buffer " << buffer
+                      << " is fully replicated (not a reducer). "
+                      << "Cannot use it to infer loop layout. "
+                      << "loop_var_to_thread = " << loop_var_to_thread
+                      << " (only depends on rep)";
+        throw LayoutConflictException("Buffer is fully replicated and cannot "
+                                      "be used for layout inference");
+      }
+
       PostOrderVisit(loop_var_to_thread, [&](const ObjectRef &objref) {
         if (auto opt_var = objref.as<Var>();
             opt_var && inner_vars_.count(*opt_var)) {
@@ -334,6 +393,13 @@ LayoutMap ParallelOpNode::InferLayout(const LayoutInferArgs &T,
                << result->DebugOutput() << '\n';
     return result;
   };
+
+  // Try to infer loop layout from buffers in order of preference:
+  // 1. Non-replicated write buffer (most reliable)
+  // 2. Non-replicated read buffer
+  // 3. Fully replicated write buffer (backup, may cause issues)
+  // 4. Free inference mode (no source buffer)
+
   if (source_buffer.defined()) {
     loop_layout_ = compute_loop_layout_from_buffer(source_buffer);
   } else if (level == InferLevel::kFree) {
@@ -388,7 +454,25 @@ LayoutMap ParallelOpNode::InferLayout(const LayoutInferArgs &T,
         auto rep = inv->Forward(fwd).back();
         AddPredicate(EQ(rep, 0));
       }
-    } else {
+    } else if (replicated_write_buffer.defined()) {
+      // Backup: try to use fully replicated write buffer
+      // This may cause replication propagation, but it's better than failing
+      DLOG(WARNING) << "Using fully replicated buffer "
+                    << replicated_write_buffer
+                    << " for loop layout inference as no other source buffer "
+                       "is available";
+      try {
+        loop_layout_ = compute_loop_layout_from_buffer(replicated_write_buffer);
+      } catch (const LayoutConflictException &e) {
+        // If fails, fall back to free mode
+        DLOG(WARNING) << "Failed to infer from replicated buffer: " << e.what()
+                      << ". Falling back to free mode";
+        replicated_write_buffer = Buffer(); // Clear to trigger free mode below
+      }
+    }
+
+    if (!loop_layout_.defined()) {
+      // No source buffer available, use free mode inference
       // Vectorize Size must be aware of the buffer_remap
       // As the pass will do post processing to the layout
       auto maybe_remapped_root_ =