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

1358

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)?

1359–1362

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

2501

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

1358

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

1359–1362

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

mlir/lib/IR/AsmPrinter.cpp

2501

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

1358

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.

1359–1362

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

2501

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

1367–1375

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

2501

+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

2190

sp. "dictionary"

[mehdi_amini]

mlir/include/mlir/IR/OpImplementation.h

1367–1375

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

2501

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

This should probably now be part of config section

mlir/include/mlir/IR/OpImplementation.h

1367

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

2501

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

[mehdi_amini]

mlir/include/mlir/IR/OpImplementation.h

1367

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

1367

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

1367

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/docs/LangRef.md

 attribute values.
 
 ### Builtin Attribute Values
 
 The [builtin dialect](Dialects/Builtin.md) defines a set of attribute values
 that are directly usable by any other dialect in MLIR. These types cover a range
 from primitive integer and floating-point values, attribute dictionaries, dense
 multi-dimensional arrays, and more.
+
+### IR Versionning
+
+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.

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;
 
+  /// Return the attribute describing the version of the provided dialect name,
+  /// if any. The version is provided by the `dialect_versions` directive at the
+  /// very beginning of the parsing.
+  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;
   }
 
+  /// Hook provided by the dialect to emit a version when printing. The
+  /// attribute will be available when parsing back, and the dialect
+  /// implementation will be able to use it to load previous known version.
+  /// This management is entirely under the responsibility of the individual
+  /// dialects, and limited by possible changes in the MLIR syntax.
+  /// There is no restriction on what kind of attribute a dialect is using to
+  /// model its versioning.
+  virtual Attribute getProducerVersion() const { return {}; }
+
+  /// Hook invoked after parsing completed, if a version directive was present
+  /// and included an entry for the current dialect. This hook offers the
+  /// opportunity to the dialect to visit the IR and upgrades constructs emitted
+  /// by the version of the dialect corresponding to the provided version
+  /// attribute.
+  virtual LogicalResult upgradeFromVersion(Operation *topLevelOp,
+                                           Attribute producerVersion) const {
+    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) {
+
+  // Retrieve the version to print for this dialect, if any.
+  auto getVersion = [&](Dialect *dialect) -> Attribute {
+    const OpAsmDialectInterface *asmIface = state->getOpAsmInterface(dialect);
+    if (!asmIface)
+      return {};
+    return asmIface->getProducerVersion();
+  };
+
+  // If any dialect has a version to print, print the `dialect_versions`
+  // directive. This directive has to be the very first directive in the file
+  // and introduce a dictionary 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.
+  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 OperationParser::finalize() {
   });
   if (walkRes.wasInterrupted())
     return failure();
 
   // Pop the top level name scope.
   if (failed(popSSANameScope()))
     return failure();
 
+  // Parsing is complete, give an opportunity to each dialect to visit the
+  // IR and perform upgrades.
+  if (state.dialectVersions) {
+    for (NamedAttribute dialectVersion : state.dialectVersions) {
+      Dialect *dialect =
+          getContext()->getLoadedDialect(dialectVersion.getName());
+      if (!dialect)
+        continue;
+      auto *asmIface = dialect->getRegisteredInterface<OpAsmDialectInterface>();
+      if (!asmIface)
+        continue;
+      if (failed(asmIface->upgradeFromVersion(topLevelOp,
+                                              dialectVersion.getValue())))
+        return failure();
+    }
+  }
+
   // Verify that the parsed operations are valid.
   if (failed(verify(topLevelOp)))
     return failure();
 
   // If we are populating the parser state, finalize the top-level operation.
   if (state.asmState)
     state.asmState->finalize(topLevelOp);
   return success();
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 the input starts with a `dialect_versions` keyword, we expect a
+  // dictionnary representing the version of the dialect at the time the IR was
+  // produced. This will be used for possibly upgrading the IR when parsing
+  // completes.
+  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());
+
       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/IR/ir_upgrade.mlir

