[mlir][bytecode] Avoid recording null arglocs & realloc opnames.

[jpienaar]

Summary

For block arg locs a common case is no/uknown location (where the producer
signifies they don't care about blockarg location). Also avoid needing to
dynamically resize opnames during parsing.

Assumed to be post lazy loading change, so chose version 3.

[mehdi_amini]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

2162–2163

This initialization seems unusual to me: isn't it the default?

2164

Can't this hasLoc be nested in the version >= 3 branch?

2167

I'm a bit puzzled: seems like this would be true is argType is valid? I don't quite get this condition.

2172

I don't quite get this "hasLoc" test here?

[jpienaar]

Update to head & make logic a bit clearer for parsing.

[rriddle]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

2171–2178

[jpienaar]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

2162–2163

I made these a bit more explicit and used more local vars.

2172

I first misunderstood your question, then noticed it can be folded in.

[jpienaar]

Simplify following suggestion

mlir/docs/BytecodeFormat.md

 The dialect section of the bytecode contains all of the dialects referenced
 within the encoded IR, and some information about the components of those
 dialects that were also referenced.
 
 ```
 dialect_section {
   numDialects: varint,
   dialectNames: varint[],
+  numTotalOpNames: varint,
   opNames: op_name_group[]
 }
 
 op_name_group {
   dialect: varint // (dialectID << 1) | (hasVersion),
   version : dialect_version_section
   numOpNames: varint,
   opNames: varint[]

 block_arguments {
   numArgs: varint?,
   args: block_argument[]
   numUseListOrders: varint?,
   useListOrders: uselist[],
 }
 
 block_argument {
-  typeIndex: varint,
-  location: varint
+  typeAndLocation: varint, // (type << 1) | (hasLocation)
+  location: varint? // Optional, else unknown location
 }
 ```
 
 A block is encoded with an array of operations and block arguments. The first
 field is an encoding that combines the number of operations in the block, with a
 flag indicating if the block has arguments.
 
 Use-list orders are attached to block arguments similarly to how they are
 attached to operation results.

mlir/include/mlir/Bytecode/Encoding.h

 // General constants
 //===----------------------------------------------------------------------===//
 
 enum {
   /// The minimum supported version of the bytecode.
   kMinSupportedVersion = 0,
 
   /// The current bytecode version.
-  kVersion = 3,
+  kVersion = 4,
 
   /// An arbitrary value used to fill alignment padding.
   kAlignmentByte = 0xCB,
 };
 
 //===----------------------------------------------------------------------===//
 // Sections
 //===----------------------------------------------------------------------===//

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

BytecodeReader::Impl::parseDialectSection(ArrayRef<uint8_t> sectionData) {
   // Parse the operation names, which are grouped by dialect.
   auto parseOpName = [&](BytecodeDialect *dialect) {
     StringRef opName;
     if (failed(stringReader.parseString(sectionReader, opName)))
       return failure();
     opNames.emplace_back(dialect, opName);
     return success();
   };
+  // Avoid re-allocation in bytecode version > 3 where the number of ops are
+  // known.
+  if (version > 3) {
+    uint64_t numOps;
+    if (failed(sectionReader.parseVarInt(numOps)))
+      return failure();
+    opNames.reserve(numOps);
+  }
   while (!sectionReader.empty())
     if (failed(parseDialectGrouping(sectionReader, dialects, parseOpName)))
       return failure();
   return success();
 }
 
 FailureOr<OperationName>
 BytecodeReader::Impl::parseOpName(EncodingReader &reader) {
LogicalResult BytecodeReader::Impl::parseBlockArguments(EncodingReader &reader,
   if (failed(reader.parseVarInt(numArgs)))
     return failure();
 
   SmallVector<Type> argTypes;
   SmallVector<Location> argLocs;
   argTypes.reserve(numArgs);
   argLocs.reserve(numArgs);
 
+  Location unknownLoc = UnknownLoc::get(config.getContext());
   while (numArgs--) {
     Type argType;
-    LocationAttr argLoc;
+    LocationAttr argLoc = unknownLoc;
+    if (version > 3) {
+      // Parse the type with hasLoc flag to determine if it has type.
+      uint64_t typeIdx;
+      bool hasLoc;
+      if (failed(reader.parseVarIntWithFlag(typeIdx, hasLoc)) ||
+          !(argType = attrTypeReader.resolveType(typeIdx)))
+        return failure();
+      if (hasLoc && failed(parseAttribute(reader, argLoc)))
+        return failure();
+    } else {
+      // All args has type and location.
       if (failed(parseType(reader, argType)) ||
           failed(parseAttribute(reader, argLoc)))
         return failure();
-
+    }
     argTypes.push_back(argType);
     argLocs.push_back(argLoc);
   }
   block->addArguments(argTypes, argLocs);
   return defineValues(reader, block->getArguments());
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

for (DialectNumbering &dialect : dialects) {
     // section.
     size_t versionAvailable = versionEmitter.size() > 0;
     dialectEmitter.emitVarIntWithFlag(nameID, versionAvailable);
     if (versionAvailable)
       dialectEmitter.emitSection(bytecode::Section::kDialectVersions,
                                  std::move(versionEmitter));
   }
 
+  if (config.bytecodeVersion > 3)
+    dialectEmitter.emitVarInt(size(numberingState.getOpNames()));
+
   // Emit the referenced operation names grouped by dialect.
   auto emitOpName = [&](OpNameNumbering &name) {
     dialectEmitter.emitVarInt(stringSection.insert(name.name.stripDialect()));
   };
   writeDialectGrouping(dialectEmitter, numberingState.getOpNames(), emitOpName);
 
   emitter.emitSection(bytecode::Section::kDialect, std::move(dialectEmitter));
 }
void BytecodeWriter::writeBlock(EncodingEmitter &emitter, Block *block) {
   // arguments.
   unsigned numOps = numberingState.getOperationCount(block);
   emitter.emitVarIntWithFlag(numOps, hasArgs);
 
   // Emit the arguments of the block.
   if (hasArgs) {
     emitter.emitVarInt(args.size());
     for (BlockArgument arg : args) {
+      Location argLoc = arg.getLoc();
+      if (config.bytecodeVersion > 3) {
+        emitter.emitVarIntWithFlag(numberingState.getNumber(arg.getType()),
+                                   !isa<UnknownLoc>(argLoc));
+        if (!isa<UnknownLoc>(argLoc))
+          emitter.emitVarInt(numberingState.getNumber(argLoc));
+      } else {
         emitter.emitVarInt(numberingState.getNumber(arg.getType()));
-      emitter.emitVarInt(numberingState.getNumber(arg.getLoc()));
+        emitter.emitVarInt(numberingState.getNumber(argLoc));
+      }
     }
     if (config.bytecodeVersion > 2) {
       uint64_t maskOffset = emitter.size();
       uint8_t encodingMask = 0;
       emitter.emitByte(0);
       writeUseListOrders(emitter, encodingMask, args);
       if (encodingMask)
         emitter.patchByte(maskOffset, encodingMask);
void BytecodeWriter::writeOp(EncodingEmitter &emitter, Operation *op) {
   // 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()) {
       // If the region is not isolated from above, or we are emitting bytecode
-      // targetting version <2, we don't use a section.
+      // targeting version <2, we don't use a section.
       if (!isIsolatedFromAbove || config.bytecodeVersion < 2) {
         writeRegion(emitter, &region);
         continue;
       }
 
       EncodingEmitter regionEmitter;
       writeRegion(regionEmitter, &region);
       emitter.emitSection(bytecode::Section::kIR, std::move(regionEmitter));

mlir/test/Bytecode/general.mlir

"bytecode.test1"() ({
   "bytecode.attributes"() {attra = 10, attrb = #bytecode.attr} : () -> ()
   %cst = "bytecode.sparse"() {value = sparse<[[2, 1], [1, 1], [1, 2]], [1.0, 5.0, 6.0]> : tensor<8x7xf32>} : () -> (tensor<8x7xf32>)
   test.graph_region {
     "bytecode.operands"(%results#0, %results#1, %results#2) : (i32, i64, i32) -> ()
     %results:3 = "bytecode.results"() : () -> (i32, i64, i32)
   }
   "bytecode.branch"()[^secondBlock] : () -> ()
 
-^secondBlock(%arg1: i32, %arg2: !bytecode.int, %arg3: !pdl.operation):
+^secondBlock(%arg1: i32 loc(unknown), %arg2: !bytecode.int, %arg3: !pdl.operation loc(unknown)):
   "bytecode.regions"() ({
     "bytecode.operands"(%arg1, %arg2, %arg3) : (i32, !bytecode.int, !pdl.operation) -> ()
     "bytecode.return"() : () -> ()
   }) : () -> ()
   "bytecode.return"() : () -> ()
 }) : () -> ()

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

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