Page MenuHomePhabricator

[mlir] Add `UnregisteredOp` to help traversing unregistered ops
Needs ReviewPublic

Authored by Chia-hungDuan on Nov 23 2021, 3:03 PM.

Details

Reviewers
rriddle
Summary

Add a helper class UnregisteredOp for the unregistered ops to get the
similar ability as registered ops while using some MLIR utilities. For
example, you can do things like,

  • op->walk([](UnregisteredOp<opName> op) { ... };

Diff Detail

Event Timeline

Chia-hungDuan created this revision.Nov 23 2021, 3:03 PM
Chia-hungDuan requested review of this revision.Nov 23 2021, 3:03 PM

River@,

We had discussion about this and we were thinking things like UnregisteredOp<"test.op">. It seems that we don't have a way to pass string literal for non-type template argument (literal class type in C++20 may be help?). If we want something like template<char *> then it needs to have somewhere to put the string literal in global scope. We can discuss this later.

Thus I think we can extend the walk for specifying the operation name, so that we don't need to do the string comparison all the time(by leveraging the OperationName)

River@,

We had discussion about this and we were thinking things like UnregisteredOp<"test.op">. It seems that we don't have a way to pass string literal for non-type template argument (literal class type in C++20 may be help?). If we want something like template<char *> then it needs to have somewhere to put the string literal in global scope. We can discuss this later.

Thus I think we can extend the walk for specifying the operation name, so that we don't need to do the string comparison all the time(by leveraging the OperationName)

Could we do something like

op->clone();

cloned->walk<WalkOrder::PostOrder>(WithName("dummy"), ...)

where WithName (or somesuch) is the constructor of a struct with (say) a WalkResult operator()(Operation *)? and so the string stored for the lifetime of the struct, internal to the struct one can then cache OperationName without walk needing to know anything more, we get documentation at callsite that filtering is by name, and it could be extended to match attributes and the like too (not sure if that's needed, just thinking more that string name matching here feels like it serves as an and on top of a general filtering mechanism)

River@,

We had discussion about this and we were thinking things like UnregisteredOp<"test.op">. It seems that we don't have a way to pass string literal for non-type template argument (literal class type in C++20 may be help?). If we want something like template<char *> then it needs to have somewhere to put the string literal in global scope. We can discuss this later.

Thus I think we can extend the walk for specifying the operation name, so that we don't need to do the string comparison all the time(by leveraging the OperationName)

Could we do something like

op->clone();

cloned->walk<WalkOrder::PostOrder>(WithName("dummy"), ...)

where WithName (or somesuch) is the constructor of a struct with (say) a WalkResult operator()(Operation *)? and so the string stored for the lifetime of the struct, internal to the struct one can then cache OperationName without walk needing to know anything more, we get documentation at callsite that filtering is by name, and it could be extended to match attributes and the like too (not sure if that's needed, just thinking more that string name matching here feels like it serves as an and on top of a general filtering mechanism)

Yes, I was thinking to add a callback like shouldVisit to determine which operation we want to visit. I was not sure if we want to support it generically. Thinking the case like, the user may not know they can use OperationName so they may do the string comparison as well (And yes, here's the question, who should have the knowledge of OperationName). Then they may not need shouldVisit, they can do it in original callback. The attribute may have the same case.

We can do both I think, i.e., have the generic shouldVisit callback and specialized way for certain filters like operation name. What do you think?

River@,

We had discussion about this and we were thinking things like UnregisteredOp<"test.op">. It seems that we don't have a way to pass string literal for non-type template argument (literal class type in C++20 may be help?). If we want something like template<char *> then it needs to have somewhere to put the string literal in global scope. We can discuss this later.

Thus I think we can extend the walk for specifying the operation name, so that we don't need to do the string comparison all the time(by leveraging the OperationName)

Could we do something like

op->clone();

cloned->walk<WalkOrder::PostOrder>(WithName("dummy"), ...)

where WithName (or somesuch) is the constructor of a struct with (say) a WalkResult operator()(Operation *)? and so the string stored for the lifetime of the struct, internal to the struct one can then cache OperationName without walk needing to know anything more, we get documentation at callsite that filtering is by name, and it could be extended to match attributes and the like too (not sure if that's needed, just thinking more that string name matching here feels like it serves as an and on top of a general filtering mechanism)

That's a good point about caching the OperationName. In the context of unregistered operations, I was also thinking about adding a class that was like template <const char *name> class UnregisteredOp : public Op {,, }; that could be used in-place of the raw operation name in certain cases. We could even consider adding some build methods to it, so that it could be used over OperationState/checking the raw operation name when possible.

Turn back to the UnregisteredOp idea.

Chia-hungDuan retitled this revision from [mlir] Support visiting operation with given name to [mlir] Add `UnregisteredOp` to help traversing unregistered ops.Dec 15 2021, 5:52 PM
Chia-hungDuan edited the summary of this revision. (Show Details)
mehdi_amini added inline comments.Dec 15 2021, 8:50 PM
mlir/include/mlir/IR/OpDefinition.h
1854 ↗(On Diff #394698)

Why template this class instead of taking the string in the constructor?

OperationName name;

UnregisteredOp(StringRef name, MLIRContext *ctx) : name(name, ctx) {}
UnregisteredOp(StrAttr name) : name(name, name.getContext()) {}
UnregisteredOp(OperationName name) : name(name) {}
UnregisteredOp(Operation *op) : name(op->getName())
1861 ↗(On Diff #394698)

That TODO is a hard blocker to me: we can't land this as-is.

1863 ↗(On Diff #394698)

There shouldn't be any different between this and what op->getName() will return.

Chia-hungDuan added inline comments.Dec 15 2021, 9:32 PM
mlir/include/mlir/IR/OpDefinition.h
1854 ↗(On Diff #394698)

This is supposed to be a specific type of an unregistered op and we need to use the 'type' in the op->walk() for filtering the op.

1861 ↗(On Diff #394698)

Agree...
The OperationName is bound to MLIRContext but it doesn't have a way to get the MLIRContext associated to itself. If we had it, maybe we can update it if we see the context change (Or fallback to string comparison?). TBH, I'm not sure if this will be acceptable.

1863 ↗(On Diff #394698)

The opName is constructed with template argument name and it'll be initialized once. Here we want to check if the op has the same operation name as the unregistered one. To avoid the string comparsion, we use OperationName comparison instead.

mehdi_amini added inline comments.Dec 15 2021, 9:54 PM
mlir/include/mlir/IR/OpDefinition.h
1854 ↗(On Diff #394698)

Isn't this just a matter of override the comparison by specializing a template in the walk implementation to handle UnregisteredOp?

1863 ↗(On Diff #394698)

I was asking what is the difference between static OperationName opName(name, op->getContext()); and static OperationName opName = op->getName(); ?
Either you misunderstood the question, or I'm not understand your answer :)

Chia-hungDuan added inline comments.Dec 15 2021, 10:55 PM
mlir/include/mlir/IR/OpDefinition.h
1854 ↗(On Diff #394698)

For registered op, like FuncOp, we can do walk([](FuncOp){} and walk() has specialization for this. It's specialized by recognizing the FuncOp type.

If we use UnregisteredOp without taking template paramter const char *, then we can't distinguish two different UnregisteredOps by only checking the type. Which means if we do walk([](UnregisteredOp) {}), we don't have a way to tell what UnregisteredOp we want to visit. We need a type as identification rather than a value.

Add few context about this change. Currently, the usual way to visit unregistered op may be,

op->walk([&](Operation *op) {
    if (op->getName().getStringRef().equals("FooOp"))
        ...
});

But it may not be ideal to do string comparison, we should use OperationName instead.

Sorry, I'm not sure if this is clear enough or is there any tricks I missed?

1863 ↗(On Diff #394698)

Sorry, I didn't explain it clearly.

in opName(name, op->getContext()), the first argument name is the template argument, the argument opof classof is passed while doing the casting, which means op->getName().getStringRef().equals(name) may not be true while init the opName.

For example, if the first call to this classof is from isa<UnregisteredOp<FooOpName>(FuncOp.getOperation()), then the op->getName().getStringRef() will be "func".

Or am I misunderstanding something?

mehdi_amini added inline comments.Dec 15 2021, 10:58 PM
mlir/include/mlir/IR/OpDefinition.h
1854 ↗(On Diff #394698)

What about: op->walk("test.foo", [&](Operation *op) { ?

Chia-hungDuan added inline comments.Dec 15 2021, 11:26 PM
mlir/include/mlir/IR/OpDefinition.h
1854 ↗(On Diff #394698)

The approach was used in the first version of this patch. Jacques has some suggestions that maybe we can do something like op->walk(WithName("test.foo"), [&](...)), WithName is the predicate to see if we want to visit the operation. But I think the predicate thing can be done in callback so it may not be so useful to separate them out. (There's a brief discussion at the beginning)

After considering it for a while, I reverted that approach and turn to UnregisteredOp one. I expect this can be used other cases like OpConversionPattern.

mehdi_amini added inline comments.Dec 15 2021, 11:57 PM
mlir/include/mlir/IR/OpDefinition.h
1854 ↗(On Diff #394698)

OK but there are some things that just aren't great with the current approach:

  • You can't cache the OperationName
  • The need for a static const char dummyOpName[] global isn't really nice.