[mlir][CallOpInterface] Add `setCalleeFromCallable` method

[Whitney]

Summary

Currently CallOpInterface has a method getCallableForCallee to have a consistent way to get the callee from an operation with CallOpInterface, but missing a consistent way to set a callee for an operation with CallOpInterface.

A set callee method is useful for transformations that operate on CallOpInterface, and change the callee, e.g., a pass that specialize function, which clone the callee, and change the CallOpInterface's callee to the cloned version. Without such method, transformation would need to understand the implementation for every operations with CallOpInterface, and have a type switch to handle them.

This review adds a method to set callee for operation with CallOpInterface.

[gysit]

Can you explain the motivation for the change (ideally in the commit message)? Could you achieve your goal also by changing the specific call operation you are working with?

Also note that the unit tests are currently failing.

[Whitney]

In D149763#4317844, @gysit wrote:

Can you explain the motivation for the change (ideally in the commit message)? Could you achieve your goal also by changing the specific call operation you are working with?

Added in commit message. I am working with a CallOpInterface, so it would need to have a big type switch to handle every possible call operation, which is not ideal.

Also note that the unit tests are currently failing.

Thanks for letting me know, updated the diff.

[gysit]

Thanks for the longer commit message. It seems there is still a compilation error though.

Any chance you can add a test to exercise the new functionality upstream? It seems there is at least a test operation that implements the interface already but probably no pass or unit test where a test could be added easily.

[zero9178]

mlir/include/mlir/Interfaces/CallInterfaces.td

47

Could you document here what the expectected preconditions and post conditions of the method are?

E.g. if I have a func.call op is it a precondition of the method that $callee must be a FlatSymbolRefAttr or is it a post condition that the implementation should then simply do nothing if its not. Vice versa with ops like func.call_indirect and Value.

