diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h b/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h --- a/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h +++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h @@ -43,6 +43,11 @@ /// type. Type getTensorTypeFromMemRefType(Type type); +/// Finds a single dealloc operation for the given allocated value. If there +/// are > 1 deallocates for `allocValue`, returns None, else returns the single +/// deallocate if it exists or nullptr. +Optional findDealloc(Value allocValue); + } // namespace memref } // namespace mlir diff --git a/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h b/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h --- a/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h +++ b/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h @@ -16,14 +16,4 @@ #ifndef MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H #define MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H -#include "mlir/Dialect/MemRef/IR/MemRef.h" - -namespace mlir { - -/// Finds a single dealloc operation for the given allocated value. If there -/// are > 1 deallocates for `allocValue`, returns None, else returns the single -/// deallocate if it exists or nullptr. -llvm::Optional findDealloc(Value allocValue); -} // namespace mlir - #endif // MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp --- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp +++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp @@ -118,11 +118,12 @@ // also consider aliases. That would also make the safety check below // redundant. llvm::Optional maybeCloneDeallocOp = - findDealloc(cloneOp.output()); + memref::findDealloc(cloneOp.output()); // Skip if either of them has > 1 deallocate operations. if (!maybeCloneDeallocOp.hasValue()) return failure(); - llvm::Optional maybeSourceDeallocOp = findDealloc(source); + llvm::Optional maybeSourceDeallocOp = + memref::findDealloc(source); if (!maybeSourceDeallocOp.hasValue()) return failure(); Operation *cloneDeallocOp = *maybeCloneDeallocOp; diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferUtils.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferUtils.cpp --- a/mlir/lib/Dialect/Bufferization/Transforms/BufferUtils.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferUtils.cpp @@ -78,7 +78,7 @@ // Get allocation result. Value allocValue = allocateResultEffects[0].getValue(); // Find the associated dealloc value and register the allocation entry. - llvm::Optional dealloc = findDealloc(allocValue); + llvm::Optional dealloc = memref::findDealloc(allocValue); // If the allocation has > 1 dealloc associated with it, skip handling it. if (!dealloc.hasValue()) return; diff --git a/mlir/lib/Dialect/MemRef/IR/CMakeLists.txt b/mlir/lib/Dialect/MemRef/IR/CMakeLists.txt --- a/mlir/lib/Dialect/MemRef/IR/CMakeLists.txt +++ b/mlir/lib/Dialect/MemRef/IR/CMakeLists.txt @@ -19,6 +19,6 @@ MLIRDialectUtils MLIRInferTypeOpInterface MLIRIR - MLIRMemRefUtils + MLIRSideEffectInterfaces MLIRViewLikeInterface ) diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp --- a/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp +++ b/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp @@ -7,6 +7,7 @@ //===----------------------------------------------------------------------===// #include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" #include "mlir/Transforms/InliningUtils.h" using namespace mlir; @@ -39,3 +40,28 @@ >(); addInterfaces(); } + +/// Finds a single dealloc operation for the given allocated value. +llvm::Optional mlir::memref::findDealloc(Value allocValue) { + Operation *dealloc = nullptr; + for (Operation *user : allocValue.getUsers()) { + auto effectInterface = dyn_cast(user); + if (!effectInterface) + continue; + // Try to find a free effect that is applied to one of our values + // that will be automatically freed by our pass. + SmallVector effects; + effectInterface.getEffectsOnValue(allocValue, effects); + const bool isFree = + llvm::any_of(effects, [&](MemoryEffects::EffectInstance &it) { + return isa(it.getEffect()); + }); + if (!isFree) + continue; + // If we found > 1 dealloc, return None. + if (dealloc) + return llvm::None; + dealloc = user; + } + return dealloc; +} diff --git a/mlir/lib/Dialect/MemRef/Utils/CMakeLists.txt b/mlir/lib/Dialect/MemRef/Utils/CMakeLists.txt --- a/mlir/lib/Dialect/MemRef/Utils/CMakeLists.txt +++ b/mlir/lib/Dialect/MemRef/Utils/CMakeLists.txt @@ -3,9 +3,5 @@ ADDITIONAL_HEADER_DIRS ${PROJECT_SOURCE_DIR}/inlude/mlir/Dialect/MemRefDialect - - LINK_LIBS PUBLIC - MLIRIR - MLIRSideEffectInterfaces ) diff --git a/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp b/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp --- a/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp +++ b/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp @@ -11,31 +11,3 @@ //===----------------------------------------------------------------------===// #include "mlir/Dialect/MemRef/Utils/MemRefUtils.h" -#include "mlir/Interfaces/SideEffectInterfaces.h" - -using namespace mlir; - -/// Finds a single dealloc operation for the given allocated value. -llvm::Optional mlir::findDealloc(Value allocValue) { - Operation *dealloc = nullptr; - for (Operation *user : allocValue.getUsers()) { - auto effectInterface = dyn_cast(user); - if (!effectInterface) - continue; - // Try to find a free effect that is applied to one of our values - // that will be automatically freed by our pass. - SmallVector effects; - effectInterface.getEffectsOnValue(allocValue, effects); - const bool isFree = - llvm::any_of(effects, [&](MemoryEffects::EffectInstance &it) { - return isa(it.getEffect()); - }); - if (!isFree) - continue; - // If we found > 1 dealloc, return None. - if (dealloc) - return llvm::None; - dealloc = user; - } - return dealloc; -}