[Mlir-commits] [mlir] [mlir][BytecodeReader] Fix FusedLoc bytecode builder (PR #99632)

[Hideto Ueno]

https://github.com/uenoku created https://github.com/llvm/llvm-project/pull/99632

[`FusedLoc::get(MLIRContext*, ArrayRef<Location>)`](https://github.com/llvm/llvm-project/blob/098bd842a7e50853fa231f8b73c24ec5006fe063/mlir/include/mlir/IR/BuiltinLocationAttributes.td#L138) is not guaranteed to return `FusedLoc` but it is used when reading bytecode. This PR changes to use a different getter `FusedLoc::get(MLIRContext*, ArrayRef<Location>, Attribute attribute)` which always returns `FusedLoc`.  

Fix https://github.com/llvm/llvm-project/issues/99626. 

>From a3daf2895be2df7722b7ef7dad7ec39d7fb6d41d Mon Sep 17 00:00:00 2001
From: Hideto Ueno <uenoku.tokotoko at gmail.com>
Date: Fri, 19 Jul 2024 03:53:44 -0700
Subject: [PATCH] [mlir][BytecodeReader] Fix FusedLoc bytecode builder

---
 .../include/mlir/IR/BuiltinDialectBytecode.td |  5 +++--
 mlir/unittests/Bytecode/BytecodeTest.cpp      | 22 +++++++++++++++++++
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/IR/BuiltinDialectBytecode.td b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
index f50b5dd7ad822..0062bcbf6ea11 100644
--- a/mlir/include/mlir/IR/BuiltinDialectBytecode.td
+++ b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
@@ -101,12 +101,12 @@ def FileLineColLoc : DialectAttribute<(attr
   VarInt:$column
 )>;
 
-let cType = "FusedLoc",
-    cBuilder = "cast<FusedLoc>(get<FusedLoc>(context, $_args))" in {
+let cType = "FusedLoc" in {
 def FusedLoc : DialectAttribute<(attr
   Array<Location>:$locations
 )> {
   let printerPredicate = "!$_val.getMetadata()";
+  let cBuilder = "cast<FusedLoc>(get<FusedLoc>(context, $_args, Attribute()))";
 }
 
 def FusedLocWithMetadata : DialectAttribute<(attr
@@ -114,6 +114,7 @@ def FusedLocWithMetadata : DialectAttribute<(attr
   Attribute:$metadata
 )> {
   let printerPredicate = "$_val.getMetadata()";
+  let cBuilder = "cast<FusedLoc>(get<FusedLoc>(context, $_args))";
 }
 }
 
diff --git a/mlir/unittests/Bytecode/BytecodeTest.cpp b/mlir/unittests/Bytecode/BytecodeTest.cpp
index a37a2afc22645..43649819057d0 100644
--- a/mlir/unittests/Bytecode/BytecodeTest.cpp
+++ b/mlir/unittests/Bytecode/BytecodeTest.cpp
@@ -10,6 +10,7 @@
 #include "mlir/Bytecode/BytecodeWriter.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/OwningOpRef.h"
 #include "mlir/Parser/Parser.h"
@@ -148,3 +149,24 @@ TEST(Bytecode, OpWithoutProperties) {
   EXPECT_TRUE(OperationEquivalence::computeHash(op.get()) ==
               OperationEquivalence::computeHash(roundtripped));
 }
+
+TEST(Bytecode, FusedLocCrash) {
+  MLIRContext context;
+  OpBuilder builder(&context);
+  SmallVector<Location> locs;
+  FusedLoc fusedLoc = FusedLoc::get(&context, locs, {});
+
+  auto module = builder.create<mlir::ModuleOp>(fusedLoc, "test");
+
+  // Write the module to bytecode
+  std::string buffer;
+  llvm::raw_string_ostream ostream(buffer);
+  ASSERT_TRUE(succeeded(writeBytecodeToFile(module, ostream)));
+  ostream.flush();
+
+  // Parse it back
+  ParserConfig parseConfig(&context);
+  OwningOpRef<Operation *> roundTripModule =
+      parseSourceString<Operation *>(buffer, parseConfig);
+  ASSERT_TRUE(roundTripModule);
+}