I also think that describing the meaning of what the SSA value represents should not be described here, given the meaning is dependent upon the Op that creates the value and/or the call op that uses it (could e.g. also be a function pointer, doesn't have to be a lambda-like operation with a region).

[Whitney]

Added mlir:: in FIROps.td.

[Whitney]

changed comment.

[Whitney]

In D149763#4318835, @gysit wrote:

Thanks for the longer commit message. It seems there is still a compilation error though.

The compilation error should be fixed now.

Any chance you can add a test to exercise the new functionality upstream? It seems there is at least a test operation that implements the interface already but probably no pass or unit test where a test could be added easily.

There seems to have no extra testing for getCallableForCallee on top of the test operation. As you said, there seems to be no easy way to add extra testing. Do you have any suggestions?

mlir/include/mlir/Interfaces/CallInterfaces.td

47

Updated comment.

[gysit]

LGTM

There seems to have no extra testing for getCallableForCallee on top of the test operation. As you said, there seems to be no easy way to add extra testing. Do you have any suggestions?

I think ideally the pass exercising the new functionality would live upstream. But let's proceed with the revision as is then.

[zero9178]

LGTM

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

let extraClassDeclaration = [{
 
     /// Return the callee of this operation.
     mlir::CallInterfaceCallable getCallableForCallee() {
       if (auto calling =
           (*this)->getAttrOfType<mlir::SymbolRefAttr>(getCalleeAttrName()))
         return calling;
       return getOperand(0);
     }
+
+    /// Set the callee for this operation.
+    void setCalleeFromCallable(mlir::CallInterfaceCallable callee) {
+      if (auto calling =
+          (*this)->getAttrOfType<mlir::SymbolRefAttr>(getCalleeAttrName()))
+        (*this)->setAttr(getCalleeAttrName(), callee.get<mlir::SymbolRefAttr>());
+      setOperand(0, callee.get<mlir::Value>());
+    }
   }];
 }
 
 def fir_DispatchOp : fir_Op<"dispatch", []> {
   let summary = "call a type-bound procedure";
 
   let description = [{
     Perform a dynamic dispatch on the method name via the dispatch table

mlir/docs/Interfaces.md

 
 *   `Interface class name`
     -   (`C++ class` -- `ODS class`(if applicable))
 
 ##### CallInterfaces
 
 *   `CallOpInterface` - Used to represent operations like 'call'
     -   `CallInterfaceCallable getCallableForCallee()`
+    -   `void setCalleeFromCallable(CallInterfaceCallable)`
 *   `CallableOpInterface` - Used to represent the target callee of call.
     -   `Region * getCallableRegion()`
     -   `ArrayRef<Type> getCallableResults()`
     -   `ArrayAttr getCallableArgAttrs()`
     -   `ArrayAttr getCallableResAttrs()`
 
 ##### RegionKindInterfaces
 

mlir/docs/Tutorials/Toy/Ch-4.md

 // ....
 
 /// Return the callee of the generic call operation, this is required by the
 /// call interface.
 CallInterfaceCallable GenericCallOp::getCallableForCallee() {
   return getAttrOfType<SymbolRefAttr>("callee");
 }
 
+/// Set the callee for the generic call operation, this is required by the call
+/// interface.
+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 inputs(); }
 ```
 
 Now that the inliner has been informed about the Toy dialect, we can add the
 inliner pass to the pass manager for Toy:
 

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

 }
 
 /// Return the callee of the generic call operation, this is required by the
 /// call interface.
 CallInterfaceCallable GenericCallOp::getCallableForCallee() {
   return (*this)->getAttrOfType<SymbolRefAttr>("callee");
 }
 
+/// Set the callee for the generic call operation, this is required by the call
+/// interface.
+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(); }
 
 //===----------------------------------------------------------------------===//
 // MulOp
 //===----------------------------------------------------------------------===//
 

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

 }
 
 /// Return the callee of the generic call operation, this is required by the
 /// call interface.
 CallInterfaceCallable GenericCallOp::getCallableForCallee() {
   return (*this)->getAttrOfType<SymbolRefAttr>("callee");
 }
 
+/// Set the callee for the generic call operation, this is required by the call
+/// interface.
+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(); }
 
 //===----------------------------------------------------------------------===//
 // MulOp
 //===----------------------------------------------------------------------===//
 

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

 }
 
 /// Return the callee of the generic call operation, this is required by the
 /// call interface.
 CallInterfaceCallable GenericCallOp::getCallableForCallee() {
   return (*this)->getAttrOfType<SymbolRefAttr>("callee");
 }
 
+/// Set the callee for the generic call operation, this is required by the call
+/// interface.
+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(); }
 
 //===----------------------------------------------------------------------===//
 // MulOp
 //===----------------------------------------------------------------------===//
 

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

 }
 
 /// Return the callee of the generic call operation, this is required by the
 /// call interface.
 CallInterfaceCallable GenericCallOp::getCallableForCallee() {
   return (*this)->getAttrOfType<SymbolRefAttr>("callee");
 }
 
+/// Set the callee for the generic call operation, this is required by the call
+/// interface.
+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(); }
 
 //===----------------------------------------------------------------------===//
 // MulOp
 //===----------------------------------------------------------------------===//
 

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

let extraClassDeclaration = [{
 
     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");
     }
+
+    /// Set the callee for this operation.
+    void setCalleeFromCallable(CallInterfaceCallable callee) {
+      (*this)->setAttr("callee", callee.get<SymbolRefAttr>());
+    }
   }];
 
   let assemblyFormat = [{
     $callee `(` $operands `)` attr-dict `:` functional-type($operands, results)
   }];
 }
 
 def Async_ReturnOp : Async_Op<"return",

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

let extraClassDeclaration = [{
 
     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");
     }
+
+    /// Set the callee for this operation.
+    void setCalleeFromCallable(CallInterfaceCallable callee) {
+      (*this)->setAttr("callee", callee.get<SymbolRefAttr>());
+    }
   }];
 
   let assemblyFormat = [{
     $callee `(` $operands `)` attr-dict `:` functional-type($operands, results)
   }];
 }
 
 //===----------------------------------------------------------------------===//
operand_range getArgOperands() {
       return {arg_operand_begin(), arg_operand_end()};
     }
 
     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) {
+      setOperand(0, callee.get<Value>());
+    }
   }];
 
   let hasCanonicalizeMethod = 1;
   let assemblyFormat = [{
     $callee `(` $callee_operands `)` attr-dict `:` type($callee)
   }];
 }
 

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

let assemblyFormat =
       "$target `failures` `(` $failure_propagation_mode `)`"
       "`(` $operands `)` attr-dict `:` functional-type($operands, $results)";
 
   let extraClassDeclaration = [{
     ::mlir::CallInterfaceCallable getCallableForCallee() {
       return getTarget();
     }
 
+    void setCalleeFromCallable(::mlir::CallInterfaceCallable callee) {
+      setTargetAttr(callee.get<SymbolRefAttr>());
+    }
+
     ::mlir::Operation::operand_range getArgOperands() {
       return getOperands();
     }
   }];
 }
 
 def MatchOperationNameOp : TransformDialectOp<"match.operation_name",
     [SingleOpMatcher,

mlir/include/mlir/Interfaces/CallInterfaces.td

InterfaceMethod<[{
         Returns the callee of this call-like operation. A `callee` is either a
         reference to a symbol, via SymbolRefAttr, or a reference to a defined
         SSA value. If the reference is an SSA value, the SSA value corresponds
         to a region of a lambda-like operation.
       }],
       "::mlir::CallInterfaceCallable", "getCallableForCallee"
     >,
     InterfaceMethod<[{
+        Sets the callee of this call-like operation. A `callee` is either a
+        reference to a symbol, via SymbolRefAttr, or a reference to a defined
+        SSA value. The type of the `callee` is expected to be the same as the
+        return type of `getCallableForCallee`, e.g., `callee` should be
+        SymbolRefAttr for `func.call`.
+      }],
+      "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"
     >,
   ];
 
   let extraClassDeclaration = [{

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

 CallInterfaceCallable CallOp::getCallableForCallee() {
   // Direct call.
   if (FlatSymbolRefAttr calleeAttr = getCalleeAttr())
     return calleeAttr;
   // Indirect call, callee Value is the first operand.
   return getOperand(0);
 }
 
+void CallOp::setCalleeFromCallable(CallInterfaceCallable callee) {
+  // Direct call.
+  if (FlatSymbolRefAttr calleeAttr = getCalleeAttr()) {
+    auto symRef = callee.get<SymbolRefAttr>();
+    return setCalleeAttr(cast<FlatSymbolRefAttr>(symRef));
+  }
+  // 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);
 }
 
 LogicalResult CallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
   if (getNumResults() > 1)
     return emitOpError("must have 0 or 1 result");
 

 CallInterfaceCallable InvokeOp::getCallableForCallee() {
   // Direct call.
   if (FlatSymbolRefAttr calleeAttr = getCalleeAttr())
     return calleeAttr;
   // Indirect call, callee Value is the first operand.
   return getOperand(0);
 }
 
+void InvokeOp::setCalleeFromCallable(CallInterfaceCallable callee) {
+  // Direct call.
+  if (FlatSymbolRefAttr calleeAttr = getCalleeAttr()) {
+    auto symRef = callee.get<SymbolRefAttr>();
+    return setCalleeAttr(cast<FlatSymbolRefAttr>(symRef));
+  }
+  // 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);
 }
 
 LogicalResult InvokeOp::verify() {
   if (getNumResults() > 1)
     return emitOpError("must have 0 or 1 result");
 

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

LogicalResult spirv::FunctionCallOp::verify() {
 
   return success();
 }
 
 CallInterfaceCallable spirv::FunctionCallOp::getCallableForCallee() {
   return (*this)->getAttrOfType<SymbolRefAttr>(kCallee);
 }
 
+void spirv::FunctionCallOp::setCalleeFromCallable(
+    CallInterfaceCallable callee) {
+  (*this)->setAttr(kCallee, callee.get<SymbolRefAttr>());
+}
+
 Operation::operand_range spirv::FunctionCallOp::getArgOperands() {
   return getArguments();
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.GLFClampOp
 //===----------------------------------------------------------------------===//
 

mlir/test/lib/Dialect/Test/TestOps.td

 def ConversionCallOp : TEST_Op<"conversion_call_op",
     [CallOpInterface]> {
   let arguments = (ins Variadic<AnyType>:$arg_operands, SymbolRefAttr:$callee);
   let results = (outs Variadic<AnyType>);
 
   let extraClassDeclaration = [{
     /// Return the callee of this operation.
     ::mlir::CallInterfaceCallable getCallableForCallee();
+
+    /// Set the callee for this operation.
+    void setCalleeFromCallable(::mlir::CallInterfaceCallable);
   }];
   let extraClassDefinition = [{
     ::mlir::CallInterfaceCallable $cppClass::getCallableForCallee() {
       return (*this)->getAttrOfType<::mlir::SymbolRefAttr>("callee");
     }
+
+    void $cppClass::setCalleeFromCallable(::mlir::CallInterfaceCallable callee) {
+      (*this)->setAttr("callee", callee.get<SymbolRefAttr>());
+    }
   }];
 }
 
 def ConversionFuncOp : TEST_Op<"conversion_func_op", [CallableOpInterface,
                                                       FunctionOpInterface]> {
   let arguments = (ins SymbolNameAttr:$sym_name,
                        TypeAttrOf<FunctionType>:$function_type,
                        OptionalAttr<DictArrayAttr>:$arg_attrs,