From c715f1ae2e871297b50dc164f1341f012d125533 Mon Sep 17 00:00:00 2001 From: David Stone Date: Wed, 19 Nov 2025 18:13:55 -0700 Subject: [PATCH] [mlir] Replace `llvm::OwningArrayRef` with `std::vector` There are several places where we use `llvm::OwningArrayRef`. The interface to this requires us to first construct temporary storage, then allocate space and set the allocated memory to 0, then copy the values we actually want into that memory, then move the array into place. Instead we can just do it all inline in a single pass by using `std::vector`. In one case we actually allocate a completely separate container and then allocate + copy the data over because `llvm::OwningArrayRef` does not (and can't) support `push_back`. Note that `llvm::SmallVector` is not a suitable replacement here because we rely on reference stability on move construction: when the outer container reallocates, we need the the contents of the inner containers to be fixed in memory, and `llvm::SmallVector` does not give us that guarantee. --- mlir/include/mlir/IR/PDLPatternMatch.h.inc | 12 ++--- mlir/lib/Rewrite/ByteCode.cpp | 60 ++++++++++------------ mlir/lib/Rewrite/ByteCode.h | 7 ++- 3 files changed, 33 insertions(+), 46 deletions(-) diff --git a/mlir/include/mlir/IR/PDLPatternMatch.h.inc b/mlir/include/mlir/IR/PDLPatternMatch.h.inc index d5fb57d7c360d..4afbcf2924965 100644 --- a/mlir/include/mlir/IR/PDLPatternMatch.h.inc +++ b/mlir/include/mlir/IR/PDLPatternMatch.h.inc @@ -152,9 +152,7 @@ public: void push_back(TypeRange value) { // The lifetime of a TypeRange can't be guaranteed, so we'll need to // allocate a storage for it. - llvm::OwningArrayRef storage(value.size()); - llvm::copy(value, storage.begin()); - allocatedTypeRanges.emplace_back(std::move(storage)); + allocatedTypeRanges.emplace_back(value.begin(), value.end()); typeRanges.push_back(allocatedTypeRanges.back()); results.push_back(&typeRanges.back()); } @@ -174,9 +172,7 @@ public: void push_back(ValueRange value) { // The lifetime of a ValueRange can't be guaranteed, so we'll need to // allocate a storage for it. - llvm::OwningArrayRef storage(value.size()); - llvm::copy(value, storage.begin()); - allocatedValueRanges.emplace_back(std::move(storage)); + allocatedValueRanges.emplace_back(value.begin(), value.end()); valueRanges.push_back(allocatedValueRanges.back()); results.push_back(&valueRanges.back()); } @@ -206,8 +202,8 @@ protected: SmallVector valueRanges; /// Memory allocated to store ranges in the result list whose lifetime was /// generated in the native function. - SmallVector> allocatedTypeRanges; - SmallVector> allocatedValueRanges; + SmallVector> allocatedTypeRanges; + SmallVector> allocatedValueRanges; }; //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Rewrite/ByteCode.cpp b/mlir/lib/Rewrite/ByteCode.cpp index 42843ea1780c4..159aa54686034 100644 --- a/mlir/lib/Rewrite/ByteCode.cpp +++ b/mlir/lib/Rewrite/ByteCode.cpp @@ -1099,12 +1099,12 @@ class ByteCodeRewriteResultList : public PDLResultList { MutableArrayRef getResults() { return results; } /// Return the type ranges allocated by this list. - MutableArrayRef> getAllocatedTypeRanges() { + MutableArrayRef> getAllocatedTypeRanges() { return allocatedTypeRanges; } /// Return the value ranges allocated by this list. - MutableArrayRef> getAllocatedValueRanges() { + MutableArrayRef> getAllocatedValueRanges() { return allocatedValueRanges; } }; @@ -1112,19 +1112,20 @@ class ByteCodeRewriteResultList : public PDLResultList { /// This class provides support for executing a bytecode stream. class ByteCodeExecutor { public: - ByteCodeExecutor( - const ByteCodeField *curCodeIt, MutableArrayRef memory, - MutableArrayRef> opRangeMemory, - MutableArrayRef typeRangeMemory, - std::vector> &allocatedTypeRangeMemory, - MutableArrayRef valueRangeMemory, - std::vector> &allocatedValueRangeMemory, - MutableArrayRef loopIndex, ArrayRef uniquedMemory, - ArrayRef code, - ArrayRef currentPatternBenefits, - ArrayRef patterns, - ArrayRef constraintFunctions, - ArrayRef rewriteFunctions) + ByteCodeExecutor(const ByteCodeField *curCodeIt, + MutableArrayRef memory, + MutableArrayRef> opRangeMemory, + MutableArrayRef typeRangeMemory, + std::vector> &allocatedTypeRangeMemory, + MutableArrayRef valueRangeMemory, + std::vector> &allocatedValueRangeMemory, + MutableArrayRef loopIndex, + ArrayRef uniquedMemory, + ArrayRef code, + ArrayRef currentPatternBenefits, + ArrayRef patterns, + ArrayRef constraintFunctions, + ArrayRef rewriteFunctions) : curCodeIt(curCodeIt), memory(memory), opRangeMemory(opRangeMemory), typeRangeMemory(typeRangeMemory), allocatedTypeRangeMemory(allocatedTypeRangeMemory), @@ -1367,13 +1368,9 @@ class ByteCodeExecutor { if (range.empty()) { rangeMemory[rangeIndex] = {}; } else { - // Allocate a buffer for this type range. - llvm::OwningArrayRef storage(llvm::size(range)); - llvm::copy(range, storage.begin()); - // Assign this to the range slot and use the range as the value for the // memory index. - allocatedRangeMemory.emplace_back(std::move(storage)); + allocatedRangeMemory.emplace_back(range.begin(), range.end()); rangeMemory[rangeIndex] = allocatedRangeMemory.back(); } memory[memIndex] = &rangeMemory[rangeIndex]; @@ -1397,11 +1394,11 @@ class ByteCodeExecutor { /// The current execution memory. MutableArrayRef memory; - MutableArrayRef opRangeMemory; + MutableArrayRef> opRangeMemory; MutableArrayRef typeRangeMemory; - std::vector> &allocatedTypeRangeMemory; + std::vector> &allocatedTypeRangeMemory; MutableArrayRef valueRangeMemory; - std::vector> &allocatedValueRangeMemory; + std::vector> &allocatedValueRangeMemory; /// The current loop indices. MutableArrayRef loopIndex; @@ -1907,10 +1904,10 @@ void ByteCodeExecutor::executeGetUsers() { LDBG() << "Executing GetUsers:"; unsigned memIndex = read(); unsigned rangeIndex = read(); - OwningOpRange &range = opRangeMemory[rangeIndex]; + std::vector &range = opRangeMemory[rangeIndex]; memory[memIndex] = ⦥ - range = OwningOpRange(); + range.clear(); if (read() == PDLValue::Kind::Value) { // Read the value. Value value = read(); @@ -1918,9 +1915,7 @@ void ByteCodeExecutor::executeGetUsers() { return; LDBG() << " * Value: " << value; - // Extract the users of a single value. - range = OwningOpRange(std::distance(value.user_begin(), value.user_end())); - llvm::copy(value.getUsers(), range.begin()); + range.assign(value.user_begin(), value.user_end()); } else { // Read a range of values. ValueRange *values = read(); @@ -1929,12 +1924,8 @@ void ByteCodeExecutor::executeGetUsers() { LDBG() << " * Values (" << values->size() << "): " << llvm::interleaved(*values); - // Extract all the users of a range of values. - SmallVector users; for (Value value : *values) - users.append(value.user_begin(), value.user_end()); - range = OwningOpRange(users.size()); - llvm::copy(users, range.begin()); + range.insert(range.end(), value.user_begin(), value.user_end()); } LDBG() << " * Result: " << range.size() << " operations"; @@ -2174,7 +2165,8 @@ ByteCodeExecutor::execute(PatternRewriter &rewriter, executeEraseOp(rewriter); break; case ExtractOp: - executeExtract(); + executeExtract, + PDLValue::Kind::Operation>(); break; case ExtractType: executeExtract(); diff --git a/mlir/lib/Rewrite/ByteCode.h b/mlir/lib/Rewrite/ByteCode.h index 4aceac7ed3a4c..566c1cb66bdbd 100644 --- a/mlir/lib/Rewrite/ByteCode.h +++ b/mlir/lib/Rewrite/ByteCode.h @@ -30,7 +30,6 @@ class PDLByteCode; /// entries. ByteCodeAddr refers to size of indices into the bytecode. using ByteCodeField = uint16_t; using ByteCodeAddr = uint32_t; -using OwningOpRange = llvm::OwningArrayRef; //===----------------------------------------------------------------------===// // PDLByteCodePattern @@ -94,21 +93,21 @@ class PDLByteCodeMutableState { /// the bytecode to store ranges of operations. These are always stored by /// owning references, because at no point in the execution of the byte code /// we get an indexed range (view) of operations. - std::vector opRangeMemory; + std::vector> opRangeMemory; /// A mutable block of memory used during the matching and rewriting phase of /// the bytecode to store ranges of types. std::vector typeRangeMemory; /// A set of type ranges that have been allocated by the byte code interpreter /// to provide a guaranteed lifetime. - std::vector> allocatedTypeRangeMemory; + std::vector> allocatedTypeRangeMemory; /// A mutable block of memory used during the matching and rewriting phase of /// the bytecode to store ranges of values. std::vector valueRangeMemory; /// A set of value ranges that have been allocated by the byte code /// interpreter to provide a guaranteed lifetime. - std::vector> allocatedValueRangeMemory; + std::vector> allocatedValueRangeMemory; /// The current index of ranges being iterated over for each level of nesting. /// These are always maintained at 0 for the loops that are not active, so we