diff --git a/example/extensions/lib_subgraph/README.md b/example/extensions/lib_subgraph/README.md index b113be267fd3..83c823676f18 100644 --- a/example/extensions/lib_subgraph/README.md +++ b/example/extensions/lib_subgraph/README.md @@ -53,9 +53,11 @@ You can start getting familiar with custom partitioners by running an example pr * **lib_subgraph/test_subgraph.py**: This file calls `mx.library.load(‘libsubgraph_lib.so’)` to load the library containing the custom components, partitions the model using the `optimize_for` API, and prints outputs of the forward passes. The outputs should be the same as the regular MXNet forward pass without partitioning. +* **include/mxnet/lib_api.h**: This file from MXNet source code is the single header file needed to include all necessary data types and function prototypes for writing a custom operator library. You can either specify the include path in the `Makefile`, or copy the header file over to `example/extensions/lib_subgraph` folder. Note that apart from this header, the custom operator library is independent of MXNet source. + ## Writing Custom Partitioner Library -For building a library containing your own custom partitioner, compose a C++ source file like `mypart_lib.cc`, include `lib_api.h` header file, and write your custom partitioner with these essential functions: +To build your own library containing a custom partitioner, compose a C++ source file like `mypart_lib.cc`, include `lib_api.h` header file, and write your custom partitioner with these essential functions: - `initialize` - Library Initialization Function - `REGISTER_PARTITIONER ` - Partitioner Registration Macro - `mySupportedOps ` - Operator Support @@ -76,34 +78,60 @@ sym, _, _ = mx.model.load_checkpoint('mymodel', 0) # Symbol/Module flow sym2 = sym.optimize_for("myPart") -# Gluon flow +# Gluon flow 1 sym_block = nn.SymbolBlock(sym, inputs) sym_block.hybridize(backend='myPart') + +# Gluon flow 2 +sym_block = nn.SymbolBlock(sym, inputs) +sym_block.optimize_for(x, backend='myPart') ``` +In the Gluon hybridize flow, the model is actually hybridized during the first inference, rather than immediately when calling `hybridize`. This hybridize-based flow is useful if a user expects to run inference immediately after hybridizing. But for users than just want to partition but not run a whole forward pass, the `optimize_for` API combines the hybrdize/forward APIs but does not run a forward pass. After calling `optimize_for` users can `export` their model immediately without running a forward pass. + ### Using a Custom Partitioner Library Partitioning APIs in MXNet are available in both Symbol and Gluon APIs. For the Symbol API, the `optimize_for` API can be called on Symbol objects to return a partitioned Symbol. ``` -optimize_for(backend, args=None, ctx=None, **kwargs) +optimize_for(backend, args=None, aux=None, ctx=None, **kwargs) ``` -The `optimize_for` API takes at least 1 argument, `backend` which is a string that identifies which backend to partition the model for. The `args` argument is optional and takes a list of NDArray or dict of str to NDArray. It is used to infer shapes and types and before partitioning. The `ctx` argument is optional and takes a device context to infer storage types. It also take any other user-specified options that will be passed to the backend partitioning APIs. +The `optimize_for` API takes at least 1 argument, `backend` which is a string that identifies which backend to partition the model for. The `args` and `aux` arguments are optional and take a list of NDArray or dict of str to NDArray. They are used to infer shapes and types and before partitioning, and passed to the backend to use during compilation. The `ctx` argument is optional and takes a device context to infer storage types. It also takes any other user-specified options that will be passed to the backend partitioning APIs. For the Gluon API, the `hybridize` API can be called on HybridBlocks to partition the internal CachedOp Symbol. ``` -hybridize(backend=None, backend_opts=None) +hybridize(backend=None, backend_opts=None, **kwargs) +``` + +The `hybridize` function prepares the HybridBlock to be converted into a backend symbol. The `backend` argument is a string that identifies which backend that will partition the model. The `backend_opts` takes other user-specified options that will be passed to the backend partitioning APIs. The actual partitioning takes place during the forward pass. + +If you just want to partition the HybridBlock but not run a complete forward pass, you can use the `optimize_for` API that combines the work done in the `hybridize` API with part of the work done in the forward pass. + +``` +optimize_for(x, backend=None, backend_opts=None, **kwargs) +``` + +When the `optimize_for` API is called on a HybridBlock it partitions immediately. This lets users export the partitioned model without running a complete forward pass. + +``` +block.optimize_for(x, backend='myPart') +block.export('partitioned') ``` -When the `hybridize` function is called, Gluon will convert the program’s execution into the style used in symbolic programming. The `backend` argument is a string that identifies which backend to partition the model for. The `backend_opts` takes other user-specified options that will be passed to the backend partitioning APIs. +But you can also use `optimize_for` in place of `hybridize` and run inference immediately after too. + +``` +block.optimize_for(x, backend='myPart') +block(x) +``` ### Writing A Custom Partitioner There are several essential building blocks for making a custom partitioner: -* [initialize](./subgraph_lib.cc#L242): +* [initialize](./subgraph_lib.cc#L261): * This function is the library initialization function necessary for any dynamic libraries. It lets you check if the user is using a compatible version of MXNet. Note that this `version` parameter is passed from MXNet when library is loaded. MXReturnValue initialize(int version) @@ -116,40 +144,37 @@ There are several essential building blocks for making a custom partitioner: std::vector& ids, std::unordered_map& options) -* [REGISTER_PARTITIONER(my_part_name)](./subgraph_lib.cc#L238): +* [REGISTER_PARTITIONER(my_part_name)](./subgraph_lib.cc#L257): * This macro registers the custom partitioner and its properties to MXNet by its name. Notice that a partitioner can have multiple partitioning strategies. This enables multiple *passes* to be run in a single partitioning call from the user. The first argument to `addStrategy` is a user-specified name. The second argument is the `supportedOps` function. The third argument is the name of the subgraph operator to create for each subgraph created during partitioning (see below for more info about subgraph operators). The `setReviewSubgraph` API registers a callback function that is called for each subgraph created during partitioning (more on this below). Notice that the first argument to this function is the strategy to associate with and the second argument is the `reviewSubgraph` function. REGISTER_PARTITIONER(my_part_name) - .addStrategy("strategy1", - supportedOps, - "_custom_subgraph_op") - .setReviewSubgraph("strategy1", - reviewSubgraph); + .addStrategy("strategy1", supportedOps, "_custom_subgraph_op") + .setReviewSubgraph("strategy1", reviewSubgraph); Also there are some optional functions you can specify: -* [reviewSubgraph](./subgraph_lib.cc#L220): +* [reviewSubgraph](./subgraph_lib.cc#L219): * This function provides an opportunity to accept/reject a subgraph after MXNet partitions it. It also allows specifying custom attributes on the subgraph (ie. user-generated IDs). If you do not register this function, subgraphs will be accepted by default. MXReturnValue reviewSubgraph( std::string json, - int subraph_id, + int subgraph_id, bool* accept, - std::unordered_map& options, - std::unordered_map& attrs) + std::unordered_map& options, + std::unordered_map& attrs, + std::map& args, + std::map& aux) Let’s take a closer look at those registry functions: -* **supportedOps**: This function takes four arguments. The 1st argument is a JSON string of the model architecture graph, where nodes are inputs/params/weights and edges are data dependencies. The graph is pre-sorted in topological order. The 2nd argument is an array of booleans, one for each operator in the model. When traversing the graph, operators to be partitioned into subgraphs are identified and an entry is set to `true` for the node ID in the `ids` array. The last argument is the map of options specified by the user. Users can pass custom options to the partitioner and they are passed to this function in the `options` map. +* **supportedOps**: This function takes four arguments. The 1st argument is a JSON string of the model architecture graph, where nodes are inputs/params/weights and edges are data dependencies. The graph is pre-sorted in topological order. The 2nd argument is an array of booleans, one for each operator in the model. When traversing the graph, operators to be partitioned into subgraphs are identified and an entry is set to `true` for the index in the `ids` array corresponding to the node ID. The last argument is the map of options specified by the user. Users can pass custom options to the partitioner and they are passed to this function in the `options` map. -* **reviewSubgraph**: This function takes five arguments. The 1st argument is a JSON string of the newly partitioned subgraph. The 2nd argument is the subgraph ID, this is just a number MXNet uses to identify this particular subgraph (it starts at zero and increments). The 3rd argument is an output to be set in this function to tell MXNet whether to accept (value: `true`) or reject (value: `false`) the subgraph. The 4th argument is the map of options specified by the user. The last argument is a map of attributes that should be set on the created subgraph. These attributes will be available later at runtime, and provides a mechanisn to pass info from partition-time to runtime. You might want to reject a subgraph if it doesnt include all the operators you want, for example. The `options` map is the same one passed to the `supportedOps` API. +* **reviewSubgraph**: This function takes five arguments. The 1st argument is a JSON string of the newly partitioned subgraph. The 2nd argument is the subgraph ID, this is just a number MXNet uses to identify this particular subgraph (it starts at zero and increments, unique for each subgraph in the model). The 3rd argument is an output to be set in this function to tell MXNet whether to accept (value: `true`) or reject (value: `false`) the subgraph. You might want to reject a subgraph if it doesnt include all the operators you want, for example. The `options` map is the same one passed to the `supportedOps` API. The 4th argument is the map of options specified by the user. The 5th argument is a map of attributes that should be set on the created subgraph. These attributes will be available later at runtime, and provides a mechanisn to pass info from partition-time to runtime. The last argument is the map of params/weights/args to the model and the associated names. For inputs the the subgraph that come directly from the params/weights of the model, you can look up the name of the input in this map to get the actual tensor values. ### Writing A Custom Subgraph Operator -A partitioning strategy specifies how to partition a model and isolate operators into subgraphs. In MXNet, subgraphs are just a [stateful operator](../lib_custom_op#writing-stateful-custom-operator). Subgraph operators have an extra attribute called `SUBGRAPH_SYM_JSON` that maps to a JSON string of the subgraph. The expectation is that when a subgraph operator executes a forward/backward call, it executes all of the operators in the subgraph. +A partitioning strategy specifies how to partition a model and isolate operators into subgraphs. In MXNet, subgraphs are just a [stateful operator](../lib_custom_op#writing-stateful-custom-operator). Subgraph operators have an extra attribute called `MX_STR_SUBGRAPH_SYM_JSON` that maps to a JSON string of the subgraph. The expectation is that when a subgraph operator executes a forward/backward call, it executes all of the operators in the subgraph. When registering a custom subgraph operator, all thats needed is to register a `createOpState` function and to set that the operator is a subgraph operator by calling the `setIsSubgraphOp` API like: diff --git a/example/extensions/lib_subgraph/subgraph_lib.cc b/example/extensions/lib_subgraph/subgraph_lib.cc index da888fd10383..8c24dd880f72 100644 --- a/example/extensions/lib_subgraph/subgraph_lib.cc +++ b/example/extensions/lib_subgraph/subgraph_lib.cc @@ -160,11 +160,11 @@ MXReturnValue createOpState(std::map attrs, std::string serialized_subgraph = "[empty]"; // MXNet subgraph is stored as Symbol in operator node attrs subgraphs field // custom subgraph is stored as json string in custom operator attrs map entry - if (attrs.count(SUBGRAPH_SYM_JSON)) { + if (attrs.count(MX_STR_SUBGRAPH_SYM_JSON)) { // user can now parse json and run other custom ops inside subgraph - serialized_subgraph = attrs[SUBGRAPH_SYM_JSON]; + serialized_subgraph = attrs[MX_STR_SUBGRAPH_SYM_JSON]; } - attrs.erase(SUBGRAPH_SYM_JSON); + attrs.erase(MX_STR_SUBGRAPH_SYM_JSON); *op_inst = new MyStatefulOp(serialized_subgraph, attrs); std::cout << "Info: stateful operator created" << std::endl; return MX_SUCCESS; @@ -177,7 +177,7 @@ REGISTER_OP(_custom_subgraph_op) const std::vector op_names({"exp","log"}); MXReturnValue mySupportedOps(std::string json, - std::vector ids, + std::vector& ids, std::unordered_map& options) { for (auto kv : options) { std::cout << "option: " << kv.first << " ==> " << kv.second << std::endl; @@ -204,8 +204,8 @@ MXReturnValue mySupportedOps(std::string json, dtype = std::stoi(attrs.map[JsonVal("dtype")].str); } - //check if op dtype is float - if(dtype == kFloat32) { + //check if op dtype is float, and if option was specified to require float types + if((dtype == kFloat32 && options.count("reqFloat") > 0) || options.count("reqFloat") == 0) { //check if op is in whitelist if(std::find(op_names.begin(),op_names.end(),op.str.c_str()) != op_names.end()) { // found op in whitelist, set value to 1 to include op in subgraph @@ -216,14 +216,34 @@ MXReturnValue mySupportedOps(std::string json, return MX_SUCCESS; } -MXReturnValue myReviewSubgraph(std::string json, int subraph_id, bool* accept, +MXReturnValue myReviewSubgraph(std::string json, int subgraph_id, bool* accept, std::unordered_map& options, - std::unordered_map& attrs) { + std::unordered_map& attrs, + std::map& args, + std::map& aux) { for (auto kv : options) { std::cout << "option: " << kv.first << " ==> " << kv.second << std::endl; } - if(options.find("reject") != options.end() && - options["reject"].compare("True") == 0) { + for (auto kv : args) { + std::cout << "arg: " << kv.first << " ==> ("; + for (auto s : kv.second.shape) + std::cout << s << ","; + std::cout << ") ["; + for (int i=0; i()[i] << ", "; + std::cout << "]" << std::endl; + } + + // check if option `reqArgs` was specified, and if so check if args were provided + if(options.count("reqArgs") > 0 && args.size() == 0) { + *accept = false; + std::cout << "rejecting subgraph since args were not provided" << std::endl; + return MX_SUCCESS; + } + + // check if option `reject` was specified, and if so check if value is 'True' + if(options.count("reject") > 0 && options["reject"].compare("True") == 0) { + // if specified, reject the subgraph. this is only used for testing *accept = false; std::cout << "rejecting subgraph" << std::endl; } else { @@ -231,7 +251,6 @@ MXReturnValue myReviewSubgraph(std::string json, int subraph_id, bool* accept, std::cout << "accepting subgraph" << std::endl; attrs["myKey"] = "myVal"; } - std::cout << json << std::endl; return MX_SUCCESS; } diff --git a/example/extensions/lib_subgraph/test_subgraph.py b/example/extensions/lib_subgraph/test_subgraph.py index 1bcecae3e21b..55a40514f105 100644 --- a/example/extensions/lib_subgraph/test_subgraph.py +++ b/example/extensions/lib_subgraph/test_subgraph.py @@ -23,8 +23,10 @@ # This test checks if dynamic loading of library into MXNet is successful # and checks the end of end computation of custom operator -import mxnet as mx import os, ctypes +import mxnet as mx +from mxnet.gluon import nn +from mxnet import nd from mxnet.base import _LIB, check_call, mx_uint, c_str, c_str_array, SymbolHandle # load library @@ -35,6 +37,10 @@ path = os.path.abspath('libsubgraph_lib.dll') mx.library.load(path) +############################################### +# Test with subgraph not consuming params +############################################### +# example model, ops to be partitioned do not have args (use outputs from other ops as inputs) a = mx.sym.var('a') b = mx.sym.var('b') c = a + b @@ -75,9 +81,6 @@ out3 = exe3.forward() print(out3) -from mxnet.gluon import nn -from mxnet import nd - # Gluon Hybridize partitioning with shapes/types print('-------------------------------') print('Testing Gluon Hybridize partitioning with shapes/types') @@ -88,3 +91,54 @@ out4 = sym_block(mx.nd.ones((3,2)),mx.nd.ones((3,2))) print(out4) +# Gluon Hybridize partitioning with shapes/types without inference +print('-------------------------------') +print('Testing Gluon Hybridize partitioning with shapes/types without inference') +inputs = [a,b] +sym_block2 = nn.SymbolBlock(sym, inputs) +sym_block2.initialize() +sym_block2.optimize_for(mx.nd.ones((3,2)), mx.nd.ones((3,2)), backend='myProp') +sym_block2.export('partitioned') + + +############################################### +# Test with subgraph directly consuming params +############################################### +# example model, ops to be partitioned have args +d2 = mx.sym.exp(a) +sym2 = mx.sym.log(d2) + +#execute in MXNet +print('-------------------------------') +print('Testing regular MXNet execution') +exe5 = sym2.bind(ctx=mx.cpu(), args={'a':mx.nd.ones((3,2))}) +out5 = exe5.forward() +print(out5) + +# with propogating shapes/types +print('-------------------------------') +print('Testing partitioning with shapes/types') +arg_array = [mx.nd.ones((3,2),dtype='float32')] +mysym6 = sym2.optimize_for("myProp", arg_array, reqArgs=True) +print(mysym6.tojson()) +exe6 = mysym6.bind(ctx=mx.cpu(), args={'a':mx.nd.ones((3,2))}) +out6 = exe6.forward() +print(out6) + +# without propogating shapes/types +print('-------------------------------') +print('Testing partitioning without shapes/types') +mysym7 = sym2.optimize_for("myProp", reqArgs=True) +exe7 = mysym7.bind(ctx=mx.cpu(), args={'a':mx.nd.ones((3,2))}) +out7 = exe7.forward() +print(out7) + +# Gluon Hybridize partitioning with shapes/types +print('-------------------------------') +print('Testing Gluon Hybridize partitioning with shapes/types') +inputs = [a] +sym2_block = nn.SymbolBlock(sym2, inputs) +sym2_block.initialize() +sym2_block.hybridize(backend='myProp') +out8 = sym2_block(mx.nd.ones((3,2))) +print(out8) diff --git a/include/mxnet/c_api.h b/include/mxnet/c_api.h index efa00330cee9..637b31dc5f55 100644 --- a/include/mxnet/c_api.h +++ b/include/mxnet/c_api.h @@ -2177,8 +2177,10 @@ MXNET_DLL int MXOptimizeForBackend(SymbolHandle sym_handle, const char* backend_name, const int dev_type, SymbolHandle* ret_sym_handle, - const mx_uint len, + const mx_uint args_len, NDArrayHandle* in_args_handle, + const mx_uint aux_len, + NDArrayHandle* in_aux_handle, const mx_uint num_options, const char** keys, const char** vals); diff --git a/include/mxnet/lib_api.h b/include/mxnet/lib_api.h index d59f2b12da37..9b32122c7d7a 100644 --- a/include/mxnet/lib_api.h +++ b/include/mxnet/lib_api.h @@ -39,7 +39,7 @@ #include #include -#define MX_LIBRARY_VERSION 3 +#define MX_LIBRARY_VERSION 4 /*! * \brief For loading multiple custom op libraries in Linux, exporting same symbol multiple @@ -234,10 +234,15 @@ enum MXReturnValue { */ struct MXTensor { MXTensor() : data_ptr(nullptr), dtype(kUNSET), verID(0) {} - + MXTensor(const MXTensor& oth) : data_ptr(oth.data_ptr), shape(oth.shape), + dtype(oth.dtype), verID(oth.verID), ctx(oth.ctx) { + setDLTensor(); + } MXTensor(void *data_ptr, const std::vector &shape, MXDType dtype, size_t vID, MXContext mx_ctx) - : data_ptr(data_ptr), shape(shape), dtype(dtype), verID(vID), ctx(mx_ctx) {} + : data_ptr(data_ptr), shape(shape), dtype(dtype), verID(vID), ctx(mx_ctx) { + setDLTensor(); + } /*! \brief populate internal tensor fields */ void setTensor(void *dptr, MXDType type, const int64_t* dims, int ndims, @@ -406,7 +411,35 @@ class OpResource { * \brief Json utility to parse serialized subgraph symbol */ /*! \brief Macro to help passing serialized subgraph through attribute dict */ -#define SUBGRAPH_SYM_JSON "subgraph_sym_json" +#define MX_STR_SUBGRAPH_SYM_JSON "subgraph_sym_json" +#define MX_STR_DTYPE "__dtype__" +#define MX_STR_SHAPE "__shape__" + +/* \brief get shape value from list of shapes string + * format: [[1]] or [[1],[2]] + */ +std::string getShapeAt(const std::string& shape, unsigned index) { + int idx = 1; // start at 1 to skip the first square bracket [ + // find the beginning of the output shape for the particular output index + for (unsigned x=0; x < index; x++) + idx = shape.find("[", idx+1); + int stop = shape.find("]", idx); // find stop index for this output shape + // add this shape to the list + return shape.substr(idx, stop-idx+1); +} + +/* \brief get dtype value from list of dtypes string + * format: [1] or [1,2] + */ +std::string getDtypeAt(const std::string& dtype, unsigned index) { + // find the beginning of the output dtype for the particular output index + int idx = 0; + for (unsigned x=0; x < index; x++) + idx = dtype.find(",", idx+1); + int stop = dtype.find(",", idx+1); // find stop index for this output dtype + if (stop == -1) stop = dtype.find("]", idx+1); + return dtype.substr(idx+1, stop-idx-1); +} /*! \brief Types of JSON objects */ enum JsonType {ERR, STR, NUM, LIST, MAP}; @@ -713,11 +746,13 @@ class CustomOp { }; /*! \brief Custom Subgraph Create function template */ -typedef MXReturnValue (*supportedOps_t)(std::string, std::vector, +typedef MXReturnValue (*supportedOps_t)(std::string, std::vector&, std::unordered_map&); typedef MXReturnValue (*reviewSubgraph_t)(std::string, int, bool*, std::unordered_map&, - std::unordered_map&); + std::unordered_map&, + std::map&, + std::map&); /*! * \brief An abstract class for subgraph property @@ -920,7 +955,17 @@ typedef int (*partCallSupportedOps_t)(supportedOps_t supportedOps, const char *j typedef int (*partCallReviewSubgraph_t)(reviewSubgraph_t reviewSubgraph, const char *json, int subgraph_id, int *accept, const char* const* opt_keys, const char* const* opt_vals, int num_opts, - char*** attr_keys, char*** attr_vals, int *num_attrs); + char*** attr_keys, char*** attr_vals, int *num_attrs, + const char* const* arg_names, int num_args, + void* const* arg_data, const int64_t* const* arg_shapes, + const int* arg_dims, const int* arg_types, + const size_t* arg_IDs, const char* const* arg_dev_type, + const int* arg_dev_id, + const char* const* aux_names, int num_aux, + void* const* aux_data, const int64_t* const* aux_shapes, + const int* aux_dims, const int* aux_types, + const size_t* aux_IDs, const char* const* aux_dev_type, + const int* aux_dev_id); #define MXLIB_INITIALIZE_STR "initialize" typedef int (*initialize_t)(int version); @@ -1266,11 +1311,11 @@ extern "C" { int num_ids, int *ids, const char* const* opt_keys, const char* const* opt_vals, int num_opts) { std::string subgraph_json(json); - // create map of attributes from list + // create map of options from list std::unordered_map opts; - for (int i = 0; i < num_opts; i++) { + for (int i = 0; i < num_opts; i++) opts[std::string(opt_keys[i])] = std::string(opt_vals[i]); - } + // create array of bools for operator support std::vector _ids(num_ids, false); // call user's supportedOps function @@ -1293,19 +1338,55 @@ extern "C" { _partCallReviewSubgraph(reviewSubgraph_t reviewSubgraph, const char *json, int subgraph_id, int *accept, const char* const* opt_keys, const char* const* opt_vals, int num_opts, - char*** attr_keys, char*** attr_vals, int *num_attrs) { + char*** attr_keys, char*** attr_vals, int *num_attrs, + const char* const* arg_names, int num_args, + void* const* arg_data, const int64_t* const* arg_shapes, + const int* arg_dims, const int* arg_types, + const size_t* arg_IDs, const char* const* arg_dev_type, + const int* arg_dev_id, + const char* const* aux_names, int num_aux, + void* const* aux_data, const int64_t* const* aux_shapes, + const int* aux_dims, const int* aux_types, + const size_t* aux_IDs, const char* const* aux_dev_type, + const int* aux_dev_id) { std::string subgraph_json(json); bool accept_bool = false; // create map of attributes from list std::unordered_map opts; - for (int i = 0; i < num_opts; i++) { + for (int i = 0; i < num_opts; i++) opts[std::string(opt_keys[i])] = std::string(opt_vals[i]); + + // create a map of named tensors for args + std::map args; + for (int i = 0; i < num_args; i++) { + std::vector shapes; + for (int j = 0; j < arg_dims[i]; j++) + shapes.push_back(arg_shapes[i][j]); + + MXTensor tensor(arg_data[i], shapes, (MXDType)arg_types[i], + arg_IDs[i], {arg_dev_type[i], arg_dev_id[i]}); + args[arg_names[i]] = tensor; + } + // create a map of named tensors for aux + std::map aux; + for (int i = 0; i < num_aux; i++) { + std::vector shapes; + for (int j = 0; j < aux_dims[i]; j++) + shapes.push_back(aux_shapes[i][j]); + + MXTensor tensor(aux_data[i], shapes, (MXDType)aux_types[i], + aux_IDs[i], {aux_dev_type[i], aux_dev_id[i]}); + aux[aux_names[i]] = tensor; } + // attributes to set on subgraph node std::unordered_map attrs; - MXReturnValue retval = reviewSubgraph(subgraph_json, subgraph_id, &accept_bool, opts, attrs); + MXReturnValue retval = reviewSubgraph(subgraph_json, subgraph_id, &accept_bool, + opts, attrs, args, aux); + if (!retval) return retval; + *accept = accept_bool; if (attrs.size() > 0) { diff --git a/perl-package/AI-MXNetCAPI/mxnet.i b/perl-package/AI-MXNetCAPI/mxnet.i index 3bc53d6442c1..846b28ff0e34 100644 --- a/perl-package/AI-MXNetCAPI/mxnet.i +++ b/perl-package/AI-MXNetCAPI/mxnet.i @@ -1633,6 +1633,8 @@ int MXOptimizeForBackend(SymbolHandle sym_handle, const mx_uint in, NDArrayHandle* in, const mx_uint in, + NDArrayHandle* in, + const mx_uint in, const char** keys, const char** vals); diff --git a/python/mxnet/gluon/block.py b/python/mxnet/gluon/block.py index da76b3efcd87..312358c8f5c5 100644 --- a/python/mxnet/gluon/block.py +++ b/python/mxnet/gluon/block.py @@ -989,11 +989,15 @@ def _build_cache(self, *args): # get list of params in the order of out.list_arguments arg_array = [args[data_names[name]] if name in data_names.keys() else params[name].data() for name in out.list_arguments()] + aux_array = [args[data_names[name]] if name in data_names.keys() else params[name].data() + for name in out.list_auxiliary_states()] # Partition the graph. - out = out.optimize_for(self._backend, arg_array, ctx, **self._backend_opts) - + out = out.optimize_for(self._backend, arg_array, aux_array, ctx, **self._backend_opts) + #update cached graph with partitioned graph + self._cached_graph = data, out self._cached_op = ndarray.CachedOp(out, flags) + def _deferred_infer_shape(self, *args): try: self.infer_shape(*args) @@ -1042,6 +1046,69 @@ def _call_cached_op(self, *args): out = [out] return _regroup(out, self._out_format) + def optimize_for(self, x, *args, backend=None, backend_opts=None, **kwargs): + """Partitions the current HybridBlock and optimizes it for a given backend + without executing a forward pass. Modifies the HybridBlock in-place. + + Immediately partitions a HybridBlock using the specified backend. Combines + the work done in the hybridize API with part of the work done in the forward + pass without calling the CachedOp. Can be used in place of hybridize, + afterwards `export` can be called or inference can be run. See README.md in + example/extensions/lib_subgraph/README.md for more details. + + Examples + -------- + # partition and then export to file + block.optimize_for(x, backend='myPart') + block.export('partitioned') + + # partition and then run inference + block.optimize_for(x, backend='myPart') + block(x) + + Parameters + ---------- + x : NDArray + first input to model + *args : NDArray + other inputs to model + backend : str + The name of backend, as registered in `SubgraphBackendRegistry`, default None + backend_opts : dict of user-specified options to pass to the backend for partitioning, optional + Passed on to `PrePartition` and `PostPartition` functions of `SubgraphProperty` + static_alloc : bool, default False + Statically allocate memory to improve speed. Memory usage may increase. + static_shape : bool, default False + Optimize for invariant input shapes between iterations. Must also + set static_alloc to True. Change of input shapes is still allowed + but slower. + """ + + # do hybrize API call + self.hybridize(True, backend, backend_opts, **kwargs) + + # do part of forward API call + has_symbol, has_ndarray, ctx_set, _ = _gather_type_ctx_info([x] + list(args)) + if has_symbol: + raise ValueError('Inputs must be NDArrays for the optimize_for API' + ' Please check the type of the args.\n') + if not has_symbol and not has_ndarray: + raise ValueError('In HybridBlock, there must be one NDArray as input.' + ' Please check the type of the args.\n') + if len(ctx_set) > 1: + raise ValueError('Find multiple contexts in the input, ' + 'After hybridized, the HybridBlock only supports one input ' + 'context. You can print the ele.ctx in the ' + 'input arguments to inspect their contexts. ' + 'Find all contexts = {}'.format(ctx_set)) + + self._build_cache(x, *args) + assert self._cached_op, "Gluon failed to build the cache. " \ + "This should never happen. " \ + "Please submit an issue on Github" \ + " https://github.com/apache/incubator-mxnet." + # do not actually call the cached_op + def _clear_cached_op(self): self._cached_graph = () self._cached_op = None diff --git a/python/mxnet/symbol/symbol.py b/python/mxnet/symbol/symbol.py index 706152f4b172..8e53a1a6b779 100644 --- a/python/mxnet/symbol/symbol.py +++ b/python/mxnet/symbol/symbol.py @@ -1446,7 +1446,7 @@ def _gen_atomic_symbol(self): return Symbol(handle) - def optimize_for(self, backend, args=None, ctx=None, **kwargs): + def optimize_for(self, backend, args=None, aux=None, ctx=None, **kwargs): """Partitions current symbol and optimizes it for a given backend, returns new partitioned symbol. @@ -1462,6 +1462,13 @@ def optimize_for(self, backend, args=None, ctx=None, **kwargs): - If type is a dict of str to `NDArray`, then it maps the name of arguments to the corresponding `NDArray`. + aux : list of NDArray or dict of str to NDArray, optional + Input auxiliary arguments to the symbol + + - If type is a list of `NDArray`, the order is the same as that of `list_arguments()`. + - If type is a dict of str to `NDArray`, then it maps the name of arguments + to the corresponding `NDArray`. + ctx : Context, optional Device context, used to infer stypes @@ -1476,13 +1483,19 @@ def optimize_for(self, backend, args=None, ctx=None, **kwargs): out = SymbolHandle() assert isinstance(backend, str) - if args is None: + if args is None or len(args) == 0: args = [] args_handle = c_array(NDArrayHandle, []) else: - listed_arguments = self.list_arguments() - args_handle, args = self._get_ndarray_inputs('args', args, listed_arguments, False) + args_handle, args = self._get_ndarray_inputs('args', args, + self.list_arguments(), False) + if aux is None or len(aux) == 0: + aux = [] + aux_handle = c_array(NDArrayHandle, []) + else: + aux_handle, aux = self._get_ndarray_inputs('aux_states', aux, + self.list_auxiliary_states(), False) if ctx is None: ctx = current_context() assert isinstance(ctx, Context) @@ -1498,6 +1511,8 @@ def optimize_for(self, backend, args=None, ctx=None, **kwargs): ctypes.byref(out), mx_uint(len(args)), args_handle, + mx_uint(len(aux)), + aux_handle, mx_uint(len(key_list)), c_str_array(key_list), c_str_array(val_list))) diff --git a/src/c_api/c_api.cc b/src/c_api/c_api.cc index ad303502b211..db0e2629a5df 100644 --- a/src/c_api/c_api.cc +++ b/src/c_api/c_api.cc @@ -120,17 +120,28 @@ void CustomFComputeDispatcher(const std::string op_name, std::vector in_verIDs, out_verIDs; std::vector in_dev_type, out_dev_type; std::vector in_dev_id, out_dev_id; + std::vector conv_mkl; // converted NDArrays from MKLDNN format // convert inputs/outpus NDArray to C types to be passed to lib_api.h for (size_t i = 0; i < inputs.size(); i++) { - in_data.push_back(inputs[i].data().dptr_); - in_shapes.push_back(inputs[i].shape().data()); - in_dims.push_back(inputs[i].shape().ndim()); - in_types.push_back(inputs[i].dtype()); - in_verIDs.push_back(inputs[i].version()); - const char* ctx_str = inputs[i].ctx().dev_mask() == Context::kCPU ? "cpu" : "gpu"; + NDArray const* in_nd = &(inputs[i]); +#if MXNET_USE_MKLDNN == 1 + // reorder data if in MKLDNN format + if (in_nd->IsMKLDNNData()) { + // convert from MKLDNN + conv_mkl.push_back(in_nd->Reorder2Default()); + in_nd = &(conv_mkl.back()); + } +#endif + // pull out parts to pass over to library + in_data.push_back(in_nd->data().dptr_); + in_shapes.push_back(in_nd->shape().data()); + in_dims.push_back(in_nd->shape().ndim()); + in_types.push_back(in_nd->dtype()); + in_verIDs.push_back(in_nd->version()); + const char* ctx_str = in_nd->ctx().dev_mask() == Context::kCPU ? "cpu" : "gpu"; in_dev_type.push_back(ctx_str); - in_dev_id.push_back(inputs[i].ctx().real_dev_id()); + in_dev_id.push_back(in_nd->ctx().real_dev_id()); } for (size_t i = 0; i < outputs.size(); i++) { @@ -193,7 +204,7 @@ void CustomFComputeDispatcher(const std::string op_name, if (fcomp_fp != nullptr) { // convert attributes to vector of char* std::vector attr_keys, attr_vals; - for (auto kv : attrs->dict) { + for (auto &kv : attrs->dict) { attr_keys.push_back(kv.first.c_str()); attr_vals.push_back(kv.second.c_str()); } @@ -361,7 +372,7 @@ int MXLoadLib(const char *path) { auto attr_parser = [=](const NodeAttrs* attrs) { // convert attributes to vector of char std::vector attr_keys, attr_vals; - for (auto kv : attrs->dict) { + for (auto &kv : attrs->dict) { attr_keys.push_back(kv.first.c_str()); attr_vals.push_back(kv.second.c_str()); } @@ -371,7 +382,7 @@ int MXLoadLib(const char *path) { nnvm::Graph g; g.outputs = attrs->subgraphs[0].get()->outputs; subgraph_json = nnvm::pass::SaveJSON(g); - attr_keys.push_back(SUBGRAPH_SYM_JSON); + attr_keys.push_back(MX_STR_SUBGRAPH_SYM_JSON); attr_vals.push_back(subgraph_json.c_str()); } @@ -388,7 +399,7 @@ int MXLoadLib(const char *path) { auto num_inputs = [=](const NodeAttrs& attrs) { // convert attributes to vector of char std::vector attr_keys, attr_vals; - for (auto kv : attrs.dict) { + for (auto &kv : attrs.dict) { attr_keys.push_back(kv.first.c_str()); attr_vals.push_back(kv.second.c_str()); } @@ -406,7 +417,7 @@ int MXLoadLib(const char *path) { auto num_outputs = [=](const NodeAttrs& attrs) { // convert attributes to vector of char* std::vector attr_keys, attr_vals; - for (auto kv : attrs.dict) { + for (auto &kv : attrs.dict) { attr_keys.push_back(kv.first.c_str()); attr_vals.push_back(kv.second.c_str()); } @@ -425,7 +436,7 @@ int MXLoadLib(const char *path) { auto num_inouts = [=](const NodeAttrs& attrs) { // convert attributes to vector of char* std::vector attr_keys, attr_vals; - for (auto kv : attrs.dict) { + for (auto &kv : attrs.dict) { attr_keys.push_back(kv.first.c_str()); attr_vals.push_back(kv.second.c_str()); } @@ -445,7 +456,7 @@ int MXLoadLib(const char *path) { mxnet::ShapeVector *out_shape) { // convert attributes to vector of char* std::vector attr_keys, attr_vals; - for (auto kv : attrs.dict) { + for (auto &kv : attrs.dict) { attr_keys.push_back(kv.first.c_str()); attr_vals.push_back(kv.second.c_str()); } @@ -516,7 +527,7 @@ int MXLoadLib(const char *path) { std::vector *out_type) { // convert attributes to vector of char* std::vector attr_keys, attr_vals; - for (auto kv : attrs.dict) { + for (auto &kv : attrs.dict) { attr_keys.push_back(kv.first.c_str()); attr_vals.push_back(kv.second.c_str()); } @@ -544,7 +555,7 @@ int MXLoadLib(const char *path) { auto mutate_inputs = [=](const nnvm::NodeAttrs& attrs) { // convert attributes to vector of char* std::vector attr_keys, attr_vals; - for (auto kv : attrs.dict) { + for (auto &kv : attrs.dict) { attr_keys.push_back(kv.first.c_str()); attr_vals.push_back(kv.second.c_str()); } @@ -629,7 +640,7 @@ int MXLoadLib(const char *path) { const std::vector& in_types) { // convert attributes to vector of char* std::vector attr_keys, attr_vals; - for (auto kv : attrs.dict) { + for (auto &kv : attrs.dict) { attr_keys.push_back(kv.first.c_str()); attr_vals.push_back(kv.second.c_str()); } @@ -640,7 +651,7 @@ int MXLoadLib(const char *path) { nnvm::Graph g; g.outputs = attrs.subgraphs[0].get()->outputs; subgraph_json = nnvm::pass::SaveJSON(g); - attr_keys.push_back(SUBGRAPH_SYM_JSON); + attr_keys.push_back(MX_STR_SUBGRAPH_SYM_JSON); attr_vals.push_back(subgraph_json.c_str()); } @@ -858,12 +869,10 @@ int MXLoadLib(const char *path) { std::string op_name_str(op_name); LOG(INFO) << "\t\tStrategy[" << j << "] " << strategy_str << " subgraphOp: '" << op_name_str << "'"; - - // MXNET_REGISTER_SUBGRAPH_PROPERTY(customBackend, CustomSubgraphProperty); - mxnet::op::SubgraphBackendRegistry::Get()->__REGISTER_CUSTOM_PROPERTY__(name_str, - std::make_shared( - strategy_str, callSupportedOps, supportedOps_fp, - callReviewSubgraph, reviewSubgraph_fp, callFree, op_name_str)); + mxnet::op::SubgraphBackendRegistry::Get()->__REGISTER_CUSTOM_PROPERTY__ + (name_str, std::make_shared + (strategy_str, callSupportedOps, supportedOps_fp, + callReviewSubgraph, reviewSubgraph_fp, callFree, op_name_str)); } } API_END(); diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc index 9042dfa4c266..4ec916da3375 100644 --- a/src/c_api/c_api_symbolic.cc +++ b/src/c_api/c_api_symbolic.cc @@ -1353,32 +1353,54 @@ int MXOptimizeForBackend(SymbolHandle sym_handle, const char* backend_name, const int dev_type, SymbolHandle* ret_sym_handle, - const mx_uint len, + const mx_uint args_len, NDArrayHandle* in_args_handle, + const mx_uint aux_len, + NDArrayHandle* in_aux_handle, const mx_uint num_options, const char** keys, const char** vals) { + // create copy of input symbol nnvm::Symbol *s = new nnvm::Symbol(); API_BEGIN(); nnvm::Symbol *sym = static_cast(sym_handle); *s = sym->Copy(); nnvm::Graph g = Symbol2Graph(*s); - if (len) { + const auto& indexed_graph = g.indexed_graph(); + const auto& mutable_nodes = indexed_graph.mutable_input_nodes(); + std::vector input_names = sym->ListInputNames(nnvm::Symbol::kAll); + size_t num_forward_inputs = input_names.size(); + if (args_len || aux_len) { NDArray **in_args_ptr = reinterpret_cast(in_args_handle); + NDArray **in_aux_ptr = reinterpret_cast(in_aux_handle); Context default_ctx = Context::Create(static_cast(dev_type), 0); - mxnet::ShapeVector arg_shapes(len); - nnvm::DTypeVector arg_dtypes(len); - StorageTypeVector arg_stypes(len); - for (mx_uint i = 0; i < len; i++) { - const auto &in_arg = *(in_args_ptr[i]); - arg_shapes[i] = in_arg.shape(); - arg_dtypes[i] = in_arg.dtype(); - arg_stypes[i] = in_arg.storage_type(); + mxnet::ShapeVector arg_shapes(args_len + aux_len); + nnvm::DTypeVector arg_dtypes(args_len + aux_len); + StorageTypeVector arg_stypes(args_len + aux_len); + size_t args_top = 0, aux_top = 0; + // loop over inputs to symbol in order and add to args/aux if mutable + for (size_t i = 0; i < num_forward_inputs; ++i) { + const uint32_t nid = indexed_graph.input_nodes().at(i); + if (mutable_nodes.count(nid)) { + CHECK_LT(aux_top, aux_len) + << "Cannot find aux '" << input_names[i] << "' in provided aux to optimize_for"; + const auto &in_arg = *(in_aux_ptr[aux_top++]); + arg_shapes[i] = in_arg.shape(); + arg_dtypes[i] = in_arg.dtype(); + arg_stypes[i] = in_arg.storage_type(); + } else { + CHECK_LT(args_top, args_len) + << "Cannot find arg '" << input_names[i] << "' in provided args to optimize_for"; + const auto &in_arg = *(in_args_ptr[args_top++]); + arg_shapes[i] = in_arg.shape(); + arg_dtypes[i] = in_arg.dtype(); + arg_stypes[i] = in_arg.storage_type(); + } } - const auto& indexed_graph = g.indexed_graph(); - const auto num_forward_inputs = indexed_graph.input_nodes().size(); + g.attrs["context"] = std::make_shared( exec::ContextVector(indexed_graph.num_nodes(), default_ctx)); + // infer shapes g = exec::InferShape(std::move(g), std::move(arg_shapes), "__shape__"); // infer dtypes @@ -1393,11 +1415,31 @@ int MXOptimizeForBackend(SymbolHandle sym_handle, common::HandleInferStorageTypeError(num_forward_inputs, indexed_graph, g.GetAttr("storage_type")); } + // set args/aux as attributes on graph so that subgraph property can use them + std::vector arg_names = sym->ListInputNames(nnvm::Symbol::kReadOnlyArgs); + g.attrs["in_args"] = std::make_shared(in_args_ptr); + g.attrs["in_arg_names"] = std::make_shared(arg_names); + + std::vector aux_names = sym->ListInputNames(nnvm::Symbol::kAuxiliaryStates); + g.attrs["in_aux"] = std::make_shared(in_aux_ptr); + g.attrs["in_aux_names"] = std::make_shared(aux_names); + } else { + // args/aux were not specified, so set nullptr/empty-lists + NDArray **in_args_ptr = static_cast(nullptr); + std::vector arg_names; + g.attrs["in_args"] = std::make_shared(in_args_ptr); + g.attrs["in_arg_names"] = std::make_shared(arg_names); + + NDArray **in_aux_ptr = static_cast(nullptr); + std::vector aux_names; + g.attrs["in_aux"] = std::make_shared(in_aux_ptr); + g.attrs["in_aux_names"] = std::make_shared(aux_names); } + // create a data structure from pointer array std::vector> options_map; - for (mx_uint i = 0; i < num_options; ++i) { + for (mx_uint i = 0; i < num_options; ++i) options_map.emplace_back(keys[i], vals[i]); - } + const auto backend = mxnet::op::SubgraphBackendRegistry::Get()->GetSubgraphBackend(backend_name); const auto& subgraph_prop_list = backend->GetSubgraphProperties(); for (auto property : subgraph_prop_list) { diff --git a/src/operator/subgraph/build_subgraph.cc b/src/operator/subgraph/build_subgraph.cc index dc0c1425cef1..413395c3b74f 100644 --- a/src/operator/subgraph/build_subgraph.cc +++ b/src/operator/subgraph/build_subgraph.cc @@ -560,11 +560,7 @@ void CutGraphInputs(const std::vector &input_entries, } nnvm::ObjectPtr n = nnvm::CreateVariableNode( var_name + std::to_string(name_count_map[var_name])); - // set attribute for subgraph input to indicate if it is from an arg/param to model - if (e->node->is_variable()) - n->attrs.dict["isArg"] = "True"; - else - n->attrs.dict["isArg"] = "False"; + *e = nnvm::NodeEntry{n, 0, 0}; } } @@ -583,7 +579,7 @@ void ReattachGraphInputs(const std::vector &input_entries, } /*! - * \brief Replace a set of nodes belonging to the same subgraph with a subgrpah node + * \brief Replace a set of nodes belonging to the same subgraph with a subgraph node * and keep the subgraph in the subgraph node. */ void CreateSubgraphNode(nnvm::Graph* g, @@ -613,6 +609,7 @@ void CreateSubgraphNode(nnvm::Graph* g, sym.outputs[i] = *output_entries[i]; } const SubgraphPropertyPtr& subg_prop = g->GetAttr("subgraph_property"); + subg_prop->InitSubgraphInputs(&input_entries, &orig_input_entries); nnvm::ObjectPtr n = subg_prop->CreateSubgraphNode(sym, subgraph_selector, subgraph_id); // CreateSubgraphNode returns NULL if subgraph property determines that subgraph is sub-optimal // In that case, subgraph node is not created and graph is not modified diff --git a/src/operator/subgraph/partitioner/custom_subgraph_property.h b/src/operator/subgraph/partitioner/custom_subgraph_property.h index 410d983fa591..b7f2cc2d0fef 100644 --- a/src/operator/subgraph/partitioner/custom_subgraph_property.h +++ b/src/operator/subgraph/partitioner/custom_subgraph_property.h @@ -33,6 +33,7 @@ #include #include #include +#include #include "../common.h" #include "../subgraph_property.h" #include "../../include/mxnet/lib_api.h" @@ -99,6 +100,75 @@ class CustomSubgraphProperty: public SubgraphProperty { const std::vector>& options_map) { // clear supported_nodes to remove state from previous calls supported_nodes.clear(); + // get input args and arg names + in_arg_names = g.GetAttr>("in_arg_names"); + in_args_ptr = g.GetAttr("in_args"); + in_aux_names = g.GetAttr>("in_aux_names"); + in_aux_ptr = g.GetAttr("in_aux"); + + // convert input args + arg_names.clear(); + arg_data.clear(); + arg_shapes.clear(); + arg_dims.clear(); + arg_types.clear(); + arg_verIDs.clear(); + arg_dev_type.clear(); + arg_dev_id.clear(); + for (size_t i=0; i < in_arg_names.size(); i++) { + arg_names.push_back(in_arg_names[i].c_str()); + const NDArray &in_arg = *(in_args_ptr[i]); + +#if MXNET_USE_MKLDNN == 1 + // reorder data if in MKLDNN format + if (in_arg.IsMKLDNNData()) { + in_arg.Reorder2DefaultAsync(); + in_arg.WaitToRead(); + } +#endif + + // pull out parts of NDArray to send to backend + arg_data.push_back(in_arg.data().dptr_); + arg_shapes.push_back(in_arg.shape().data()); + arg_dims.push_back(in_arg.shape().ndim()); + arg_types.push_back(in_arg.dtype()); + arg_verIDs.push_back(in_arg.version()); + const char* arg_ctx_str = in_arg.ctx().dev_mask() == Context::kCPU ? "cpu" : "gpu"; + arg_dev_type.push_back(arg_ctx_str); + arg_dev_id.push_back(in_arg.ctx().real_dev_id()); + } + + // convert input aux + aux_names.clear(); + aux_data.clear(); + aux_shapes.clear(); + aux_dims.clear(); + aux_types.clear(); + aux_verIDs.clear(); + aux_dev_type.clear(); + aux_dev_id.clear(); + for (size_t i=0; i < in_aux_names.size(); i++) { + aux_names.push_back(in_aux_names[i].c_str()); + const auto &in_aux = *(in_aux_ptr[i]); + +#if MXNET_USE_MKLDNN == 1 + // reorder data if in MKLDNN format + if (in_aux.IsMKLDNNData()) { + in_aux.Reorder2DefaultAsync(); + in_aux.WaitToRead(); + } +#endif + + // pull out parts of NDArray to send to backend + aux_data.push_back(in_aux.data().dptr_); + aux_shapes.push_back(in_aux.shape().data()); + aux_dims.push_back(in_aux.shape().ndim()); + aux_types.push_back(in_aux.dtype()); + aux_verIDs.push_back(in_aux.version()); + const char* aux_ctx_str = in_aux.ctx().dev_mask() == Context::kCPU ? "cpu" : "gpu"; + aux_dev_type.push_back(aux_ctx_str); + aux_dev_id.push_back(in_aux.ctx().real_dev_id()); + } // remove all graph attrs, some cannot be saved to json nnvm::Graph graph = std::move(g); @@ -108,23 +178,37 @@ class CustomSubgraphProperty: public SubgraphProperty { // set shape attrs for each node in the graph if (g.HasAttr("shape")) { mxnet::ShapeVector shapes = g.GetAttr("shape"); - for (unsigned i = 0; i < indexed_graph.num_nodes(); i++) { - nnvm::Node* node = const_cast(indexed_graph[i].source); - mxnet::TShape shape = shapes[i]; + for (unsigned nid = 0; nid < indexed_graph.num_nodes(); nid++) { + nnvm::Node* node = const_cast(indexed_graph[nid].source); std::stringstream ss; - ss << shape; - node->attrs.dict["shape"] = ss.str(); + ss << "["; + // set the output shapes for this node + for (unsigned oid = 0; oid < node->num_outputs(); oid++) { + const uint32_t out_entry_id = indexed_graph.entry_id(nid, oid); + mxnet::TShape& shape = shapes[out_entry_id]; + ss << shape; + if (oid < node->num_outputs()-1) ss << ","; + } + ss << "]"; + node->attrs.dict[MX_STR_SHAPE] = ss.str(); } } // set dtype attrs for each node in the graph if (g.HasAttr("dtype")) { std::vector dtypes = g.GetAttr >("dtype"); - for (unsigned i = 0; i < indexed_graph.num_nodes(); i++) { - nnvm::Node* node = const_cast(indexed_graph[i].source); - int dtype = dtypes[i]; + for (unsigned nid = 0; nid < indexed_graph.num_nodes(); nid++) { + nnvm::Node* node = const_cast(indexed_graph[nid].source); std::stringstream ss; - ss << dtype; - node->attrs.dict["dtype"] = ss.str(); + ss << "["; + // set the output dtypes for this node + for (unsigned oid = 0; oid < node->num_outputs(); oid++) { + const uint32_t out_entry_id = indexed_graph.entry_id(nid, oid); + int dtype = dtypes[out_entry_id]; + ss << dtype; + if (oid < node->num_outputs()-1) ss << ","; + } + ss << "]"; + node->attrs.dict[MX_STR_DTYPE] = ss.str(); } } @@ -142,10 +226,14 @@ class CustomSubgraphProperty: public SubgraphProperty { opt_keys_.clear(); opt_vals_.clear(); options_map_.clear(); - for (auto kv : options_map) { + // store options in map in subgraph property to re-use later for reviewSubgraph + for (auto& kv : options_map) { options_map_.push_back(kv); - opt_keys_.push_back(options_map_.back().first.c_str()); - opt_vals_.push_back(options_map_.back().second.c_str()); + } + // convert options_map_ to char* to pass to backend library + for (auto& kv : options_map_) { + opt_keys_.push_back(kv.first.c_str()); + opt_vals_.push_back(kv.second.c_str()); } CHECK(call_supported_ops_(supported_ops_, json, supported_node_IDs.size(), ids, @@ -162,9 +250,10 @@ class CustomSubgraphProperty: public SubgraphProperty { } // override CreateSubgraphNode virtual nnvm::ObjectPtr CreateSubgraphNode(const nnvm::Symbol &sym, - const int subgraph_id = 0) const { + const int subgraph_id = 0) const { int accept = 1; int num_attr = 0; + std::map user_attrs; char** attr_keys = nullptr; char** attr_vals = nullptr; if (review_subgraph_) { @@ -173,8 +262,9 @@ class CustomSubgraphProperty: public SubgraphProperty { const auto& idx = g.indexed_graph(); // set isArg/isAux for each null op/param in the graph - const std::vector aux_names = sym.ListInputNames(nnvm::Symbol::kAuxiliaryStates); - std::unordered_set aux_set(aux_names.begin(), aux_names.end()); + const std::vector aux_state_names = + sym.ListInputNames(nnvm::Symbol::kAuxiliaryStates); + std::unordered_set aux_set(aux_state_names.begin(), aux_state_names.end()); for (unsigned i = 0; i < idx.num_nodes(); i++) { nnvm::Node* node = const_cast(idx[i].source); // check if this node is input to subgraph @@ -188,31 +278,121 @@ class CustomSubgraphProperty: public SubgraphProperty { } std::string subgraph_json = nnvm::pass::SaveJSON(g); - CHECK(call_review_subgraph_(review_subgraph_, subgraph_json.c_str(), - subgraph_id, &accept, opt_keys_.data(), - opt_vals_.data(), opt_keys_.size(), - &attr_keys, &attr_vals, &num_attr)) + CHECK(call_review_subgraph_(review_subgraph_, subgraph_json.c_str(), subgraph_id, + &accept, opt_keys_.data(), opt_vals_.data(), + opt_keys_.size(), &attr_keys, &attr_vals, &num_attr, + arg_names.data(), arg_names.size(), arg_data.data(), + arg_shapes.data(), arg_dims.data(), arg_types.data(), + arg_verIDs.data(), arg_dev_type.data(), + arg_dev_id.data(), aux_names.data(), aux_names.size(), + aux_data.data(), aux_shapes.data(), aux_dims.data(), + aux_types.data(), aux_verIDs.data(), + aux_dev_type.data(), aux_dev_id.data())) << "Error calling review_subgraph for '" << subgraph_prop << "'"; + + if (num_attr > 0) { + // set user specified attributes + for (int i=0; i < num_attr; i++) { + user_attrs[attr_keys[i]] = attr_vals[i]; + call_free_(attr_vals[i]); + call_free_(attr_keys[i]); + } + // free memory used by custom op to allocate attributes + call_free_(attr_vals); + call_free_(attr_keys); + } } + if (accept) { nnvm::ObjectPtr n = nnvm::Node::Create(); n->attrs.op = Op::Get(subgraph_op_name); n->attrs.name = "_op" + std::to_string(subgraph_id); n->attrs.subgraphs.push_back(std::make_shared(sym)); - // set user specified attributes - for (int i=0; i < num_attr; i++) { - n->attrs.dict[attr_keys[i]] = attr_vals[i]; - call_free_(attr_vals[i]); - call_free_(attr_keys[i]); + + // set shapes + { + std::stringstream ss; + ss << "["; + for (unsigned i=0; i < sym.outputs.size(); i++) { + const nnvm::NodeEntry& e = sym.outputs[i]; + if (e.node->attrs.dict.count("__shape__") > 0) { + std::string& shape = e.node->attrs.dict["__shape__"]; + // add this shape to the list + ss << getShapeAt(shape, e.index); + } + if (i < sym.outputs.size()-1) + ss << ","; + } + ss << "]"; + n->attrs.dict["__shape__"] = ss.str(); + } + // set dtypes + { + std::stringstream ss; + ss << "["; + for (unsigned i=0; i < sym.outputs.size(); i++) { + const nnvm::NodeEntry& e = sym.outputs[i]; + if (e.node->attrs.dict.count("__dtype__") > 0) { + std::string& dtype = e.node->attrs.dict["__dtype__"]; + // add this dtype to the list + ss << getDtypeAt(dtype, e.index); + } + if (i < sym.outputs.size()-1) + ss << ","; + } + ss << "]"; + n->attrs.dict["__dtype__"] = ss.str(); } - // free memory used by custom op to allocate attributes - call_free_(attr_vals); - call_free_(attr_keys); + // set user specified attributes + for (auto attr : user_attrs) + n->attrs.dict[attr.first] = attr.second; return n; } else { return nullptr; } } + + virtual void InitSubgraphInputs(std::vector* input_entries, + std::vector* orig_input_entries) const { + for (size_t i = 0; i < input_entries->size(); ++i) { + nnvm::NodeEntry *e = input_entries->at(i); + nnvm::NodeEntry& orig = orig_input_entries->at(i); + + // set attribute for subgraph input to indicate if it is from an arg/param to model + if (orig.node->is_variable()) { + // get name of original output entry + nnvm::Symbol sym; + sym.outputs.push_back(orig); + const auto output_names = sym.ListOutputNames(); + CHECK_EQ(output_names.size(), 1U); + const std::string& var_name = output_names[0]; + + e->node->attrs.dict["isArg"] = "True"; + e->node->attrs.dict["argName"] = var_name; + } else { + e->node->attrs.dict["isArg"] = "False"; + } + + // pass down other attributes if available + if (orig.node->attrs.dict.count("__dtype__") > 0) { + // get dtype string from other node + std::string& dtype = orig.node->attrs.dict["__dtype__"]; + std::stringstream ss; + ss << "[" << getDtypeAt(dtype, orig.index) << "]"; + e->node->attrs.dict["__dtype__"] = ss.str(); + } + + if (orig.node->attrs.dict.count("__shape__") > 0) { + // get shape string from other node + std::string& shape = orig.node->attrs.dict["__shape__"]; + // create new shape string for this node + std::stringstream ss; + ss << "[" << getShapeAt(shape, orig.index) << "]"; + e->node->attrs.dict["__shape__"] = ss.str(); + } + } + } + // override CreateSubgraphSelector virtual SubgraphSelectorPtr CreateSubgraphSelector() const { return std::make_shared(supported_nodes); @@ -228,6 +408,17 @@ class CustomSubgraphProperty: public SubgraphProperty { std::string subgraph_op_name; std::vector> options_map_; std::vector opt_keys_, opt_vals_; + std::vector in_arg_names, in_aux_names; + NDArray **in_args_ptr; + NDArray **in_aux_ptr; + std::vector arg_names, aux_names; + std::vector arg_data, aux_data; + std::vector arg_shapes, aux_shapes; + std::vector arg_dims, aux_dims; + std::vector arg_types, aux_types; + std::vector arg_verIDs, aux_verIDs; + std::vector arg_dev_type, aux_dev_type; + std::vector arg_dev_id, aux_dev_id; }; } // namespace op } // namespace mxnet diff --git a/src/operator/subgraph/subgraph_property.h b/src/operator/subgraph/subgraph_property.h index e68fc6877202..5e87626659bf 100644 --- a/src/operator/subgraph/subgraph_property.h +++ b/src/operator/subgraph/subgraph_property.h @@ -358,6 +358,14 @@ class SubgraphProperty { std::vector* orig_input_entries) const { subgraph_node->inputs = *orig_input_entries; } + /*! + * \brief Initialize subgraph internal inputs with external input entries. + * Called before CreateSubgraphNode, optional + * \param input_entries input entries inside subgraph + * \param orig_input_entries input entries outside subgraph + */ + virtual void InitSubgraphInputs(std::vector* input_entries, + std::vector* orig_input_entries) const {} /*! * \brief Set an attr with name in the attr map. */ diff --git a/tests/python/unittest/test_extensions.py b/tests/python/unittest/test_extensions.py index 799615b60aa2..d00f1494e4d5 100644 --- a/tests/python/unittest/test_extensions.py +++ b/tests/python/unittest/test_extensions.py @@ -167,3 +167,17 @@ def test_subgraph(): out4 = sym_block(mx.nd.ones((3,2)),mx.nd.ones((3,2))) # check that result matches one executed by MXNet assert_almost_equal(out[0].asnumpy(), out4[0].asnumpy(), rtol=1e-3, atol=1e-3) + + # Gluon Hybridize partitioning with shapes/types + sym_block2 = nn.SymbolBlock(sym, [a,b]) + sym_block2.initialize() + a_data = mx.nd.ones((3,2)) + b_data = mx.nd.ones((3,2)) + sym_block2.optimize_for(a_data, b_data, backend='myProp') + sym_block2.export('optimized') + sym_block3 = nn.SymbolBlock.imports('optimized-symbol.json',['a','b'], + 'optimized-0000.params') + + out5 = sym_block3(a_data, b_data) + # check that result matches one executed by MXNet + assert_almost_equal(out[0].asnumpy(), out5[0].asnumpy(), rtol=1e-3, atol=1e-3) diff --git a/tests/python/unittest/test_subgraph_op.py b/tests/python/unittest/test_subgraph_op.py index f1572e71f128..e414a9836ccb 100644 --- a/tests/python/unittest/test_subgraph_op.py +++ b/tests/python/unittest/test_subgraph_op.py @@ -282,7 +282,7 @@ def check_subgraph_exe6(sym, subgraph_backend, op_names): # infer shape/type before partition before simple_bind check_call(_LIB.MXSetSubgraphPropertyOpNamesV2(c_str(subgraph_backend), mx_uint(len(op_names)), c_str_array(op_names))) - part_sym = sym.optimize_for(subgraph_backend, exe1.arg_dict) + part_sym = sym.optimize_for(subgraph_backend, exe1.arg_dict, exe1.aux_dict) check_call(_LIB.MXRemoveSubgraphPropertyOpNamesV2(c_str(subgraph_backend))) exe2 = part_sym.simple_bind(ctx=mx.current_context(), grad_req='null') @@ -335,7 +335,7 @@ def check_subgraph_exe8(sym, subgraph_backend, op_names): # infer shape/type before partition before bind check_call(_LIB.MXSetSubgraphPropertyOpNamesV2(c_str(subgraph_backend), mx_uint(len(op_names)), c_str_array(op_names))) - part_sym = sym.optimize_for(subgraph_backend, arg_array) + part_sym = sym.optimize_for(subgraph_backend, arg_array, aux_array) check_call(_LIB.MXRemoveSubgraphPropertyOpNamesV2(c_str(subgraph_backend))) exe2 = part_sym.bind(ctx=mx.current_context(), args=arg_array, aux_states=aux_array, grad_req='null')