[mlir] Rewrite the internal representation of OpResult to be optimized for memory.

Summary:
This changes the implementation of OpResult to have some of the results be represented inline in Value, via a pointer int pair of Operation*+result number, and the rest being trailing objects on the main operation. The full details of the new representation is detailed in the proposal here:
https://groups.google.com/a/tensorflow.org/g/mlir/c/XXzzKhqqF_0/m/v6bKb08WCgAJ

The only difference between here and the above proposal is that we only steal 2-bits for the Value kind instead of 3. This means that we can only fit 2-results inline instead of 6. This allows for other users to steal the final bit for PointerUnion/etc. If necessary, we can always steal this bit back in the future to save more space if 3-6 results are common enough.

Reviewed By: jpienaar

Differential Revision: https://reviews.llvm.org/D72020

1 files changed, 97 insertions, 0 deletions

diff --git a/mlir/lib/IR/Value.cpp b/mlir/lib/IR/Value.cpp
index fd9fda53e68..13bd0bcd8f3 100644
--- a/mlir/lib/IR/Value.cpp
+++ b/mlir/lib/IR/Value.cpp
@@ -9,8 +9,63 @@
 #include "mlir/IR/Value.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/IR/StandardTypes.h"
 using namespace mlir;
 
+/// Construct a value.
+Value::Value(detail::BlockArgumentImpl *impl)
+    : ownerAndKind(impl, Kind::BlockArgument) {}
+Value::Value(Operation *op, unsigned resultNo) {
+  assert(op->getNumResults() > resultNo && "invalid result number");
+  if (LLVM_LIKELY(canPackResultInline(resultNo))) {
+    ownerAndKind = {op, static_cast<Kind>(resultNo)};
+    return;
+  }
+
+  // If we can't pack the result directly, we need to represent this as a
+  // trailing result.
+  unsigned trailingResultNo =
+      resultNo - static_cast<unsigned>(Kind::TrailingOpResult);
+  ownerAndKind = {op->getTrailingResult(trailingResultNo),
+                  Kind::TrailingOpResult};
+}
+
+/// Return the type of this value.
+Type Value::getType() const {
+  if (BlockArgument arg = dyn_cast<BlockArgument>())
+    return arg.getType();
+
+  // If this is an operation result, query the parent operation.
+  OpResult result = cast<OpResult>();
+  Operation *owner = result.getOwner();
+  if (owner->hasSingleResult)
+    return owner->resultType;
+  return owner->resultType.cast<TupleType>().getType(result.getResultNumber());
+}
+
+/// Mutate the type of this Value to be of the specified type.
+void Value::setType(Type newType) {
+  if (BlockArgument arg = dyn_cast<BlockArgument>())
+    return arg.setType(newType);
+  OpResult result = cast<OpResult>();
+
+  // If the owner has a single result, simply update it directly.
+  Operation *owner = result.getOwner();
+  if (owner->hasSingleResult) {
+    owner->resultType = newType;
+    return;
+  }
+  unsigned resultNo = result.getResultNumber();
+
+  // Otherwise, rebuild the tuple if the new type is different from the current.
+  auto curTypes = owner->resultType.cast<TupleType>().getTypes();
+  if (curTypes[resultNo] == newType)
+    return;
+  auto newTypes = llvm::to_vector<4>(curTypes);
+  newTypes[resultNo] = newType;
+  owner->resultType = TupleType::get(newTypes, newType.getContext());
+}
+
 /// If this value is the result of an Operation, return the operation that
 /// defines it.
 Operation *Value::getDefiningOp() const {
@@ -84,6 +139,48 @@ bool Value::use_empty() const {
 }
 
 //===----------------------------------------------------------------------===//
+// OpResult
+//===----------------------------------------------------------------------===//
+
+/// Returns the operation that owns this result.
+Operation *OpResult::getOwner() const {
+  // If the result is in-place, the `owner` is the operation.
+  if (LLVM_LIKELY(getKind() != Kind::TrailingOpResult))
+    return reinterpret_cast<Operation *>(ownerAndKind.getPointer());
+
+  // Otherwise, we need to do some arithmetic to get the operation pointer.
+  // Move the trailing owner to the start of the array.
+  auto *trailingIt =
+      static_cast<detail::TrailingOpResult *>(ownerAndKind.getPointer());
+  trailingIt -= trailingIt->trailingResultNumber;
+
+  // This point is the first trailing object after the operation. So all we need
+  // to do here is adjust for the operation size.
+  return reinterpret_cast<Operation *>(trailingIt) - 1;
+}
+
+/// Return the result number of this result.
+unsigned OpResult::getResultNumber() const {
+  // If the result is in-place, we can use the kind directly.
+  if (LLVM_LIKELY(getKind() != Kind::TrailingOpResult))
+    return static_cast<unsigned>(ownerAndKind.getInt());
+  // Otherwise, we add the number of inline results to the trailing owner.
+  auto *trailingIt =
+      static_cast<detail::TrailingOpResult *>(ownerAndKind.getPointer());
+  unsigned trailingNumber = trailingIt->trailingResultNumber;
+  return trailingNumber + static_cast<unsigned>(Kind::TrailingOpResult);
+}
+
+/// Given a number of operation results, returns the number that need to be
+/// stored as trailing.
+unsigned OpResult::getNumTrailing(unsigned numResults) {
+  // If we can pack all of the results, there is no need for additional storage.
+  if (numResults <= static_cast<unsigned>(Kind::TrailingOpResult))
+    return 0;
+  return numResults - static_cast<unsigned>(Kind::TrailingOpResult);
+}
+
+//===----------------------------------------------------------------------===//
 // BlockOperand
 //===----------------------------------------------------------------------===//