diff options
| -rw-r--r-- | mlir/include/mlir/IR/SubElementInterfaces.h | 6 | ||||
| -rw-r--r-- | mlir/lib/IR/SubElementInterfaces.cpp | 9 |
2 files changed, 15 insertions, 0 deletions
diff --git a/mlir/include/mlir/IR/SubElementInterfaces.h b/mlir/include/mlir/IR/SubElementInterfaces.h index 5c362c1..9b14b17 100644 --- a/mlir/include/mlir/IR/SubElementInterfaces.h +++ b/mlir/include/mlir/IR/SubElementInterfaces.h @@ -40,6 +40,12 @@ public: void replaceElementsIn(Operation *op, bool replaceAttrs = true, bool replaceLocs = false, bool replaceTypes = false); + /// Replace the elements within the given operation, and all nested + /// operations. + void recursivelyReplaceElementsIn(Operation *op, bool replaceAttrs = true, + bool replaceLocs = false, + bool replaceTypes = false); + /// Replace the given attribute/type, and recursively replace any sub /// elements. Returns either the new attribute/type, or nullptr in the case of /// failure. diff --git a/mlir/lib/IR/SubElementInterfaces.cpp b/mlir/lib/IR/SubElementInterfaces.cpp index d6130a6..606b3c6 100644 --- a/mlir/lib/IR/SubElementInterfaces.cpp +++ b/mlir/lib/IR/SubElementInterfaces.cpp @@ -145,6 +145,15 @@ void AttrTypeReplacer::replaceElementsIn(Operation *op, bool replaceAttrs, } } +void AttrTypeReplacer::recursivelyReplaceElementsIn(Operation *op, + bool replaceAttrs, + bool replaceLocs, + bool replaceTypes) { + op->walk([&](Operation *nestedOp) { + replaceElementsIn(nestedOp, replaceAttrs, replaceLocs, replaceTypes); + }); +} + template <typename T> static void updateSubElementImpl(T element, AttrTypeReplacer &replacer, DenseMap<T, T> &elementMap, |