xref: /llvm-project/mlir/lib/Bytecode/Writer/IRNumbering.cpp (revision f1ac7725e4fd5afa21fb244f9bcc33de654ed80c)
1f3acb54cSRiver Riddle //===- IRNumbering.cpp - MLIR Bytecode IR numbering -----------------------===//
2f3acb54cSRiver Riddle //
3f3acb54cSRiver Riddle // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4f3acb54cSRiver Riddle // See https://llvm.org/LICENSE.txt for license information.
5f3acb54cSRiver Riddle // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6f3acb54cSRiver Riddle //
7f3acb54cSRiver Riddle //===----------------------------------------------------------------------===//
8f3acb54cSRiver Riddle 
9f3acb54cSRiver Riddle #include "IRNumbering.h"
1002c2ecb9SRiver Riddle #include "mlir/Bytecode/BytecodeImplementation.h"
11660f714eSMehdi Amini #include "mlir/Bytecode/BytecodeOpInterface.h"
12*f1ac7725SJacques Pienaar #include "mlir/Bytecode/BytecodeWriter.h"
135ed11e76SAlex Zinenko #include "mlir/Bytecode/Encoding.h"
146ab2bcffSRiver Riddle #include "mlir/IR/AsmState.h"
15f3acb54cSRiver Riddle #include "mlir/IR/BuiltinTypes.h"
16f3acb54cSRiver Riddle #include "mlir/IR/OpDefinition.h"
17f3acb54cSRiver Riddle 
18f3acb54cSRiver Riddle using namespace mlir;
19f3acb54cSRiver Riddle using namespace mlir::bytecode::detail;
20f3acb54cSRiver Riddle 
21f3acb54cSRiver Riddle //===----------------------------------------------------------------------===//
2202c2ecb9SRiver Riddle // NumberingDialectWriter
2302c2ecb9SRiver Riddle //===----------------------------------------------------------------------===//
2402c2ecb9SRiver Riddle 
2502c2ecb9SRiver Riddle struct IRNumberingState::NumberingDialectWriter : public DialectBytecodeWriter {
267ad9e9dcSMatteo Franciolini   NumberingDialectWriter(
277ad9e9dcSMatteo Franciolini       IRNumberingState &state,
287ad9e9dcSMatteo Franciolini       llvm::StringMap<std::unique_ptr<DialectVersion>> &dialectVersionMap)
297ad9e9dcSMatteo Franciolini       : state(state), dialectVersionMap(dialectVersionMap) {}
3002c2ecb9SRiver Riddle 
3102c2ecb9SRiver Riddle   void writeAttribute(Attribute attr) override { state.number(attr); }
32660f714eSMehdi Amini   void writeOptionalAttribute(Attribute attr) override {
33660f714eSMehdi Amini     if (attr)
34660f714eSMehdi Amini       state.number(attr);
35660f714eSMehdi Amini   }
3602c2ecb9SRiver Riddle   void writeType(Type type) override { state.number(type); }
376ab2bcffSRiver Riddle   void writeResourceHandle(const AsmDialectResourceHandle &resource) override {
386ab2bcffSRiver Riddle     state.number(resource.getDialect(), resource);
396ab2bcffSRiver Riddle   }
4002c2ecb9SRiver Riddle 
4102c2ecb9SRiver Riddle   /// Stubbed out methods that are not used for numbering.
4202c2ecb9SRiver Riddle   void writeVarInt(uint64_t) override {}
432f90764cSRiver Riddle   void writeSignedVarInt(int64_t value) override {}
442f90764cSRiver Riddle   void writeAPIntWithKnownWidth(const APInt &value) override {}
452f90764cSRiver Riddle   void writeAPFloatWithKnownSemantics(const APFloat &value) override {}
4602c2ecb9SRiver Riddle   void writeOwnedString(StringRef) override {
4702c2ecb9SRiver Riddle     // TODO: It might be nice to prenumber strings and sort by the number of
4802c2ecb9SRiver Riddle     // references. This could potentially be useful for optimizing things like
4902c2ecb9SRiver Riddle     // file locations.
5002c2ecb9SRiver Riddle   }
515fb1bbe6SRiver Riddle   void writeOwnedBlob(ArrayRef<char> blob) override {}
5279c83e12SAndrzej Warzynski   void writeOwnedBool(bool value) override {}
5302c2ecb9SRiver Riddle 
540610e2f6SJacques Pienaar   int64_t getBytecodeVersion() const override {
559ea6b30aSMehdi Amini     return state.getDesiredBytecodeVersion();
560610e2f6SJacques Pienaar   }
570610e2f6SJacques Pienaar 
587ad9e9dcSMatteo Franciolini   FailureOr<const DialectVersion *>
597ad9e9dcSMatteo Franciolini   getDialectVersion(StringRef dialectName) const override {
607ad9e9dcSMatteo Franciolini     auto dialectEntry = dialectVersionMap.find(dialectName);
617ad9e9dcSMatteo Franciolini     if (dialectEntry == dialectVersionMap.end())
627ad9e9dcSMatteo Franciolini       return failure();
637ad9e9dcSMatteo Franciolini     return dialectEntry->getValue().get();
647ad9e9dcSMatteo Franciolini   }
657ad9e9dcSMatteo Franciolini 
6602c2ecb9SRiver Riddle   /// The parent numbering state that is populated by this writer.
6702c2ecb9SRiver Riddle   IRNumberingState &state;
687ad9e9dcSMatteo Franciolini 
697ad9e9dcSMatteo Franciolini   /// A map containing dialect version information for each dialect to emit.
707ad9e9dcSMatteo Franciolini   llvm::StringMap<std::unique_ptr<DialectVersion>> &dialectVersionMap;
7102c2ecb9SRiver Riddle };
7202c2ecb9SRiver Riddle 
7302c2ecb9SRiver Riddle //===----------------------------------------------------------------------===//
74f3acb54cSRiver Riddle // IR Numbering
75f3acb54cSRiver Riddle //===----------------------------------------------------------------------===//
76f3acb54cSRiver Riddle 
77f3acb54cSRiver Riddle /// Group and sort the elements of the given range by their parent dialect. This
78f3acb54cSRiver Riddle /// grouping is applied to sub-sections of the ranged defined by how many bytes
79f3acb54cSRiver Riddle /// it takes to encode a varint index to that sub-section.
80f3acb54cSRiver Riddle template <typename T>
81f3acb54cSRiver Riddle static void groupByDialectPerByte(T range) {
82f3acb54cSRiver Riddle   if (range.empty())
83f3acb54cSRiver Riddle     return;
84f3acb54cSRiver Riddle 
85f3acb54cSRiver Riddle   // A functor used to sort by a given dialect, with a desired dialect to be
86f3acb54cSRiver Riddle   // ordered first (to better enable sharing of dialects across byte groups).
87f3acb54cSRiver Riddle   auto sortByDialect = [](unsigned dialectToOrderFirst, const auto &lhs,
88f3acb54cSRiver Riddle                           const auto &rhs) {
89f3acb54cSRiver Riddle     if (lhs->dialect->number == dialectToOrderFirst)
90f3acb54cSRiver Riddle       return rhs->dialect->number != dialectToOrderFirst;
911fbccfd6SRiver Riddle     if (rhs->dialect->number == dialectToOrderFirst)
921fbccfd6SRiver Riddle       return false;
93f3acb54cSRiver Riddle     return lhs->dialect->number < rhs->dialect->number;
94f3acb54cSRiver Riddle   };
95f3acb54cSRiver Riddle 
96f3acb54cSRiver Riddle   unsigned dialectToOrderFirst = 0;
97f3acb54cSRiver Riddle   size_t elementsInByteGroup = 0;
98f3acb54cSRiver Riddle   auto iterRange = range;
99f3acb54cSRiver Riddle   for (unsigned i = 1; i < 9; ++i) {
100f3acb54cSRiver Riddle     // Update the number of elements in the current byte grouping. Reminder
101f3acb54cSRiver Riddle     // that varint encodes 7-bits per byte, so that's how we compute the
102f3acb54cSRiver Riddle     // number of elements in each byte grouping.
10393cf0e8aSRiver Riddle     elementsInByteGroup = (1ULL << (7ULL * i)) - elementsInByteGroup;
104f3acb54cSRiver Riddle 
105f3acb54cSRiver Riddle     // Slice out the sub-set of elements that are in the current byte grouping
106f3acb54cSRiver Riddle     // to be sorted.
107f3acb54cSRiver Riddle     auto byteSubRange = iterRange.take_front(elementsInByteGroup);
108f3acb54cSRiver Riddle     iterRange = iterRange.drop_front(byteSubRange.size());
109f3acb54cSRiver Riddle 
110f3acb54cSRiver Riddle     // Sort the sub range for this byte.
111f3acb54cSRiver Riddle     llvm::stable_sort(byteSubRange, [&](const auto &lhs, const auto &rhs) {
112f3acb54cSRiver Riddle       return sortByDialect(dialectToOrderFirst, lhs, rhs);
113f3acb54cSRiver Riddle     });
114f3acb54cSRiver Riddle 
115f3acb54cSRiver Riddle     // Update the dialect to order first to be the dialect at the end of the
116f3acb54cSRiver Riddle     // current grouping. This seeks to allow larger dialect groupings across
117f3acb54cSRiver Riddle     // byte boundaries.
118f3acb54cSRiver Riddle     dialectToOrderFirst = byteSubRange.back()->dialect->number;
119f3acb54cSRiver Riddle 
120f3acb54cSRiver Riddle     // If the data range is now empty, we are done.
121f3acb54cSRiver Riddle     if (iterRange.empty())
122f3acb54cSRiver Riddle       break;
123f3acb54cSRiver Riddle   }
124f3acb54cSRiver Riddle 
125f3acb54cSRiver Riddle   // Assign the entry numbers based on the sort order.
1268c258fdaSJakub Kuderski   for (auto [idx, value] : llvm::enumerate(range))
1278c258fdaSJakub Kuderski     value->number = idx;
128f3acb54cSRiver Riddle }
129f3acb54cSRiver Riddle 
130660f714eSMehdi Amini IRNumberingState::IRNumberingState(Operation *op,
131660f714eSMehdi Amini                                    const BytecodeWriterConfig &config)
132660f714eSMehdi Amini     : config(config) {
1334af01bf9SRiver Riddle   computeGlobalNumberingState(op);
13461278191SMatteo Franciolini 
135f3acb54cSRiver Riddle   // Number the root operation.
136f3acb54cSRiver Riddle   number(*op);
137f3acb54cSRiver Riddle 
1384af01bf9SRiver Riddle   // A worklist of region contexts to number and the next value id before that
1394af01bf9SRiver Riddle   // region.
140f3acb54cSRiver Riddle   SmallVector<std::pair<Region *, unsigned>, 8> numberContext;
1414af01bf9SRiver Riddle 
1424af01bf9SRiver Riddle   // Functor to push the regions of the given operation onto the numbering
1434af01bf9SRiver Riddle   // context.
1444af01bf9SRiver Riddle   auto addOpRegionsToNumber = [&](Operation *op) {
1454af01bf9SRiver Riddle     MutableArrayRef<Region> regions = op->getRegions();
1464af01bf9SRiver Riddle     if (regions.empty())
1474af01bf9SRiver Riddle       return;
1484af01bf9SRiver Riddle 
1494af01bf9SRiver Riddle     // Isolated regions don't share value numbers with their parent, so we can
1504af01bf9SRiver Riddle     // start numbering these regions at zero.
1514af01bf9SRiver Riddle     unsigned opFirstValueID = isIsolatedFromAbove(op) ? 0 : nextValueID;
1524af01bf9SRiver Riddle     for (Region &region : regions)
1534af01bf9SRiver Riddle       numberContext.emplace_back(&region, opFirstValueID);
1544af01bf9SRiver Riddle   };
1554af01bf9SRiver Riddle   addOpRegionsToNumber(op);
156f3acb54cSRiver Riddle 
157f3acb54cSRiver Riddle   // Iteratively process each of the nested regions.
158f3acb54cSRiver Riddle   while (!numberContext.empty()) {
159f3acb54cSRiver Riddle     Region *region;
160f3acb54cSRiver Riddle     std::tie(region, nextValueID) = numberContext.pop_back_val();
161f3acb54cSRiver Riddle     number(*region);
162f3acb54cSRiver Riddle 
163f3acb54cSRiver Riddle     // Traverse into nested regions.
1644af01bf9SRiver Riddle     for (Operation &op : region->getOps())
1654af01bf9SRiver Riddle       addOpRegionsToNumber(&op);
166f3acb54cSRiver Riddle   }
167f3acb54cSRiver Riddle 
168f3acb54cSRiver Riddle   // Number each of the dialects. For now this is just in the order they were
169f3acb54cSRiver Riddle   // found, given that the number of dialects on average is small enough to fit
170f3acb54cSRiver Riddle   // within a singly byte (128). If we ever have real world use cases that have
171f3acb54cSRiver Riddle   // a huge number of dialects, this could be made more intelligent.
1728c258fdaSJakub Kuderski   for (auto [idx, dialect] : llvm::enumerate(dialects))
1738c258fdaSJakub Kuderski     dialect.second->number = idx;
174f3acb54cSRiver Riddle 
175f3acb54cSRiver Riddle   // Number each of the recorded components within each dialect.
176f3acb54cSRiver Riddle 
177f3acb54cSRiver Riddle   // First sort by ref count so that the most referenced elements are first. We
178f3acb54cSRiver Riddle   // try to bias more heavily used elements to the front. This allows for more
179f3acb54cSRiver Riddle   // frequently referenced things to be encoded using smaller varints.
180f3acb54cSRiver Riddle   auto sortByRefCountFn = [](const auto &lhs, const auto &rhs) {
181f3acb54cSRiver Riddle     return lhs->refCount > rhs->refCount;
182f3acb54cSRiver Riddle   };
183f3acb54cSRiver Riddle   llvm::stable_sort(orderedAttrs, sortByRefCountFn);
184f3acb54cSRiver Riddle   llvm::stable_sort(orderedOpNames, sortByRefCountFn);
185f3acb54cSRiver Riddle   llvm::stable_sort(orderedTypes, sortByRefCountFn);
186f3acb54cSRiver Riddle 
187f3acb54cSRiver Riddle   // After that, we apply a secondary ordering based on the parent dialect. This
188f3acb54cSRiver Riddle   // ordering is applied to sub-sections of the element list defined by how many
189f3acb54cSRiver Riddle   // bytes it takes to encode a varint index to that sub-section. This allows
190f3acb54cSRiver Riddle   // for more efficiently encoding components of the same dialect (e.g. we only
191f3acb54cSRiver Riddle   // have to encode the dialect reference once).
192a288d7f9SJoe Loser   groupByDialectPerByte(llvm::MutableArrayRef(orderedAttrs));
193a288d7f9SJoe Loser   groupByDialectPerByte(llvm::MutableArrayRef(orderedOpNames));
194a288d7f9SJoe Loser   groupByDialectPerByte(llvm::MutableArrayRef(orderedTypes));
1956ab2bcffSRiver Riddle 
1966ab2bcffSRiver Riddle   // Finalize the numbering of the dialect resources.
1976ab2bcffSRiver Riddle   finalizeDialectResourceNumberings(op);
198f3acb54cSRiver Riddle }
199f3acb54cSRiver Riddle 
2004af01bf9SRiver Riddle void IRNumberingState::computeGlobalNumberingState(Operation *rootOp) {
2014af01bf9SRiver Riddle   // A simple state struct tracking data used when walking operations.
2024af01bf9SRiver Riddle   struct StackState {
2034af01bf9SRiver Riddle     /// The operation currently being walked.
2044af01bf9SRiver Riddle     Operation *op;
2054af01bf9SRiver Riddle 
2064af01bf9SRiver Riddle     /// The numbering of the operation.
2074af01bf9SRiver Riddle     OperationNumbering *numbering;
2084af01bf9SRiver Riddle 
2094af01bf9SRiver Riddle     /// A flag indicating if the current state or one of its parents has
2104af01bf9SRiver Riddle     /// unresolved isolation status. This is tracked separately from the
2114af01bf9SRiver Riddle     /// isIsolatedFromAbove bit on `numbering` because we need to be able to
2124af01bf9SRiver Riddle     /// handle the given case:
2134af01bf9SRiver Riddle     ///   top.op {
2144af01bf9SRiver Riddle     ///     %value = ...
2154af01bf9SRiver Riddle     ///     middle.op {
2164af01bf9SRiver Riddle     ///       %value2 = ...
2174af01bf9SRiver Riddle     ///       inner.op {
2184af01bf9SRiver Riddle     ///         // Here we mark `inner.op` as not isolated. Note `middle.op`
2194af01bf9SRiver Riddle     ///         // isn't known not isolated yet.
2204af01bf9SRiver Riddle     ///         use.op %value2
2214af01bf9SRiver Riddle     ///
2224af01bf9SRiver Riddle     ///         // Here inner.op is already known to be non-isolated, but
2234af01bf9SRiver Riddle     ///         // `middle.op` is now also discovered to be non-isolated.
2244af01bf9SRiver Riddle     ///         use.op %value
2254af01bf9SRiver Riddle     ///       }
2264af01bf9SRiver Riddle     ///     }
2274af01bf9SRiver Riddle     ///   }
2284af01bf9SRiver Riddle     bool hasUnresolvedIsolation;
2294af01bf9SRiver Riddle   };
2304af01bf9SRiver Riddle 
2314af01bf9SRiver Riddle   // Compute a global operation ID numbering according to the pre-order walk of
2324af01bf9SRiver Riddle   // the IR. This is used as reference to construct use-list orders.
2334af01bf9SRiver Riddle   unsigned operationID = 0;
2344af01bf9SRiver Riddle 
2354af01bf9SRiver Riddle   // Walk each of the operations within the IR, tracking a stack of operations
2364af01bf9SRiver Riddle   // as we recurse into nested regions. This walk method hooks in at two stages
2374af01bf9SRiver Riddle   // during the walk:
2384af01bf9SRiver Riddle   //
2394af01bf9SRiver Riddle   //   BeforeAllRegions:
2404af01bf9SRiver Riddle   //     Here we generate a numbering for the operation and push it onto the
2414af01bf9SRiver Riddle   //     stack if it has regions. We also compute the isolation status of parent
2424af01bf9SRiver Riddle   //     regions at this stage. This is done by checking the parent regions of
2434af01bf9SRiver Riddle   //     operands used by the operation, and marking each region between the
2444af01bf9SRiver Riddle   //     the operand region and the current as not isolated. See
2454af01bf9SRiver Riddle   //     StackState::hasUnresolvedIsolation above for an example.
2464af01bf9SRiver Riddle   //
2474af01bf9SRiver Riddle   //   AfterAllRegions:
2484af01bf9SRiver Riddle   //     Here we pop the operation from the stack, and if it hasn't been marked
2494af01bf9SRiver Riddle   //     as non-isolated, we mark it as so. A non-isolated use would have been
2504af01bf9SRiver Riddle   //     found while walking the regions, so it is safe to mark the operation at
2514af01bf9SRiver Riddle   //     this point.
2524af01bf9SRiver Riddle   //
2534af01bf9SRiver Riddle   SmallVector<StackState> opStack;
2544af01bf9SRiver Riddle   rootOp->walk([&](Operation *op, const WalkStage &stage) {
2554af01bf9SRiver Riddle     // After visiting all nested regions, we pop the operation from the stack.
2560e6000f6SRiver Riddle     if (op->getNumRegions() && stage.isAfterAllRegions()) {
2574af01bf9SRiver Riddle       // If no non-isolated uses were found, we can safely mark this operation
2584af01bf9SRiver Riddle       // as isolated from above.
2594af01bf9SRiver Riddle       OperationNumbering *numbering = opStack.pop_back_val().numbering;
2604af01bf9SRiver Riddle       if (!numbering->isIsolatedFromAbove.has_value())
2614af01bf9SRiver Riddle         numbering->isIsolatedFromAbove = true;
2624af01bf9SRiver Riddle       return;
2634af01bf9SRiver Riddle     }
2644af01bf9SRiver Riddle 
2654af01bf9SRiver Riddle     // When visiting before nested regions, we process "IsolatedFromAbove"
2664af01bf9SRiver Riddle     // checks and compute the number for this operation.
2674af01bf9SRiver Riddle     if (!stage.isBeforeAllRegions())
2684af01bf9SRiver Riddle       return;
2694af01bf9SRiver Riddle     // Update the isolation status of parent regions if any have yet to be
2704af01bf9SRiver Riddle     // resolved.
2714af01bf9SRiver Riddle     if (!opStack.empty() && opStack.back().hasUnresolvedIsolation) {
2724af01bf9SRiver Riddle       Region *parentRegion = op->getParentRegion();
2734af01bf9SRiver Riddle       for (Value operand : op->getOperands()) {
2744af01bf9SRiver Riddle         Region *operandRegion = operand.getParentRegion();
2754af01bf9SRiver Riddle         if (operandRegion == parentRegion)
2764af01bf9SRiver Riddle           continue;
2774af01bf9SRiver Riddle         // We've found a use of an operand outside of the current region,
2784af01bf9SRiver Riddle         // walk the operation stack searching for the parent operation,
2794af01bf9SRiver Riddle         // marking every region on the way as not isolated.
2804af01bf9SRiver Riddle         Operation *operandContainerOp = operandRegion->getParentOp();
2814af01bf9SRiver Riddle         auto it = std::find_if(
2824af01bf9SRiver Riddle             opStack.rbegin(), opStack.rend(), [=](const StackState &it) {
2834af01bf9SRiver Riddle               // We only need to mark up to the container region, or the first
2844af01bf9SRiver Riddle               // that has an unresolved status.
2854af01bf9SRiver Riddle               return !it.hasUnresolvedIsolation || it.op == operandContainerOp;
2864af01bf9SRiver Riddle             });
2874af01bf9SRiver Riddle         assert(it != opStack.rend() && "expected to find the container");
2884af01bf9SRiver Riddle         for (auto &state : llvm::make_range(opStack.rbegin(), it)) {
2894af01bf9SRiver Riddle           // If we stopped at a region that knows its isolation status, we can
2904af01bf9SRiver Riddle           // stop updating the isolation status for the parent regions.
2914af01bf9SRiver Riddle           state.hasUnresolvedIsolation = it->hasUnresolvedIsolation;
2924af01bf9SRiver Riddle           state.numbering->isIsolatedFromAbove = false;
2934af01bf9SRiver Riddle         }
2944af01bf9SRiver Riddle       }
2954af01bf9SRiver Riddle     }
2964af01bf9SRiver Riddle 
2974af01bf9SRiver Riddle     // Compute the number for this op and push it onto the stack.
2984af01bf9SRiver Riddle     auto *numbering =
2994af01bf9SRiver Riddle         new (opAllocator.Allocate()) OperationNumbering(operationID++);
3004af01bf9SRiver Riddle     if (op->hasTrait<OpTrait::IsIsolatedFromAbove>())
3014af01bf9SRiver Riddle       numbering->isIsolatedFromAbove = true;
3024af01bf9SRiver Riddle     operations.try_emplace(op, numbering);
3034af01bf9SRiver Riddle     if (op->getNumRegions()) {
3044af01bf9SRiver Riddle       opStack.emplace_back(StackState{
3054af01bf9SRiver Riddle           op, numbering, !numbering->isIsolatedFromAbove.has_value()});
3064af01bf9SRiver Riddle     }
3074af01bf9SRiver Riddle   });
3084af01bf9SRiver Riddle }
3094af01bf9SRiver Riddle 
310f3acb54cSRiver Riddle void IRNumberingState::number(Attribute attr) {
311f3acb54cSRiver Riddle   auto it = attrs.insert({attr, nullptr});
312f3acb54cSRiver Riddle   if (!it.second) {
313f3acb54cSRiver Riddle     ++it.first->second->refCount;
314f3acb54cSRiver Riddle     return;
315f3acb54cSRiver Riddle   }
316f3acb54cSRiver Riddle   auto *numbering = new (attrAllocator.Allocate()) AttributeNumbering(attr);
317f3acb54cSRiver Riddle   it.first->second = numbering;
318f3acb54cSRiver Riddle   orderedAttrs.push_back(numbering);
319f3acb54cSRiver Riddle 
320f3acb54cSRiver Riddle   // Check for OpaqueAttr, which is a dialect-specific attribute that didn't
321f3acb54cSRiver Riddle   // have a registered dialect when it got created. We don't want to encode this
322f3acb54cSRiver Riddle   // as the builtin OpaqueAttr, we want to encode it as if the dialect was
323f3acb54cSRiver Riddle   // actually loaded.
3245550c821STres Popp   if (OpaqueAttr opaqueAttr = dyn_cast<OpaqueAttr>(attr)) {
325f3acb54cSRiver Riddle     numbering->dialect = &numberDialect(opaqueAttr.getDialectNamespace());
32602c2ecb9SRiver Riddle     return;
32702c2ecb9SRiver Riddle   }
328f3acb54cSRiver Riddle   numbering->dialect = &numberDialect(&attr.getDialect());
32902c2ecb9SRiver Riddle 
33002c2ecb9SRiver Riddle   // If this attribute will be emitted using the bytecode format, perform a
33102c2ecb9SRiver Riddle   // dummy writing to number any nested components.
33202c2ecb9SRiver Riddle   // TODO: We don't allow custom encodings for mutable attributes right now.
3336ab2bcffSRiver Riddle   if (!attr.hasTrait<AttributeTrait::IsMutable>()) {
334bff6a429SMatteo Franciolini     // Try overriding emission with callbacks.
335bff6a429SMatteo Franciolini     for (const auto &callback : config.getAttributeWriterCallbacks()) {
3367ad9e9dcSMatteo Franciolini       NumberingDialectWriter writer(*this, config.getDialectVersionMap());
337bff6a429SMatteo Franciolini       // The client has the ability to override the group name through the
338bff6a429SMatteo Franciolini       // callback.
339bff6a429SMatteo Franciolini       std::optional<StringRef> groupNameOverride;
340bff6a429SMatteo Franciolini       if (succeeded(callback->write(attr, groupNameOverride, writer))) {
341bff6a429SMatteo Franciolini         if (groupNameOverride.has_value())
342bff6a429SMatteo Franciolini           numbering->dialect = &numberDialect(*groupNameOverride);
343bff6a429SMatteo Franciolini         return;
344bff6a429SMatteo Franciolini       }
345bff6a429SMatteo Franciolini     }
346bff6a429SMatteo Franciolini 
347bff6a429SMatteo Franciolini     if (const auto *interface = numbering->dialect->interface) {
3487ad9e9dcSMatteo Franciolini       NumberingDialectWriter writer(*this, config.getDialectVersionMap());
3496ab2bcffSRiver Riddle       if (succeeded(interface->writeAttribute(attr, writer)))
3506ab2bcffSRiver Riddle         return;
35102c2ecb9SRiver Riddle     }
352f3acb54cSRiver Riddle   }
3536ab2bcffSRiver Riddle   // If this attribute will be emitted using the fallback, number the nested
3546ab2bcffSRiver Riddle   // dialect resources. We don't number everything (e.g. no nested
3556ab2bcffSRiver Riddle   // attributes/types), because we don't want to encode things we won't decode
3566ab2bcffSRiver Riddle   // (the textual format can't really share much).
3576ab2bcffSRiver Riddle   AsmState tempState(attr.getContext());
3586ab2bcffSRiver Riddle   llvm::raw_null_ostream dummyOS;
3596ab2bcffSRiver Riddle   attr.print(dummyOS, tempState);
3606ab2bcffSRiver Riddle 
3616ab2bcffSRiver Riddle   // Number the used dialect resources.
3626ab2bcffSRiver Riddle   for (const auto &it : tempState.getDialectResources())
3636ab2bcffSRiver Riddle     number(it.getFirst(), it.getSecond().getArrayRef());
3646ab2bcffSRiver Riddle }
365f3acb54cSRiver Riddle 
366f3acb54cSRiver Riddle void IRNumberingState::number(Block &block) {
367f3acb54cSRiver Riddle   // Number the arguments of the block.
368f3acb54cSRiver Riddle   for (BlockArgument arg : block.getArguments()) {
369f3acb54cSRiver Riddle     valueIDs.try_emplace(arg, nextValueID++);
370f3acb54cSRiver Riddle     number(arg.getLoc());
371f3acb54cSRiver Riddle     number(arg.getType());
372f3acb54cSRiver Riddle   }
373f3acb54cSRiver Riddle 
374f3acb54cSRiver Riddle   // Number the operations in this block.
375f3acb54cSRiver Riddle   unsigned &numOps = blockOperationCounts[&block];
376f3acb54cSRiver Riddle   for (Operation &op : block) {
377f3acb54cSRiver Riddle     number(op);
378f3acb54cSRiver Riddle     ++numOps;
379f3acb54cSRiver Riddle   }
380f3acb54cSRiver Riddle }
381f3acb54cSRiver Riddle 
382f3acb54cSRiver Riddle auto IRNumberingState::numberDialect(Dialect *dialect) -> DialectNumbering & {
383f3acb54cSRiver Riddle   DialectNumbering *&numbering = registeredDialects[dialect];
384f3acb54cSRiver Riddle   if (!numbering) {
385f3acb54cSRiver Riddle     numbering = &numberDialect(dialect->getNamespace());
38602c2ecb9SRiver Riddle     numbering->interface = dyn_cast<BytecodeDialectInterface>(dialect);
3876ab2bcffSRiver Riddle     numbering->asmInterface = dyn_cast<OpAsmDialectInterface>(dialect);
388f3acb54cSRiver Riddle   }
389f3acb54cSRiver Riddle   return *numbering;
390f3acb54cSRiver Riddle }
391f3acb54cSRiver Riddle 
392f3acb54cSRiver Riddle auto IRNumberingState::numberDialect(StringRef dialect) -> DialectNumbering & {
393f3acb54cSRiver Riddle   DialectNumbering *&numbering = dialects[dialect];
394f3acb54cSRiver Riddle   if (!numbering) {
395f3acb54cSRiver Riddle     numbering = new (dialectAllocator.Allocate())
396f3acb54cSRiver Riddle         DialectNumbering(dialect, dialects.size() - 1);
397f3acb54cSRiver Riddle   }
398f3acb54cSRiver Riddle   return *numbering;
399f3acb54cSRiver Riddle }
400f3acb54cSRiver Riddle 
401f3acb54cSRiver Riddle void IRNumberingState::number(Region &region) {
402f3acb54cSRiver Riddle   if (region.empty())
403f3acb54cSRiver Riddle     return;
404f3acb54cSRiver Riddle   size_t firstValueID = nextValueID;
405f3acb54cSRiver Riddle 
406f3acb54cSRiver Riddle   // Number the blocks within this region.
407f3acb54cSRiver Riddle   size_t blockCount = 0;
4088c258fdaSJakub Kuderski   for (auto it : llvm::enumerate(region)) {
409f3acb54cSRiver Riddle     blockIDs.try_emplace(&it.value(), it.index());
410f3acb54cSRiver Riddle     number(it.value());
411f3acb54cSRiver Riddle     ++blockCount;
412f3acb54cSRiver Riddle   }
413f3acb54cSRiver Riddle 
414f3acb54cSRiver Riddle   // Remember the number of blocks and values in this region.
415f3acb54cSRiver Riddle   regionBlockValueCounts.try_emplace(&region, blockCount,
416f3acb54cSRiver Riddle                                      nextValueID - firstValueID);
417f3acb54cSRiver Riddle }
418f3acb54cSRiver Riddle 
419f3acb54cSRiver Riddle void IRNumberingState::number(Operation &op) {
420f3acb54cSRiver Riddle   // Number the components of an operation that won't be numbered elsewhere
421f3acb54cSRiver Riddle   // (e.g. we don't number operands, regions, or successors here).
422f3acb54cSRiver Riddle   number(op.getName());
423f3acb54cSRiver Riddle   for (OpResult result : op.getResults()) {
424f3acb54cSRiver Riddle     valueIDs.try_emplace(result, nextValueID++);
425f3acb54cSRiver Riddle     number(result.getType());
426f3acb54cSRiver Riddle   }
427f3acb54cSRiver Riddle 
4285ed11e76SAlex Zinenko   // Prior to a version with native property encoding, or when properties are
4295ed11e76SAlex Zinenko   // not used, we need to number also the merged dictionary containing both the
4305ed11e76SAlex Zinenko   // inherent and discardable attribute.
431fb771fe3SJeff Niu   DictionaryAttr dictAttr;
432fb771fe3SJeff Niu   if (config.getDesiredBytecodeVersion() >= bytecode::kNativePropertiesEncoding)
433fb771fe3SJeff Niu     dictAttr = op.getRawDictionaryAttrs();
434fb771fe3SJeff Niu   else
435660f714eSMehdi Amini     dictAttr = op.getAttrDictionary();
436fb771fe3SJeff Niu   // Only number the operation's dictionary if it isn't empty.
437f3acb54cSRiver Riddle   if (!dictAttr.empty())
438f3acb54cSRiver Riddle     number(dictAttr);
439f3acb54cSRiver Riddle 
440660f714eSMehdi Amini   // Visit the operation properties (if any) to make sure referenced attributes
441660f714eSMehdi Amini   // are numbered.
442fb771fe3SJeff Niu   if (config.getDesiredBytecodeVersion() >=
443fb771fe3SJeff Niu           bytecode::kNativePropertiesEncoding &&
444660f714eSMehdi Amini       op.getPropertiesStorageSize()) {
445660f714eSMehdi Amini     if (op.isRegistered()) {
446660f714eSMehdi Amini       // Operation that have properties *must* implement this interface.
447660f714eSMehdi Amini       auto iface = cast<BytecodeOpInterface>(op);
4487ad9e9dcSMatteo Franciolini       NumberingDialectWriter writer(*this, config.getDialectVersionMap());
449660f714eSMehdi Amini       iface.writeProperties(writer);
450660f714eSMehdi Amini     } else {
451660f714eSMehdi Amini       // Unregistered op are storing properties as an optional attribute.
452660f714eSMehdi Amini       if (Attribute prop = *op.getPropertiesStorage().as<Attribute *>())
453660f714eSMehdi Amini         number(prop);
454660f714eSMehdi Amini     }
455660f714eSMehdi Amini   }
456660f714eSMehdi Amini 
457f3acb54cSRiver Riddle   number(op.getLoc());
458f3acb54cSRiver Riddle }
459f3acb54cSRiver Riddle 
460f3acb54cSRiver Riddle void IRNumberingState::number(OperationName opName) {
461f3acb54cSRiver Riddle   OpNameNumbering *&numbering = opNames[opName];
462f3acb54cSRiver Riddle   if (numbering) {
463f3acb54cSRiver Riddle     ++numbering->refCount;
464f3acb54cSRiver Riddle     return;
465f3acb54cSRiver Riddle   }
466f3acb54cSRiver Riddle   DialectNumbering *dialectNumber = nullptr;
467f3acb54cSRiver Riddle   if (Dialect *dialect = opName.getDialect())
468f3acb54cSRiver Riddle     dialectNumber = &numberDialect(dialect);
469f3acb54cSRiver Riddle   else
470f3acb54cSRiver Riddle     dialectNumber = &numberDialect(opName.getDialectNamespace());
471f3acb54cSRiver Riddle 
472f3acb54cSRiver Riddle   numbering =
473f3acb54cSRiver Riddle       new (opNameAllocator.Allocate()) OpNameNumbering(dialectNumber, opName);
474f3acb54cSRiver Riddle   orderedOpNames.push_back(numbering);
475f3acb54cSRiver Riddle }
476f3acb54cSRiver Riddle 
477f3acb54cSRiver Riddle void IRNumberingState::number(Type type) {
478f3acb54cSRiver Riddle   auto it = types.insert({type, nullptr});
479f3acb54cSRiver Riddle   if (!it.second) {
480f3acb54cSRiver Riddle     ++it.first->second->refCount;
481f3acb54cSRiver Riddle     return;
482f3acb54cSRiver Riddle   }
483f3acb54cSRiver Riddle   auto *numbering = new (typeAllocator.Allocate()) TypeNumbering(type);
484f3acb54cSRiver Riddle   it.first->second = numbering;
485f3acb54cSRiver Riddle   orderedTypes.push_back(numbering);
486f3acb54cSRiver Riddle 
487f3acb54cSRiver Riddle   // Check for OpaqueType, which is a dialect-specific type that didn't have a
488f3acb54cSRiver Riddle   // registered dialect when it got created. We don't want to encode this as the
489f3acb54cSRiver Riddle   // builtin OpaqueType, we want to encode it as if the dialect was actually
490f3acb54cSRiver Riddle   // loaded.
4915550c821STres Popp   if (OpaqueType opaqueType = dyn_cast<OpaqueType>(type)) {
492f3acb54cSRiver Riddle     numbering->dialect = &numberDialect(opaqueType.getDialectNamespace());
49302c2ecb9SRiver Riddle     return;
49402c2ecb9SRiver Riddle   }
495f3acb54cSRiver Riddle   numbering->dialect = &numberDialect(&type.getDialect());
49602c2ecb9SRiver Riddle 
49702c2ecb9SRiver Riddle   // If this type will be emitted using the bytecode format, perform a dummy
49802c2ecb9SRiver Riddle   // writing to number any nested components.
49902c2ecb9SRiver Riddle   // TODO: We don't allow custom encodings for mutable types right now.
5006ab2bcffSRiver Riddle   if (!type.hasTrait<TypeTrait::IsMutable>()) {
501bff6a429SMatteo Franciolini     // Try overriding emission with callbacks.
502bff6a429SMatteo Franciolini     for (const auto &callback : config.getTypeWriterCallbacks()) {
5037ad9e9dcSMatteo Franciolini       NumberingDialectWriter writer(*this, config.getDialectVersionMap());
504bff6a429SMatteo Franciolini       // The client has the ability to override the group name through the
505bff6a429SMatteo Franciolini       // callback.
506bff6a429SMatteo Franciolini       std::optional<StringRef> groupNameOverride;
507bff6a429SMatteo Franciolini       if (succeeded(callback->write(type, groupNameOverride, writer))) {
508bff6a429SMatteo Franciolini         if (groupNameOverride.has_value())
509bff6a429SMatteo Franciolini           numbering->dialect = &numberDialect(*groupNameOverride);
510bff6a429SMatteo Franciolini         return;
511bff6a429SMatteo Franciolini       }
512bff6a429SMatteo Franciolini     }
513bff6a429SMatteo Franciolini 
514bff6a429SMatteo Franciolini     // If this attribute will be emitted using the bytecode format, perform a
515bff6a429SMatteo Franciolini     // dummy writing to number any nested components.
516bff6a429SMatteo Franciolini     if (const auto *interface = numbering->dialect->interface) {
5177ad9e9dcSMatteo Franciolini       NumberingDialectWriter writer(*this, config.getDialectVersionMap());
5186ab2bcffSRiver Riddle       if (succeeded(interface->writeType(type, writer)))
5196ab2bcffSRiver Riddle         return;
5206ab2bcffSRiver Riddle     }
5216ab2bcffSRiver Riddle   }
5226ab2bcffSRiver Riddle   // If this type will be emitted using the fallback, number the nested dialect
5236ab2bcffSRiver Riddle   // resources. We don't number everything (e.g. no nested attributes/types),
5246ab2bcffSRiver Riddle   // because we don't want to encode things we won't decode (the textual format
5256ab2bcffSRiver Riddle   // can't really share much).
5266ab2bcffSRiver Riddle   AsmState tempState(type.getContext());
5276ab2bcffSRiver Riddle   llvm::raw_null_ostream dummyOS;
5286ab2bcffSRiver Riddle   type.print(dummyOS, tempState);
5296ab2bcffSRiver Riddle 
5306ab2bcffSRiver Riddle   // Number the used dialect resources.
5316ab2bcffSRiver Riddle   for (const auto &it : tempState.getDialectResources())
5326ab2bcffSRiver Riddle     number(it.getFirst(), it.getSecond().getArrayRef());
5336ab2bcffSRiver Riddle }
5346ab2bcffSRiver Riddle 
5356ab2bcffSRiver Riddle void IRNumberingState::number(Dialect *dialect,
5366ab2bcffSRiver Riddle                               ArrayRef<AsmDialectResourceHandle> resources) {
5376ab2bcffSRiver Riddle   DialectNumbering &dialectNumber = numberDialect(dialect);
5386ab2bcffSRiver Riddle   assert(
5396ab2bcffSRiver Riddle       dialectNumber.asmInterface &&
5406ab2bcffSRiver Riddle       "expected dialect owning a resource to implement OpAsmDialectInterface");
5416ab2bcffSRiver Riddle 
5426ab2bcffSRiver Riddle   for (const auto &resource : resources) {
5436ab2bcffSRiver Riddle     // Check if this is a newly seen resource.
5446ab2bcffSRiver Riddle     if (!dialectNumber.resources.insert(resource))
54502c2ecb9SRiver Riddle       return;
54602c2ecb9SRiver Riddle 
5476ab2bcffSRiver Riddle     auto *numbering =
5486ab2bcffSRiver Riddle         new (resourceAllocator.Allocate()) DialectResourceNumbering(
5496ab2bcffSRiver Riddle             dialectNumber.asmInterface->getResourceKey(resource));
5506ab2bcffSRiver Riddle     dialectNumber.resourceMap.insert({numbering->key, numbering});
5516ab2bcffSRiver Riddle     dialectResources.try_emplace(resource, numbering);
5526ab2bcffSRiver Riddle   }
5536ab2bcffSRiver Riddle }
5546ab2bcffSRiver Riddle 
5559ea6b30aSMehdi Amini int64_t IRNumberingState::getDesiredBytecodeVersion() const {
5569ea6b30aSMehdi Amini   return config.getDesiredBytecodeVersion();
5579ea6b30aSMehdi Amini }
5589ea6b30aSMehdi Amini 
5596ab2bcffSRiver Riddle namespace {
5606ab2bcffSRiver Riddle /// A dummy resource builder used to number dialect resources.
5616ab2bcffSRiver Riddle struct NumberingResourceBuilder : public AsmResourceBuilder {
5626ab2bcffSRiver Riddle   NumberingResourceBuilder(DialectNumbering *dialect, unsigned &nextResourceID)
5636ab2bcffSRiver Riddle       : dialect(dialect), nextResourceID(nextResourceID) {}
5646ab2bcffSRiver Riddle   ~NumberingResourceBuilder() override = default;
5656ab2bcffSRiver Riddle 
5666ab2bcffSRiver Riddle   void buildBlob(StringRef key, ArrayRef<char>, uint32_t) final {
5676ab2bcffSRiver Riddle     numberEntry(key);
5686ab2bcffSRiver Riddle   }
5696ab2bcffSRiver Riddle   void buildBool(StringRef key, bool) final { numberEntry(key); }
5706ab2bcffSRiver Riddle   void buildString(StringRef key, StringRef) final {
5716ab2bcffSRiver Riddle     // TODO: We could pre-number the value string here as well.
5726ab2bcffSRiver Riddle     numberEntry(key);
5736ab2bcffSRiver Riddle   }
5746ab2bcffSRiver Riddle 
5756ab2bcffSRiver Riddle   /// Number the dialect entry for the given key.
5766ab2bcffSRiver Riddle   void numberEntry(StringRef key) {
5776ab2bcffSRiver Riddle     // TODO: We could pre-number resource key strings here as well.
5786ab2bcffSRiver Riddle 
5795e458f5aSMehdi Amini     auto *it = dialect->resourceMap.find(key);
5806ab2bcffSRiver Riddle     if (it != dialect->resourceMap.end()) {
5816ab2bcffSRiver Riddle       it->second->number = nextResourceID++;
5826ab2bcffSRiver Riddle       it->second->isDeclaration = false;
5836ab2bcffSRiver Riddle     }
5846ab2bcffSRiver Riddle   }
5856ab2bcffSRiver Riddle 
5866ab2bcffSRiver Riddle   DialectNumbering *dialect;
5876ab2bcffSRiver Riddle   unsigned &nextResourceID;
5886ab2bcffSRiver Riddle };
5896ab2bcffSRiver Riddle } // namespace
5906ab2bcffSRiver Riddle 
5916ab2bcffSRiver Riddle void IRNumberingState::finalizeDialectResourceNumberings(Operation *rootOp) {
5926ab2bcffSRiver Riddle   unsigned nextResourceID = 0;
5936ab2bcffSRiver Riddle   for (DialectNumbering &dialect : getDialects()) {
5946ab2bcffSRiver Riddle     if (!dialect.asmInterface)
5956ab2bcffSRiver Riddle       continue;
5966ab2bcffSRiver Riddle     NumberingResourceBuilder entryBuilder(&dialect, nextResourceID);
5976ab2bcffSRiver Riddle     dialect.asmInterface->buildResources(rootOp, dialect.resources,
5986ab2bcffSRiver Riddle                                          entryBuilder);
5996ab2bcffSRiver Riddle 
6006ab2bcffSRiver Riddle     // Number any resources that weren't added by the dialect. This can happen
6016ab2bcffSRiver Riddle     // if there was no backing data to the resource, but we still want these
6026ab2bcffSRiver Riddle     // resource references to roundtrip, so we number them and indicate that the
6036ab2bcffSRiver Riddle     // data is missing.
6046ab2bcffSRiver Riddle     for (const auto &it : dialect.resourceMap)
6056ab2bcffSRiver Riddle       if (it.second->isDeclaration)
6066ab2bcffSRiver Riddle         it.second->number = nextResourceID++;
60702c2ecb9SRiver Riddle   }
608f3acb54cSRiver Riddle }
609