[mlir] Fix allocateAndCopyWithAlign for immutable (#108679)

Previously this would assert when attempting to getMutableData.

Author: jpienaar

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>

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