[Refactor] Enhance layout inference logic in ParallelOp (#420)

* Updated the layout inference in ParallelOp to improve the selection of source buffers for layout accuracy.
* Introduced logic to choose the read source buffer based on the number of indices, ensuring more precise layout inference.
* Refactored the loop handling to maintain clarity and improve the overall robustness of the layout inference process.

src/op/parallel.cc

@@ -131,13 +131,21 @@ LayoutMap ParallelOp::InferLayout(const LayoutInferArgs &T, InferLevel level) {
   auto block_size = T.thread_bounds->extent - T.thread_bounds->min;
   // Step 1: try to infer loop's partition from a source fragment
   Buffer source_buffer, read_source_buffer;
-  for (const auto &[buffer, _] : indice_map_) {
+  for (const auto &[buffer, indices] : indice_map_) {
     if (T.layout_map.count(buffer)) {
       auto frag = T.layout_map[buffer].as<Fragment>().value();
-      if (buffer_is_write_.count(buffer))
+      if (buffer_is_write_.count(buffer)) {
         source_buffer = buffer;
-      else
-        read_source_buffer = buffer;
+      } else {
+        // Keep the buffer with largest number of indices
+        // (which means the inference based on that buffer is more accurate)
+        // as read_source_buffer to get more accurate layout
+        if (!read_source_buffer.defined() ||
+            indice_map_[buffer].size() >
+                indice_map_[read_source_buffer].size()) {
+          read_source_buffer = buffer;
+        }
+      }
     }
   }
   auto compute_loop_layout_from_buffer = [&](const Buffer &buffer) {