From 8dfc6042e93fe05b5936e24f734cd77d9906e85b Mon Sep 17 00:00:00 2001
From: Jeff Niu <jeff@modular.com>
Date: Fri, 21 Oct 2022 15:00:20 -0700
Subject: [PATCH] [mlir][llvm] Allow literal structs to
 replaceImmediateSubElements

SubElementInterfaces forbids all mutable types and attributes from
implementing `replaceImmediateSubElements`. However, this prohibits
literal structs, which are immutable, from implementing that function.
This patch defers the decision on whether to support
`replaceImmediateSubElements` to the individual types/attributes.

Depends on D136505

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D136507
---
 mlir/include/mlir/IR/SubElementInterfaces.td |  7 ++++++-
 mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp     |  9 ++++++---
 mlir/lib/IR/SubElementInterfaces.cpp         | 14 --------------
 3 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/mlir/include/mlir/IR/SubElementInterfaces.td b/mlir/include/mlir/IR/SubElementInterfaces.td
index e857beb..3718b38 100644
--- a/mlir/include/mlir/IR/SubElementInterfaces.td
+++ b/mlir/include/mlir/IR/SubElementInterfaces.td
@@ -42,6 +42,11 @@ class SubElementInterfaceBase<string interfaceName, string attrOrType,
         would replace the very first attribute given by `walkImmediateSubElements`.
         On success, the new instance with the values replaced is returned. If replacement
         fails, nullptr is returned.
+
+        Note that replacing the sub-elements of mutable types or attributes is
+        not currently supported by the interface. If an implementing type or
+        attribute is mutable, it should return `nullptr` if it has no mechanism
+        for replacing sub elements.
       }], attrOrType, "replaceImmediateSubElements", (ins
         "::llvm::ArrayRef<::mlir::Attribute>":$replAttrs,
         "::llvm::ArrayRef<::mlir::Type>":$replTypes
@@ -106,7 +111,7 @@ class SubElementInterfaceBase<string interfaceName, string attrOrType,
     void walkSubTypes(llvm::function_ref<void(mlir::Type)> walkFn) {
       walkSubElements(/*walkAttrsFn=*/[](mlir::Attribute) {}, walkFn);
     }
-    
+
     /// Recursively replace all of the nested sub-attributes using the provided
     /// map function. Returns nullptr in the case of failure.
     }] # attrOrType # [{ replaceSubElements(
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
index 9187814..99fa193 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
@@ -698,9 +698,12 @@ void LLVMStructType::walkImmediateSubElements(
 
 Type LLVMStructType::replaceImmediateSubElements(
     ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
-  // TODO: It's not clear how we support replacing sub-elements of mutable
-  // types.
-  return nullptr;
+  if (isIdentified()) {
+    // TODO: It's not clear how we support replacing sub-elements of mutable
+    // types.
+    return nullptr;
+  }
+  return getLiteral(getContext(), replTypes, isPacked());
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/SubElementInterfaces.cpp b/mlir/lib/IR/SubElementInterfaces.cpp
index a362479..fd05b9d 100644
--- a/mlir/lib/IR/SubElementInterfaces.cpp
+++ b/mlir/lib/IR/SubElementInterfaces.cpp
@@ -93,14 +93,6 @@ void SubElementTypeInterface::walkSubElements(
 //===----------------------------------------------------------------------===//
 // ReplaceSubElements
 
-/// Return if the given element is mutable.
-static bool isMutable(Attribute attr) {
-  return attr.hasTrait<AttributeTrait::IsMutable>();
-}
-static bool isMutable(Type type) {
-  return type.hasTrait<TypeTrait::IsMutable>();
-}
-
 template <typename InterfaceT, typename T, typename ReplaceSubElementFnT>
 static void updateSubElementImpl(
     T element, function_ref<std::pair<T, WalkResult>(T)> walkFn,
@@ -187,12 +179,6 @@ replaceSubElementsImpl(InterfaceT interface,
   if (!*changed)
     return interface;
 
-  // If this element is mutable, we don't support changing its sub elements, the
-  // sub element walk doesn't give us a valid ordering for what we need here. If
-  // we want to support mutable elements, we'll need something more.
-  if (isMutable(interface))
-    return {};
-
   // Use the new elements during the replacement.
   return interface.replaceImmediateSubElements(newAttrs, newTypes);
 }
-- 
2.7.4