[Mlir-commits] [mlir] b96ebee - [mlir] Fix allocateAndCopyWithAlign for immutable (#108679)

[llvmlistbot at llvm.org]

Author: Jacques Pienaar
Date: 2024-10-11T09:10:01-07:00
New Revision: b96ebee1fab2b281c97deb54f3d61c469fe07d01

URL: https://github.com/llvm/llvm-project/commit/b96ebee1fab2b281c97deb54f3d61c469fe07d01
DIFF: https://github.com/llvm/llvm-project/commit/b96ebee1fab2b281c97deb54f3d61c469fe07d01.diff

LOG: [mlir] Fix allocateAndCopyWithAlign for immutable (#108679)

Previously this would assert when attempting to getMutableData.

Added: 
    

Modified: 
    mlir/include/mlir/IR/AsmState.h
    mlir/unittests/IR/AttributeTest.cpp

Removed: 
    


################################################################################
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 981e919289c306..2b2ec056739431 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);