shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

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

Browse Source 
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.
This commit is contained in:
David Stone 2025-11-20 11:23:12 -07:00 committed by GitHub
parent 01e5e4fd00
commit bfbd191f35
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 33 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
mlir/include/mlir/IR/PDLPatternMatch.h.inc
View File

@ -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());
}
@ -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());
}
@ -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;
};
//===----------------------------------------------------------------------===//

60
mlir/lib/Rewrite/ByteCode.cpp
View File

@ -1099,12 +1099,12 @@ public:
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;
}
};
@ -1112,19 +1112,20 @@ public:
/// 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),
@ -1367,13 +1368,9 @@ private:
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];
@ -1397,11 +1394,11 @@ private:
/// 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;
@ -1907,10 +1904,10 @@ 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();
if (read<PDLValue::Kind>() == PDLValue::Kind::Value) {
// Read the value.
Value value = read<Value>();
@ -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<ValueRange *>();
@ -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";
@ -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>();

7
mlir/lib/Rewrite/ByteCode.h
View File

@ -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
@ -94,21 +93,21 @@ private:
/// 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