[StableHLO] Improve the PJRT parsing error message to include StableH…

…LO version and portable artifact version.

This uses a hypothetical "future" portable artifact, serialized targeting StableHLO v2.0.0 and using an op `vhlo.constant_v99`.

```
/tmp/t2.mlir:3:10: error: unregistered operation 'vhlo.constant_v99' found in dialect ('vhlo') that does not allow unknown operations
/tmp/t2.mlir:3:10: note: see current operation: %0 = "vhlo.constant_v99"() <#vhlo.tensor_v1<dense<1.000000e+00> : tensor<f32>>> : () -> !vhlo.tensor_v1<!vhlo.f32_v1>
/tmp/t2.mlir:3:10: note: in bytecode version 6 produced by: StableHLO_v2.0.0
<unknown>:0: error: Failed to parse using StableHLO v1.7.5, this could indicate forward incompatibility, >12w old unsupported plugin, or a portable artifact that needs to be further downgraded.
```

PiperOrigin-RevId: 681117697

xla/pjrt/BUILD

@@ -640,6 +640,7 @@ xla_cc_test(
         "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest_main",
         "@llvm-project//mlir:IR",
+        "@stablehlo//:stablehlo_portable_api",
         "@tsl//tsl/platform:statusor",
     ],
 )

xla/pjrt/mlir_to_hlo.cc

@@ -42,6 +42,8 @@ limitations under the License.
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Location.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OwningOpRef.h"
 #include "mlir/IR/Visitors.h"
@@ -136,27 +138,18 @@ absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ParseMlirModuleString(
   mlir::OwningOpRef<mlir::ModuleOp> module =
       mlir::parseSourceString<mlir::ModuleOp>(
           llvm::StringRef(mlir_module_str.data(), mlir_module_str.size()),
-          // IR may be invalid because some fields may be using DenseElements
-          // instead of DenseArray. We rectify that below and verify after.
-          mlir::ParserConfig{&context, /*verifyAfterParse=*/false});
+          mlir::ParserConfig{&context});
   if (!module) {
+    mlir::emitError(mlir::UnknownLoc::get(&context))
+        << "Failed to parse using StableHLO v"
+        << mlir::vhlo::Version::getCurrentVersion() << ", "
+        << "this could indicate forward incompatibility, >12w old "
+           "unsupported plugin, or a portable artifact that needs to be "
+           "further downgraded.";
     return diagnostic_handler.ConsumeStatus();
   }
 
-  // In
-  // https://github.com/google/jax/commit/184e3a88004680dbf34328b05c5fc0d869cc4a93,
-  // fields on some ops were changed to use Dense{Bool,I64}ArrayAttr instead of
-  // I64DenseElementsAttr (DenseIntElementsAttr). Some clients still expect
-  // dense elements, not dense arrays, so when serializing we always convert the
-  // arrays to elements. The elements need to be converted back to arrays when
-  // deserializing.
-  // TODO: b/320507168 - Remove the conversion code, and verifyAfterParse.
   TF_RETURN_IF_ERROR(UpgradeVersionedStablehlo(*module));
-  if (failed(module->verifyInvariants())) {
-    VLOG(1) << "MLIR verification failed.";
-    module->dump();
-    return diagnostic_handler.ConsumeStatus();
-  }
   return std::move(module);
 }
 

xla/pjrt/mlir_to_hlo_test.cc

@@ -21,6 +21,7 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OwningOpRef.h"
+#include "stablehlo/api/PortableApi.h"
 #include "xla/test.h"
 #include "tsl/platform/statusor.h"
 
@@ -109,5 +110,50 @@ TEST(MlirToHloTest, MhloTest) {
   EXPECT_THAT(blob, Not(IsVhloArtifact("1.0.0")));
 }
 
