Implement IR versioning through post-parsing upgrade through OpAsmDialectInterface

[mehdi_amini]

Summary

Dialect can opt-in to handle versioning in a custom way. The input IR may start
with the keyword dialect_versions, followed by a dictionary attribute where
the keys are dialect names and the value are attributes representing the
producer version for the given dialect. There is no restriction on what kind of
attribute a dialect is using to model its versioning.

On parsing the version is made available through the parser, and a dialect for
which a version is present has the ability to use the stored version to upgrade
the IR.

[rriddle]

One thing that this doesn't really cover right now is how the assembly syntax itself can differ between versions. Realistically, if we do allow versioning the IR I would also assume that we may also want to be able to tag versions of the IR syntax. You wouldn't want to ship all old versions of the IR syntax with your current dialect obviously, but we could conceptually have mechanisms of linking in extensions that support parsing old versions? The alternative there obviously is that we just say "you only get versioned output when using the generic form". I see all of that as separable though, upgrading an in-memory form of an operation is separable from parsing and upgrading an old syntax; though they could be hooked into via the same versioning mechanism I assume.

mlir/include/mlir/IR/OpImplementation.h

1352

I like using an attribute for this as it leaves it all up to the dialect to define. Do we impose any kind of constraints on this though? i.e., if these were dialect attributes, would they need to be upgraded as well (if the dialect changed)?

1353–1356

What are the constraints on what this is allowed to modify? Can it touch anything? I guess in some cases it kind of should be able to (e.g. if old argument attributes are now represented via some other contract, the dialect would want to be able to update the function).

mlir/lib/IR/AsmPrinter.cpp

2505

Not important for now, but it'd be nice to have some reserved syntax for directives (helps prevent confusing it for an operation or some other IR construct).

[mehdi_amini]

One thing that this doesn't really cover right now is how the assembly syntax itself can differ between versions

Oh right, there is something I haven't plumbed yet: exposing an accessor for the version through the AsmParser, so that custom parsing function can query it before parsing a given operation.

mlir/include/mlir/IR/OpImplementation.h

1352

I'd say it's up to the dialect to figure it out?

1353–1356

I didn't think about putting any constraint right now, I couldn't foresee a reason to limit this.

mlir/lib/IR/AsmPrinter.cpp

2505

Yeah I wrote this anticipating that you would complain about it, so that I could ask you for a suggestion ;)

[rriddle]

mlir/include/mlir/IR/OpImplementation.h

1352

I would be fine leaving it up to the dialect, but it still feels like we need to understand the implications here: just so that dialect authors don't feel fancy, only for it to break in some way down the line.

1353–1356

I'm fine with leaving it open, but it feels like we need some kind of order for upgrades, otherwise we might end up some weird situations with cross dialect things. For example, with the example case I picked above: what if the new way of representing a dialect attribute requires some new construct that is only available in a new version; what happens if the dialect attribute is upgraded first? will these things even compose(i.e. will we crash?) or will the upgrade for the function potentially destroy/overwrite that? We may need some other callback that acts on dialect attributes that gets called after an operation is upgraded? I don't entirely know, but feels like we should sketch out some potential cross-interaction cases.

mlir/lib/IR/AsmPrinter.cpp

2505

Dang. I naturally would just go with #directive, but unfortunately we use that for attributes... We could go with something similar to rust metadata: #![thing = ...]/#!thing = .... I'd say if we can reuse a scheme from an existing language, that makes sense (from a perspective of reusing existing conventions).

[mehdi_amini]

In D117761#3260333, @mehdi_amini wrote:

One thing that this doesn't really cover right now is how the assembly syntax itself can differ between versions

Oh right, there is something I haven't plumbed yet: exposing an accessor for the version through the AsmParser, so that custom parsing function can query it before parsing a given operation.

Nevermind, I already exposed it: virtual Attribute getDialectVersion(StringRef dialect) = 0; is available on the AsmParser, someone really motivated can use this in their custom parser!

[mehdi_amini]

Rebase, document, and add a test.

[stephenneuendorffer]