+// RUN: mlir-opt -split-input-file --verify-diagnostics %s | FileCheck %s
+
+dialect_versions { test = 41 }
+
+// -----
+
+// CHECK: dialect_versions { test = 42 : i32 }
+// CHECK:  "test.versionedA"() {dims = 123 : i64} : () -> ()
+"test.versionedA"() {dims = 123 : i64} : () -> ()
+
+// -----
+
+// CHECK: dialect_versions { test = 42 : i32 }
+// CHECK: "test.versionedA"() {dims = 123 : i64} : () -> ()
+dialect_versions { test = 41 }
+"test.versionedA"() {dimensions = 123 : i64} : () -> ()
+
+// -----
+
+// CHECK: dialect_versions { test = 42 : i32 }
+// CHECK: "test.versionedA"() {dims = 123 : i64} : () -> ()
+dialect_versions { test = 41 }
+"test.versionedA"() {dims = 123 : i64} : () -> ()
+
+// -----
+
+// CHECK: dialect_versions { test = 42 : i32 }
+// CHECK: test.versionedB current_version
+dialect_versions { test = 41 }
+test.versionedB  deprecated_syntax
+
+// -----
+
+// CHECK: dialect_versions { test = 42 : i32 }
+// CHECK: test.versionedB current_version
+dialect_versions { test = 42 }
+test.versionedB  current_version
+
+// -----
+
+dialect_versions { test = 42 }
+// expected-error@+1{{custom op 'test.versionedB' expected 'current_version'}}
+test.versionedB  deprecated_syntax
+
+// -----
+// expected-error@-2{{Current test dialect version is 42, can't parse version: 43 : i64}}
+dialect_versions { test = 43 }
+"test.versionedA"() {dims = 123 : i64} : () -> ()
+
+

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);
+  }
+  LogicalResult upgradeFromVersion(Operation *topLevelOp,
+                                   Attribute producerVersion) const final {
+    auto version = producerVersion.dyn_cast<IntegerAttr>();
+    if (!version) {
+      return topLevelOp->emitError()
+             << "Expected an IntegerAttr for test dialect version, got: "
+             << producerVersion;
+    }
+    if (version.getInt() == 42)
+      return success();
+    if (version.getInt() > 42) {
+      return topLevelOp->emitError()
+             << "Current test dialect version is 42, can't parse version: "
+             << version;
+    }
+    topLevelOp->walk([](TestVersionedOpA op) {
+      if (auto dims = op->getAttr("dimensions")) {
+        op->removeAttr("dimensions");
+        op->setAttr("dims", dims);
+      }
+    });
+
+    return success();
+  }
 };
 
 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.

 static void print(SingleNoTerminatorCustomAsmOp op, OpAsmPrinter &printer) {
   printer.printRegion(
       op.getRegion(), /*printEntryBlockArgs=*/false,
       // This op has a single block without terminators. But explicitly mark
       // as not printing block terminators for testing.
       /*printBlockTerminators=*/false);
 }
 
+//===----------------------------------------------------------------------===//
+// Test dialect_version upgrade by supporting an old syntax
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseVersionedOp(OpAsmParser &parser,
+                                    OperationState &state) {
+  if (auto version =
+          parser.getDialectVersion("test").dyn_cast_or_null<IntegerAttr>()) {
+    if (version.getInt() < 42)
+      return parser.parseKeyword("deprecated_syntax");
+  }
+  return parser.parseKeyword("current_version");
+}
+
+static void print(TestVersionedOpB op, OpAsmPrinter &printer) {
+  printer << " current_version";
+}
+
 #include "TestOpEnums.cpp.inc"
 #include "TestOpInterfaces.cpp.inc"
 #include "TestOpStructs.cpp.inc"
 #include "TestTypeInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.cpp.inc"

mlir/test/lib/Dialect/Test/TestOps.td

let arguments = (ins
         );
 
     let results = (outs Res<AnyMemRef, "", [MemAlloc<TestResource>]>);
 }
 
 def TestEffectsOpB : TEST_Op<"op_with_effects_b",
     [MemoryEffects<[MemWrite<TestResource>]>]>;
 
+//===----------------------------------------------------------------------===//
+// Test Ops to upgrade base on the dialect versions
+//===----------------------------------------------------------------------===//
+
+def TestVersionedOpA : TEST_Op<"versionedA"> {
+    let arguments = (ins
+      // A previous version of this dialect used the name "dimensions" for this
+      // attribute, it got renamed but we support loading old IR through
+      // upgrading, see `upgradeFromVersion()` in `TestOpAsmInterface`.
+      AnyI64Attr:$dims
+    );
+}
+
+// This op will be able to parse based on an old syntax and "auto-upgrade".
+def TestVersionedOpB : TEST_Op<"versionedB"> {
+  let arguments = (ins);
+  let parser = [{ return ::parseVersionedOp(parser, result); }];
+  let printer = [{ return ::print(*this, p); }];
+}
+
+
 #endif // TEST_OPS