Skip to content

[mlir] Replace llvm::OwningArrayRef with std::vector #168803

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Nov 20, 2025
+33 −46
Merged

[mlir] Replace llvm::OwningArrayRef with std::vector #168803

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 4 additions & 8 deletions mlir/include/mlir/IR/PDLPatternMatch.h.inc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Type> 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());
}
Expand All @@ -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<Value> 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());
}
Expand Down Expand Up @@ -206,8 +202,8 @@ protected:
SmallVector<ValueRange> valueRanges;
/// Memory allocated to store ranges in the result list whose lifetime was
/// generated in the native function.
SmallVector<llvm::OwningArrayRef<Type>> allocatedTypeRanges;
SmallVector<llvm::OwningArrayRef<Value>> allocatedValueRanges;
SmallVector<std::vector<Type>> allocatedTypeRanges;
SmallVector<std::vector<Value>> allocatedValueRanges;
};

//===----------------------------------------------------------------------===//
Expand Down
60 changes: 26 additions & 34 deletions mlir/lib/Rewrite/ByteCode.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1099,32 +1099,33 @@ class ByteCodeRewriteResultList : public PDLResultList {
MutableArrayRef<PDLValue> getResults() { return results; }

/// Return the type ranges allocated by this list.
MutableArrayRef<llvm::OwningArrayRef<Type>> getAllocatedTypeRanges() {
MutableArrayRef<std::vector<Type>> getAllocatedTypeRanges() {
return allocatedTypeRanges;
}

/// Return the value ranges allocated by this list.
MutableArrayRef<llvm::OwningArrayRef<Value>> getAllocatedValueRanges() {
MutableArrayRef<std::vector<Value>> getAllocatedValueRanges() {
return allocatedValueRanges;
}
};

/// This class provides support for executing a bytecode stream.
class ByteCodeExecutor {
public:
ByteCodeExecutor(
const ByteCodeField *curCodeIt, MutableArrayRef<const void *> memory,
MutableArrayRef<llvm::OwningArrayRef<Operation *>> opRangeMemory,
MutableArrayRef<TypeRange> typeRangeMemory,
std::vector<llvm::OwningArrayRef<Type>> &allocatedTypeRangeMemory,
MutableArrayRef<ValueRange> valueRangeMemory,
std::vector<llvm::OwningArrayRef<Value>> &allocatedValueRangeMemory,
MutableArrayRef<unsigned> loopIndex, ArrayRef<const void *> uniquedMemory,
ArrayRef<ByteCodeField> code,
ArrayRef<PatternBenefit> currentPatternBenefits,
ArrayRef<PDLByteCodePattern> patterns,
ArrayRef<PDLConstraintFunction> constraintFunctions,
ArrayRef<PDLRewriteFunction> rewriteFunctions)
ByteCodeExecutor(const ByteCodeField *curCodeIt,
MutableArrayRef<const void *> memory,
MutableArrayRef<std::vector<Operation *>> opRangeMemory,
MutableArrayRef<TypeRange> typeRangeMemory,
std::vector<std::vector<Type>> &allocatedTypeRangeMemory,
MutableArrayRef<ValueRange> valueRangeMemory,
std::vector<std::vector<Value>> &allocatedValueRangeMemory,
MutableArrayRef<unsigned> loopIndex,
ArrayRef<const void *> uniquedMemory,
ArrayRef<ByteCodeField> code,
ArrayRef<PatternBenefit> currentPatternBenefits,
ArrayRef<PDLByteCodePattern> patterns,
ArrayRef<PDLConstraintFunction> constraintFunctions,
ArrayRef<PDLRewriteFunction> rewriteFunctions)
: curCodeIt(curCodeIt), memory(memory), opRangeMemory(opRangeMemory),
typeRangeMemory(typeRangeMemory),
allocatedTypeRangeMemory(allocatedTypeRangeMemory),
Expand Down Expand Up @@ -1367,13 +1368,9 @@ class ByteCodeExecutor {
if (range.empty()) {
rangeMemory[rangeIndex] = {};
} else {
// Allocate a buffer for this type range.
llvm::OwningArrayRef<T> 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];
Expand All @@ -1397,11 +1394,11 @@ class ByteCodeExecutor {

/// The current execution memory.
MutableArrayRef<const void *> memory;
MutableArrayRef<OwningOpRange> opRangeMemory;
MutableArrayRef<std::vector<Operation *>> opRangeMemory;
MutableArrayRef<TypeRange> typeRangeMemory;
std::vector<llvm::OwningArrayRef<Type>> &allocatedTypeRangeMemory;
std::vector<std::vector<Type>> &allocatedTypeRangeMemory;
MutableArrayRef<ValueRange> valueRangeMemory;
std::vector<llvm::OwningArrayRef<Value>> &allocatedValueRangeMemory;
std::vector<std::vector<Value>> &allocatedValueRangeMemory;

/// The current loop indices.
MutableArrayRef<unsigned> loopIndex;
Expand Down Expand Up @@ -1907,20 +1904,18 @@ void ByteCodeExecutor::executeGetUsers() {
LDBG() << "Executing GetUsers:";
unsigned memIndex = read();
unsigned rangeIndex = read();
OwningOpRange &range = opRangeMemory[rangeIndex];
std::vector<Operation *> &range = opRangeMemory[rangeIndex];
memory[memIndex] = &range;

range = OwningOpRange();
range.clear();
Copy link
Contributor Author

@davidstone davidstone Nov 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One note on this line. The old behavior always deallocated memory at this line and then immediately reallocated memory later on. That is much worse for performance but ensures we always use exactly the amount of memory we need. This version holds onto the old memory and allocates only if the new memory requirements are greater than the previous allocation. This avoids memory allocations, but if we have one very large allocation followed by much smaller allocations this will use more memory. I don't understand the ByteCode use case well enough to know if this is an issue. I suspect the old behavior was chosen because that is the only choice you have with llvm::OwningArrayRef, not out of a deliberate design decision, but I want to call out this change. It's a trivial change to go back to the old memory usage pattern.

if (read<PDLValue::Kind>() == PDLValue::Kind::Value) {
// Read the value.
Value value = read<Value>();
if (!value)
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<ValueRange *>();
Expand All @@ -1929,12 +1924,8 @@ void ByteCodeExecutor::executeGetUsers() {
LDBG() << " * Values (" << values->size()
<< "): " << llvm::interleaved(*values);

// Extract all the users of a range of values.
SmallVector<Operation *> 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";
Expand Down Expand Up @@ -2174,7 +2165,8 @@ ByteCodeExecutor::execute(PatternRewriter &rewriter,
executeEraseOp(rewriter);
break;
case ExtractOp:
executeExtract<Operation *, OwningOpRange, PDLValue::Kind::Operation>();
executeExtract<Operation *, std::vector<Operation *>,
PDLValue::Kind::Operation>();
break;
case ExtractType:
executeExtract<Type, TypeRange, PDLValue::Kind::Type>();
Expand Down
7 changes: 3 additions & 4 deletions mlir/lib/Rewrite/ByteCode.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Operation *>;

//===----------------------------------------------------------------------===//
// PDLByteCodePattern
Expand Down Expand Up @@ -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<OwningOpRange> opRangeMemory;
std::vector<std::vector<Operation *>> opRangeMemory;

/// A mutable block of memory used during the matching and rewriting phase of
/// the bytecode to store ranges of types.
std::vector<TypeRange> typeRangeMemory;
/// A set of type ranges that have been allocated by the byte code interpreter
/// to provide a guaranteed lifetime.
std::vector<llvm::OwningArrayRef<Type>> allocatedTypeRangeMemory;
std::vector<std::vector<Type>> allocatedTypeRangeMemory;

/// A mutable block of memory used during the matching and rewriting phase of
/// the bytecode to store ranges of values.
std::vector<ValueRange> valueRangeMemory;
/// A set of value ranges that have been allocated by the byte code
/// interpreter to provide a guaranteed lifetime.
std::vector<llvm::OwningArrayRef<Value>> allocatedValueRangeMemory;
std::vector<std::vector<Value>> 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
Expand Down