[mlir] Don't allocate an operand storage if the operation is known to never have operands

[rriddle]

Summary

Certain classes of operations, such as FuncOp, are known to never have operands. This revision adds a bit to operation to detect this case and avoid allocating the unnecessary operand storage. This saves 1 word for each instance of these operations.

Depends On D78875

[mehdi_amini]

mlir/lib/IR/Operation.cpp

113

Nice!

mlir/include/mlir/IR/Operation.h

public:
   //===--------------------------------------------------------------------===//
   // Operands
   //===--------------------------------------------------------------------===//
 
   /// Replace the current operands of this operation with the ones provided in
   /// 'operands'.
   void setOperands(ValueRange operands);
 
-  unsigned getNumOperands() { return getOperandStorage().size(); }
+  unsigned getNumOperands() {
+    return LLVM_LIKELY(hasOperandStorage) ? getOperandStorage().size() : 0;
+  }
 
   Value getOperand(unsigned idx) { return getOpOperand(idx).get(); }
   void setOperand(unsigned idx, Value value) {
     return getOpOperand(idx).set(value);
   }
 
   // Support operand iteration.
   using operand_range = OperandRange;
   using operand_iterator = operand_range::iterator;
 
   operand_iterator operand_begin() { return getOperands().begin(); }
   operand_iterator operand_end() { return getOperands().end(); }
 
   /// Returns an iterator on the underlying Value's (Value ).
   operand_range getOperands() { return operand_range(this); }
 
   /// Erase the operand at position `idx`.
   void eraseOperand(unsigned idx) { getOperandStorage().eraseOperand(idx); }
 
   MutableArrayRef<OpOperand> getOpOperands() {
-    return getOperandStorage().getOperands();
+    return LLVM_LIKELY(hasOperandStorage) ? getOperandStorage().getOperands()
+                                          : MutableArrayRef<OpOperand>();
   }
 
   OpOperand &getOpOperand(unsigned idx) { return getOpOperands()[idx]; }
 
   // Support operand type iteration.
   using operand_type_iterator = operand_range::type_iterator;
   using operand_type_range = operand_range::type_range;
   operand_type_iterator operand_type_begin() { return operand_begin(); }
private:
   void updateOrderIfNecessary();
 
   /// Returns true if this operation has a valid order.
   bool hasValidOrder() { return orderIndex != kInvalidOrderIdx; }
 
 private:
   Operation(Location location, OperationName name, ArrayRef<Type> resultTypes,
             unsigned numSuccessors, unsigned numRegions,
-            const NamedAttributeList &attributes);
+            const NamedAttributeList &attributes, bool hasOperandStorage);
 
   // Operations are deleted through the destroy() member because they are
   // allocated with malloc.
   ~Operation();
 
   /// Returns the operand storage object.
   detail::OperandStorage &getOperandStorage() {
+    assert(hasOperandStorage && "expected operation to have operand storage");
     return *getTrailingObjects<detail::OperandStorage>();
   }
 
   /// Returns a pointer to the use list for the given trailing result.
   detail::TrailingOpResult *getTrailingResult(unsigned resultNumber) {
     return getTrailingObjects<detail::TrailingOpResult>() + resultNumber;
   }
 
private:
   /// or derived from.
   Location location;
 
   /// Relative order of this operation in its parent block. Used for
   /// O(1) local dominance checks between operations.
   mutable unsigned orderIndex = 0;
 
   const unsigned numSuccs;
-  const unsigned numRegions : 31;
+  const unsigned numRegions : 30;
+
+  /// This bit signals whether this operation has an operand storage or not. The
+  /// operand storage may be elided for operations that are known to never have
+  /// operands.
+  bool hasOperandStorage : 1;
 
   /// This holds the result types of the operation. There are three different
   /// states recorded here:
   /// - 0 results : The type below is null.
   /// - 1 result  : The single result type is held here.
   /// - N results : The type here is a tuple holding the result types.
   /// Note: We steal a bit for 'hasSingleResult' from somewhere else so that we
   /// can use 'resultType` in an ArrayRef<Type>.

mlir/lib/IR/Operation.cpp