mlir/include/mlir/IR/OpImplementation.h

1361–1369

This seems awkward to me, since it assumes that the syntax can always be parsed. Why not just rely on the parser to do this without a separate hook?

mlir/lib/IR/AsmPrinter.cpp

2505

+1 for some distinguished syntax here.

[stellaraccident]

In D117761#3260322, @rriddle wrote:

One thing that this doesn't really cover right now is how the assembly syntax itself can differ between versions. Realistically, if we do allow versioning the IR I would also assume that we may also want to be able to tag versions of the IR syntax. You wouldn't want to ship all old versions of the IR syntax with your current dialect obviously, but we could conceptually have mechanisms of linking in extensions that support parsing old versions? The alternative there obviously is that we just say "you only get versioned output when using the generic form". I see all of that as separable though, upgrading an in-memory form of an operation is separable from parsing and upgrading an old syntax; though they could be hooked into via the same versioning mechanism I assume.

I was thinking that there could be an an addition version like "assembly_version" or something. Could just be hard-coded in the source code and we bump it on breaking printer/parser changes? The naming of "dialect_versions" somewhat precludes putting such a thing in the dictionary, but maybe one special exception spelled to not conflict with a dialect name?

[stellaraccident]

mlir/lib/Parser/Parser.cpp

2159

sp. "dictionary"

[mehdi_amini]

mlir/include/mlir/IR/OpImplementation.h

1361–1369

Because the main point is to support the generic IR format. Also, you can rename operation, or even delete an operation from a dialect, while being able to "upgrade it" to the new form.

[jpienaar]

mlir/lib/IR/AsmPrinter.cpp

2505

May I suggest https://downloads.haskell.org/~ghc/7.0.4/docs/html/users_guide/pragmas.html , it is "foreign" enough that it sticks out (and can't be confused with IR), these naturally feel like file scope pragmas, it is familiar to some (but not as general (i think :-)) as Rust), and i don't think any parsing ambiguity at file scope.

[jpienaar]

mlir/docs/LangRef.md

836 ↗

(On Diff #401883)

This should probably now be part of config section

mlir/include/mlir/IR/OpImplementation.h

1361

This covers the main use case but makes the current version of dialect implicit/the hook always gets fed latest. Is that sufficient? E.g., if one wanted to upgrade to an older version (say i want to use this as part of upgrade to certain version for deployment) it means you need to be synced to that version when building and so if you have multiple target versions, you'd need binary per?

mlir/lib/IR/AsmPrinter.cpp

2505

External resources uses this syntax, so if this is moved there then perfect

[mehdi_amini]

mlir/include/mlir/IR/OpImplementation.h

1361

Yes: this only offers the ability for the compiler to support the "latest version of the dialect" and upgrade older versions to the current one supported by the compiler.

It seems that what you're describing is a bit more involved: wouldn't that account to include the notion of versioning to the ODS definitions directly somehow?

[jpienaar]

mlir/include/mlir/IR/OpImplementation.h

1361

Would it? E.g., in the simple case I see this as a sequence of if statements in upgrade hook which incrementally updates. So initially no hooks , then at version Foo one needs to introduce upgrades, so upgrade from < Foo to Foo. Now at version Baz another bump is needed so add upgrade from Foo to Baz. Of course one could rewrite upgrade hooks for every version and so have more efficient ones and that means you only have head as target possible. These should mostly not be on the critical path though (if you are actually serving, run update offline and reserialize is option rather than adding any cost). But I don't want to restrict being able to do optimized hooks and for those one can't guarantee it'll be anything other than input or head.

[mehdi_amini]

mlir/include/mlir/IR/OpImplementation.h

1361

But regardless of the input, we still upgrade to the current latest version of the dialect right?

I guess I don't get what you meant with "one wanted to upgrade to an older version"?
I thought you meant I want my 4.x compiler to be able to load version 1.x and decide at runtime with a flag whether it'll upgrade to 3.x or 4.x : that would require to support two versions of the dialect in the current compiler.

mlir/include/mlir/IR/OpImplementation.h

ParseResult getCurrentLocation(llvm::SMLoc *loc) {
     return success();
   }
 
   /// Re-encode the given source location as an MLIR location and return it.
   /// Note: This method should only be used when a `Location` is necessary, as
   /// the encoding process is not efficient.
   virtual Location getEncodedSourceLoc(llvm::SMLoc loc) = 0;
 
+  virtual Attribute getDialectVersion(StringRef dialect) = 0;
+
   //===--------------------------------------------------------------------===//
   // Token Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse a '->' token.
   virtual ParseResult parseArrow() = 0;
 
   /// Parse a '->' token if present
public:
   /// end with a numeric digit([0-9]+).
   virtual AliasResult getAlias(Attribute attr, raw_ostream &os) const {
     return AliasResult::NoAlias;
   }
   virtual AliasResult getAlias(Type type, raw_ostream &os) const {
     return AliasResult::NoAlias;
   }
 
+  virtual Attribute getProducerVersion() const { return {}; }
+  virtual LogicalResult upgradeFromVersion(Operation *topLevelOp,
+                                           Attribute attr) {
+    return success();
+  }
 };
 } // namespace mlir
 
 //===--------------------------------------------------------------------===//
 // Operation OpAsm interface.
 //===--------------------------------------------------------------------===//
 
 /// The OpAsmOpInterface, see OpAsmInterface.td for more details.
 #include "mlir/IR/OpAsmInterface.h.inc"
 
 #endif

