Enhancements for custom subgraph op

[samskalicky]

Description

Enhancements for custom subgraph operators in external libraries.

• Adds API to set boolean flag to denote a subgraph operator in an external library
• Removes the requirement to manually implement infer shape/type/storage/mutate/request
• Uses default functions from subgraph/common.h

Design

Its difficult to implement the full shape/type propagation in an external library for a whole subgraph. Rather that re-implement these complex pieces, we'll use the default functions in the subgraph API.

Starting from the user in the custom library, the registration is now only:

REGISTER_OP(_custom_subgraph_op)
.setParseAttrs(parseAttrs)
.isSubgraphOp(true)
.setCreateOpState(createOpState);

Here, setting isSubgraphOp means we dont have to register the infer shape/type functions. When re-registering the operator in MXNet, if isSubgraphOp is true, then we re-use the functions from the subgraph API:

if(isSubgraphOp) {
      using namespace mxnet::op;
      regOp.set_attr<nnvm::FInferType>("FInferType",
                                       DefaultSubgraphOpType, plevel);
      regOp.set_attr<mxnet::FInferShape>("FInferShape",
                                         DefaultSubgraphOpShape, plevel);
      regOp.set_attr<FInferStorageType>("FInferStorageType",
                                        DefaultSubgraphOpStorageType, plevel);
      regOp.set_attr<FResourceRequest>("FResourceRequest",
                                       DefaultSubgraphOpResourceRequest, plevel);
      regOp.set_attr<nnvm::FMutateInputs>("FMutateInputs",
                                          DefaultSubgraphOpMutableInputs, plevel);
}

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage:
• Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
• Nightly tests are added for complicated/long-running ones (e.g. changing distributed kvstore)
• Build tests will be added for build configuration changes (e.g. adding a new build option with NCCL)
• Code is well-documented:
• For user-facing API changes, API doc string has been updated.
• For new C++ functions in header files, their functionalities and arguments are documented.
• For new examples, README.md is added to explain the what the example does, the source of the dataset, expected performance on test set and reference to the original paper if applicable
• Check the API doc at https://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
• To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

[rondogency]

LGTM!

[samskalicky]

@mxnet-label-bot add [pr-awaiting-review]

[samskalicky]

@TaoLv @ZhennanQin would you please review since you're familiar with the subgraph API?

[TaoLv]

Sounds this function should return true or false. How about using setXXX?

[samskalicky]

done! changed it to setIsSubgraphOp() and removed argument

[mseth10]

should we check if (mutate_fp != nullptr) here?

[samskalicky]

as discussed offline, we'll only use default functions when the user sets the setIsSubgraphOp flag. Later we can revisit if the user wants to use custom functions mixed with default functions

[mseth10]

looks good to me.

[samskalicky]

Hey @wkcn this PR is finished and ready for review, would you please take a look into this enhancement for custom subgraph ops? Happy New Year!

[wkcn]

@samskalicky Happy New Year!

This PR looks good to me : )
Is it available to copy a custom subgraph op test into tests/python/unitest/test_extensions.py?

[samskalicky]

@wkcn

Is it available to copy a custom subgraph op test into tests/python/unitest/test_extensions.py?

The current subgraph op has an example tests here, but is SUPER hacky:
https://github.com/apache/incubator-mxnet/blob/e65fc4bee5b8602b38ba2e419a4ea2e5ce5b7e9a/example/extensions/lib_custom_op/test_subgraph.py#L47-L56
In this example we're literally doing a string replace to change the op in the symbol json string. This isnt something we want to test in the CI. Typically subgraph ops are inserted when the graph is partitioned using subgraph properties.

It will be officially tested in test_extensions.py in the other PR that adds support for custom subgraph properties and actually uses custom subgraph ops in #17034 here are the code changes:
https://github.com/apache/incubator-mxnet/pull/17034/files#diff-460d06cb3af2af996c5ec259b086712fR87-R150

This PR (#17194) modifies the custom op that is used in that test with the changes in this PR:
https://github.com/apache/incubator-mxnet/pull/17194/files#diff-4347bc72a85f439003da183493ba2089R87

Would be great to get your feedback on #17034 too.

If this is acceptable would you please approve/merge this PR?

[wkcn]

LGTM. Thank you!