Operation *Operation::create(Location location, OperationName name,
                              ArrayRef<Block *> successors,
                              unsigned numRegions) {
   // We only need to allocate additional memory for a subset of results.
   unsigned numTrailingResults = OpResult::getNumTrailing(resultTypes.size());
   unsigned numInlineResults = OpResult::getNumInline(resultTypes.size());
   unsigned numSuccessors = successors.size();
   unsigned numOperands = operands.size();
 
+  // If the operation is known to have no operands, don't allocate an operand
+  // storage.
+  bool needsOperandStorage = true;
+  if (operands.empty()) {
+    if (const AbstractOperation *abstractOp = name.getAbstractOperation())
+      needsOperandStorage = !abstractOp->hasTrait<OpTrait::ZeroOperands>();
+  }
+
   // Compute the byte size for the operation and the operand storage.
   auto byteSize =
       totalSizeToAlloc<detail::InLineOpResult, detail::TrailingOpResult,
                        BlockOperand, Region, detail::OperandStorage>(
           numInlineResults, numTrailingResults, numSuccessors, numRegions,
-          /*detail::OperandStorage*/ 1);
+          needsOperandStorage ? 1 : 0);
   byteSize +=
       llvm::alignTo(detail::OperandStorage::additionalAllocSize(numOperands),
                     alignof(Operation));
   void *rawMem = malloc(byteSize);
 
   // Create the new Operation.
-  auto op = ::new (rawMem) Operation(location, name, resultTypes, numSuccessors,
-                                     numRegions, attributes);
+  Operation *op =
+      ::new (rawMem) Operation(location, name, resultTypes, numSuccessors,
+                               numRegions, attributes, needsOperandStorage);
 
   assert((numSuccessors == 0 || !op->isKnownNonTerminator()) &&
          "unexpected successors in a non-terminator operation");
 
   // Initialize the results.
   for (unsigned i = 0; i < numInlineResults; ++i)
     new (op->getInlineResult(i)) detail::InLineOpResult();
   for (unsigned i = 0; i < numTrailingResults; ++i)
     new (op->getTrailingResult(i)) detail::TrailingOpResult(i);
 
   // Initialize the regions.
   for (unsigned i = 0; i != numRegions; ++i)
     new (&op->getRegion(i)) Region(op);
 
   // Initialize the operands.
+  if (needsOperandStorage)
     new (&op->getOperandStorage()) detail::OperandStorage(op, operands);
 
   // Initialize the successors.
   auto blockOperands = op->getBlockOperands();
   for (unsigned i = 0; i != numSuccessors; ++i)
     new (&blockOperands[i]) BlockOperand(op, successors[i]);
 
   return op;
 }
 
 Operation::Operation(Location location, OperationName name,
                      ArrayRef<Type> resultTypes, unsigned numSuccessors,
-                     unsigned numRegions, const NamedAttributeList &attributes)
+                     unsigned numRegions, const NamedAttributeList &attributes,
+                     bool hasOperandStorage)
     : location(location), numSuccs(numSuccessors), numRegions(numRegions),
-      hasSingleResult(false), name(name), attrs(attributes) {
+      hasOperandStorage(hasOperandStorage), hasSingleResult(false), name(name),
+      attrs(attributes) {
   if (!resultTypes.empty()) {
     // If there is a single result it is stored in-place, otherwise use a tuple.
     hasSingleResult = resultTypes.size() == 1;
     if (hasSingleResult)
       resultType = resultTypes.front();
     else
       resultType = TupleType::get(resultTypes, location->getContext());
   }
 }
 
 // Operations are deleted through the destroy() member because they are
 // allocated via malloc.
 Operation::~Operation() {
   assert(block == nullptr && "operation destroyed but still in a block");
 
-  // Explicitly run the destructors for the operands and results.
+  // Explicitly run the destructors for the operands.
+  if (hasOperandStorage)
     getOperandStorage().~OperandStorage();
 
   // Explicitly run the destructors for the successors.
   for (auto &successor : getBlockOperands())
     successor.~BlockOperand();
 
   // Explicitly destroy the regions.
   for (auto &region : getRegions())
     region.~Region();
void Operation::replaceUsesOfWith(Value from, Value to) {
   for (auto &operand : getOpOperands())
     if (operand.get() == from)
       operand.set(to);
 }
 
 /// Replace the current operands of this operation with the ones provided in
 /// 'operands'.
 void Operation::setOperands(ValueRange operands) {
-  getOperandStorage().setOperands(this, operands);
+  if (LLVM_LIKELY(hasOperandStorage))
+    return getOperandStorage().setOperands(this, operands);
+  assert(operands.empty() && "setting operands without an operand storage");
 }
 
 //===----------------------------------------------------------------------===//
 // Diagnostics
 //===----------------------------------------------------------------------===//
 
 /// Emit an error about fatal conditions with this operation, reporting up to
 /// any diagnostic handlers that may be listening.