[Mlir-commits] [mlir] [mlir] fix Operation::getDiscardableAttrs in absence of properties (PR #76816)

Oleksandr Alex Zinenko llvmlistbot at llvm.org
Wed Jan 3 07:32:59 PST 2024


https://github.com/ftynse updated https://github.com/llvm/llvm-project/pull/76816

>From ab612deb4ededc11e538c34adf96ab9ca1aa7029 Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko at google.com>
Date: Wed, 3 Jan 2024 13:53:50 +0000
Subject: [PATCH 1/3] [mlir] fix Operation::getDiscardableAttrs in absence of
 properties

When properties are not enabled in an operation, inherent attributes are
stored in the common dictionary with discardable attributes. However,
`getDiscardableAttrs` and `getDiscardableAttrDictionary` were returning
the entire dictionary, making the caller mistakenly believe that all
inherent attributes are discardable. Fix this by filtering out
attributes whose names are registered with the operation, i.e., inherent
attributes. This requires an API change so `getDiscardableAttrs` returns
a filter range.
---
 mlir/include/mlir/IR/Operation.h       | 28 ++++++++++++------
 mlir/lib/CAPI/IR/IR.cpp                |  6 ++--
 mlir/lib/IR/AsmPrinter.cpp             |  2 +-
 mlir/unittests/IR/OpPropertiesTest.cpp | 40 +++++++++++++++++++++++++-
 4 files changed, 64 insertions(+), 12 deletions(-)

diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index 3ba0daeca0258a..d4b0c439bd70ae 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -475,19 +475,31 @@ class alignas(8) Operation final
     return removeDiscardableAttr(StringAttr::get(getContext(), name));
   }
 
-  /// Return all of the discardable attributes on this operation.
-  ArrayRef<NamedAttribute> getDiscardableAttrs() { return attrs.getValue(); }
+  /// Return a range of all of discardable attributes on this operation. Note
+  /// that for unregistered operations, all attributes are considered
+  /// discardable.
+  auto getDiscardableAttrs() {
+    std::optional<RegisteredOperationName> opName = getRegisteredInfo();
+    ArrayRef<StringAttr> attributeNames =
+        opName ? getRegisteredInfo()->getAttributeNames()
+               : ArrayRef<StringAttr>();
+    return llvm::make_filter_range(
+        attrs.getValue(),
+        [this, attributeNames](const NamedAttribute attribute) {
+          return getPropertiesStorage() ||
+                 !llvm::is_contained(attributeNames, attribute.getName());
+        });
+  }
 
   /// Return all of the discardable attributes on this operation as a
   /// DictionaryAttr.
-  DictionaryAttr getDiscardableAttrDictionary() { return attrs; }
+  DictionaryAttr getDiscardableAttrDictionary() {
+    return DictionaryAttr::get(getContext(),
+                               llvm::to_vector(getDiscardableAttrs()));
+  }
 
   /// Return all of the attributes on this operation.
-  ArrayRef<NamedAttribute> getAttrs() {
-    if (!getPropertiesStorage())
-      return getDiscardableAttrs();
-    return getAttrDictionary().getValue();
-  }
+  ArrayRef<NamedAttribute> getAttrs() { return getAttrDictionary().getValue(); }
 
   /// Return all of the attributes on this operation as a DictionaryAttr.
   DictionaryAttr getAttrDictionary();
diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index ac9889df11f80d..a97cfe5b0ea364 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -613,12 +613,14 @@ void mlirOperationSetInherentAttributeByName(MlirOperation op,
 }
 
 intptr_t mlirOperationGetNumDiscardableAttributes(MlirOperation op) {
-  return static_cast<intptr_t>(unwrap(op)->getDiscardableAttrs().size());
+  return static_cast<intptr_t>(
+      llvm::range_size(unwrap(op)->getDiscardableAttrs()));
 }
 
 MlirNamedAttribute mlirOperationGetDiscardableAttribute(MlirOperation op,
                                                         intptr_t pos) {
-  NamedAttribute attr = unwrap(op)->getDiscardableAttrs()[pos];
+  NamedAttribute attr =
+      *std::next(unwrap(op)->getDiscardableAttrs().begin(), pos);
   return MlirNamedAttribute{wrap(attr.getName()), wrap(attr.getValue())};
 }
 
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 1f7cbf349255d5..0cc6d11ddc576c 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -3543,7 +3543,7 @@ void OperationPrinter::printGenericOp(Operation *op, bool printOpName) {
   }
 
   auto attrs = op->getDiscardableAttrs();
