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

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

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

783

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

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

mlir/include/mlir/IR/AsmState.h

463

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

1198–1199

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

1235

Please document the struct.

1240

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

1637–1640

[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

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

mlir/include/mlir/IR/AsmState.h

463

Removed. Thanks.

mlir/include/mlir/IR/Value.h

192

Should be fixed now.

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

1198–1199

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

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

783

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

199 ↗

(On Diff #521408)

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

450 ↗

(On Diff #521408)

(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

199 ↗

(On Diff #521408)

Agreed, added!

[mehdi_amini]

Can you rebase? The version was just bumped.

[mehdi_amini]

mlir/docs/BytecodeFormat.md

399 ↗

(On Diff #524069)

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

448 ↗

(On Diff #524069)

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/include/mlir/Bytecode/BytecodeWriter.h

 /// This class contains the configuration used for the bytecode writer. It
 /// controls various aspects of bytecode generation, and contains all of the
 /// various bytecode writer hooks.
 class BytecodeWriterConfig {
 public:
   /// `producer` is an optional string that can be used to identify the producer
   /// of the bytecode when reading. It has no functional effect on the bytecode
   /// serialization.
-  BytecodeWriterConfig(StringRef producer = "MLIR" LLVM_VERSION_STRING);
+  /// `shouldEmitUseListOrder` is an optional boolean that controls whether to
+  /// emit use-list order info to the bytecode to preserve walk behaviours after
+  /// a bytecode roundtrip.
+  BytecodeWriterConfig(bool shouldEmitUseListOrder = false,
+                       StringRef producer = "MLIR" LLVM_VERSION_STRING);
   /// `map` is a fallback resource map, which when provided will attach resource
   /// printers for the fallback resources within the map.
   BytecodeWriterConfig(FallbackAsmResourceMap &map,
+                       bool shouldEmitUseListOrder = false,
                        StringRef producer = "MLIR" LLVM_VERSION_STRING);
   ~BytecodeWriterConfig();
 
   /// An internal implementation class that contains the state of the
   /// configuration.
   struct Impl;
 
   /// Return an instance of the internal implementation.

mlir/include/mlir/IR/AsmState.h

 //===----------------------------------------------------------------------===//
 
 /// This class represents a configuration for the MLIR assembly parser. It
 /// contains all of the necessary state to parse a MLIR source file.
 class ParserConfig {
 public:
   /// Construct a parser configuration with the given context.
   /// `verifyAfterParse` indicates if the IR should be verified after parsing.
+  /// `preserveUseListOrders` indicates if the in memory use-list orders should
+  /// be preserved when interacting with bytecode.
   /// `fallbackResourceMap` is an optional fallback handler that can be used to
   /// parse external resources not explicitly handled by another parser.
   ParserConfig(MLIRContext *context, bool verifyAfterParse = true,
+               bool preserveUseListOrders = false,
                FallbackAsmResourceMap *fallbackResourceMap = nullptr)
       : context(context), verifyAfterParse(verifyAfterParse),
+        preserveUseListOrders(preserveUseListOrders),
         fallbackResourceMap(fallbackResourceMap) {
     assert(context && "expected valid MLIR context");
   }
 
   /// Return the MLIRContext to be used when parsing.
   MLIRContext *getContext() const { return context; }
 
   /// Returns if the parser should verify the IR after parsing.
   bool shouldVerifyAfterParse() const { return verifyAfterParse; }
 
+  /// Returns if the parser should preserve the use-list if available in
+  /// bytecode.
+  bool shouldPreserveUseListOrders() const { return preserveUseListOrders; }
+
   /// Return the resource parser registered to the given name, or nullptr if no
   /// parser with `name` is registered.
   AsmResourceParser *getResourceParser(StringRef name) const {
     auto it = resourceParsers.find(name);
     if (it != resourceParsers.end())
       return it->second.get();
     if (fallbackResourceMap)
       return &fallbackResourceMap->getParserFor(name);
public:
   attachResourceParser(StringRef name, CallableT &&parserFn) {
     attachResourceParser(AsmResourceParser::fromCallable(
         name, std::forward<CallableT>(parserFn)));
   }
 
 private:
   MLIRContext *context;
   bool verifyAfterParse;
+  bool preserveUseListOrders;
   DenseMap<StringRef, std::unique_ptr<AsmResourceParser>> resourceParsers;
   FallbackAsmResourceMap *fallbackResourceMap;
 };
 
 //===----------------------------------------------------------------------===//
 // AsmState
 //===----------------------------------------------------------------------===//
 

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);
 
+  /// Set the use list order according to the provided indices. The indices need
+  /// to map the current use-list order chain to another valid use-list chain.
+  void setUseListOrder(ArrayRef<unsigned>order);
+
   //===--------------------------------------------------------------------===//
   // 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/include/mlir/Tools/mlir-opt/MlirOptMain.h

public:
 
   /// Set whether to run the verifier after each transformation pass.
   MlirOptMainConfig &verifyPasses(bool verify) {
     verifyPassesFlag = verify;
     return *this;
   }
   bool shouldVerifyPasses() const { return verifyPassesFlag; }
 
+  bool shouldPreserveUseListOrders() const { return preserveUseListOrdersFlag; }
+
 protected:
   /// Allow operation with no registered dialects.
   /// This option is for convenience during testing only and discouraged in
   /// general.
   bool allowUnregisteredDialectsFlag = false;
 
   /// Configuration for the debugging hooks.
   tracing::DebugConfig debugConfig;
protected:
   bool useExplicitModuleFlag = false;
 
   /// Set whether to check that emitted diagnostics match `expected-*` lines on
   /// the corresponding line. This is meant for implementing diagnostic tests.
   bool verifyDiagnosticsFlag = false;
 
   /// Run the verifier after each transformation pass.
   bool verifyPassesFlag = true;
+
+  /// Preserve uselist orders when interacting with bytecode.
+  bool preserveUseListOrdersFlag = false;
 };
 
 /// This defines the function type used to setup the pass manager. This can be
 /// used to pass in a callback to setup a default pass pipeline to be applied on
 /// the loaded IR.
 using PassPipelineFn = llvm::function_ref<LogicalResult(PassManager &pm)>;
 
 /// Perform the core processing behind `mlir-opt`.

mlir/lib/Bytecode/Encoding.h

 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
 
 } // namespace bytecode
 } // namespace mlir
 
 #endif

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

 #include "mlir/IR/Verifier.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 <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:
   Value parseOperand(EncodingReader &reader);
 
   /// Sequentially define the given value range.
   LogicalResult defineValues(EncodingReader &reader, ValueRange values);
 
   /// Create a value to use for a forward reference.
   Value createForwardRef();
 
+  /// Parse use-list orders from bytecode if available. On success, return a map
+  /// between the op result number and the use-list encoding.
+  using UseListMapT =
+      DenseMap<unsigned, std::pair<bool, SmallVector<unsigned, 4>>>;
+  FailureOr<UseListMapT> parseOpResultsUseListOrder(EncodingReader &reader,
+                                                    uint64_t numResults);
+
+  /// Process the values with custom use-list and modify the use-chain according
+  /// to the order parsed.
+  LogicalResult processUseListOrderWorklist();
+
   //===--------------------------------------------------------------------===//
   // 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.
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;
   };
 
+  struct UseListOrder {
+    UseListOrder(Value val, bool isIndexPairEncoding,
+                 SmallVector<unsigned, 4> &&indices)
+        : value(val), isIndexPairEncoding(isIndexPairEncoding),
+          indices(std::move(indices)){};
+    // The Value owning the custom  use-list the uselist.
+    Value value;
+
+    // Whether the order is encoded as an explicit pair of indices `(src, dst)`
+    // or it is a direct indexing, such as `dst = order[src]`.
+    bool isIndexPairEncoding;
+
+    // The vector containing the order information required to reorder the
+    // use-list of value.
+    SmallVector<unsigned, 4> indices;
+  };
+
   /// The configuration of the parser.
   const ParserConfig &config;
 
   /// A location to use when emitting errors.
   Location fileLoc;
 
   /// The reader used to process attribute and types within the bytecode.
   AttrTypeReader attrTypeReader;
 
   /// The version of the bytecode being read.
   uint64_t version = 0;
 
   /// The producer of the bytecode being read.
   StringRef producer;
 
   /// 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.
+  SmallVector<UseListOrder> useListOrderWorklist;
+
   /// The table of strings referenced within the bytecode file.
   StringSectionReader stringReader;
 
   /// The current set of available IR value scopes.
   std::vector<ValueScope> valueScopes;
+
   /// 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;
 
   /// The optional owning source manager, which when present may be used to
   /// extend the lifetime of the input buffer.
   const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef;
 };
 } // namespace
LogicalResult BytecodeReader::parseResourceSection(
   DialectReader dialectReader(attrTypeReader, stringReader, resourceReader,
                               reader);
   return resourceReader.initialize(fileLoc, config, dialects, stringReader,
                                    *resourceData, *resourceOffsetData,
                                    dialectReader, bufferOwnerRef);
 }
 
 //===----------------------------------------------------------------------===//
+// UseListOrder Helpers
+
+FailureOr<BytecodeReader::UseListMapT>
+BytecodeReader::parseOpResultsUseListOrder(EncodingReader &reader,
+                                           uint64_t numResults) {
+  BytecodeReader::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,
+                    std::make_pair(indexPairEncoding, useListOrders));
+  }
+
+  return map;
+}
+
+LogicalResult BytecodeReader::processUseListOrderWorklist() {
+  while (!useListOrderWorklist.empty()) {
+    UseListOrder entry = useListOrderWorklist.pop_back_val();
+    uint64_t numUses = std::distance(entry.value.getUses().begin(),
+                                     entry.value.getUses().end());
+    SmallVector<unsigned, 4> shuffle = std::move(entry.indices);
+
+    // 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 (entry.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();
+
+    entry.value.setUseListOrder(shuffle);
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // IR Section
 
 LogicalResult BytecodeReader::parseIRSection(ArrayRef<uint8_t> sectionData,
                                              Block *block) {
   EncodingReader reader(sectionData, fileLoc);
 
   // A stack of operation regions currently being read from the bytecode.
   std::vector<RegionReadState> regionStack;
LogicalResult BytecodeReader::parseIRSection(ArrayRef<uint8_t> sectionData,
   while (!regionStack.empty())
     if (failed(parseRegions(reader, regionStack, regionStack.back())))
       return failure();
   if (!forwardRefOps.empty()) {
     return reader.emitError(
         "not all forward unresolved forward operand references");
   }
 
+  if (config.shouldPreserveUseListOrders())
+    if (failed(processUseListOrderWorklist()))
+      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.
+  std::optional<UseListMapT> useListOrderMap = std::nullopt;
+  if (opMask & bytecode::OpEncodingMask::kHasUseListOrders) {
+    size_t numResults = opState.types.size();
+    auto parseResult = parseOpResultsUseListOrder(reader, numResults);
+    if (failed(parseResult))
+      return failure();
+    useListOrderMap = 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();
 
+  // We can't update the uselist at this point, since the uses of the results
+  // don't yet exist.
+  if (useListOrderMap.has_value()) {
+    for (size_t idx = 0; idx < op->getNumResults(); idx++) {
+      if (useListOrderMap->contains(idx)) {
+        bool isIndexPairEncoding = useListOrderMap->at(idx).first;
+        auto order = useListOrderMap->at(idx).second;
+        useListOrderWorklist.emplace_back(UseListOrder(
+            op->getResult(idx), isIndexPairEncoding, std::move(order)));
+      }
+    }
+  }
   return op;
 }
 
 LogicalResult BytecodeReader::parseRegion(EncodingReader &reader,
                                           RegionReadState &readState) {
   // Parse the number of blocks in the region.
   uint64_t numBlocks;
   if (failed(reader.parseVarInt(numBlocks)))

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

 using namespace mlir;
 using namespace mlir::bytecode::detail;
 
 //===----------------------------------------------------------------------===//
 // BytecodeWriterConfig
 //===----------------------------------------------------------------------===//
 
 struct BytecodeWriterConfig::Impl {
-  Impl(StringRef producer) : producer(producer) {}
+  Impl(bool shouldEmitUseListOrder, StringRef producer)
+      : producer(producer), shouldEmitUseListOrder(shouldEmitUseListOrder) {}
 
   /// Version to use when writing.
   /// Note: This only differs from kVersion if a specific version is set.
   int64_t bytecodeVersion = bytecode::kVersion;
 
   /// The producer of the bytecode.
   StringRef producer;
 
+  /// Flag indicating if we should emit use-list order for operations.
+  bool shouldEmitUseListOrder;
+
   /// A collection of non-dialect resource printers.
   SmallVector<std::unique_ptr<AsmResourcePrinter>> externalResourcePrinters;
 };
 
-BytecodeWriterConfig::BytecodeWriterConfig(StringRef producer)
-    : impl(std::make_unique<Impl>(producer)) {}
+BytecodeWriterConfig::BytecodeWriterConfig(bool shouldEmitUseListOrder,
+                                           StringRef producer)
+    : impl(std::make_unique<Impl>(shouldEmitUseListOrder, producer)) {}
 BytecodeWriterConfig::BytecodeWriterConfig(FallbackAsmResourceMap &map,
+                                           bool shouldEmitUseListOrder,
                                            StringRef producer)
-    : BytecodeWriterConfig(producer) {
+    : BytecodeWriterConfig(shouldEmitUseListOrder, producer) {
   attachFallbackResourcePrinter(map);
 }
 BytecodeWriterConfig::~BytecodeWriterConfig() = default;
 
 void BytecodeWriterConfig::attachResourcePrinter(
     std::unique_ptr<AsmResourcePrinter> printer) {
   impl->externalResourcePrinters.emplace_back(std::move(printer));
 }
private:
   void writeResourceSection(Operation *op, EncodingEmitter &emitter);
 
   //===--------------------------------------------------------------------===//
   // Strings
 
   void writeStringSection(EncodingEmitter &emitter);
 
   //===--------------------------------------------------------------------===//
+  // Helpers
+
+  void writeUseListOrders(EncodingEmitter &emitter, uint8_t &opEncodingMask,
+                          Operation *op);
+
+  //===--------------------------------------------------------------------===//
   // Fields
 
   /// The builder used for the string section.
   StringSectionBuilder stringSection;
 
   /// The IR numbering state generated for the root operation.
   IRNumberingState numberingState;
 
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));
   }
 
+  // If uselist order preservation is requested, record the indices that map the
+  // current use-list order of the op with respect to the order printed to the
+  // IR. We can use this information to rebuild the correct uselist when parsing
+  // the bytecode.
+  if (config.shouldEmitUseListOrder)
+    writeUseListOrders(emitter, opEncodingMask, op);
+
   // Check for regions.
   unsigned numRegions = op->getNumRegions();
   if (numRegions)
     opEncodingMask |= bytecode::OpEncodingMask::kHasInlineRegions;
 
   // Update the mask for the operation.
   emitter.patchByte(maskOffset, opEncodingMask);
 
   // With the mask emitted, we can now emit the regions of the operation. We do
   // this after mask emission to avoid offset complications that may arise by
   // emitting the regions first (e.g. if the regions are huge, backpatching the
   // op encoding mask is more annoying).
   if (numRegions) {
     bool isIsolatedFromAbove = op->hasTrait<OpTrait::IsIsolatedFromAbove>();
     emitter.emitVarIntWithFlag(numRegions, isIsolatedFromAbove);
 
     for (Region &region : op->getRegions())
       writeRegion(emitter, &region);
   }
 }
 
+void BytecodeWriter::writeUseListOrders(EncodingEmitter &emitter,
+                                        uint8_t &opEncodingMask,
+                                        Operation *op) {
+  /// Loop over the results and store the use-list order per result index.
+  DenseMap<unsigned, llvm::SmallVector<std::pair<unsigned, unsigned>>> map;
+  for (auto resultIterator : llvm::enumerate(op->getResults())) {
+    auto result = resultIterator.value();
+    // No need to store a custom use-list order if the result does not have
+    // multiple uses.
+    if (result.use_empty() || result.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;
+    unsigned prevIndex =
+        numberingState.getNumber(result.getUses().begin()->getOwner());
+    llvm::SmallVector<std::pair<unsigned, unsigned>> useListPairs(
+        {{0, prevIndex}});
+
+    for (auto item : llvm::drop_begin(llvm::enumerate(result.getUses()))) {
+      unsigned currentIdx = numberingState.getNumber(item.value().getOwner());
+      alreadyOrdered &= (prevIndex >= currentIdx);
+      useListPairs.push_back({item.index(), currentIdx});
+      prevIndex = currentIdx;
+    }
+
+    // Do not emit if the order is already sorted.
+    if (alreadyOrdered)
+      continue;
+
+    map.try_emplace(resultIterator.index(), useListPairs);
+  }
+
+  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 (op->getNumResults() != 1)
+    emitter.emitVarInt(map.size());
+
+  for (const auto &elem : map) {
+    auto resultIdx = elem.getFirst();
+    auto useListPairs = elem.getSecond();
+    // Sort the use indices by the IR operation ordering and insert them to the
+    // builder for emission.
+    std::sort(
+        useListPairs.begin(), useListPairs.end(),
+        [](auto elem1, auto elem2) { return elem1.second > elem2.second; });
+
+    // 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(useListPairs), [](auto item) {
+          return item.index() != item.value().first;
+        });
+    bool indexPairEncoding = shuffledElements < (useListPairs.size() / 2);
+
+    // For single result, we don't need to store the result index.
+    if (op->getNumResults() != 1)
+      emitter.emitVarInt(resultIdx);
+
+    if (indexPairEncoding) {
+      emitter.emitVarIntWithFlag(shuffledElements * 2, indexPairEncoding);
+      for (auto pair : llvm::enumerate(useListPairs)) {
+        if (pair.index() != pair.value().first) {
+          emitter.emitVarInt(pair.value().first);
+          emitter.emitVarInt(pair.index());
+        }
+      }
+    } else {
+      emitter.emitVarIntWithFlag(useListPairs.size(), indexPairEncoding);
+      for (auto elem : useListPairs)
+        emitter.emitVarInt(elem.first);
+    }
+  }
+}
+
 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) && "block 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 Block, Operation 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;
 
   /// The next value ID to assign when numbering.
   unsigned nextValueID = 0;