mlir/lib/IR/AsmPrinter.cpp

private:
   const static unsigned indentWidth = 2;
 
   // This is the current indentation level for nested structures.
   unsigned currentIndent = 0;
 };
 } // namespace
 
 void OperationPrinter::printTopLevelOperation(Operation *op) {
+  auto getVersion = [&](Dialect *dialect) -> Attribute {
+    const OpAsmDialectInterface *asmIface = state->getOpAsmInterface(dialect);
+    if (!asmIface)
+      return {};
+    return asmIface->getProducerVersion();
+  };
+
+  std::vector<Dialect *> loadedDialects = op->getContext()->getLoadedDialects();
+  llvm::erase_if(loadedDialects,
+                 [&](Dialect *dialect) { return !getVersion(dialect); });
+  if (llvm::any_of(loadedDialects, getVersion)) {
+    os << "dialect_versions { ";
+    interleaveComma(loadedDialects, [&](Dialect *dialect) {
+      if (Attribute version = getVersion(dialect))
+        os << dialect->getNamespace() << " = " << version;
+    });
+    os << " }" << newLine;
+  }
+
   // Output the aliases at the top level that can't be deferred.
   state->getAliasState().printNonDeferredAliases(os, newLine);
 
   // Print the module.
   print(op);
   os << newLine;
 
   // Output the aliases at the top level that can be deferred.

mlir/lib/Parser/AsmParserImpl.h

llvm::SMLoc getCurrentLocation() override {
     return parser.getToken().getLoc();
   }
 
   /// Re-encode the given source location as an MLIR location and return it.
   Location getEncodedSourceLoc(llvm::SMLoc loc) override {
     return parser.getEncodedSourceLocation(loc);
   }
 
+  Attribute getDialectVersion(StringRef dialect) override {
+    return parser.getDialectVersion(dialect);
+  }
+
   //===--------------------------------------------------------------------===//
   // Token Parsing
   //===--------------------------------------------------------------------===//
 
   using Delimiter = AsmParser::Delimiter;
 
   /// Parse a `->` token.
   ParseResult parseArrow() override {

mlir/lib/Parser/Parser.h

public:
   Builder builder;
 
   Parser(ParserState &state) : builder(state.context), state(state) {}
 
   // Helper methods to get stuff from the parser-global state.
   ParserState &getState() const { return state; }
   MLIRContext *getContext() const { return state.context; }
   const llvm::SourceMgr &getSourceMgr() { return state.lex.getSourceMgr(); }
+  Attribute getDialectVersion(StringRef dialect) const {
+    return state.dialectVersions.get(dialect);
+  }
 
   /// Parse a comma-separated list of elements up until the specified end token.
   ParseResult
   parseCommaSeparatedListUntil(Token::Kind rightToken,
                                function_ref<ParseResult()> parseElement,
                                bool allowEmptyList = true);
 
   /// Parse a list of comma-separated items with an optional delimiter.  If a

mlir/lib/Parser/Parser.cpp

ParseResult TopLevelOperationParser::parseTypeAliasDef() {
   return success();
 }
 
 ParseResult TopLevelOperationParser::parse(Block *topLevelBlock,
                                            Location parserLoc) {
   // Create a top-level operation to contain the parsed state.
   OwningOpRef<ModuleOp> topLevelOp(ModuleOp::create(parserLoc));
   OperationParser opParser(state, topLevelOp.get());
+
+  if (getToken().getSpelling() == "dialect_versions") {
+    consumeToken();
+    Attribute versions = opParser.parseAttribute();
+    if (!versions)
+      return failure();
+    state.dialectVersions = versions.dyn_cast<DictionaryAttr>();
+    if (!state.dialectVersions) {
+      emitError("Expects Dictionary attribute for dialect versions");
+      return failure();
+    }
+  }
+
   while (true) {
     switch (getToken().getKind()) {
     default:
       // Parse a top-level operation.
       if (opParser.parseOperation())
         return failure();
       break;
 
     // If we got to the end of the file, then we're done.
     case Token::eof: {
       if (opParser.finalize())
         return failure();
 
       // Splice the blocks of the parsed operation over to the provided
       // top-level block.
       auto &parsedOps = topLevelOp->getBody()->getOperations();
       auto &destOps = topLevelBlock->getOperations();
       destOps.splice(destOps.empty() ? destOps.end() : std::prev(destOps.end()),
                      parsedOps, parsedOps.begin(), parsedOps.end());
+
+      for (NamedAttribute dialectVersion : state.dialectVersions) {
+        Dialect *dialect =
+            getContext()->getLoadedDialect(dialectVersion.getName());
+        if (!dialect)
+          continue;
+        auto *asmIface = dyn_cast<OpAsmDialectInterface>(dialect);
+        if (!asmIface)
+          continue;
+        if (failed(asmIface->upgradeFromVersion(topLevelOp.get(),
+                                                dialectVersion.getValue())))
+          return failure();
+      }
       return success();
     }
 
     // If we got an error token, then the lexer already emitted an error, just
     // stop.  Someday we could introduce error recovery if there was demand
     // for it.
     case Token::error:
       return failure();

mlir/lib/Parser/ParserState.h

struct ParserState {
   AsmParserState *asmState;
 
   // Contains the stack of default dialect to use when parsing regions.
   // A new dialect get pushed to the stack before parsing regions nested
   // under an operation implementing `OpAsmOpInterface`, and
   // popped when done. At the top-level we start with "builtin" as the
   // default, so that the top-level `module` operation parses as-is.
   SmallVector<StringRef> defaultDialectStack{"builtin"};
+
+  DictionaryAttr dialectVersions;
 };
 
 } // namespace detail
 } // namespace mlir
 
 #endif // MLIR_LIB_PARSER_PARSERSTATE_H

mlir/test/lib/Dialect/Test/TestDialect.cpp

if (auto intType = type.dyn_cast<TestIntegerType>()) {
               TestIntegerType::SignednessSemantics::Unsigned &&
           intType.getWidth() == 8) {
         os << "test_ui8";
         return AliasResult::FinalAlias;
       }
     }
     return AliasResult::NoAlias;
   }
+
+  Attribute getProducerVersion() const final {
+    return Builder(getDialect()->getContext()).getI32IntegerAttr(42);
+  }
 };
 
 struct TestDialectFoldInterface : public DialectFoldInterface {
   using DialectFoldInterface::DialectFoldInterface;
 
   /// Registered hook to check if the given region, which is attached to an
   /// operation that is *not* isolated from above, should be used when
   /// materializing constants.