+TEST(MlirToHloTest, InvalidBytecodeTest) {
+  // MLIR bytecode format has full compatibility.
+  // Program using StableHLO v2.0.0 with op vhlo.constant_v99.
+  // TODO: Once this file is exposed via the StableHLO repo, replace this
+  // bytecode string with a read of the StableHLO file.
+  unsigned char invalid_future_vhlo_mlirbc[] = {
+      0x4d, 0x4c, 0xef, 0x52, 0x0d, 0x53, 0x74, 0x61, 0x62, 0x6c, 0x65, 0x48,
+      0x4c, 0x4f, 0x5f, 0x76, 0x32, 0x2e, 0x30, 0x2e, 0x30, 0x00, 0x01, 0x19,
+      0x05, 0x01, 0x05, 0x09, 0x01, 0x03, 0x0b, 0x03, 0x07, 0x0f, 0x13, 0x17,
+      0x03, 0x2b, 0x15, 0x07, 0x01, 0x0b, 0x0b, 0x13, 0x13, 0x13, 0x13, 0x03,
+      0x0b, 0x0b, 0x0b, 0x0b, 0x0b, 0x1f, 0x03, 0x07, 0x0f, 0x13, 0x07, 0x02,
+      0x53, 0x05, 0x0d, 0x17, 0x01, 0x03, 0x03, 0x17, 0x01, 0x05, 0x07, 0x17,
+      0x01, 0x07, 0x15, 0x17, 0x01, 0x09, 0x0b, 0x03, 0x01, 0x23, 0x03, 0x1d,
+      0x0f, 0x1d, 0x11, 0x1f, 0x01, 0x09, 0x00, 0x00, 0x80, 0x3f, 0x29, 0x01,
+      0x05, 0x11, 0x01, 0x03, 0x01, 0x09, 0x04, 0x41, 0x05, 0x01, 0x50, 0x03,
+      0x01, 0x07, 0x04, 0x31, 0x03, 0x01, 0x05, 0x03, 0x50, 0x05, 0x03, 0x07,
+      0x04, 0x1d, 0x03, 0x03, 0x09, 0x05, 0x42, 0x07, 0x05, 0x03, 0x01, 0x07,
+      0x04, 0x09, 0x03, 0x01, 0x06, 0x03, 0x01, 0x05, 0x01, 0x00, 0xad, 0x13,
+      0x0f, 0x0b, 0x1b, 0x15, 0x1b, 0x11, 0x0f, 0x0b, 0x11, 0x62, 0x75, 0x69,
+      0x6c, 0x74, 0x69, 0x6e, 0x00, 0x76, 0x68, 0x6c, 0x6f, 0x00, 0x6d, 0x6f,
+      0x64, 0x75, 0x6c, 0x65, 0x00, 0x66, 0x75, 0x6e, 0x63, 0x5f, 0x76, 0x31,
+      0x00, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x76, 0x39,
+      0x39, 0x00, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x5f, 0x76, 0x31, 0x00,
+      0x2f, 0x74, 0x6d, 0x70, 0x2f, 0x74, 0x32, 0x2e, 0x6d, 0x6c, 0x69, 0x72,
+      0x00, 0x6d, 0x61, 0x69, 0x6e, 0x00, 0x70, 0x75, 0x62, 0x6c, 0x69, 0x63,
+      0x00, 0x08, 0x19, 0x07, 0x05, 0x01, 0x01, 0x0b, 0x0b, 0x0d, 0x0b, 0x0f,
+      0x11, 0x03, 0x13};
+  unsigned int invalid_future_vhlo_mlirbc_len = 243;
+
+  std::string buffer(reinterpret_cast<char*>(invalid_future_vhlo_mlirbc),
+                     invalid_future_vhlo_mlirbc_len);
+
+  mlir::MLIRContext context;
+  auto status = ParseMlirModuleString(buffer, context);
+  ASSERT_FALSE(status.ok());
+  // Check that the error message contains:
+  //   - The name of the op that is not supported (vhlo.constant_v99)
+  //   - The version that the StableHLO portable artifact was emit for (v2.0.0)
+  //   - The current version of StableHLO (v1.X.Y)
+  EXPECT_THAT(status.status().message(), HasSubstr("vhlo.constant_v99"));
+  EXPECT_THAT(status.status().message(), HasSubstr("StableHLO_v2.0.0"));
+  EXPECT_THAT(status.status().message(),
+              HasSubstr(mlir::stablehlo::getCurrentVersion()));
+}
+
 }  // namespace
 }  // namespace xla