Preserve use-list orders in mlir bytecode

[mfrancio]

Summary

This patch implements a mechanism to read/write use-list orders from/to the mlir bytecode format. When producing bytecode, use-list orders are appended to each value of the IR. When reading bytecode, use-lists orders are loaded in memory and used at the end of parsing to sort the existing use-list chains.

[mehdi_amini]

You need to update the Bytecode documentation still I think.

mlir/include/mlir/Bytecode/BytecodeWriter.h

37 ↗

(On Diff #519093)

Please keep the default to true: I'd really want the default to always be the "safe" way. We already got it wrong in other places and it's impossible to fix later.

Also what about a fluent API like we do for other config? (OpPrintingFlags for example)

mlir/include/mlir/IR/AsmState.h

463 ↗

(On Diff #519093)

I don't quite get why we need this? I would just honor what's in the bytecode without config here.

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

793

How do you manager the order of two uses by the same operation?

mlir/lib/IR/Value.cpp

103

Seems like you could just zip order with the uses here instead of enumerate?

107

Does C++17 allows for (const auto &[use, newUse] : llvm::zip(getUses(), newUseList)) {?
Structured bindings might look nicer than the std::get

[rriddle]

mlir/include/mlir/Bytecode/BytecodeWriter.h

37 ↗

(On Diff #519093)

+1. Let's focus on making the default correct and efficient.

mlir/include/mlir/IR/AsmState.h

463 ↗

(On Diff #519093)

Yeah, I don't see an initial need for this functionality. Let's start simple, this can always be added if we really really need it later.

mlir/include/mlir/IR/Value.h

192

The formatting here looks off.

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

1286–1287

Can you turn the value of this into a struct? It's not obvious what the bool means.

1348

Please document the struct.

1353

// -> /// for each of these.

1835–1838

[mfrancio]

The patch contains the following updates:

1. removes the command line arguments to optionally print/read use-lists to bytecode, which are now printed and parsed by default;
2. extends the use-lists orders info to block arguments and increases the bytecode version to 2 (this was a breaking change);
3. implements a testing pass to do somewhat of a rigorous functional check (essentially, the pass takes any module, applies a random shuffle to the use-list of every value/blockarg, does a roundtrip to bytecode and verifies that the use-list order is preserved).
4. Fixes few functionality issues of the initial implementation:
   • the use-list was not shuffled correctly, the implementation is now updated and tested rigorously within the testing pass;
   • the case with multiple uses in the same operand was not handled correctly - now I introduced a new enumeration that takes into account the operand number of the use;
   • I found few corner cases where the op enumeration I was doing was not consistent with what I needed (pre-order walk). Now this is done separately before the recursive enumeration of regions

I experimented a little bit with the testing pass and run it through the whole set of mlir file check tests that exist in the repo, and saw no failures.

I hope you will find this work interesting. Looking forward for further feedback.

[mfrancio]

mlir/include/mlir/Bytecode/BytecodeWriter.h

37 ↗

(On Diff #519093)

Ok, I removed this. We could add this later if needed.

mlir/include/mlir/IR/AsmState.h

463 ↗

(On Diff #519093)

Removed. Thanks.

mlir/include/mlir/IR/Value.h

192

Should be fixed now.

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

1286–1287

I simplified the implementation a little bit and added a struct with documentation.

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

793

It was indeed not done correctly - I updated the implementation and now this is done by sorting with keys that contain operand number and operation number.

mlir/lib/IR/Value.cpp

107

Apologies, but this wasn't done correctly - I thought I could be able to change pointee without changing pointers, but the API wasn't really doing that. I added a couple of helpers to shuffle the uselist explicitly.

[mfrancio]

Adds a test case that exercises preserving use-lists on IR without dominance.

[mehdi_amini]

LG, with a nit.

mlir/test/lib/IR/TestUseListOrders.cpp

200

I'm concerned that we'll have reproducibility issues here. Can we explicitly initialize the engine with a seed that would be a pass option?
Even if you run the test with a random seed, we'll at least see it in the invocation and have reproducibility.

[mehdi_amini]

mlir/docs/BytecodeFormat.md

458

(newline)

[mfrancio]

• Adds random number generator seed as pass input to allow test reproducibility
• Removes the do{}while(is_sorted) from permutation testing - essentially, this allows to verify that sorted uselists remain sorted after a bytecode roundtrip
• Improves comments of few helper functions

[mfrancio]

mlir/test/lib/IR/TestUseListOrders.cpp

200

Agreed, added!

[mehdi_amini]

Can you rebase? The version was just bumped.

[mehdi_amini]

mlir/docs/BytecodeFormat.md

407

Nit: seems like you have spaces at the end of every line here

456

Nit: space at the end of line here as well

[mfrancio]

Removed trailing whitespaces and rebased.

[vitalybuka]

Breaks bots https://lab.llvm.org/buildbot/#/builders/5/builds/33857

[mehdi_amini]

Thanks for the report, seems like a leak in the test, I'll have a fix shortly (building/testing locally)

[mehdi_amini]

f6fb639c76ce may be enough to fix it, this bot seems really slow though: any reason ccache isn't enabled?

[vitalybuka]

It's enabled. but it spends most of the time in "check" steps.
Also splitting it into 3 bots, 1 per sanitizer, could be useful.

[mehdi_amini]

There is something fishy, this is the build-time:

real	8m51.140s
user	228m58.976s
sys	16m24.143s

That seems far too much for just this patch alone: either the ccache cache is too small, or it isn't setup correctly (maybe wiped between builds?).

[vitalybuka]

From which build/step is this from?

[vitalybuka]

In D149755#4359548, @vitalybuka wrote:

From which build/step is this from?

Found that: asan build from https://lab.llvm.org/buildbot/#/builders/5/builds/33857

[vitalybuka]

In D149755#4359550, @vitalybuka wrote:

In D149755#4359548, @vitalybuka wrote:

From which build/step is this from?

Found that: asan build from https://lab.llvm.org/buildbot/#/builders/5/builds/33857

You can find ccach stat in Info of the bot, hit rate is quite high.
The step is 8 min, usually, it's less, one before is 4. Failed step is "asan check" and it's 30m.

E.g. here ccache was whipped (settings changed) https://lab.llvm.org/buildbot/#/builders/5/builds/33859 and the build is 39m.

mlir/docs/BytecodeFormat.md

op {
   resultTypes: varint[],
 
   numOperands: varint?,
   operands: varint[],
 
   numSuccessors: varint?,
   successors: varint[],
 
+  numUseListOrders: varint?,
+  useListOrders: uselist[],
+
   regionEncoding: varint?, // (numRegions << 1) | (isIsolatedFromAbove)
 
   // regions are stored in a section if isIsolatedFromAbove
   regions: (region | region_section)[]
 }
+
+uselist {
+  indexInRange: varint?,
+  useListEncoding: varint, // (numIndices << 1) | (isIndexPairEncoding)
+  indices: varint[]
+}
 ```
 
 The encoding of an operation is important because this is generally the most
 commonly appearing structure in the bytecode. A single encoding is used for
 every type of operation. Given this prevelance, many of the fields of an
 operation are optional. The `encodingMask` field is a bitmask which indicates
 which of the components of the operation are present.
 

 each operand is encoded. This value index is the relative ordering of the
 definition of that value from the start of the first ancestor isolated region.
 
 ##### Successors
 
 If the operation has successors, the number of successors and the indexes of the
 successor blocks within the parent region are encoded.
 
+##### Use-list orders
+
+The reference use-list order is assumed to be the reverse of the global
+enumeration of all the op operands that one would obtain with a pre-order walk
+of the IR. This order is naturally obtained by building blocks of operations
+op-by-op. However, some transformations may shuffle the use-lists with respect
+to this reference ordering. If any of the results of the operation have a
+use-list order that is not sorted with respect to the reference use-list order,
+an encoding is emitted such that it is possible to reconstruct such order after
+parsing the bytecode. The encoding represents an index map from the reference
+operand order to the current use-list order. A bit flag is used to detect if
+this encoding is of type index-pair or not. When the bit flag is set to zero,
+the element at `i` represent the position of the use `i` of the reference list
+into the current use-list. When the bit flag is set to `1`, the encoding
+represent index pairs `(i, j)`, which indicate that the use at position `i` of
+the reference list is mapped to position `j` in the current use-list. When only
+less than half of the elements in the current use-list are shuffled with respect
+to the reference use-list, the index-pair encoding is used to reduce the
+bytecode memory requirements.
+
 ##### Regions
 
 If the operation has regions, the number of regions and if the regions are
 isolated from above are encoded together in a single varint. Afterwards, each
 region is encoded inline.
 
 #### Region Encoding
 
block {
   encoding: varint, // (numOps << 1) | (hasBlockArgs)
   arguments: block_arguments?, // Optional based on encoding
   ops : op[]
 }
 
 block_arguments {
   numArgs: varint?,
   args: block_argument[]
+  numUseListOrders: varint?,
+  useListOrders: uselist[],
 }
 
 block_argument {
   typeIndex: varint,
   location: varint
 }
 ```
 
 A block is encoded with an array of operations and block arguments. The first
 field is an encoding that combines the number of operations in the block, with a
 flag indicating if the block has arguments.
+
+Use-list orders are attached to block arguments similarly to how they are
+attached to operation results.

mlir/include/mlir/Bytecode/Encoding.h

 //===- Encoding.h - MLIR binary format encoding information -----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This header defines enum values describing the structure of MLIR bytecode
 // files.
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef LIB_MLIR_BYTECODE_ENCODING_H
-#define LIB_MLIR_BYTECODE_ENCODING_H
+#ifndef MLIR_BYTECODE_ENCODING_H
+#define MLIR_BYTECODE_ENCODING_H
 
 #include <cstdint>
 
 namespace mlir {
 namespace bytecode {
 //===----------------------------------------------------------------------===//
 // General constants
 //===----------------------------------------------------------------------===//
 
 enum {
   /// The minimum supported version of the bytecode.
   kMinSupportedVersion = 0,
 
   /// The current bytecode version.
-  kVersion = 2,
+  kVersion = 3,
 
   /// An arbitrary value used to fill alignment padding.
   kAlignmentByte = 0xCB,
 };
 
 //===----------------------------------------------------------------------===//
 // Sections
 //===----------------------------------------------------------------------===//

 namespace OpEncodingMask {
 enum : uint8_t {
   // clang-format off
   kHasAttrs         = 0b00000001,
   kHasResults       = 0b00000010,
   kHasOperands      = 0b00000100,
   kHasSuccessors    = 0b00001000,
   kHasInlineRegions = 0b00010000,
+  kHasUseListOrders = 0b00100000,
   // clang-format on
 };
 } // namespace OpEncodingMask
 
+/// Get the unique ID of a value use. We encode the unique ID combining an owner
+/// number and the argument number such as if ownerID(op1) < ownerID(op2), then
+/// useID(op1) < useID(op2). If uses have the same owner, then argNumber(op1) <
+/// argNumber(op2) implies useID(op1) < useID(op2).
+template <typename OperandT>
+static inline uint64_t getUseID(OperandT &val, unsigned ownerID) {
+  uint32_t operandNumberID;
+  if constexpr (std::is_same_v<OpOperand, OperandT>)
+    operandNumberID = val.getOperandNumber();
+  else if constexpr (std::is_same_v<BlockArgument, OperandT>)
+    operandNumberID = val.getArgNumber();
+  else
+    llvm_unreachable("unexpected operand type");
+  return (static_cast<uint64_t>(ownerID) << 32) | operandNumberID;
+}
+
 } // namespace bytecode
 } // namespace mlir
 
 #endif

mlir/include/mlir/IR/UseDefLists.h

public:
   /// Return the owner of this operand.
   Operation *getOwner() const { return owner; }
 
   /// Return the next operand on the use-list of the value we are referring to.
   /// This should generally only be used by the internal implementation details
   /// of the SSA machinery.
   IROperandBase *getNextOperandUsingThisValue() { return nextUse; }
 
+  /// Initialize the use-def chain by setting the back address to self and
+  /// nextUse to nullptr.
+  void initChainWithUse(IROperandBase **self) {
+    assert(this == *self);
+    back = self;
+    nextUse = nullptr;
+  }
+
+  /// Link the current node to next.
+  void linkTo(IROperandBase *next) {
+    nextUse = next;
+    if (nextUse)
+      nextUse->back = &nextUse;
+  }
+
 protected:
   IROperandBase(Operation *owner) : owner(owner) {}
   IROperandBase(IROperandBase &&other) : owner(other.owner) {
     *this = std::move(other);
   }
   IROperandBase &operator=(IROperandBase &&other) {
     removeFromCurrent();
     other.removeFromCurrent();
public:
   template <typename ValueT>
   void replaceAllUsesWith(ValueT &&newValue) {
     assert((!newValue || this != OperandType::getUseList(newValue)) &&
            "cannot RAUW a value with itself");
     while (!use_empty())
       use_begin()->set(newValue);
   }
 
+  /// Shuffle the use-list chain according to the provided indices vector, which
+  /// need to represent a valid shuffle. That is, a vector of unique integers in
+  /// range [0, numUses - 1]. Users of this function need to guarantee the
+  /// validity of the indices vector.
+  void shuffleUseList(ArrayRef<unsigned> indices) {
+    assert((size_t)std::distance(getUses().begin(), getUses().end()) ==
+               indices.size() &&
+           "indices vector expected to have a number of elements equal to the "
+           "number of uses");
+    SmallVector<detail::IROperandBase *> shuffled(indices.size());
+    detail::IROperandBase *ptr = firstUse;
+    for (size_t idx = 0; idx < indices.size();
+         idx++, ptr = ptr->getNextOperandUsingThisValue())
+      shuffled[indices[idx]] = ptr;
+
+    initFirstUse(shuffled.front());
+    auto *current = firstUse;
+    for (auto &next : llvm::drop_begin(shuffled)) {
+      current->linkTo(next);
+      current = next;
+    }
+    current->linkTo(nullptr);
+  }
+
   //===--------------------------------------------------------------------===//
   // Uses
   //===--------------------------------------------------------------------===//
 
   using use_iterator = ValueUseIterator<OperandType>;
   using use_range = iterator_range<use_iterator>;
 
   use_iterator use_begin() const { return use_iterator(firstUse); }

 protected:
   IRObjectWithUseList() = default;
 
   /// Return the first operand that is using this value, for use by custom
   /// use/def iterators.
   OperandType *getFirstUse() const { return (OperandType *)firstUse; }
 
 private:
+  /// Set use as the first use of the chain.
+  void initFirstUse(detail::IROperandBase *use) {
+    firstUse = use;
+    firstUse->initChainWithUse(&firstUse);
+  }
+
   detail::IROperandBase *firstUse = nullptr;
 
   /// Allow access to `firstUse`.
   friend detail::IROperandBase;
 };
 
 //===----------------------------------------------------------------------===//
 // ValueUseIterator

mlir/include/mlir/IR/Value.h

public:
   /// Replace all uses of 'this' value with 'newValue' if the given callback
   /// returns true.
   void replaceUsesWithIf(Value newValue,
                          function_ref<bool(OpOperand &)> shouldReplace);
 
   /// Returns true if the value is used outside of the given block.
   bool isUsedOutsideOfBlock(Block *block);
 
+  /// Shuffle the use list order according to the provided indices. It is
+  /// responsibility of the caller to make sure that the indices map the current
+  /// use-list chain to another valid use-list chain.
+  void shuffleUseList(ArrayRef<unsigned> indices);
+
   //===--------------------------------------------------------------------===//
   // Uses
 
   /// This class implements an iterator over the uses of a value.
   using use_iterator = ValueUseIterator<OpOperand>;
   using use_range = iterator_range<use_iterator>;
 
   use_iterator use_begin() const { return impl->use_begin(); }

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

 //===- BytecodeReader.cpp - MLIR Bytecode Reader --------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 // TODO: Support for big-endian architectures.
-// TODO: Properly preserve use lists of values.
 
 #include "mlir/Bytecode/BytecodeReader.h"
-#include "../Encoding.h"
 #include "mlir/AsmParser/AsmParser.h"
 #include "mlir/Bytecode/BytecodeImplementation.h"
+#include "mlir/Bytecode/Encoding.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/IR/Visitors.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/Support/SaveAndRestore.h"
 #include "llvm/Support/SourceMgr.h"
 #include <list>
 #include <memory>
+#include <numeric>
 #include <optional>
 
 #define DEBUG_TYPE "mlir-bytecode-reader"
 
 using namespace mlir;
 
 /// Stringify the given section ID.
 static std::string toString(bytecode::Section::ID sectionID) {
private:
 
   /// Sequentially define the given value range.
   LogicalResult defineValues(EncodingReader &reader, ValueRange values);
 
   /// Create a value to use for a forward reference.
   Value createForwardRef();
 
   //===--------------------------------------------------------------------===//
+  // Use-list order helpers
+
+  /// This struct is a simple storage that contains information required to
+  /// reorder the use-list of a value with respect to the pre-order traversal
+  /// ordering.
+  struct UseListOrderStorage {
+    UseListOrderStorage(bool isIndexPairEncoding,
+                        SmallVector<unsigned, 4> &&indices)
+        : indices(std::move(indices)),
+          isIndexPairEncoding(isIndexPairEncoding){};
+    /// The vector containing the information required to reorder the
+    /// use-list of a value.
+    SmallVector<unsigned, 4> indices;
+
+    /// Whether indices represent a pair of type `(src, dst)` or it is a direct
+    /// indexing, such as `dst = order[src]`.
+    bool isIndexPairEncoding;
+  };
+
+  /// Parse use-list order from bytecode for a range of values if available. The
+  /// range is expected to be either a block argument or an op result range. On
+  /// success, return a map of the position in the range and the use-list order
+  /// encoding. The function assumes to know the size of the range it is
+  /// processing.
+  using UseListMapT = DenseMap<unsigned, UseListOrderStorage>;
+  FailureOr<UseListMapT> parseUseListOrderForRange(EncodingReader &reader,
+                                                   uint64_t rangeSize);
+
+  /// Shuffle the use-chain according to the order parsed.
+  LogicalResult sortUseListOrder(Value value);
+
+  /// Recursively visit all the values defined within topLevelOp and sort the
+  /// use-list orders according to the indices parsed.
+  LogicalResult processUseLists(Operation *topLevelOp);
+
+  //===--------------------------------------------------------------------===//
   // Fields
 
   /// This class represents a single value scope, in which a value scope is
   /// delimited by isolated from above regions.
   struct ValueScope {
     /// Push a new region state onto this scope, reserving enough values for
     /// those defined within the current region of the provided state.
     void push(RegionReadState &readState) {
struct ValueScope {
     /// The set of values defined in this scope.
     std::vector<Value> values;
 
     /// The ID for the next defined value for each region current being
     /// processed in this scope.
     SmallVector<unsigned, 4> nextValueIDs;
   };
 
   /// The configuration of the parser.
   const ParserConfig &config;
 
   /// A location to use when emitting errors.
   Location fileLoc;
 
   /// Flag that indicates if lazyloading is enabled.
   bool lazyLoading;
 
   /// Keep track of operations that have been lazy loaded (their regions haven't
   /// been materialized), along with the `RegionReadState` that allows to
   /// lazy-load the regions nested under the operation.
   LazyLoadableOpsInfo lazyLoadableOps;
   LazyLoadableOpsMap lazyLoadableOpsMap;
private:
 
   /// The table of IR units referenced within the bytecode file.
   SmallVector<BytecodeDialect> dialects;
   SmallVector<BytecodeOperationName> opNames;
 
   /// The reader used to process resources within the bytecode.
   ResourceSectionReader resourceReader;
 
+  /// Worklist of values with custom use-list orders to process before the end
+  /// of the parsing.
+  DenseMap<void *, UseListOrderStorage> valueToUseListMap;
+
   /// The table of strings referenced within the bytecode file.
   StringSectionReader stringReader;
 
   /// The current set of available IR value scopes.
   std::vector<ValueScope> valueScopes;
+
+  /// The global pre-order operation ordering.
+  DenseMap<Operation *, unsigned> operationIDs;
+
   /// A block containing the set of operations defined to create forward
   /// references.
   Block forwardRefOps;
+
   /// A block containing previously created, and no longer used, forward
   /// reference operations.
   Block openForwardRefOps;
+
   /// An operation state used when instantiating forward references.
   OperationState forwardRefOpState;
 
   /// Reference to the input buffer.
   llvm::MemoryBufferRef buffer;
 
   /// The optional owning source manager, which when present may be used to
   /// extend the lifetime of the input buffer.
LogicalResult BytecodeReader::Impl::parseResourceSection(
   DialectReader dialectReader(attrTypeReader, stringReader, resourceReader,
                               reader);
   return resourceReader.initialize(fileLoc, config, dialects, stringReader,
                                    *resourceData, *resourceOffsetData,
                                    dialectReader, bufferOwnerRef);
 }
 
 //===----------------------------------------------------------------------===//
+// UseListOrder Helpers
+
+FailureOr<BytecodeReader::Impl::UseListMapT>
+BytecodeReader::Impl::parseUseListOrderForRange(EncodingReader &reader,
+                                                uint64_t numResults) {
+  BytecodeReader::Impl::UseListMapT map;
+  uint64_t numValuesToRead = 1;
+  if (numResults > 1 && failed(reader.parseVarInt(numValuesToRead)))
+    return failure();
+
+  for (size_t valueIdx = 0; valueIdx < numValuesToRead; valueIdx++) {
+    uint64_t resultIdx = 0;
+    if (numResults > 1 && failed(reader.parseVarInt(resultIdx)))
+      return failure();
+
+    uint64_t numValues;
+    bool indexPairEncoding;
+    if (failed(reader.parseVarIntWithFlag(numValues, indexPairEncoding)))
+      return failure();
+
+    SmallVector<unsigned, 4> useListOrders;
+    for (size_t idx = 0; idx < numValues; idx++) {
+      uint64_t index;
+      if (failed(reader.parseVarInt(index)))
+        return failure();
+      useListOrders.push_back(index);
+    }
+
+    // Store in a map the result index
+    map.try_emplace(resultIdx, UseListOrderStorage(indexPairEncoding,
+                                                   std::move(useListOrders)));
+  }
+
+  return map;
+}
+
+/// Sorts each use according to the order specified in the use-list parsed. If
+/// the custom use-list is not found, this means that the order needs to be
+/// consistent with the reverse pre-order walk of the IR. If multiple uses lie
+/// on the same operation, the order will follow the reverse operand number
+/// ordering.
+LogicalResult BytecodeReader::Impl::sortUseListOrder(Value value) {
+  // Early return for trivial use-lists.
+  if (value.use_empty() || value.hasOneUse())
+    return success();
+
+  bool hasIncomingOrder =
+      valueToUseListMap.contains(value.getAsOpaquePointer());
+
+  // Compute the current order of the use-list with respect to the global
+  // ordering. Detect if the order is already sorted while doing so.
+  bool alreadySorted = true;
+  auto &firstUse = *value.use_begin();
+  uint64_t prevID =
+      bytecode::getUseID(firstUse, operationIDs.at(firstUse.getOwner()));
+  llvm::SmallVector<std::pair<unsigned, uint64_t>> currentOrder = {{0, prevID}};
+  for (auto item : llvm::drop_begin(llvm::enumerate(value.getUses()))) {
+    uint64_t currentID = bytecode::getUseID(
+        item.value(), operationIDs.at(item.value().getOwner()));
+    alreadySorted &= prevID > currentID;
+    currentOrder.push_back({item.index(), currentID});
+    prevID = currentID;
+  }
+
+  // If the order is already sorted, and there wasn't a custom order to apply
+  // from the bytecode file, we are done.
+  if (alreadySorted && !hasIncomingOrder)
+    return success();
+
+  // If not already sorted, sort the indices of the current order by descending
+  // useIDs.
+  if (!alreadySorted)
+    std::sort(
+        currentOrder.begin(), currentOrder.end(),
+        [](auto elem1, auto elem2) { return elem1.second > elem2.second; });
+
+  if (!hasIncomingOrder) {
+    // If the bytecode file did not contain any custom use-list order, it means
+    // that the order was descending useID. Hence, shuffle by the first index
+    // of the `currentOrder` pair.
+    SmallVector<unsigned> shuffle = SmallVector<unsigned>(
+        llvm::map_range(currentOrder, [&](auto item) { return item.first; }));
+    value.shuffleUseList(shuffle);
+    return success();
+  }
+
+  // Pull the custom order info from the map.
+  UseListOrderStorage customOrder =
+      valueToUseListMap.at(value.getAsOpaquePointer());
+  SmallVector<unsigned, 4> shuffle = std::move(customOrder.indices);
+  uint64_t numUses =
+      std::distance(value.getUses().begin(), value.getUses().end());
+
+  // 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.
+  if (customOrder.isIndexPairEncoding) {
+    // Return failure if the number of indices was not representing pairs.
+    if (shuffle.size() & 1)
+      return failure();
+
+    SmallVector<unsigned, 4> newShuffle(numUses);
+    size_t idx = 0;
+    std::iota(newShuffle.begin(), newShuffle.end(), idx);
+    for (idx = 0; idx < shuffle.size(); idx += 2)
+      newShuffle[shuffle[idx]] = shuffle[idx + 1];
+
+    shuffle = std::move(newShuffle);
+  }
+
+  // Make sure that the indices represent a valid mapping. That is, the sum of
+  // all the values needs to be equal to (numUses - 1) * numUses / 2, and no
+  // duplicates are allowed in the list.
+  DenseSet<unsigned> set;
+  uint64_t accumulator = 0;
+  for (const auto &elem : shuffle) {
+    if (set.contains(elem))
+      return failure();
+    accumulator += elem;
+    set.insert(elem);
+  }
+  if (numUses != shuffle.size() ||
+      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]; }));
+  value.shuffleUseList(shuffle);
+  return success();
+}
+
+LogicalResult BytecodeReader::Impl::processUseLists(Operation *topLevelOp) {
+  // Precompute operation IDs according to the pre-order walk of the IR. We
+  // can't do this while parsing since parseRegions ordering is not strictly
+  // equal to the pre-order walk.
+  unsigned operationID = 0;
+  topLevelOp->walk<mlir::WalkOrder::PreOrder>(
+      [&](Operation *op) { operationIDs.try_emplace(op, operationID++); });
+
+  auto blockWalk = topLevelOp->walk([this](Block *block) {
+    for (auto arg : block->getArguments())
+      if (failed(sortUseListOrder(arg)))
+        return WalkResult::interrupt();
+    return WalkResult::advance();
+  });
+
+  auto resultWalk = topLevelOp->walk([this](Operation *op) {
+    for (auto result : op->getResults())
+      if (failed(sortUseListOrder(result)))
+        return WalkResult::interrupt();
+    return WalkResult::advance();
+  });
+
+  return failure(blockWalk.wasInterrupted() || resultWalk.wasInterrupted());
+}
+
+//===----------------------------------------------------------------------===//
 // IR Section
 
 LogicalResult
 BytecodeReader::Impl::parseIRSection(ArrayRef<uint8_t> sectionData,
                                      Block *block) {
   EncodingReader reader(sectionData, fileLoc);
 
   // A stack of operation regions currently being read from the bytecode.
BytecodeReader::Impl::parseIRSection(ArrayRef<uint8_t> sectionData,
   while (!regionStack.empty())
     if (failed(parseRegions(regionStack, regionStack.back())))
       return failure();
   if (!forwardRefOps.empty()) {
     return reader.emitError(
         "not all forward unresolved forward operand references");
   }
 
+  // Sort use-lists according to what specified in bytecode.
+  if (failed(processUseLists(*moduleOp)))
+    return reader.emitError(
+        "parsed use-list orders were invalid and could not be applied");
+
   // Resolve dialect version.
   for (const BytecodeDialect &byteCodeDialect : dialects) {
     // Parsing is complete, give an opportunity to each dialect to visit the
     // IR and perform upgrades.
     if (!byteCodeDialect.loadedVersion)
       continue;
     if (byteCodeDialect.interface &&
         failed(byteCodeDialect.interface->upgradeFromVersion(
if (opMask & bytecode::OpEncodingMask::kHasSuccessors) {
     opState.successors.resize(numSuccs);
     for (int i = 0, e = numSuccs; i < e; ++i) {
       if (failed(parseEntry(reader, readState.curBlocks, opState.successors[i],
                             "successor")))
         return failure();
     }
   }
 
+  /// Parse the use-list orders for the results of the operation. Use-list
+  /// orders are available since version 3 of the bytecode.
+  std::optional<UseListMapT> resultIdxToUseListMap = std::nullopt;
+  if (version > 2 && (opMask & bytecode::OpEncodingMask::kHasUseListOrders)) {
+    size_t numResults = opState.types.size();
+    auto parseResult = parseUseListOrderForRange(reader, numResults);
+    if (failed(parseResult))
+      return failure();
+    resultIdxToUseListMap = std::move(*parseResult);
+  }
+
   /// Parse the regions of the operation.
   if (opMask & bytecode::OpEncodingMask::kHasInlineRegions) {
     uint64_t numRegions;
     if (failed(reader.parseVarIntWithFlag(numRegions, isIsolatedFromAbove)))
       return failure();
 
     opState.regions.reserve(numRegions);
     for (int i = 0, e = numRegions; i < e; ++i)
       opState.regions.push_back(std::make_unique<Region>());
   }
 
   // Create the operation at the back of the current block.
   Operation *op = Operation::create(opState);
   readState.curBlock->push_back(op);
 
   // If the operation had results, update the value references.
   if (op->getNumResults() && failed(defineValues(reader, op->getResults())))
     return failure();
 
+  /// Store a map for every value that received a custom use-list order from the
+  /// bytecode file.
+  if (resultIdxToUseListMap.has_value()) {
+    for (size_t idx = 0; idx < op->getNumResults(); idx++) {
+      if (resultIdxToUseListMap->contains(idx)) {
+        valueToUseListMap.try_emplace(op->getResult(idx).getAsOpaquePointer(),
+                                      resultIdxToUseListMap->at(idx));
+      }
+    }
+  }
   return op;
 }
 
 LogicalResult BytecodeReader::Impl::parseRegion(RegionReadState &readState) {
   EncodingReader &reader = *readState.reader;
 
   // Parse the number of blocks in the region.
   uint64_t numBlocks;
BytecodeReader::Impl::parseBlockHeader(EncodingReader &reader,
   bool hasArgs;
   if (failed(reader.parseVarIntWithFlag(readState.numOpsRemaining, hasArgs)))
     return failure();
 
   // Parse the arguments of the block.
   if (hasArgs && failed(parseBlockArguments(reader, &*readState.curBlock)))
     return failure();
 
+  // Uselist orders are available since version 3 of the bytecode.
+  if (version < 3)
+    return success();
+
+  uint8_t hasUseListOrders = 0;
+  if (hasArgs && failed(reader.parseByte(hasUseListOrders)))
+    return failure();
+
+  if (!hasUseListOrders)
+    return success();
+
+  Block &blk = *readState.curBlock;
+  auto argIdxToUseListMap =
+      parseUseListOrderForRange(reader, blk.getNumArguments());
+  if (failed(argIdxToUseListMap) || argIdxToUseListMap->empty())
+    return failure();
+
+  for (size_t idx = 0; idx < blk.getNumArguments(); idx++)
+    if (argIdxToUseListMap->contains(idx))
+      valueToUseListMap.try_emplace(blk.getArgument(idx).getAsOpaquePointer(),
+                                    argIdxToUseListMap->at(idx));
+
   // We don't parse the operations of the block here, that's done elsewhere.
   return success();
 }
 
 LogicalResult BytecodeReader::Impl::parseBlockArguments(EncodingReader &reader,
                                                         Block *block) {
   // Parse the value ID for the first argument, and the number of arguments.
   uint64_t numArgs;

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

 //===- BytecodeWriter.cpp - MLIR Bytecode Writer --------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Bytecode/BytecodeWriter.h"
-#include "../Encoding.h"
 #include "IRNumbering.h"
 #include "mlir/Bytecode/BytecodeImplementation.h"
+#include "mlir/Bytecode/Encoding.h"
 #include "mlir/IR/OpImplementation.h"
 #include "llvm/ADT/CachedHashString.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallString.h"
 
 #define DEBUG_TYPE "mlir-bytecode-writer"
 
 using namespace mlir;
private:
   void writeResourceSection(Operation *op, EncodingEmitter &emitter);
 
   //===--------------------------------------------------------------------===//
   // Strings
 
   void writeStringSection(EncodingEmitter &emitter);
 
   //===--------------------------------------------------------------------===//
+  // Helpers
+
+  void writeUseListOrders(EncodingEmitter &emitter, uint8_t &opEncodingMask,
+                          ValueRange range);
+
+  //===--------------------------------------------------------------------===//
   // Fields
 
   /// The builder used for the string section.
   StringSectionBuilder stringSection;
 
   /// The IR numbering state generated for the root operation.
   IRNumberingState numberingState;
 
void BytecodeWriter::writeBlock(EncodingEmitter &emitter, Block *block) {
 
   // Emit the arguments of the block.
   if (hasArgs) {
     emitter.emitVarInt(args.size());
     for (BlockArgument arg : args) {
       emitter.emitVarInt(numberingState.getNumber(arg.getType()));
       emitter.emitVarInt(numberingState.getNumber(arg.getLoc()));
     }
+    if (config.bytecodeVersion > 2) {
+      uint64_t maskOffset = emitter.size();
+      uint8_t encodingMask = 0;
+      emitter.emitByte(0);
+      writeUseListOrders(emitter, encodingMask, args);
+      if (encodingMask)
+        emitter.patchByte(maskOffset, encodingMask);
+    }
   }
 
   // Emit the operations within the block.
   for (Operation &op : *block)
     writeOp(emitter, &op);
 }
 
 void BytecodeWriter::writeOp(EncodingEmitter &emitter, Operation *op) {
void BytecodeWriter::writeOp(EncodingEmitter &emitter, Operation *op) {
   // Emit the successors of the operation.
   if (unsigned numSuccessors = op->getNumSuccessors()) {
     opEncodingMask |= bytecode::OpEncodingMask::kHasSuccessors;
     emitter.emitVarInt(numSuccessors);
     for (Block *successor : op->getSuccessors())
       emitter.emitVarInt(numberingState.getNumber(successor));
   }
 
+  // Emit the use-list orders to bytecode, so we can reconstruct the same order
+  // at parsing.
+  if (config.bytecodeVersion > 2)
+    writeUseListOrders(emitter, opEncodingMask, ValueRange(op->getResults()));
+
   // Check for regions.
   unsigned numRegions = op->getNumRegions();
   if (numRegions)
     opEncodingMask |= bytecode::OpEncodingMask::kHasInlineRegions;
 
   // Update the mask for the operation.
   emitter.patchByte(maskOffset, opEncodingMask);
 
for (Region &region : op->getRegions()) {
 
       EncodingEmitter regionEmitter;
       writeRegion(regionEmitter, &region);
       emitter.emitSection(bytecode::Section::kIR, std::move(regionEmitter));
     }
   }
 }
 
+void BytecodeWriter::writeUseListOrders(EncodingEmitter &emitter,
+                                        uint8_t &opEncodingMask,
+                                        ValueRange range) {
+  // Loop over the results and store the use-list order per result index.
+  DenseMap<unsigned, llvm::SmallVector<unsigned>> map;
+  for (auto item : llvm::enumerate(range)) {
+    auto value = item.value();
+    // No need to store a custom use-list order if the result does not have
+    // multiple uses.
+    if (value.use_empty() || value.hasOneUse())
+      continue;
+
+    // For each result, assemble the list of pairs (use-list-index,
+    // global-value-index). While doing so, detect if the global-value-index is
+    // already ordered with respect to the use-list-index.
+    bool alreadyOrdered = true;
+    auto &firstUse = *value.use_begin();
+    uint64_t prevID = bytecode::getUseID(
+        firstUse, numberingState.getNumber(firstUse.getOwner()));
+    llvm::SmallVector<std::pair<unsigned, uint64_t>> useListPairs(
+        {{0, prevID}});
+
+    for (auto use : llvm::drop_begin(llvm::enumerate(value.getUses()))) {
+      uint64_t currentID = bytecode::getUseID(
+          use.value(), numberingState.getNumber(use.value().getOwner()));
+      // The use-list order achieved when building the IR at parsing always
+      // pushes new uses on front. Hence, if the order by unique ID is
+      // monotonically decreasing, a roundtrip to bytecode preserves such order.
+      alreadyOrdered &= (prevID > currentID);
+      useListPairs.push_back({use.index(), currentID});
+      prevID = currentID;
+    }
+
+    // Do not emit if the order is already sorted.
+    if (alreadyOrdered)
+      continue;
+
+    // Sort the use indices by the unique ID indices in descending order.
+    std::sort(
+        useListPairs.begin(), useListPairs.end(),
+        [](auto elem1, auto elem2) { return elem1.second > elem2.second; });
+
+    map.try_emplace(item.index(), llvm::map_range(useListPairs, [](auto elem) {
+                      return elem.first;
+                    }));
+  }
+
+  if (map.empty())
+    return;
+
+  opEncodingMask |= bytecode::OpEncodingMask::kHasUseListOrders;
+  // Emit the number of results that have a custom use-list order if the number
+  // of results is greater than one.
+  if (range.size() != 1)
+    emitter.emitVarInt(map.size());
+
+  for (const auto &item : map) {
+    auto resultIdx = item.getFirst();
+    auto useListOrder = item.getSecond();
+
+    // Compute the number of uses that are actually shuffled. If those are less
+    // than half of the total uses, encoding the index pair `(src, dst)` is more
+    // space efficient.
+    size_t shuffledElements =
+        llvm::count_if(llvm::enumerate(useListOrder),
+                       [](auto item) { return item.index() != item.value(); });
+    bool indexPairEncoding = shuffledElements < (useListOrder.size() / 2);
+
+    // For single result, we don't need to store the result index.
+    if (range.size() != 1)
+      emitter.emitVarInt(resultIdx);
+
+    if (indexPairEncoding) {
+      emitter.emitVarIntWithFlag(shuffledElements * 2, indexPairEncoding);
+      for (auto pair : llvm::enumerate(useListOrder)) {
+        if (pair.index() != pair.value()) {
+          emitter.emitVarInt(pair.value());
+          emitter.emitVarInt(pair.index());
+        }
+      }
+    } else {
+      emitter.emitVarIntWithFlag(useListOrder.size(), indexPairEncoding);
+      for (const auto &index : useListOrder)
+        emitter.emitVarInt(index);
+    }
+  }
+}
+
 void BytecodeWriter::writeRegion(EncodingEmitter &emitter, Region *region) {
   // If the region is empty, we only need to emit the number of blocks (which is
   // zero).
   if (region->empty())
     return emitter.emitVarInt(/*numBlocks*/ 0);
 
   // Emit the number of blocks and values within the region.
   unsigned numBlocks, numValues;

mlir/lib/Bytecode/Writer/IRNumbering.h

public:
   unsigned getNumber(Attribute attr) {
     assert(attrs.count(attr) && "attribute not numbered");
     return attrs[attr]->number;
   }
   unsigned getNumber(Block *block) {
     assert(blockIDs.count(block) && "block not numbered");
     return blockIDs[block];
   }
+  unsigned getNumber(Operation *op) {
+    assert(operationIDs.count(op) && "operation not numbered");
+    return operationIDs[op];
+  }
   unsigned getNumber(OperationName opName) {
     assert(opNames.count(opName) && "opName not numbered");
     return opNames[opName]->number;
   }
   unsigned getNumber(Type type) {
     assert(types.count(type) && "type not numbered");
     return types[type]->number;
   }
private:
 
   /// Allocators used for the various numbering entries.
   llvm::SpecificBumpPtrAllocator<AttributeNumbering> attrAllocator;
   llvm::SpecificBumpPtrAllocator<DialectNumbering> dialectAllocator;
   llvm::SpecificBumpPtrAllocator<OpNameNumbering> opNameAllocator;
   llvm::SpecificBumpPtrAllocator<DialectResourceNumbering> resourceAllocator;
   llvm::SpecificBumpPtrAllocator<TypeNumbering> typeAllocator;
 
-  /// The value ID for each Block and Value.
+  /// The value ID for each Operation, Block and Value.
+  DenseMap<Operation *, unsigned> operationIDs;
   DenseMap<Block *, unsigned> blockIDs;
   DenseMap<Value, unsigned> valueIDs;
 
   /// The number of operations in each block.
   DenseMap<Block *, unsigned> blockOperationCounts;
 
   /// A map from region to the number of blocks and values within that region.
   DenseMap<Region *, std::pair<unsigned, unsigned>> regionBlockValueCounts;

mlir/lib/Bytecode/Writer/IRNumbering.cpp

 //===- IRNumbering.cpp - MLIR Bytecode IR numbering -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "IRNumbering.h"
-#include "../Encoding.h"
 #include "mlir/Bytecode/BytecodeImplementation.h"
-#include "mlir/Bytecode/BytecodeWriter.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
 #include "llvm/Support/ErrorHandling.h"
 
 using namespace mlir;
 using namespace mlir::bytecode::detail;
 
static void groupByDialectPerByte(T range) {
   }
 
   // Assign the entry numbers based on the sort order.
   for (auto [idx, value] : llvm::enumerate(range))
     value->number = idx;
 }
 
 IRNumberingState::IRNumberingState(Operation *op) {
+  // Compute a global operation ID numbering according to the pre-order walk of
+  // the IR. This is used as reference to construct use-list orders.
+  unsigned operationID = 0;
+  op->walk<WalkOrder::PreOrder>(
+      [&](Operation *op) { operationIDs.try_emplace(op, operationID++); });
+
   // Number the root operation.
   number(*op);
 
   // Push all of the regions of the root operation onto the worklist.
   SmallVector<std::pair<Region *, unsigned>, 8> numberContext;
   for (Region &region : op->getRegions())
     numberContext.emplace_back(&region, nextValueID);
 

mlir/lib/IR/Value.cpp

 
 /// Returns true if the value is used outside of the given block.
 bool Value::isUsedOutsideOfBlock(Block *block) {
   return llvm::any_of(getUsers(), [block](Operation *user) {
     return user->getBlock() != block;
   });
 }
 
+/// Shuffles the use-list order according to the provided indices.
+void Value::shuffleUseList(ArrayRef<unsigned> indices) {
+  getImpl()->shuffleUseList(indices);
+}
+
 //===----------------------------------------------------------------------===//
 // OpResult
 //===----------------------------------------------------------------------===//
 
 /// Returns the parent operation of this trailing result.
 Operation *OpResultImpl::getOwner() const {
   // We need to do some arithmetic to get the operation pointer. Results are
   // stored in reverse order before the operation, so move the trailing owner up
   // to the start of the array. A rough diagram of the memory layout is:
   //
   // | Out-of-Line results | Inline results | Operation |
   //
   // Given that the results are reverse order we use the result number to know
   // how far to jump to get to the operation. So if we are currently the 0th
   // result, the layout would be:

mlir/test/Bytecode/invalid/invalid-structure.mlir

 // This file contains various failure test cases related to the structure of
 // a bytecode file.
 
 // Bytecode currently does not support big-endian platforms
 // UNSUPPORTED: target=s390x-{{.*}}
 
 //===--------------------------------------------------------------------===//
 // Version
 //===--------------------------------------------------------------------===//
 
 // RUN: not mlir-opt %S/invalid-structure-version.mlirbc 2>&1 | FileCheck %s --check-prefix=VERSION
-// VERSION: bytecode version 127 is newer than the current version 2
+// VERSION: bytecode version 127 is newer than the current version 3
 
 //===--------------------------------------------------------------------===//
 // Producer
 //===--------------------------------------------------------------------===//
 
 // RUN: not mlir-opt %S/invalid-structure-producer.mlirbc 2>&1 | FileCheck %s --check-prefix=PRODUCER
 // PRODUCER: malformed null-terminated string, no null character found
 

mlir/test/Bytecode/uselist_orders.mlir

+// RUN: mlir-opt %s -split-input-file --test-verify-uselistorder -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
+//      uselist order are not preserved when doing a roundtrip to bytecode. The
+//      test needs to verify diagnostics to be functional.
+
+func.func @base_test(%arg0 : i32) -> i32 {
+  %0 = arith.constant 45 : i32
+  %1 = arith.constant 46 : i32
+  %2 = "test.addi"(%arg0, %arg0) : (i32, i32) -> i32
+  %3 = "test.addi"(%2, %0) : (i32, i32) -> i32
+  %4 = "test.addi"(%2, %1) : (i32, i32) -> i32
+  %5 = "test.addi"(%3, %4) : (i32, i32) -> i32
+  %6 = "test.addi"(%5, %4) : (i32, i32) -> i32
+  %7 = "test.addi"(%6, %4) : (i32, i32) -> i32
+  return %7 : i32
+}
+
+// -----
+
+func.func @test_with_multiple_uses_in_same_op(%arg0 : i32) -> i32 {
+  %0 = arith.constant 45 : i32
+  %1 = arith.constant 46 : i32
+  %2 = "test.addi"(%arg0, %arg0) : (i32, i32) -> i32
+  %3 = "test.addi"(%2, %0) : (i32, i32) -> i32
+  %4 = "test.addi"(%2, %1) : (i32, i32) -> i32
+  %5 = "test.addi"(%2, %2) : (i32, i32) -> i32
+  %6 = "test.addi"(%3, %4) : (i32, i32) -> i32
+  %7 = "test.addi"(%6, %5) : (i32, i32) -> i32
+  %8 = "test.addi"(%7, %4) : (i32, i32) -> i32
+  %9 = "test.addi"(%8, %4) : (i32, i32) -> i32
+  return %9 : i32
+}
+
+// -----
+
+func.func @test_with_multiple_block_arg_uses(%arg0 : i32) -> i32 {
+  %0 = arith.constant 45 : i32
+  %1 = arith.constant 46 : i32
+  %2 = "test.addi"(%arg0, %arg0) : (i32, i32) -> i32
+  %3 = "test.addi"(%2, %arg0) : (i32, i32) -> i32
+  %4 = "test.addi"(%2, %1) : (i32, i32) -> i32
+  %5 = "test.addi"(%2, %2) : (i32, i32) -> i32
+  %6 = "test.addi"(%3, %4) : (i32, i32) -> i32
+  %7 = "test.addi"(%6, %5) : (i32, i32) -> i32
+  %8 = "test.addi"(%7, %4) : (i32, i32) -> i32
+  %9 = "test.addi"(%8, %4) : (i32, i32) -> i32
+  return %9 : i32
+}
+
+// -----
+
+// Test that use-lists in region with no dominance are preserved
+test.graph_region {
+  %0 = "test.foo"(%1) : (i32) -> i32
+  test.graph_region attributes {a} {
+    %a = "test.a"(%b) : (i32) -> i32
+    %b = "test.b"(%2) : (i32) -> i32
+  }
+  %1 = "test.bar"(%2) : (i32) -> i32
+  %2 = "test.baz"() : () -> i32
+}

mlir/test/lib/IR/CMakeLists.txt

add_mlir_library(MLIRTestIR
   TestPrintDefUse.cpp
   TestPrintInvalid.cpp
   TestPrintNesting.cpp
   TestSideEffects.cpp
   TestSlicing.cpp
   TestSymbolUses.cpp
   TestRegions.cpp
   TestTypes.cpp
+  TestUseListOrders.cpp
   TestVisitors.cpp
   TestVisitorsGeneric.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRPass
   MLIRBytecodeReader

mlir/test/lib/IR/TestUseListOrders.cpp

+//===- TestPrintDefUse.cpp - Passes to illustrate the IR def-use chains ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Bytecode/BytecodeWriter.h"
+#include "mlir/Bytecode/Encoding.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Parser/Parser.h"
+#include "mlir/Pass/Pass.h"
+
+#include <numeric>
+#include <random>
+
+using namespace mlir;
+
+namespace {
+/// This pass tests that:
+/// 1) we can shuffle use-lists correctly;
+/// 2) use-list orders are preserved after a roundtrip to bytecode.
+class TestPreserveUseListOrders
+    : public PassWrapper<TestPreserveUseListOrders, OperationPass<ModuleOp>> {
+public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPreserveUseListOrders)
+
+  TestPreserveUseListOrders() = default;
+  TestPreserveUseListOrders(const TestPreserveUseListOrders &pass)
+      : PassWrapper(pass) {}
+  StringRef getArgument() const final { return "test-verify-uselistorder"; }
+  StringRef getDescription() const final {
+    return "Verify that roundtripping the IR to bytecode preserves the order "
+           "of the uselists";
+  }
+  Option<unsigned> rngSeed{*this, "rng-seed",
+                           llvm::cl::desc("Specify an input random seed"),
+                           llvm::cl::init(1)};
+  void runOnOperation() override {
+    // Clone the module so that we can plug in this pass to any other
+    // independently.
+    auto cloneModule = getOperation().clone();
+
+    // 1. Compute the op numbering of the module.
+    computeOpNumbering(cloneModule);
+
+    // 2. Loop over all the values and shuffle the uses. While doing so, check
+    // that each shuffle is correct.
+    if (failed(shuffleUses(cloneModule)))
+      return signalPassFailure();
+
+    // 3. Do a bytecode roundtrip to version 3, which supports use-list order
+    // preservation.
+    auto roundtripModuleOr = doRoundtripToBytecode(cloneModule, 3);
+    // If the bytecode roundtrip failed, try to roundtrip the original module
+    // to version 2, which does not support use-list. If this also fails, the
+    // original module had an issue unrelated to uselists.
+    if (failed(roundtripModuleOr)) {
+      auto testModuleOr = doRoundtripToBytecode(getOperation(), 2);
+      if (failed(testModuleOr))
+        return;
+
+      return signalPassFailure();
+    }
+
+    // 4. Recompute the op numbering on the new module. The numbering should be
+    // the same as (1), but on the new operation pointers.
+    computeOpNumbering(roundtripModuleOr->get());
+
+    // 5. Loop over all the values and verify that the use-list is consistent
+    // with the post-shuffle order of step (2).
+    if (failed(verifyUseListOrders(roundtripModuleOr->get())))
+      return signalPassFailure();
+  }
+
+private:
+  FailureOr<OwningOpRef<Operation *>> doRoundtripToBytecode(Operation *module,
+                                                            uint32_t version) {
+    std::string str;
+    llvm::raw_string_ostream m(str);
+    BytecodeWriterConfig config;
+    config.setDesiredBytecodeVersion(version);
+    if (failed(writeBytecodeToFile(module, m, config)))
+      return failure();
+
+    ParserConfig parseConfig(&getContext(), /*verifyAfterParse=*/true);
+    auto newModuleOp = parseSourceString(StringRef(str), parseConfig);
+    if (!newModuleOp.get())
+      return failure();
+    return newModuleOp;
+  }
+
+  /// Compute an ordered numbering for all the operations in the IR.
+  void computeOpNumbering(Operation *topLevelOp) {
+    uint32_t operationID = 0;
+    opNumbering.clear();
+    topLevelOp->walk<mlir::WalkOrder::PreOrder>(
+        [&](Operation *op) { opNumbering.try_emplace(op, operationID++); });
+  }
+
+  template <typename ValueT>
+  SmallVector<uint64_t> getUseIDs(ValueT val) {
+    return SmallVector<uint64_t>(llvm::map_range(val.getUses(), [&](auto &use) {
+      return bytecode::getUseID(use, opNumbering.at(use.getOwner()));
+    }));
+  }
+
+  LogicalResult shuffleUses(Operation *topLevelOp) {
+    uint32_t valueID = 0;
+    /// Permute randomly the use-list of each value. It is guaranteed that at
+    /// least one pair of the use list is permuted.
+    auto doShuffleForRange = [&](ValueRange range) -> LogicalResult {
+      for (auto val : range) {
+        if (val.use_empty() || val.hasOneUse())
+          continue;
+
+        /// Get a valid index permutation for the uses of value.
+        SmallVector<unsigned> permutation = getRandomPermutation(val);
+
+        /// Store original order and verify that the shuffle was applied
+        /// correctly.
+        auto useIDs = getUseIDs(val);
+
+        /// Apply shuffle to the uselist.
+        val.shuffleUseList(permutation);
+
+        /// Get the new order and verify the shuffle happened correctly.
+        auto permutedIDs = getUseIDs(val);
+        if (permutedIDs.size() != useIDs.size())
+          return failure();
+        for (size_t idx = 0; idx < permutation.size(); idx++)
+          if (useIDs[idx] != permutedIDs[permutation[idx]])
+            return failure();
+
+        referenceUseListOrder.try_emplace(
+            valueID++, llvm::map_range(val.getUses(), [&](auto &use) {
+              return bytecode::getUseID(use, opNumbering.at(use.getOwner()));
+            }));
+      }
+      return success();
+    };
+
+    return walkOverValues(topLevelOp, doShuffleForRange);
+  }
+
+  LogicalResult verifyUseListOrders(Operation *topLevelOp) {
+    uint32_t valueID = 0;
+    /// Check that the use-list for the value range matches the one stored in
+    /// the reference.
+    auto doValidationForRange = [&](ValueRange range) -> LogicalResult {
+      for (auto val : range) {
+        if (val.use_empty() || val.hasOneUse())
+          continue;
+        auto referenceOrder = referenceUseListOrder.at(valueID++);
+        for (auto [use, referenceID] :
+             llvm::zip(val.getUses(), referenceOrder)) {
+          uint64_t uniqueID =
+              bytecode::getUseID(use, opNumbering.at(use.getOwner()));
+          if (uniqueID != referenceID) {
+            use.getOwner()->emitError()
+                << "found use-list order mismatch for value: " << val;
+            return failure();
+          }
+        }
+      }
+      return success();
+    };
+
+    return walkOverValues(topLevelOp, doValidationForRange);
+  }
+
+  /// Walk over blocks and operations and execute a callable over the ranges of
+  /// operands/results respectively.
+  template <typename FuncT>
+  LogicalResult walkOverValues(Operation *topLevelOp, FuncT callable) {
+    auto blockWalk = topLevelOp->walk([&](Block *block) {
+      if (failed(callable(block->getArguments())))
+        return WalkResult::interrupt();
+      return WalkResult::advance();
+    });
+
+    if (blockWalk.wasInterrupted())
+      return failure();
+
+    auto resultsWalk = topLevelOp->walk([&](Operation *op) {
+      if (failed(callable(op->getResults())))
+        return WalkResult::interrupt();
+      return WalkResult::advance();
+    });
+
+    return failure(resultsWalk.wasInterrupted());
+  }
+
+  /// Creates a random permutation of the uselist order chain of the provided
+  /// value.
+  SmallVector<unsigned> getRandomPermutation(Value value) {
+    size_t numUses = std::distance(value.use_begin(), value.use_end());
+    SmallVector<unsigned> permutation(numUses);
+    unsigned zero = 0;
+    std::iota(permutation.begin(), permutation.end(), zero);
+    auto rng = std::default_random_engine(rngSeed);
+    std::shuffle(permutation.begin(), permutation.end(), rng);
+    return permutation;
+  }
+
+  /// Map each value to its use-list order encoded with unique use IDs.
+  DenseMap<uint32_t, SmallVector<uint64_t>> referenceUseListOrder;
+
+  /// Map each operation to its global ID.
+  DenseMap<Operation *, uint32_t> opNumbering;
+};
+} // namespace
+
+namespace mlir {
+void registerTestPreserveUseListOrders() {
+  PassRegistration<TestPreserveUseListOrders>();
+}
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestLoopPermutationPass();
 void registerTestMatchers();
 void registerTestOperationEqualPass();
 void registerTestPrintDefUsePass();
 void registerTestPrintInvalidPass();
 void registerTestPrintNestingPass();
+void registerTestPreserveUseListOrders();
 void registerTestReducer();
 void registerTestSpirvEntryPointABIPass();
 void registerTestSpirvModuleCombinerPass();
 void registerTestTraitsPass();
 void registerTosaTestQuantUtilAPIPass();
 void registerVectorizerTestPass();
 
 namespace test {
void registerTestPasses() {
   registerTestFunc();
   registerTestGpuMemoryPromotionPass();
   registerTestLoopPermutationPass();
   registerTestMatchers();
   registerTestOperationEqualPass();
   registerTestPrintDefUsePass();
   registerTestPrintInvalidPass();
   registerTestPrintNestingPass();
+  registerTestPreserveUseListOrders();
   registerTestReducer();
   registerTestSpirvEntryPointABIPass();
   registerTestSpirvModuleCombinerPass();
   registerTestTraitsPass();
   registerVectorizerTestPass();
   registerTosaTestQuantUtilAPIPass();
 
   mlir::test::registerCommutativityUtils();