samskalicky commented on a change in pull request #17270: [WIP] Dynamic custom
operator GPU support
URL: https://github.com/apache/incubator-mxnet/pull/17270#discussion_r369153621
##########
File path: src/c_api/c_api.cc
##########
@@ -164,39 +186,60 @@ int MXLoadLib(const char *path) {
for (int i = 0; i < numOps; i++) {
const char* name;
// function pointers holding implementation from custom library
- fcomp_t fcomp_fp = nullptr;
parseAttrs_t parse_fp = nullptr;
inferType_t type_fp = nullptr;
inferShape_t shape_fp = nullptr;
// optional attributes
- fcomp_t fgrad_fp = nullptr;
mutateInputs_t mutate_fp = nullptr;
- createOpState_t create_opstate_fp = nullptr;
bool isSubgraphOp = false;
int _isSubgraphOp = 0;
-
- // get custom operator implemenation from the dynamic library
- opRegGet(i, &name, &fcomp_fp, &fgrad_fp, &parse_fp, &type_fp, &shape_fp,
- &mutate_fp, &create_opstate_fp, &_isSubgraphOp);
+ // lists of forward and backward function associated with each context
+ const char **forward_ctx, **backward_ctx, **createop_ctx;
+ fcomp_t *forward_fcomp, *backward_fcomp;
+ createOpState_t *createop_fp;
+ int forward_count, backward_count, createop_count;
+
+ // main function to get custom operator implemenation from the custom
library
+ opRegGet(i, &name, &_isSubgraphOp,
+ &forward_ctx, &forward_fcomp, &forward_count,
+ &backward_ctx, &backward_fcomp, &backward_count,
+ &createop_ctx, &createop_fp, &createop_count,
+ &parse_fp, &type_fp, &shape_fp, &mutate_fp);
+
+ // construct maps of context to forward/backward custom library function
+ std::unordered_map<std::string, fcomp_t> forward_ctx_map;
+ std::unordered_map<std::string, fcomp_t> backward_ctx_map;
+ std::unordered_map<std::string, createOpState_t> createop_map;
+ for (int i=0; i<forward_count; i++) {
+ std::string ctx_str(forward_ctx[i]);
+ forward_ctx_map[ctx_str] = forward_fcomp[i];
Review comment:
should we check that `ctx_str` was already in `forward_ctx_map` and error
out if we find it (since that means the user registered the same context twice)?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
