LLVM Discussion Forums

RFC: Extending SingleBlockImplicitTerminator to support additional type checking

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.

Such regions should also probably be implementing RegionBranchOpInterface as well, which provides needed primitives for doing this check. Can you use that to implement this verification?

Let me take a look. AFAICT not all such Ops are implementing the RegionBranchOpInterface so that might be the first step. I’ll circle back once I have additional info.

I wouldn’t tie it to the SingleBlockImplicitTerminator trait. This trait is intended for ops that can omit the trailing terminator in the custom form, which only happens for terminators that don’t have operands. At which point, the check is unnecessary. Some of the ops with this trait grew to support terminators with operands, and I would be actually open to reconsider the implicitness of their terminators.

Instead, I would consider either:

  • providing a type constrait TCOpResMatchesParentRes; we have a set of TC... predicates for type constraints on ops, e.g., the types of op operands and results match, and the new proposed check makes sense there; or
  • extending the HasParent trait that many terminators have with the type constraint.

Thanks for the feedback.

SingleBlockImplicitTerminator does seems to used in cases where the terminator has operands. scf.for is an example, so looks like what the intended trait here is SingleBlockWithTerminator and nothing implicitly added. So we can either split or parameterize the trait to clearly specify which of the 2 behaviors is intended here. If you think that makes sense, I can file a bug for this issue.

Getting back to type check between region outputs and parents, for one of the reasons outlined above (that the same terminator can be used with multiple ops and for only a subset of those this type constraint is satisfied), extending HasParent or ParentOneOf may not support that use case (or at least not in a straightforward manner). So a new type constraint seems like the way to go. I still need to look into how this ties in with the RegionBranchOpInterface, I’ll look into that this week.

I think you can still write

scf.for %i = %lb to %ub step %step {
  %0 = index_cast %i : index to i64
  store %0, %memref[%i] : memref<?xi64>
  // no terminator here, but a no-operand scf.yield will be inserted for you
}

and have it accepted by the parser. That’s the implicit terminator part. The trait provides functionality for that to happen. The “single block” is there for a historical reasons: this trait predates the possibility to specify constraints on the regions with let regions = .... We can indeed separate out the “implicit terminator” part, but it would still be only applicable to single-block regions.

This just means that HasParent should not take an op name, but a constraint on the parent. This way, one can check for a set of ops, a trait, a dynamic condition based on specific ops, etc. We did the same for types, which used to be a specific type, but became a constraint instead.

Got it (about the implicit terminator part). About extending HasParent vs adding a new trait, I agree that HasParent can be modified to encode a generic constraint on the parent (which can include the Parents Op as well as any type constraints), but for the use case I am thinking of where the same terminator is used with multiple parent ops (like say scf.if and scf.for) different parent Ops may impose different constraints. And encoding all that (a set of parent ops and an optional set of per-parent op type constraints) all on the terminator seems too much. I am proposing that the parent op instead encodes these constraint. I looked at what the RegionBranchOpInterface provides, I suspect we can build this trait using that interface. The trait can be something like “OpResultsMatchRegionResults” with a list of regions and a predicate to use for type matching. The trait will be expected to be attached to the parent op and will use the RegionBranchOpInterface to extract which region results feed into the Op results and do the type matching.

Please let me know if that sounds reasonable, and I’ll work on a more concrete proposal. Also, looks like we want to revisit the SingleBlockImplicitTerminator trait and a new trait where an implicit terminator may not always make sense. I’ll file a bug for that.

Based on the feedback above, here’s take 2 on the proposal, with a standalone trait that uses the RegionBranchOpInterface:

OpResultsMatchRegionResults

MLIR has several instances of ops with attached regions where the regions compute 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). Currently, verifying this type-matching is done using custom verification code for either the terminator ops within these regions or the parent op. This RFC proposes a new trait to support this type matching. This trait could be customized in 2 ways:

  • Customize the predicate used to check if 2 types match (default will be equality).
  • 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 within these regions (common example of such ops also have the SingleBlockImplictTerminator trait attached to them) because the same terminator op 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.

MLIR abstracts how values flow from regions attached to ops to the results of the parent op using the RegionBranchOpInterface. Specifically, the getSuccessorRegion method when called with the region index returns a vector of RegionSuccessor objects, and if one of these RegionSuccessor objects represents the successor being a Parent op, then the corresponding getSuccessorInputs should type-match with the parent result types. The trait will fail if its attached to an op that does not implement the RegionBranchOpInterface (TODO: Is it possible to enforce this at compile time?)

Proposed C++ Interface

template<typename Predicate = std::equal_to<Type>, unsigned... MatchingRegions> 
class OpResultsMatchRegionResults {

  // Pseudo implementation 
  // auto regionOpInterface = dyn_cast<RegionOpInterface>(op)
  // if (!regionOpInterface)
  //   return op.emitError() << "...";
  // for (unsigned ri : MatchingRegions) {
  //    auto successors = regionOpInterface->getSuccessorInputs(ri, nullptr);
  //    for (RegionSuccessor successor : successors) {
  //        if (!successor.getSuccessor()) continue; // not a parent successor
  //        auto inputs = successor.getSuccessorInputs();
  //        if (inputs.size() != op.getNumResults())
  //           return op.emitError() << "..." ;
  //        for (auto it: llvm::zip(inputs, op->getResults()) {
  //           // check type using predicate, flag error if they do not match
  //        }
  //    }
  // }

};

Predicate will be a functor object that will be instantiated within the trait to check if types match.

Proposed ODS Interface

The ODS interface will mostly mimic the C++ interface but with the list of regions and the predicate swapped to handle the common use case of using equality check for type matching.

class OpResultsMatchRegionResults<list<int> regions, string predicate="std::equal_to<mlir::Type>"> :
		ParamNativeOpTrait<"OpResultsMatchRegionResults", preicate # ", " StrJoinInt(regions).result>

Adoption

Once the trait is defined, it will be adopted in several places like scf.if, scf.for and the custom verification code attached to the terminators for these ops can be removed (in conjunction with ParentOneOf trait and may be hoisting some remaining Op specific verification to the Op’s verification code). As an example, verification of AtomicYieldOp in the standard dialect can be completely eliminated.

1 Like

@_sean_silva, @ftynse and others, if there are no more comments, I will start on the implementation.

SGTM
(Post must be at least 20 characters so here is some filler text, this is annoying for RFCs…)

The new proposal looks good to me.

You’ll probably want to do something similar to:

SGTM, thanks for digging into this!

FYI, I have filed https://bugs.llvm.org/show_bug.cgi?id=46442 to investigate the issue that @ftynse brought up with the SingleBlockImplicitTerminator trait.