1 //===- SubElementInterfaces.cpp - Attr and Type SubElement Interfaces -----===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/IR/SubElementInterfaces.h"
10 
11 #include "llvm/ADT/DenseSet.h"
12 
13 using namespace mlir;
14 
15 //===----------------------------------------------------------------------===//
16 // SubElementInterface
17 //===----------------------------------------------------------------------===//
18 
19 //===----------------------------------------------------------------------===//
20 // WalkSubElements
21 
22 template <typename InterfaceT>
23 static void walkSubElementsImpl(InterfaceT interface,
24                                 function_ref<void(Attribute)> walkAttrsFn,
25                                 function_ref<void(Type)> walkTypesFn,
26                                 DenseSet<Attribute> &visitedAttrs,
27                                 DenseSet<Type> &visitedTypes) {
28   interface.walkImmediateSubElements(
29       [&](Attribute attr) {
30         // Guard against potentially null inputs. This removes the need for the
31         // derived attribute/type to do it.
32         if (!attr)
33           return;
34 
35         // Avoid infinite recursion when visiting sub attributes later, if this
36         // is a mutable attribute.
37         if (LLVM_UNLIKELY(attr.hasTrait<AttributeTrait::IsMutable>())) {
38           if (!visitedAttrs.insert(attr).second)
39             return;
40         }
41 
42         // Walk any sub elements first.
43         if (auto interface = attr.dyn_cast<SubElementAttrInterface>())
44           walkSubElementsImpl(interface, walkAttrsFn, walkTypesFn, visitedAttrs,
45                               visitedTypes);
46 
47         // Walk this attribute.
48         walkAttrsFn(attr);
49       },
50       [&](Type type) {
51         // Guard against potentially null inputs. This removes the need for the
52         // derived attribute/type to do it.
53         if (!type)
54           return;
55 
56         // Avoid infinite recursion when visiting sub types later, if this
57         // is a mutable type.
58         if (LLVM_UNLIKELY(type.hasTrait<TypeTrait::IsMutable>())) {
59           if (!visitedTypes.insert(type).second)
60             return;
61         }
62 
63         // Walk any sub elements first.
64         if (auto interface = type.dyn_cast<SubElementTypeInterface>())
65           walkSubElementsImpl(interface, walkAttrsFn, walkTypesFn, visitedAttrs,
66                               visitedTypes);
67 
68         // Walk this type.
69         walkTypesFn(type);
70       });
71 }
72 
73 void SubElementAttrInterface::walkSubElements(
74     function_ref<void(Attribute)> walkAttrsFn,
75     function_ref<void(Type)> walkTypesFn) {
76   assert(walkAttrsFn && walkTypesFn && "expected valid walk functions");
77   DenseSet<Attribute> visitedAttrs;
78   DenseSet<Type> visitedTypes;
79   walkSubElementsImpl(*this, walkAttrsFn, walkTypesFn, visitedAttrs,
80                       visitedTypes);
81 }
82 
83 void SubElementTypeInterface::walkSubElements(
84     function_ref<void(Attribute)> walkAttrsFn,
85     function_ref<void(Type)> walkTypesFn) {
86   assert(walkAttrsFn && walkTypesFn && "expected valid walk functions");
87   DenseSet<Attribute> visitedAttrs;
88   DenseSet<Type> visitedTypes;
89   walkSubElementsImpl(*this, walkAttrsFn, walkTypesFn, visitedAttrs,
90                       visitedTypes);
91 }
92 
93 //===----------------------------------------------------------------------===//
94 // ReplaceSubElements
95 
96 /// Return if the given element is mutable.
97 static bool isMutable(Attribute attr) {
98   return attr.hasTrait<AttributeTrait::IsMutable>();
99 }
100 static bool isMutable(Type type) {
101   return type.hasTrait<TypeTrait::IsMutable>();
102 }
103 
104 template <typename InterfaceT, typename T, typename ReplaceSubElementFnT>
105 static void updateSubElementImpl(T element, function_ref<T(T)> walkFn,
106                                  DenseMap<T, T> &visited,
107                                  SmallVectorImpl<T> &newElements,
108                                  FailureOr<bool> &changed,
109                                  ReplaceSubElementFnT &&replaceSubElementFn) {
110   // Bail early if we failed at any point.
111   if (failed(changed))
112     return;
113   newElements.push_back(element);
114 
115   // Guard against potentially null inputs. We always map null to null.
116   if (!element)
117     return;
118 
119   // Check for an existing mapping for this element, and walk it if we haven't
120   // yet.
121   T &mappedElement = visited[element];
122   if (!mappedElement) {
123     // Try walking this element.
124     if (!(mappedElement = walkFn(element))) {
125       changed = failure();
126       return;
127     }
128 
129     // Handle replacing sub-elements if this element is also a container.
130     if (auto interface = mappedElement.template dyn_cast<InterfaceT>()) {
131       if (!(mappedElement = replaceSubElementFn(interface))) {
132         changed = failure();
133         return;
134       }
135     }
136   }
137 
138   // Update to the mapped element.
139   if (mappedElement != element) {
140     newElements.back() = mappedElement;
141     changed = true;
142   }
143 }
144 
145 template <typename InterfaceT>
146 static typename InterfaceT::ValueType
147 replaceSubElementsImpl(InterfaceT interface,
148                        function_ref<Attribute(Attribute)> walkAttrsFn,
149                        function_ref<Type(Type)> walkTypesFn,
150                        DenseMap<Attribute, Attribute> &visitedAttrs,
151                        DenseMap<Type, Type> &visitedTypes) {
152   // Walk the current sub-elements, replacing them as necessary.
153   SmallVector<Attribute, 16> newAttrs;
154   SmallVector<Type, 16> newTypes;
155   FailureOr<bool> changed = false;
156   auto replaceSubElementFn = [&](auto subInterface) {
157     return replaceSubElementsImpl(subInterface, walkAttrsFn, walkTypesFn,
158                                   visitedAttrs, visitedTypes);
159   };
160   interface.walkImmediateSubElements(
161       [&](Attribute element) {
162         updateSubElementImpl<SubElementAttrInterface>(
163             element, walkAttrsFn, visitedAttrs, newAttrs, changed,
164             replaceSubElementFn);
165       },
166       [&](Type element) {
167         updateSubElementImpl<SubElementTypeInterface>(
168             element, walkTypesFn, visitedTypes, newTypes, changed,
169             replaceSubElementFn);
170       });
171   if (failed(changed))
172     return {};
173 
174   // If the sub-elements didn't change, just return the original value.
175   if (!*changed)
176     return interface;
177 
178   // If this element is mutable, we don't support changing its sub elements, the
179   // sub element walk doesn't give us a valid ordering for what we need here. If
180   // we want to support mutable elements, we'll need something more.
181   if (isMutable(interface))
182     return {};
183 
184   // Use the new elements during the replacement.
185   return interface.replaceImmediateSubElements(newAttrs, newTypes);
186 }
187 
188 Attribute SubElementAttrInterface::replaceSubElements(
189     function_ref<Attribute(Attribute)> replaceAttrFn,
190     function_ref<Type(Type)> replaceTypeFn) {
191   assert(replaceAttrFn && replaceTypeFn && "expected valid replace functions");
192   DenseMap<Attribute, Attribute> visitedAttrs;
193   DenseMap<Type, Type> visitedTypes;
194   return replaceSubElementsImpl(*this, replaceAttrFn, replaceTypeFn,
195                                 visitedAttrs, visitedTypes);
196 }
197 
198 Type SubElementTypeInterface::replaceSubElements(
199     function_ref<Attribute(Attribute)> replaceAttrFn,
200     function_ref<Type(Type)> replaceTypeFn) {
201   assert(replaceAttrFn && replaceTypeFn && "expected valid replace functions");
202   DenseMap<Attribute, Attribute> visitedAttrs;
203   DenseMap<Type, Type> visitedTypes;
204   return replaceSubElementsImpl(*this, replaceAttrFn, replaceTypeFn,
205                                 visitedAttrs, visitedTypes);
206 }
207 
208 //===----------------------------------------------------------------------===//
209 // SubElementInterface Tablegen definitions
210 //===----------------------------------------------------------------------===//
211 
212 #include "mlir/IR/SubElementAttrInterfaces.cpp.inc"
213 #include "mlir/IR/SubElementTypeInterfaces.cpp.inc"
214