[Enhancement] Improve layout inference accuracy in ParallelOp (#441) (#442)

* Added logic to use non-replicated buffers as source buffers for more accurate layout inference.
* Enhanced comments to clarify the rationale behind buffer selection in layout inference process.

src/op/parallel.cc

@@ -163,6 +163,11 @@ LayoutMap ParallelOp::InferLayout(const LayoutInferArgs &T, InferLevel level) {
                 indice_map_[read_source_buffer].size()) {
           read_source_buffer = buffer;
         }
+        // If the buffer is not replicated, use it as source_buffer
+        // because the layout inference is more accurate
+        if (is_one(frag->ReplicateExtent())) {
+          source_buffer = buffer;
+        }
       }
     }
   }