+  unsigned nextOperationID = 0;
 };
 } // namespace detail
 } // namespace bytecode
 } // namespace mlir
 
 #endif

mlir/lib/Bytecode/Writer/IRNumbering.cpp

void IRNumberingState::number(Region &region) {
   }
 
   // Remember the number of blocks and values in this region.
   regionBlockValueCounts.try_emplace(&region, blockCount,
                                      nextValueID - firstValueID);
 }
 
 void IRNumberingState::number(Operation &op) {
+  operationIDs.try_emplace(&op, nextOperationID++);
   // Number the components of an operation that won't be numbered elsewhere
   // (e.g. we don't number operands, regions, or successors here).
   number(op.getName());
   for (OpResult result : op.getResults()) {
     valueIDs.try_emplace(result, nextValueID++);
     number(result.getType());
   }
 

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;
   });
 }
 
+void Value::setUseListOrder(ArrayRef<unsigned int> order) {
+  assert((size_t)std::distance(getUses().begin(), getUses().end()) ==
+             order.size() &&
+         "order vector expected to have a number of elements equal to the "
+         "number of uses");
+  // Order uses into a vector and set them for every operand type.
+  SmallVector<OpOperand *> newUseList(order.size());
+  for (const auto &elem : llvm::enumerate(getUses()))
+    newUseList[order[elem.index()]] = &elem.value();
+
+  // For every use in the range, replace it with the one in the list.
+  for (const auto &item : llvm::zip(getUses(), newUseList)) {
+    std::get<0>(item).set((*std::get<1>(item)).get());
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // 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

mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp

MLIRDocument::MLIRDocument(MLIRContext &context, const lsp::URIForFile &uri,
   // Try to parsed the given IR string.
   auto memBuffer = llvm::MemoryBuffer::getMemBufferCopy(contents, uri.file());
   if (!memBuffer) {
     lsp::Logger::error("Failed to create memory buffer for file", uri.file());
     return;
   }
 
   ParserConfig config(&context, /*verifyAfterParse=*/true,
-                      &fallbackResourceMap);
+                      /*preserveUseListOrders=*/false, &fallbackResourceMap);
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
   if (failed(parseAsmSourceFile(sourceMgr, &parsedIR, config, &asmState))) {
     // If parsing failed, clear out any of the current state.
     parsedIR.clear();
     asmState = AsmParserState();
     fallbackResourceMap = FallbackAsmResourceMap();
     return;
   }
ScopedDiagnosticHandler diagHandler(
       &tempContext,
       [&](mlir::Diagnostic &diag) { errorMsg += diag.str() + "\n"; });
 
   // Handling for external resources, which we want to propagate up to the user.
   FallbackAsmResourceMap fallbackResourceMap;
 
   // Setup the parser config.
   ParserConfig parserConfig(&tempContext, /*verifyAfterParse=*/true,
+                            /*preserveUseListOrders=*/false,
                             &fallbackResourceMap);
 
   // Try to parse the given source file.
   Block parsedBlock;
   if (failed(parseSourceFile(uri.file(), &parsedBlock, parserConfig))) {
     return llvm::make_error<lsp::LSPError>(
         "failed to parse bytecode source file: " + errorMsg,
         lsp::ErrorCode::RequestFailed);

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

static cl::opt<bool, /*ExternalStorage=*/true> verifyDiagnostics(
                  "expected-* lines on the corresponding line"),
         cl::location(verifyDiagnosticsFlag), cl::init(false));
 
     static cl::opt<bool, /*ExternalStorage=*/true> verifyPasses(
         "verify-each",
         cl::desc("Run the verifier after each transformation pass"),
         cl::location(verifyPassesFlag), cl::init(true));
 
+    static cl::opt<bool, /*ExternalStorage=*/true> preserveUseListOrders(
+        "preserve-use-list-orders",
+        cl::desc("Preserve use-list orders when interacting with bytecode"),
+        cl::location(preserveUseListOrdersFlag), cl::init(false));
+
     static cl::list<std::string> passPlugins(
         "load-pass-plugin", cl::desc("Load passes from plugin library"));
     /// Set the callback to load a pass plugin.
     passPlugins.setCallback([&](const std::string &pluginPath) {
       auto plugin = PassPlugin::load(pluginPath);
       if (!plugin) {
         errs() << "Failed to load passes from '" << pluginPath
                << "'. Request ignored.\n";
performActions(raw_ostream &os,
 
   // Prepare the parser config, and attach any useful/necessary resource
   // handlers. Unhandled external resources are treated as passthrough, i.e.
   // they are not processed and will be emitted directly to the output
   // untouched.
   PassReproducerOptions reproOptions;
   FallbackAsmResourceMap fallbackResourceMap;
   ParserConfig parseConfig(context, /*verifyAfterParse=*/true,
+                           config.shouldPreserveUseListOrders(),
                            &fallbackResourceMap);
   reproOptions.attachResourceParser(parseConfig);
 
   // Parse the input file and reset the context threading state.
   TimingScope parserTiming = timing.nest("Parser");
   OwningOpRef<Operation *> op = parseSourceFileForTool(
       sourceMgr, parseConfig, !config.shouldUseExplicitModule());
   context->enableMultithreading(wasThreadingEnabled);
performActions(raw_ostream &os,
 
   // Run the pipeline.
   if (failed(pm.run(*op)))
     return failure();
 
   // Print the output.
   TimingScope outputTiming = timing.nest("Output");
   if (config.shouldEmitBytecode()) {
-    BytecodeWriterConfig writerConfig(fallbackResourceMap);
+    BytecodeWriterConfig writerConfig(fallbackResourceMap,
+                                      config.shouldPreserveUseListOrders());
     if (auto v = config.bytecodeVersionToEmit())
       writerConfig.setDesiredBytecodeVersion(*v);
     return writeBytecodeToFile(op.get(), os, writerConfig);
   }
 
   if (config.bytecodeVersionToEmit().has_value())
     return emitError(UnknownLoc::get(pm.getContext()))
            << "bytecode version while not emitting bytecode";