[Mlir-commits] [mlir] [mlirbc] Fix use-list ordering during deserialization (PR #191942)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Jacques Pienaar (jpienaar)

<details>
<summary>Changes</summary>

This CL fixes an issue in the MLIR bytecode reader where use-lists were correctly reconstructed when they had permutations that are not own inverse. Fixed the use-list reconstruction mapping logic in to correctly restore the stable memory order of uses, both full shuffle and index-pair encodings consistently.

Gemini/LLM assisted.

---
Full diff: https://github.com/llvm/llvm-project/pull/191942.diff


2 Files Affected:

- (modified) mlir/lib/Bytecode/Reader/BytecodeReader.cpp (+27-4) 
- (modified) mlir/test/Bytecode/uselist_orders.mlir (+20) 


``````````diff
diff --git a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
index caa17ea57b950..00c53dae83ea8 100644
--- a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
+++ b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
@@ -2309,6 +2309,7 @@ LogicalResult BytecodeReader::Impl::sortUseListOrder(Value value) {
   // If the encoding was a pair of indices `(src, dst)` for every permutation,
   // reconstruct the shuffle vector for every use. Initialize the shuffle vector
   // as identity, and then apply the mapping encoded in the indices.
+  // This produces shuffle[oldIdx] = newPos (i.e., old_index -> new_position).
   if (customOrder.isIndexPairEncoding) {
     // Return failure if the number of indices was not representing pairs.
     if (shuffle.size() & 1)
@@ -2337,10 +2338,32 @@ LogicalResult BytecodeReader::Impl::sortUseListOrder(Value value) {
       accumulator != (((numUses - 1) * numUses) >> 1))
     return failure();
 
-  // Apply the current ordering map onto the shuffle vector to get the final
-  // use-list sorting indices before shuffling.
-  shuffle = SmallVector<unsigned, 4>(llvm::map_range(
-      currentOrder, [&](auto item) { return shuffle[item.first]; }));
+  // Compose the shuffle with the current memory layout to produce the final
+  // indices for shuffleUseList. shuffleUseList(indices) places the use at
+  // current position i into position indices[i], so we need to compute
+  // finalShuffle[readerMemIdx] = writerMemIdx.
+  //
+  // The two encoding paths have different shuffle conventions:
+  //
+  // Index-pair encoding (already normalized above):
+  //   shuffle[writerMemIdx] = sortedPos  (old_index -> new_position)
+  //
+  // Full-shuffle encoding (writer's native format):
+  //   shuffle[sortedPos] = writerMemIdx  (new_position -> old_index)
+  //
+  // In both cases, currentOrder[sortedPos].first gives the readerMemIdx for
+  // a given sorted position. We fold the permutation inversion for the
+  // full-shuffle case directly into the composition to avoid an extra pass.
+  SmallVector<unsigned, 4> finalShuffle(numUses);
+  if (customOrder.isIndexPairEncoding) {
+    for (size_t writerMemIdx = 0; writerMemIdx < numUses; ++writerMemIdx)
+      finalShuffle[currentOrder[shuffle[writerMemIdx]].first] = writerMemIdx;
+  } else {
+    for (size_t sortedPos = 0; sortedPos < numUses; ++sortedPos)
+      finalShuffle[currentOrder[sortedPos].first] = shuffle[sortedPos];
+  }
+  shuffle = std::move(finalShuffle);
+
   value.shuffleUseList(shuffle);
   return success();
 }
diff --git a/mlir/test/Bytecode/uselist_orders.mlir b/mlir/test/Bytecode/uselist_orders.mlir
index b8f4c3df6542c..73049d5d66173 100644
--- a/mlir/test/Bytecode/uselist_orders.mlir
+++ b/mlir/test/Bytecode/uselist_orders.mlir
@@ -1,4 +1,5 @@
 // RUN: mlir-opt %s -split-input-file --test-verify-uselistorder -verify-diagnostics
+// RUN: mlir-opt %s -split-input-file --test-verify-uselistorder="rng-seed=54" -verify-diagnostics
 
 // COM: --test-verify-uselistorder will randomly shuffle the uselist of every
 //      value and do a roundtrip to bytecode. An error is returned if the
@@ -61,3 +62,22 @@ test.graph_region {
   %1 = "test.bar"(%2) : (i32) -> i32
   %2 = "test.baz"() : () -> i32
 }
+
+// -----
+
+// This is a reproducer test (for the fixed seed run) for use-def list
+// ordering.
+func.func @test_with_8_uses(%arg0 : i32) -> i32 {
+  %0 = arith.constant 45 : i32
+  %1 = "test.addi"(%arg0, %0) : (i32, i32) -> i32
+  %2 = "test.addi"(%1, %0) : (i32, i32) -> i32
+  %3 = "test.addi"(%1, %0) : (i32, i32) -> i32
+  %4 = "test.addi"(%1, %0) : (i32, i32) -> i32
+  %5 = "test.addi"(%1, %0) : (i32, i32) -> i32
+  %6 = "test.addi"(%1, %0) : (i32, i32) -> i32
+  %7 = "test.addi"(%1, %0) : (i32, i32) -> i32
+  %8 = "test.addi"(%1, %0) : (i32, i32) -> i32
+  %9 = "test.addi"(%1, %0) : (i32, i32) -> i32
+  return %9 : i32
+}
+

``````````

</details>


https://github.com/llvm/llvm-project/pull/191942