-  printOptionalAttrDict(attrs);
+  printOptionalAttrDict(llvm::to_vector(attrs));
 
   // Print the type signature of the operation.
   os << " : ";
diff --git a/mlir/unittests/IR/OpPropertiesTest.cpp b/mlir/unittests/IR/OpPropertiesTest.cpp
index 5c3e150f6b94c2..6dcba16e656b33 100644
--- a/mlir/unittests/IR/OpPropertiesTest.cpp
+++ b/mlir/unittests/IR/OpPropertiesTest.cpp
@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Parser/Parser.h"
 #include "gtest/gtest.h"
@@ -132,6 +133,23 @@ class OpWithProperties : public Op<OpWithProperties> {
   }
 };
 
+/// A custom operation for the purpose of showcasing how discardable attributes
+/// are handled in absence of properties.
+class OpWithoutProperties : public Op<OpWithoutProperties> {
+public:
+  // Begin boilerplate.
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpWithoutProperties)
+  using Op::Op;
+  static ArrayRef<StringRef> getAttributeNames() {
+    static StringRef attributeNames[] = {StringRef("inherent_attr")};
+    return ArrayRef(attributeNames);
+  };
+  static StringRef getOperationName() {
+    return "test_op_properties.op_without_properties";
+  }
+  // End boilerplate.
+};
+
 // A trivial supporting dialect to register the above operation.
 class TestOpPropertiesDialect : public Dialect {
 public:
@@ -142,7 +160,7 @@ class TestOpPropertiesDialect : public Dialect {
   explicit TestOpPropertiesDialect(MLIRContext *context)
       : Dialect(getDialectNamespace(), context,
                 TypeID::get<TestOpPropertiesDialect>()) {
-    addOperations<OpWithProperties>();
+    addOperations<OpWithProperties, OpWithoutProperties>();
   }
 };
 
@@ -359,4 +377,24 @@ TEST(OpPropertiesTest, getOrAddProperties) {
   op->erase();
 }
 
+constexpr StringLiteral withoutPropertiesAttrsSrc = R"mlir(
+    "test_op_properties.op_without_properties"()
+      {inherent_attr = 42, other_attr = 56} : () -> ()
+)mlir";
+
+TEST(OpPropertiesTest, withoutPropertiesDiscardableAttrs) {
+  MLIRContext context;
+  context.getOrLoadDialect<TestOpPropertiesDialect>();
+  ParserConfig config(&context);
+  OwningOpRef<Operation *> op =
+      parseSourceString(withoutPropertiesAttrsSrc, config);
+  ASSERT_EQ(llvm::range_size(op->getDiscardableAttrs()), 1u);
+  EXPECT_EQ(op->getDiscardableAttrs().begin()->getName().getValue(),
+            "other_attr");
+
+  EXPECT_EQ(op->getAttrs().size(), 2u);
+  EXPECT_TRUE(op->getInherentAttr("inherent_attr") != std::nullopt);
+  EXPECT_TRUE(op->getDiscardableAttr("other_attr") != Attribute());
+}
+
 } // namespace

>From fa4c4a53699c16c06cb837ca756955483042cc6d Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko at google.com>
Date: Wed, 3 Jan 2024 14:18:45 +0000
Subject: [PATCH 2/3] Reinstate fast path in getDiscardableAttrDictionary

---
 mlir/include/mlir/IR/Operation.h | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index d4b0c439bd70ae..d47e9d1192608e 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -494,6 +494,8 @@ class alignas(8) Operation final
   /// Return all of the discardable attributes on this operation as a
   /// DictionaryAttr.
   DictionaryAttr getDiscardableAttrDictionary() {
+    if (getPropertiesStorage())
+      return attrs;
     return DictionaryAttr::get(getContext(),
                                llvm::to_vector(getDiscardableAttrs()));
   }

>From 6ae31be005227ec1599c8adff44315361335fc12 Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko at google.com>
Date: Wed, 3 Jan 2024 15:32:42 +0000
Subject: [PATCH 3/3] update comment

---
 mlir/include/mlir/IR/Operation.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index d47e9d1192608e..d2f52cf1afee81 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -476,8 +476,8 @@ class alignas(8) Operation final
   }
 
   /// Return a range of all of discardable attributes on this operation. Note
-  /// that for unregistered operations, all attributes are considered
-  /// discardable.
+  /// that for unregistered operations that are not storing inherent attributes
+  /// as properties, all attributes are considered discardable.
   auto getDiscardableAttrs() {
     std::optional<RegisteredOperationName> opName = getRegisteredInfo();
     ArrayRef<StringAttr> attributeNames =



More information about the Mlir-commits mailing list