MLIR has several instances of Ops with single block regions where the region computes values that are expected to be of the same types as the parent Op result types. Examples includes scf.if, scf.for, generic_atomic_rmw as well as tfl.while (within Tensorflow). For such Ops, its common to have the terminator inputs of these single block regions be considered as the region results and verification needs to check that the number and types of these region results match the parent Op result types. This RFC proposes extending the existing SingleBlockImplicitTerminator trait to optionally provide this verification. This trait could be parameterized in 2 ways:
- Customize the predicate used to check if the 2 types match
- Customize the subset of regions attached to an Op that should satisfy the trait.
Note that in general this trait cannot be unconditionally attached to the terminator op, because the same terminator can be used with multiple parent Ops and not all of them satisfy this constraint. As an example, the scf.YieldOp can be used as a terminator for regions attached to scf.for, scf.if, and scf.parallel but only within the scf.if and scf.for context the yield inputs are expected to type match the parent result types. Similar use cases exist within TensorFlow as well.
The proposed C++ Interface would look like:
template <typename TerminatorOpType,
typename Predicate = std::equal_to<Type>,
unsigned... ResultMatchingRegions>
struct SingleBlockImplicitTerminator {
...
};
A custom type matcher can be provided by using a new functor type for Predicate (which the code will instantiate). ODS framework will also be extended accordingly to have 2 additional input for SingleBlockImplicitTerminator:
class SingleBlockImplicitTerminator<string op,
list<int> resultMatchingRegions = [],
string predicate="std::equal_to<mlir::Type>">
For the ODS interface, the list of regions and the predicate to use are swapped so that for the common use case of using Type equality, we don’t need to spell out the predicate explicitly. With this change, existing Op’s like generic_atomic_rmw can be extended as follows:
def GenericAtomicRMWOp : Std_Op<"generic_atomic_rmw", [
SingleBlockImplicitTerminator<"AtomicYieldOp", [0]>,
And the verification that happens on AtomicYieldOp:
static LogicalResult verify(AtomicYieldOp op) {
Type parentType = op.getParentOp()->getResultTypes().front();
Type resultType = op.result().getType();
if (parentType != resultType)
return op.emitOpError() << "types mismatch between yield op: " << resultType
<< " and its parent: " << parentType;
return success();
}
Can be removed.
An alternative way to support this is to introduce a new Trait as opposed to extending SingleBlockImplicitTerminator. However, it turns out that that trait would have to do most of the checks that SingleBlockImplicitTerminator would have to do to get to the terminator, and would have to define similar constraints (like regions having a single blocks so that the region results are well defined).
Another extension is to have a standalone OperandsMatchParentResultType trait that can be attached to terminator ops in cases where all uses of that terminator Op satisfy the type-matching constraint.