[MLIR][cf] Add branch operand edition function to cf.cond_br

[Tyker]

is there interest in having this upstream ?
I am not sure how it should be tested upstream ?

for context I needed this to build a pass that remove all cross block uses and instead pass all value by block arguments. there is not no support for switch ops because they never happen in my context so i don't know if it can be up-streamed as a test.

[rriddle]

CondBranchOp has getTrueDestOperandsMutable() and getFalseDestOperandsMutable() methods that return a MutableOperandRange. That class is what we should be using, as it has support for assignment/erasing/etc. We can extend MutableOperandRange if some necessary functionality is missing.

Switch op provides similar functions for accessing the MutableOperandRange (e.g. getCaseOperandsMutable)

[Tyker]

I didn't realize this existed. thank you for the information.

[rriddle]

In D122863#3421150, @Tyker wrote:

I didn't realize this existed. thank you for the information.

Please feel free to open a patch adding any necessary functionality to MutableOperandRange! It's possible it might be missing something useful, so patches definitely welcome.

mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td

let builders = [
     OpBuilder<(ins "Value":$condition, "Block *":$trueDest,
       "Block *":$falseDest, CArg<"ValueRange", "{}">:$falseOperands), [{
       build($_builder, $_state, condition, trueDest, ValueRange(), falseDest,
             falseOperands);
     }]>];
 
   let extraClassDeclaration = [{
     // These are the indices into the dests list.
-    enum { trueIndex = 0, falseIndex = 1 };
+    enum { trueIndex = 0, falseIndex = 1, trueODSIndex = 1, falseODSIndex = 2 };
 
     // Accessors for operands to the 'true' destination.
     Value getTrueOperand(unsigned idx) {
       assert(idx < getNumTrueOperands());
       return getOperand(getTrueDestOperandIndex() + idx);
     }
 
     void setTrueOperand(unsigned idx, Value value) {
Value getFalseOperand(unsigned idx) {
       assert(idx < getNumFalseOperands());
       return getOperand(getFalseDestOperandIndex() + idx);
     }
     void setFalseOperand(unsigned idx, Value value) {
       assert(idx < getNumFalseOperands());
       setOperand(getFalseDestOperandIndex() + idx, value);
     }
 
+    /// Replace the operands of a branch beginning at 'start' and
+    /// ending at 'start' + 'length' with the ones provided in 'operands'.
+    /// 'operands' may be smaller or larger than the range pointed to by
+    /// 'start'+'length'.
+    void setTrueOperands(unsigned start, unsigned int length, ValueRange values) {
+      setBranchOperands(trueODSIndex, start, length, values);
+    }
+    void setFalseOperands(unsigned start, unsigned int length, ValueRange values) {
+      setBranchOperands(falseODSIndex, start, length, values);
+    }
+
+    /// Insert the given operands into the operand list of a branch at
+    /// the given 'index'.
+    void insertTrueOperands(unsigned start, ValueRange values) {
+      setTrueOperands(start, 0, values);
+    }
+    void insertFalseOperands(unsigned start, ValueRange values) {
+      setFalseOperands(start, 0, values);
+    }
+
     operand_range getTrueOperands() { return getTrueDestOperands(); }
     operand_range getFalseOperands() { return getFalseDestOperands(); }
 
     unsigned getNumFalseOperands() { return getFalseOperands().size(); }
 
     /// Erase the operand at 'index' from the false operand list.
     void eraseFalseOperand(unsigned index) {
       getFalseDestOperandsMutable().erase(index);
     }
 
   private:
     /// Get the index of the first true destination operand.
     unsigned getTrueDestOperandIndex() { return 1; }
 
     /// Get the index of the first false destination operand.
     unsigned getFalseDestOperandIndex() {
       return getTrueDestOperandIndex() + getNumTrueOperands();
     }
+    /// Insert operands values on branch brIndex starting from operandIdx
+    void setBranchOperands(unsigned brODSIndex, unsigned int startIdx,
+                           unsigned int length, ValueRange values);
   }];
 
   let hasCanonicalizer = 1;
   let assemblyFormat = [{
     $condition `,`
     $trueDest (`(` $trueDestOperands^ `:` type($trueDestOperands) `)`)? `,`
     $falseDest (`(` $falseDestOperands^ `:` type($falseDestOperands) `)`)?
     attr-dict

mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp

 }
 
 Block *CondBranchOp::getSuccessorForOperands(ArrayRef<Attribute> operands) {
   if (IntegerAttr condAttr = operands.front().dyn_cast_or_null<IntegerAttr>())
     return condAttr.getValue().isOneValue() ? getTrueDest() : getFalseDest();
   return nullptr;
 }
 
+/// Replace the operands of a branch beginning at 'start' and
+/// ending at 'start' + 'length' with the ones provided in 'operands'.
+/// 'operands' may be smaller or larger than the range pointed to by
+/// 'start'+'length'.
+void CondBranchOp::setBranchOperands(unsigned brODSIndex, unsigned int startIdx,
+                                     unsigned int length, ValueRange values) {
+  auto range = getODSOperandIndexAndLength(brODSIndex);
+  assert(startIdx <= range.second && "insertion out of range");
+  auto sizeAttr = (*this)
+                      ->getAttr(getOperandSegmentSizesAttrName())
+                      .cast<::mlir::DenseIntElementsAttr>();
+
+  /// Update sizes
+  llvm::SmallVector<uint32_t> sizes(sizeAttr.value_begin<uint32_t>(),
+                                    sizeAttr.value_end<uint32_t>());
+  sizes[brODSIndex] += values.size() - length;
+  auto newAttr = DenseIntElementsAttr::get(sizeAttr.getType(), sizes);
+  (*this)->setAttr(getOperandSegmentSizesAttrName(), newAttr);
+
+  /// Update operands
+  (*this)->setOperands(range.first + startIdx, length, values);
+}
+
 //===----------------------------------------------------------------------===//
 // SwitchOp
 //===----------------------------------------------------------------------===//
 
 void SwitchOp::build(OpBuilder &builder, OperationState &result, Value value,
                      Block *defaultDestination, ValueRange defaultOperands,
                      DenseIntElementsAttr caseValues,
                      BlockRange caseDestinations,