[mlir][llvm] Allow literal structs to replaceImmediateSubElements

[Mogball]

Summary

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

[rriddle]

This LGTM with the addition of a Note on the interface definition of replaceImmediateSubElements that mutable things aren't supported (given that this now places the onus on the attribute/type to detect when it should fail).

[Mogball]

add note to interface doc

mlir/include/mlir/IR/SubElementInterfaces.td

list<InterfaceMethod> methods = [
     InterfaceMethod<
       /*desc=*/[{
         Replace the immediately nested sub-attributes and sub-types with those provided.
         The order of the provided elements is derived from the order of the elements
         returned by the callbacks of `walkImmediateSubElements`. The element at index 0
         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
       )>,
   ];
 
   code extraClassDeclaration = [{
     /// Walk all of the held sub-attributes and sub-types.
code extraSharedClassDeclaration = [{
     /// Walk all of the held sub-attributes.
     void walkSubAttrs(llvm::function_ref<void(mlir::Attribute)> walkFn) {
       walkSubElements(walkFn, /*walkTypesFn=*/[](mlir::Type) {});
     }
     /// Walk all of the held sub-types.
     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(
       mlir::SubElementReplFn<mlir::Attribute> replaceAttrFn) {
       return replaceSubElements(
         replaceAttrFn, [](mlir::Type type) { return type; });
     }
     }] # attrOrType # [{ replaceSubElements(

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

void LLVMStructType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   for (Type type : getBody())
     walkTypesFn(type);
 }
 
 Type LLVMStructType::replaceImmediateSubElements(
     ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
+  if (isIdentified()) {
     // TODO: It's not clear how we support replacing sub-elements of mutable
     // types.
     return nullptr;
   }
+  return getLiteral(getContext(), replTypes, isPacked());
+}
 
 //===----------------------------------------------------------------------===//
 // Vector types.
 //===----------------------------------------------------------------------===//
 
 /// Verifies that the type about to be constructed is well-formed.
 template <typename VecTy>
 static LogicalResult

mlir/lib/IR/SubElementInterfaces.cpp

void SubElementTypeInterface::walkSubElements(
   DenseSet<Type> visitedTypes;
   walkSubElementsImpl(*this, walkAttrsFn, walkTypesFn, visitedAttrs,
                       visitedTypes);
 }
 
 //===----------------------------------------------------------------------===//
 // 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,
     DenseMap<T, T> &visited, SmallVectorImpl<T> &newElements,
     FailureOr<bool> &changed, ReplaceSubElementFnT &&replaceSubElementFn) {
   // Bail early if we failed at any point.
   if (failed(changed))
     return;
interface.walkImmediateSubElements(
       });
   if (failed(changed))
     return {};
 
   // If the sub-elements didn't change, just return the original value.
   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);
 }
 
 Attribute SubElementAttrInterface::replaceSubElements(
     SubElementResultReplFn<Attribute> replaceAttrFn,
     SubElementResultReplFn<Type> replaceTypeFn) {
   assert(replaceAttrFn && replaceTypeFn && "expected valid replace functions");