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

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.

Author: davidstone

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

mlir/lib/Rewrite/ByteCode.cpp

@@ -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),
@@ -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];
@@ -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;
@@ -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();
   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 *>();
@@ -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>();

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<Operation *>;
 
 //===----------------------------------------------------------------------===//
 // 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<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