[mlir] Add getArgOperandsMutable method to CallOpInterface

[maerhart]

Summary

Add a method to the CallOpInterface to get a mutable operand range over
the function arguments. This allows to add, remove, or change the type
of call arguments in a generic manner without having to assume that the
argument operand range is at the end of the operand list, or having to
type switch on all supported concrete operation kinds.

Alternatively, a new OpInterface could be added which inherits from
CallOpInterface and appends it with the mutable variants of the base
interface.

There will be two users of this new function in the beginning:
(1) A few passes in the Arc dialect in CIRCT already use a downstream
implementation of the alternative case mentioned above: https://github.com/llvm/circt/blob/main/include/circt/Dialect/Arc/ArcInterfaces.td#L15
(2) The BufferDeallocation pass will be modified to be able to pass
ownership of memrefs to called private functions if the caller does not
need the memref anymore by appending the function argument list with a
boolean value per memref, thus enabling earlier deallocation of the
memref which can lead to lower peak memory usage.

[maerhart]

Fix flang build

flang/include/flang/Optimizer/Dialect/FIROps.td

let extraClassDeclaration = [{
 
     /// Get the argument operands to the called function.
     operand_range getArgOperands() {
       if ((*this)->getAttrOfType<mlir::SymbolRefAttr>(getCalleeAttrName()))
         return {arg_operand_begin(), arg_operand_end()};
       return {arg_operand_begin() + 1, arg_operand_end()};
     }
 
+    mlir::MutableOperandRange getArgOperandsMutable() {
+      if ((*this)->getAttrOfType<mlir::SymbolRefAttr>(getCalleeAttrName()))
+        return getArgsMutable();
+      return mlir::MutableOperandRange(*this, 1, getArgs().size() - 1);
+    }
+
     operand_iterator arg_operand_begin() { return operand_begin(); }
     operand_iterator arg_operand_end() { return operand_end(); }
 
     /// Return the callee of this operation.
     mlir::CallInterfaceCallable getCallableForCallee() {
       if (auto calling =
           (*this)->getAttrOfType<mlir::SymbolRefAttr>(getCalleeAttrName()))
         return calling;

mlir/examples/toy/Ch4/mlir/Dialect.cpp

 void GenericCallOp::setCalleeFromCallable(CallInterfaceCallable callee) {
   (*this)->setAttr("callee", callee.get<SymbolRefAttr>());
 }
 
 /// Get the argument operands to the called function, this is required by the
 /// call interface.
 Operation::operand_range GenericCallOp::getArgOperands() { return getInputs(); }
 
+/// Get the argument operands to the called function as a mutable range, this is
+/// required by the call interface.
+MutableOperandRange GenericCallOp::getArgOperandsMutable() {
+  return getInputsMutable();
+}
+
 //===----------------------------------------------------------------------===//
 // MulOp
 //===----------------------------------------------------------------------===//
 
 void MulOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
                   mlir::Value lhs, mlir::Value rhs) {
   state.addTypes(UnrankedTensorType::get(builder.getF64Type()));
   state.addOperands({lhs, rhs});

mlir/examples/toy/Ch5/mlir/Dialect.cpp

 void GenericCallOp::setCalleeFromCallable(CallInterfaceCallable callee) {
   (*this)->setAttr("callee", callee.get<SymbolRefAttr>());
 }
 
 /// Get the argument operands to the called function, this is required by the
 /// call interface.
 Operation::operand_range GenericCallOp::getArgOperands() { return getInputs(); }
 
+/// Get the argument operands to the called function as a mutable range, this is
+/// required by the call interface.
+MutableOperandRange GenericCallOp::getArgOperandsMutable() {
+  return getInputsMutable();
+}
+
 //===----------------------------------------------------------------------===//
 // MulOp
 //===----------------------------------------------------------------------===//
 
 void MulOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
                   mlir::Value lhs, mlir::Value rhs) {
   state.addTypes(UnrankedTensorType::get(builder.getF64Type()));
   state.addOperands({lhs, rhs});

mlir/examples/toy/Ch6/mlir/Dialect.cpp

 void GenericCallOp::setCalleeFromCallable(CallInterfaceCallable callee) {
   (*this)->setAttr("callee", callee.get<SymbolRefAttr>());
 }
 
 /// Get the argument operands to the called function, this is required by the
 /// call interface.
 Operation::operand_range GenericCallOp::getArgOperands() { return getInputs(); }
 
+/// Get the argument operands to the called function as a mutable range, this is
+/// required by the call interface.
+MutableOperandRange GenericCallOp::getArgOperandsMutable() {
+  return getInputsMutable();
+}
+
 //===----------------------------------------------------------------------===//
 // MulOp
 //===----------------------------------------------------------------------===//
 
 void MulOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
                   mlir::Value lhs, mlir::Value rhs) {
   state.addTypes(UnrankedTensorType::get(builder.getF64Type()));
   state.addOperands({lhs, rhs});

mlir/examples/toy/Ch7/mlir/Dialect.cpp

 void GenericCallOp::setCalleeFromCallable(CallInterfaceCallable callee) {
   (*this)->setAttr("callee", callee.get<SymbolRefAttr>());
 }
 
 /// Get the argument operands to the called function, this is required by the
 /// call interface.
 Operation::operand_range GenericCallOp::getArgOperands() { return getInputs(); }
 
+/// Get the argument operands to the called function as a mutable range, this is
+/// required by the call interface.
+MutableOperandRange GenericCallOp::getArgOperandsMutable() {
+  return getInputsMutable();
+}
+
 //===----------------------------------------------------------------------===//
 // MulOp
 //===----------------------------------------------------------------------===//
 
 void MulOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
                   mlir::Value lhs, mlir::Value rhs) {
   state.addTypes(UnrankedTensorType::get(builder.getF64Type()));
   state.addOperands({lhs, rhs});

mlir/include/mlir/Dialect/Async/IR/AsyncOps.td

def Async_CallOp : Async_Op<"call",
   let extraClassDeclaration = [{
     FunctionType getCalleeType();
 
     /// Get the argument operands to the called function.
     operand_range getArgOperands() {
       return {arg_operand_begin(), arg_operand_end()};
     }
 
+    MutableOperandRange getArgOperandsMutable() {
+      return getOperandsMutable();
+    }
+
     operand_iterator arg_operand_begin() { return operand_begin(); }
     operand_iterator arg_operand_end() { return operand_end(); }
 
     /// Return the callee of this operation.
     CallInterfaceCallable getCallableForCallee() {
       return (*this)->getAttrOfType<SymbolRefAttr>("callee");
     }
 

mlir/include/mlir/Dialect/Func/IR/FuncOps.td

def CallOp : Func_Op<"call",
   let extraClassDeclaration = [{
     FunctionType getCalleeType();
 
     /// Get the argument operands to the called function.
     operand_range getArgOperands() {
       return {arg_operand_begin(), arg_operand_end()};
     }
 
+    MutableOperandRange getArgOperandsMutable() {
+      return getOperandsMutable();
+    }
+
     operand_iterator arg_operand_begin() { return operand_begin(); }
     operand_iterator arg_operand_end() { return operand_end(); }
 
     /// Return the callee of this operation.
     CallInterfaceCallable getCallableForCallee() {
       return (*this)->getAttrOfType<SymbolRefAttr>("callee");
     }
 
let extraClassDeclaration = [{
     // TODO: Remove once migrated callers.
     ValueRange operands() { return getCalleeOperands(); }
 
     /// Get the argument operands to the called function.
     operand_range getArgOperands() {
       return {arg_operand_begin(), arg_operand_end()};
     }
 
+    MutableOperandRange getArgOperandsMutable() {
+      return getCalleeOperandsMutable();
+    }
+
     operand_iterator arg_operand_begin() { return ++operand_begin(); }
     operand_iterator arg_operand_end() { return operand_end(); }
 
     /// Return the callee of this operation.
     CallInterfaceCallable getCallableForCallee() { return getCallee(); }
 
     /// Set the callee for this operation.
     void setCalleeFromCallable(CallInterfaceCallable callee) {

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

let description = [{
     llvm.call @bar(%0) : (f32) -> ()
 
     // Indirect call with an argument and without a result.
     llvm.call %1(%0) : !llvm.ptr, (f32) -> ()
     ```
   }];
 
   dag args = (ins OptionalAttr<FlatSymbolRefAttr>:$callee,
-                  Variadic<LLVM_Type>,
+                  Variadic<LLVM_Type>:$callee_operands,
                   DefaultValuedAttr<LLVM_FastmathFlagsAttr,
                                    "{}">:$fastmathFlags,
                   OptionalAttr<DenseI32ArrayAttr>:$branch_weights);
   // Append the aliasing related attributes defined in LLVM_MemAccessOpBase.
   let arguments = !con(args, aliasAttrs);
   let results = (outs Optional<LLVM_Type>:$result);
   let builders = [
     OpBuilder<(ins "LLVMFuncOp":$func, "ValueRange":$args)>,

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

let extraClassDeclaration = [{
 
     void setCalleeFromCallable(::mlir::CallInterfaceCallable callee) {
       setTargetAttr(callee.get<SymbolRefAttr>());
     }
 
     ::mlir::Operation::operand_range getArgOperands() {
       return getOperands();
     }
+
+    ::mlir::MutableOperandRange getArgOperandsMutable() {
+      return getOperandsMutable();
+    }
   }];
 }
 
 def MatchOperationNameOp : TransformDialectOp<"match.operation_name",
     [SingleOpMatcher,
      MatchOpInterface,
      MemoryEffectsOpInterface]> {
   let summary = "Matches a single operation of one of the given kinds";

mlir/include/mlir/Interfaces/CallInterfaces.td

InterfaceMethod<[{
       "void", "setCalleeFromCallable", (ins "::mlir::CallInterfaceCallable":$callee)
     >,
     InterfaceMethod<[{
         Returns the operands within this call that are used as arguments to the
         callee.
       }],
       "::mlir::Operation::operand_range", "getArgOperands"
     >,
+    InterfaceMethod<[{
+        Returns the operands within this call that are used as arguments to the
+        callee as a mutable range.
+      }],
+      "::mlir::MutableOperandRange", "getArgOperandsMutable">,
   ];
 
   let extraClassDeclaration = [{
     /// Resolve the callable operation for given callee to a
     /// CallableOpInterface, or nullptr if a valid callable was not resolved.
     /// `symbolTable` is an optional parameter that will allow for using a
     /// cached symbol table for symbol lookups instead of performing an O(N)
     /// scan.

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

void CallOp::setCalleeFromCallable(CallInterfaceCallable callee) {
   // Indirect call, callee Value is the first operand.
   return setOperand(0, callee.get<Value>());
 }
 
 Operation::operand_range CallOp::getArgOperands() {
   return getOperands().drop_front(getCallee().has_value() ? 0 : 1);
 }
 
+MutableOperandRange CallOp::getArgOperandsMutable() {
+  return MutableOperandRange(*this, getCallee().has_value() ? 0 : 1,
+                             getCalleeOperands().size());
+}
+
 LogicalResult CallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
   if (getNumResults() > 1)
     return emitOpError("must have 0 or 1 result");
 
   // Type for the callee, we'll get it differently depending if it is a direct
   // or indirect call.
   Type fnType;
 
void InvokeOp::setCalleeFromCallable(CallInterfaceCallable callee) {
   // Indirect call, callee Value is the first operand.
   return setOperand(0, callee.get<Value>());
 }
 
 Operation::operand_range InvokeOp::getArgOperands() {
   return getOperands().drop_front(getCallee().has_value() ? 0 : 1);
 }
 
+MutableOperandRange InvokeOp::getArgOperandsMutable() {
+  return MutableOperandRange(*this, getCallee().has_value() ? 0 : 1,
+                             getCalleeOperands().size());
+}
+
 LogicalResult InvokeOp::verify() {
   if (getNumResults() > 1)
     return emitOpError("must have 0 or 1 result");
 
   Block *unwindDest = getUnwindDest();
   if (unwindDest->empty())
     return emitError("must have at least one operation in unwind destination");
 

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

 void FunctionCallOp::setCalleeFromCallable(CallInterfaceCallable callee) {
   (*this)->setAttr(kCallee, callee.get<SymbolRefAttr>());
 }
 
 Operation::operand_range FunctionCallOp::getArgOperands() {
   return getArguments();
 }
 
+MutableOperandRange FunctionCallOp::getArgOperandsMutable() {
+  return getArgumentsMutable();
+}
+
 //===----------------------------------------------------------------------===//
 // spirv.mlir.loop
 //===----------------------------------------------------------------------===//
 
 void LoopOp::build(OpBuilder &builder, OperationState &state) {
   state.addAttribute("loop_control", builder.getAttr<spirv::LoopControlAttr>(
                                          spirv::LoopControl::None));
   state.addRegion();

mlir/test/lib/Dialect/Test/TestDialect.cpp

 void TestCallAndStoreOp::setCalleeFromCallable(CallInterfaceCallable callee) {
   setCalleeAttr(callee.get<SymbolRefAttr>());
 }
 
 Operation::operand_range TestCallAndStoreOp::getArgOperands() {
   return getCalleeOperands();
 }
 
+MutableOperandRange TestCallAndStoreOp::getArgOperandsMutable() {
+  return getCalleeOperandsMutable();
+}
+
 void TestStoreWithARegion::getSuccessorRegions(
     std::optional<unsigned> index, ArrayRef<Attribute> operands,
     SmallVectorImpl<RegionSuccessor> &regions) {
   if (!index) {
     regions.emplace_back(&getBody(), getBody().front().getArguments());
   } else {
     regions.emplace_back();
   }