Add operand support to the Instruction base class. Add setOperand methods

to all the things.  Fill out the OneOperand trait class with support for
getting and setting operands, allowing DimOp to have a working
get/setOperand() method.

I'm not thrilled with the extra template argument on OneOperand, I'll will
investigate removing that in a follow-on patch.

PiperOrigin-RevId: 205679696

mlir/include/mlir/IR/Instructions.h

@@ -61,6 +61,59 @@ public:
   void print(raw_ostream &os) const;
   void dump() const;
 
+  //===--------------------------------------------------------------------===//
+  // Operands
+  //===--------------------------------------------------------------------===//
+
+  unsigned getNumOperands() const;
+
+  CFGValue *getOperand(unsigned idx) { return getInstOperand(idx).get(); }
+  const CFGValue *getOperand(unsigned idx) const {
+    return getInstOperand(idx).get();
+  }
+  void setOperand(unsigned idx, CFGValue *value) {
+    return getInstOperand(idx).set(value);
+  }
+
+  // Support non-const operand iteration.
+  using operand_iterator = OperandIterator<Instruction, CFGValue>;
+
+  operand_iterator operand_begin() { return operand_iterator(this, 0); }
+
+  operand_iterator operand_end() {
+    return operand_iterator(this, getNumOperands());
+  }
+
+  llvm::iterator_range<operand_iterator> getOperands() {
+    return {operand_begin(), operand_end()};
+  }
+
+  // Support const operand iteration.
+  using const_operand_iterator =
+      OperandIterator<const Instruction, const CFGValue>;
+
+  const_operand_iterator operand_begin() const {
+    return const_operand_iterator(this, 0);
+  }
+
+  const_operand_iterator operand_end() const {
+    return const_operand_iterator(this, getNumOperands());
+  }
+
+  llvm::iterator_range<const_operand_iterator> getOperands() const {
+    return {operand_begin(), operand_end()};
+  }
+
+  MutableArrayRef<InstOperand> getInstOperands();
+  ArrayRef<InstOperand> getInstOperands() const {
+    return const_cast<Instruction *>(this)->getInstOperands();
+  }
+
+  InstOperand &getInstOperand(unsigned idx) { return getInstOperands()[idx]; }
+  const InstOperand &getInstOperand(unsigned idx) const {
+    return getInstOperands()[idx];
+  }
+
 protected:
   Instruction(Kind kind) : kind(kind) {}
 
@@ -106,6 +159,9 @@ public:
   const CFGValue *getOperand(unsigned idx) const {
     return getInstOperand(idx).get();
   }
+  void setOperand(unsigned idx, CFGValue *value) {
+    return getInstOperand(idx).set(value);
+  }
 
   // Support non-const operand iteration.
   using operand_iterator = OperandIterator<OperationInst, CFGValue>;
@@ -268,6 +324,9 @@ public:
   const CFGValue *getOperand(unsigned idx) const {
     return getInstOperand(idx).get();
   }
+  void setOperand(unsigned idx, CFGValue *value) {
+    return getInstOperand(idx).set(value);
+  }
 
   ArrayRef<InstOperand> getInstOperands() const { return operands; }
   MutableArrayRef<InstOperand> getInstOperands() { return operands; }
@@ -317,6 +376,10 @@ public:
     return getInstOperand(idx).get();
   }
 
+  void setOperand(unsigned idx, CFGValue *value) {
+    return getInstOperand(idx).set(value);
+  }
+
   // Support non-const operand iteration.
   using operand_iterator = OperandIterator<ReturnInst, CFGValue>;
 

mlir/include/mlir/IR/Operation.h

@@ -54,6 +54,7 @@ public:
   const SSAValue *getOperand(unsigned idx) const {
     return const_cast<Operation *>(this)->getOperand(idx);
   }
+  void setOperand(unsigned idx, SSAValue *value);
 
   // Support non-const operand iteration.
   using operand_iterator = OperandIterator<Operation, SSAValue>;

mlir/include/mlir/IR/OperationImpl.h

@@ -147,13 +147,13 @@ private:
 
 /// This class provides the API for ops that are known to have exactly one
 /// SSA operand.
-class OneOperand {
+template <typename ConcreteType> class OneOperand {
 public:
-  void getOperand() const {
-    /// TODO.
+  SSAValue *getOperand() const {
+    return static_cast<ConcreteType *>(this)->getOperand(0);
   }
-  void setOperand() {
-    /// TODO.
+  void setOperand(SSAValue *value) {
+    static_cast<ConcreteType *>(this)->setOperand(0, value);
   }
 
   static const char *verifyBase(const Operation *op) {

mlir/include/mlir/IR/StandardOps.h

@@ -57,7 +57,7 @@ private:
 ///   %1 = dim %0, 2 : tensor<?x?x?xf32>
 ///
 class DimOp
-    : public OpImpl::Base<DimOp, OpImpl::OneOperand, OpImpl::OneResult> {
+    : public OpImpl::Base<DimOp, OpImpl::OneOperand<DimOp>, OpImpl::OneResult> {
 public:
   /// This returns the dimension number that the 'dim' is inspecting.
   unsigned getIndex() const {

mlir/lib/IR/Instructions.cpp

@@ -70,6 +70,28 @@ CFGFunction *Instruction::getFunction() const {
   return getBlock()->getFunction();
 }
 
+unsigned Instruction::getNumOperands() const {
+  switch (getKind()) {
+  case Kind::Operation:
+    return cast<OperationInst>(this)->getNumOperands();
+  case Kind::Branch:
+    return cast<BranchInst>(this)->getNumOperands();
+  case Kind::Return:
+    return cast<ReturnInst>(this)->getNumOperands();
+  }
+}
+
+MutableArrayRef<InstOperand> Instruction::getInstOperands() {
+  switch (getKind()) {
+  case Kind::Operation:
+    return cast<OperationInst>(this)->getInstOperands();
+  case Kind::Branch:
+    return cast<BranchInst>(this)->getInstOperands();
+  case Kind::Return:
+    return cast<ReturnInst>(this)->getInstOperands();
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // OperationInst
 //===----------------------------------------------------------------------===//

mlir/lib/IR/Operation.cpp

@@ -57,6 +57,17 @@ SSAValue *Operation::getOperand(unsigned idx) {
   }
 }
 
+void Operation::setOperand(unsigned idx, SSAValue *value) {
+  if (auto *inst = dyn_cast<OperationInst>(this)) {
+    inst->setOperand(idx, cast<CFGValue>(value));
+  } else {
+    auto *stmt = cast<OperationStmt>(this);
+    (void)stmt;
+    // TODO: Add operands to OperationStmt.
+    abort();
+  }
+}
+
 /// Return the number of results this operation has.
 unsigned Operation::getNumResults() const {
   if (auto *inst = dyn_cast<OperationInst>(this)) {