[mlir][bytecode] Avoid crash when test dialect version is missing

[Men-cotton]

• Guard the test bytecode roundtrip writer/reader in mlir/test/lib/IR/TestBytecodeRoundtrip.cpp so --test-bytecode-roundtrip=test-dialect-version=2.0 no longer asserts when the test dialect version isn’t present.
• Add regression coverage for the missing-version scenario in mlir/test/Bytecode/roundtrip-missing-dialect.mlir.

Fixes: #128325
Fixes: #128321

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Men-cotton (Men-cotton)

Changes

• Guard the test bytecode roundtrip writer/reader in mlir/test/lib/IR/TestBytecodeRoundtrip.cpp so --test-bytecode-roundtrip=test-dialect-version=2.0 no longer asserts when the test dialect version isn’t present.
• Add regression coverage for the missing-version scenario in mlir/test/Bytecode/roundtrip-missing-dialect.mlir.

Fixes: #128325
Fixes: #128321

---

Full diff: https://github.com/llvm/llvm-project/pull/170173.diff

2 Files Affected:

• (added) mlir/test/Bytecode/roundtrip-missing-dialect.mlir (+6)
• (modified) mlir/test/lib/IR/TestBytecodeRoundtrip.cpp (+4-4)

diff --git a/mlir/test/Bytecode/roundtrip-missing-dialect.mlir b/mlir/test/Bytecode/roundtrip-missing-dialect.mlir
new file mode 100644
index 0000000000000..8a93fcfd55695
--- /dev/null
+++ b/mlir/test/Bytecode/roundtrip-missing-dialect.mlir
@@ -0,0 +1,6 @@
+// RUN: mlir-opt %s --test-bytecode-roundtrip=test-dialect-version=2.0 | FileCheck %s
+
+// CHECK-LABEL: func.func @main
+func.func @main() {
+  return
+}
diff --git a/mlir/test/lib/IR/TestBytecodeRoundtrip.cpp b/mlir/test/lib/IR/TestBytecodeRoundtrip.cpp
index 4894ad5294990..f6ee97831c9dd 100644
--- a/mlir/test/lib/IR/TestBytecodeRoundtrip.cpp
+++ b/mlir/test/lib/IR/TestBytecodeRoundtrip.cpp
@@ -141,8 +141,8 @@ struct TestBytecodeRoundtripPass
             DialectBytecodeWriter &writer) -> LogicalResult {
           // Do not override anything if version greater than 2.0.
           auto versionOr = writer.getDialectVersion<test::TestDialect>();
-          assert(succeeded(versionOr) && "expected reader to be able to access "
-                                         "the version for test dialect");
+          if (failed(versionOr))
+            return failure();
           const auto *version =
               reinterpret_cast<const test::TestDialectVersion *>(*versionOr);
           if (version->major_ >= 2)
@@ -166,8 +166,8 @@ struct TestBytecodeRoundtripPass
             Type &entry) -> LogicalResult {
           // Get test dialect version from the version map.
           auto versionOr = reader.getDialectVersion<test::TestDialect>();
-          assert(succeeded(versionOr) && "expected reader to be able to access "
-                                         "the version for test dialect");
+          if (failed(versionOr))
+            return success();
           const auto *version =
               reinterpret_cast<const test::TestDialectVersion *>(*versionOr);
           if (version->major_ >= 2)

[ftynse]

Misclicked

[ftynse]

Should this be failure?

[Men-cotton]

Thanks for the comment!
I’d prefer to keep this as return success(); rather than failure().

My understanding is that for the Reader callbacks, returning failure() indicates a hard error and aborts the reading process immediately, whereas returning success() with entry left unset means “I didn’t handle this, please fall back to the default path”.

Since the goal here is to gracefully fallback (skip this callback) when the version is missing, I believe return success() is the correct return value for the Reader side (in contrast to the Writer side, where failure() triggers a fallback). This also matches the overall design goal described in this comment, that bytecode handling should be robust and not fail in such cases.

If I’ve misunderstood the intended reader contract, I’m happy to adjust.