[mlir] Fix allocateAndCopyWithAlign for immutable

[jpienaar]

Previously this would assert when attempting to getMutableData.

[llvmbot]

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Jacques Pienaar (jpienaar)

Changes

Previously this would assert when attempting to getMutableData.

---

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

2 Files Affected:

• (modified) mlir/include/mlir/IR/AsmState.h (+8-1)
• (modified) mlir/unittests/IR/AttributeTest.cpp (+15)

diff --git a/mlir/include/mlir/IR/AsmState.h b/mlir/include/mlir/IR/AsmState.h
index 42cbedcf9f8837..edbd3bb6fc15db 100644
--- a/mlir/include/mlir/IR/AsmState.h
+++ b/mlir/include/mlir/IR/AsmState.h
@@ -82,6 +82,8 @@ class AsmStateImpl;
 //===----------------------------------------------------------------------===//
 // Resource Entry
 
+class HeapAsmResourceBlob;
+
 /// This class represents a processed binary blob of data. A resource blob is
 /// essentially a collection of data, potentially mutable, with an associated
 /// deleter function (used if the data needs to be destroyed).
@@ -177,6 +179,8 @@ class AsmResourceBlob {
 
   /// Whether the data is mutable.
   bool dataIsMutable;
+
+  friend class HeapAsmResourceBlob;
 };
 
 /// This class provides a simple utility wrapper for creating heap allocated
@@ -196,8 +200,11 @@ class HeapAsmResourceBlob {
   static AsmResourceBlob allocateAndCopyWithAlign(ArrayRef<char> data,
                                                   size_t align,
                                                   bool dataIsMutable = true) {
-    AsmResourceBlob blob = allocate(data.size(), align, dataIsMutable);
+    // This sets the blob to be mutable initially to allow writing
+    // (getMutableData) below.
+    AsmResourceBlob blob = allocate(data.size(), align, /*dataIsMutable=*/true);
     std::memcpy(blob.getMutableData().data(), data.data(), data.size());
+    blob.dataIsMutable = dataIsMutable;
     return blob;
   }
   template <typename T>
diff --git a/mlir/unittests/IR/AttributeTest.cpp b/mlir/unittests/IR/AttributeTest.cpp
index e72bfe9d82e7cf..5a3649d7be923a 100644
--- a/mlir/unittests/IR/AttributeTest.cpp
+++ b/mlir/unittests/IR/AttributeTest.cpp
@@ -351,6 +351,21 @@ TEST(DenseResourceElementsAttrTest, CheckNoCast) {
   EXPECT_FALSE(isa<DenseBoolResourceElementsAttr>(i32ResourceAttr));
 }
 
+TEST(DenseResourceElementsAttrTest, CheckNotMutableAllocateAndCopy) {
+  MLIRContext context;
+  Builder builder(&context);
+
+  // Create a i32 attribute.
+  std::vector<int32_t> data = {10, 20, 30};
+  auto type = RankedTensorType::get(data.size(), builder.getI32Type());
+  Attribute i32ResourceAttr = DenseI32ResourceElementsAttr::get(
+      type, "resource",
+      HeapAsmResourceBlob::allocateAndCopyInferAlign<int32_t>(
+          data, /*is_mutable=*/false));
+
+  EXPECT_TRUE(isa<DenseI32ResourceElementsAttr>(i32ResourceAttr));
+}
+
 TEST(DenseResourceElementsAttrTest, CheckInvalidData) {
   MLIRContext context;
   Builder builder(&context);