src/operator/subgraph_op_common.h

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#ifndef MXNET_OPERATOR_SUBGRAPH_OP_COMMON_H_
#define MXNET_OPERATOR_SUBGRAPH_OP_COMMON_H_

#include <mxnet/io.h>
#include <mxnet/base.h>
#include <mxnet/op_attr_types.h>
#include <vector>
#include <utility>
#include <string>
#include "../imperative/cached_op.h"
#include "../imperative/imperative_utils.h"

namespace mxnet {
namespace op {

/*
 * Infer the data types of inputs and outputs of an operator that contains a
 * subgraph.
 */
bool InferSubgraphDataType(const nnvm::Symbol& subgraph,
                           std::vector<int>* in_type,
                           std::vector<int>* out_type);

/*
 * Infer the shape of inputs and outputs of an operator that contains a
 * subgraph.
 */
bool InferSubgraphShape(const nnvm::Symbol& subgraph,
                        mxnet::ShapeVector* in_shape,
                        mxnet::ShapeVector* out_shape);

/*
 * Infer the storage types of inputs and outputs of an operator that contains a
 * subgraph.
 */
bool InferSubgraphStorage(const nnvm::Symbol& subgraph,
                          const int dev_mask,
                          DispatchMode* dispatch_mode,
                          std::vector<int>* in_attrs,
                          std::vector<int>* out_attrs);

bool as_bool_scalar(const NDArray& a);

bool is_shape_udf(const mxnet::TShape& x);

bool is_stype_udf(const int& x);

bool is_type_udf(const int& x);

template <typename T>
void extract_by_loc(const std::vector<T>& array,
                    const mxnet::Tuple<dim_t> input_locs,
                    std::vector<T>* out) {
  out->clear();
  out->reserve(input_locs.ndim());
  for (dim_t i : input_locs) {
    out->push_back(array[i]);
  }
}

template <typename T>
bool fill_value(T* x, T* y, bool x_empty, bool y_empty) {
  if (*x == *y || (x_empty && y_empty)) {
    return true;
  }
  if (!x_empty && !y_empty) {
    return false;
  }
  if (x_empty) {
    *x = *y;
  }
  if (y_empty) {
    *y = *x;
  }
  return true;
}

template <typename T>
bool sync_in_in(const mxnet::Tuple<dim_t>& input_locs,
                std::vector<T>* in,
                std::vector<T>* subg_in,
                std::function<bool(const T&)> is_empty) {
  for (int i = 0; i < input_locs.ndim(); ++i) {
    T& x = in->at(input_locs[i]);
    T& y = subg_in->at(i);
    fill_value(&x, &y, is_empty(x), is_empty(y));
  }
  return true;
}

template <typename T>
bool sync_out_out(std::vector<T>* out_1,
                  std::vector<T>* out_2,
                  std::function<bool(const T&)> is_empty) {
  CHECK_EQ(out_1->size(), out_2->size());
  for (size_t i = 0; i < out_1->size(); ++i) {
    T& x = out_1->at(i);
    T& y = out_2->at(i);
    fill_value(&x, &y, is_empty(x), is_empty(y));
  }
  return true;
}

/*
 * This contains the states for running a loop and provides methods
 * of running the subgraph computation for an iteration.
 */
class LoopState {
  // These are output arrays from all iterations.
  // They also contain the Op state for each CachedOp.
  std::vector<std::vector<NDArray> > all_outputs;
  std::vector<std::vector<NDArray> > all_inputs;
  // For inference, there should be only one cached op because we
  // want to share the memory in iterations.
  // For training, each iteration has a cached op because each iteration
  // needs to maintain a set of memory buffers for all computation states,
  // which will be used in the backward.
  std::vector<OpStatePtr> all_states;
  CachedOpPtr iter_op;
  nnvm::Symbol subgraph_sym;
  nnvm::Graph subgraph;

 public:
  explicit LoopState(const nnvm::Symbol& g, bool is_dynamic = true);

  void Forward(int iter_no,
               const std::vector<NDArray>& inputs,
               const std::vector<OpReqType>& req,
               const std::vector<NDArray>& outputs,
               bool is_recording);
  void Backward(int iter_no,
                const std::vector<NDArray>& ograds,
                const std::vector<OpReqType>& req,
                const std::vector<NDArray>& igrads);
  void Cleanup() {
    all_outputs.clear();
    all_inputs.clear();
    all_states.clear();
  }
  static CachedOpPtr MakeSharedOp(const nnvm::Symbol& sym, bool is_dynamic = true) {
    // We turn on static_alloc for two reasons.
    // It avoids the overhead of unnecessary memory allocation.
    // only static_alloc supports nested call of CachedOp.
    std::vector<std::pair<std::string, std::string> > kwargs = {{"inline_limit", "0"},
                                                                {"static_alloc", "1"}};
    if (is_dynamic) {
      kwargs.push_back({"is_dynamic", "1"});
    } else {
      kwargs.push_back({"is_dynamic", "0"});
    }
    return std::make_shared<CachedOp>(sym, kwargs);
  }
};

}  // namespace op
}  // namespace mxnet

#endif  // MXNET_OPERATOR_SUBGRAPH_OP_COMMON_H_