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 ®ion : regions) 1534af01bf9SRiver Riddle numberContext.emplace_back(®ion, 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 ®ion) { 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(®ion, 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