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`.
Reviewed By: gysit, zero9178o
Differential Revision: https://reviews.llvm.org/D149763
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>());
+ }
}];
}
* `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()`
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(); }
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(); }
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(); }
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(); }
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(); }
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 = [{
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 = [{
/// 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;
return getTarget();
}
+ void setCalleeFromCallable(::mlir::CallInterfaceCallable callee) {
+ setTargetAttr(callee.get<SymbolRefAttr>());
+ }
+
::mlir::Operation::operand_range getArgOperands() {
return getOperands();
}
"::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.
}],
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);
}
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);
}
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();
}
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>());
+ }
}];
}