| /* |
| * Copyright (c) 2005-2019, NumPy Developers. |
| * |
| * All rights reserved. |
| * |
| * Redistribution and use in source and binary forms, with or without |
| * modification, are permitted provided that the following conditions are |
| * met: |
| * |
| * * Redistributions of source code must retain the above copyright |
| * notice, this list of conditions and the following disclaimer. |
| * |
| * * Redistributions in binary form must reproduce the above |
| * copyright notice, this list of conditions and the following |
| * disclaimer in the documentation and/or other materials provided |
| * with the distribution. |
| * |
| * * Neither the name of the NumPy Developers nor the names of any |
| * contributors may be used to endorse or promote products derived |
| * from this software without specific prior written permission. |
| * |
| * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS |
| * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT |
| * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR |
| * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT |
| * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, |
| * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT |
| * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, |
| * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY |
| * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT |
| * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE |
| * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. |
| */ |
| |
| /*! |
| * \file np_einsum_op-inl.h |
| * \brief Function definition of numpy-compatible einsum operator |
| * modified by Haozheng Fan(@hzfan) from: |
| * https://github.com/numpy/numpy/blob/master/numpy/core/src/multiarray/einsum.c.src |
| */ |
| |
| #ifndef MXNET_OPERATOR_NUMPY_NP_EINSUM_OP_INL_H_ |
| #define MXNET_OPERATOR_NUMPY_NP_EINSUM_OP_INL_H_ |
| |
| #include <mxnet/operator_util.h> |
| #include <string> |
| #include <vector> |
| #include <algorithm> |
| #include "./np_tensordot_op-inl.h" |
| #include "./np_einsum_path_op-inl.h" |
| #include "../../common/static_array.h" |
| #include "../mxnet_op.h" |
| #include "../operator_common.h" |
| #include "../mshadow_op.h" |
| #include "../elemwise_op_common.h" |
| |
| namespace mxnet { |
| namespace op { |
| |
| #define NPY_MAXDIMS 16 |
| #define NPY_MAXARGS 16 |
| |
| inline TShape get_stride(const TShape& shape) { |
| int ndim = shape.ndim(), prod = 1; |
| TShape stride = TShape(ndim, -1); |
| for (int i = ndim - 1; i >= 0; i--) { |
| stride[i] = shape[i] > 1 ? prod : 0; |
| prod = prod * shape[i]; |
| } |
| return stride; |
| } |
| |
| inline TShape pad(const TShape& shape, int odim) { |
| int ndim = shape.ndim(); |
| CHECK_GE(odim, ndim); |
| TShape ret(odim, 1); |
| for (int idim = 0; idim < ndim; ++idim) { |
| ret[idim] = shape[idim]; |
| } |
| return ret; |
| } |
| |
| /* |
| * Parses the subscripts for one operand into an output of 'ndim' |
| * labels. The resulting 'op_labels' array will have: |
| * - the ASCII code of the label for the first occurrence of a label; |
| * - the (negative) offset to the first occurrence of the label for |
| * repeated labels; |
| * - zero for broadcast dimensions, if subscripts has an ellipsis. |
| * For example: |
| * - subscripts="abbcbc", ndim=6 -> op_labels=[97, 98, -1, 99, -3, -2] |
| * - subscripts="ab...bc", ndim=6 -> op_labels=[97, 98, 0, 0, -3, 99] |
| */ |
| inline int parse_operand_subscripts(const char* subscripts, |
| int length, |
| int ndim, |
| int iop, |
| char* op_labels, |
| char* label_counts, |
| int* min_label, |
| int* max_label) { |
| using namespace mxnet_op; |
| int i; |
| int idim = 0; |
| int ellipsis = -1; |
| |
| /* Process all labels for this operand */ |
| for (i = 0; i < length; ++i) { |
| int label = subscripts[i]; |
| |
| /* A proper label for an axis. */ |
| if (label > 0 && isalpha(label)) { |
| /* Check we don't exceed the operator dimensions. */ |
| CHECK(idim < ndim) << "einstein sum subscripts string contains " |
| << "too many subscripts for operand " << iop; |
| |
| op_labels[idim++] = label; |
| if (label < *min_label) { |
| *min_label = label; |
| } |
| if (label > *max_label) { |
| *max_label = label; |
| } |
| label_counts[label]++; |
| } else if (label == '.') { |
| /* The beginning of the ellipsis. */ |
| /* Check it's a proper ellipsis. */ |
| CHECK( |
| !(ellipsis != -1 || i + 2 >= length || subscripts[++i] != '.' || subscripts[++i] != '.')) |
| << "einstein sum subscripts string contains a " |
| << "'.' that is not part of an ellipsis ('...') " |
| << "in operand " << iop; |
| |
| ellipsis = idim; |
| } else { |
| CHECK(label == ' ') << "invalid subscript '" << static_cast<char>(label) |
| << "' in einstein sum " |
| << "subscripts string, subscripts must " |
| << "be letters"; |
| } |
| } |
| |
| /* No ellipsis found, labels must match dimensions exactly. */ |
| if (ellipsis == -1) { |
| CHECK(idim == ndim) << "operand has more dimensions than subscripts " |
| << "given in einstein sum, but no '...' ellipsis " |
| << "provided to broadcast the extra dimensions."; |
| } else if (idim < ndim) { |
| /* Ellipsis found, may have to add broadcast dimensions. */ |
| /* Move labels after ellipsis to the end. */ |
| for (i = 0; i < idim - ellipsis; ++i) { |
| op_labels[ndim - i - 1] = op_labels[idim - i - 1]; |
| } |
| /* Set all broadcast dimensions to zero. */ |
| for (i = 0; i < ndim - idim; ++i) { |
| op_labels[ellipsis + i] = 0; |
| } |
| } |
| |
| /* |
| * Find any labels duplicated for this operand, and turn them |
| * into negative offsets to the axis to merge with. |
| * |
| * In C, the char type may be signed or unsigned, but with |
| * twos complement arithmetic the char is ok either way here, and |
| * later where it matters the char is cast to a signed char. |
| */ |
| for (idim = 0; idim < ndim - 1; ++idim) { |
| int label = op_labels[idim]; |
| /* If it is a proper label, find any duplicates of it. */ |
| if (label > 0) { |
| /* Search for the next matching label. */ |
| char* next = reinterpret_cast<char*>(memchr(op_labels + idim + 1, label, ndim - idim - 1)); |
| |
| while (next != nullptr) { |
| /* The offset from next to op_labels[idim] (negative). */ |
| *next = static_cast<char>((op_labels + idim) - next); |
| /* Search for the next matching label. */ |
| next = reinterpret_cast<char*>(memchr(next + 1, label, op_labels + ndim - 1 - next)); |
| } |
| } |
| } |
| return 0; |
| } |
| |
| /* |
| * Parses the subscripts for the output operand into an output that |
| * includes 'ndim_broadcast' unlabeled dimensions, and returns the total |
| * number of output dimensions, or -1 if there is an error. Similarly |
| * to parse_operand_subscripts, the 'out_labels' array will have, for |
| * each dimension: |
| * - the ASCII code of the corresponding label; |
| * - zero for broadcast dimensions, if subscripts has an ellipsis. |
| */ |
| inline int parse_output_subscripts(const char* subscripts, |
| int length, |
| int ndim_broadcast, |
| const char* label_counts, |
| char* out_labels) { |
| using namespace mxnet_op; |
| int i, bdim; |
| int ndim = 0; |
| int ellipsis = 0; |
| |
| /* Process all the output labels. */ |
| for (i = 0; i < length; ++i) { |
| int label = subscripts[i]; |
| |
| /* A proper label for an axis. */ |
| if (label > 0 && isalpha(label)) { |
| /* Check that it doesn't occur again. */ |
| CHECK(memchr(subscripts + i + 1, label, length - i - 1) == nullptr) |
| << "einstein sum subscripts string includes " |
| << "output subscript '" << static_cast<char>(label) << "' multiple times"; |
| |
| /* Check that it was used in the inputs. */ |
| CHECK(label_counts[label] != 0) |
| << "einstein sum subscripts string included " |
| << "output subscript '" << static_cast<char>(label) << "' which never appeared " |
| << "in an input"; |
| |
| /* Check that there is room in out_labels for this label. */ |
| CHECK(ndim < NPY_MAXDIMS) << "einstein sum subscripts string contains " |
| << "too many subscripts in the output"; |
| |
| out_labels[ndim++] = label; |
| } else if (label == '.') { |
| /* The beginning of the ellipsis. */ |
| /* Check it is a proper ellipsis. */ |
| CHECK(!(ellipsis || i + 2 >= length || subscripts[++i] != '.' || subscripts[++i] != '.')) |
| << "einstein sum subscripts string " |
| << "contains a '.' that is not part of " |
| << "an ellipsis ('...') in the output"; |
| |
| /* Check there is room in out_labels for broadcast dims. */ |
| CHECK(ndim + ndim_broadcast <= NPY_MAXDIMS) << "einstein sum subscripts string contains " |
| << "too many subscripts in the output"; |
| |
| ellipsis = 1; |
| for (bdim = 0; bdim < ndim_broadcast; ++bdim) { |
| out_labels[ndim++] = 0; |
| } |
| } else { |
| CHECK(label == ' ') << "invalid subscript '" << static_cast<char>(label) |
| << "' in einstein sum " |
| << "subscripts string, subscripts must " |
| << "be letters"; |
| } |
| } |
| |
| /* If no ellipsis was found there should be no broadcast dimensions. */ |
| CHECK(!(!ellipsis && ndim_broadcast > 0)) << "output has more dimensions than subscripts " |
| << "given in einstein sum, but no '...' ellipsis " |
| << "provided to broadcast the extra dimensions."; |
| |
| return ndim; |
| } |
| |
| inline void get_combined_dims_view(const TBlob& op, |
| int iop, |
| char* labels, |
| TShape* newshape, |
| TShape* newstride) { |
| using namespace mxnet_op; |
| int idim, ndim, icombine, combineoffset; |
| int icombinemap[NPY_MAXDIMS]; |
| int newdim; |
| |
| const TShape& shape = op.shape_; |
| TShape stride = get_stride(shape); |
| ndim = op.shape_.ndim(); |
| newdim = newshape->ndim(); |
| |
| /* Initialize the dimensions and strides to zero */ |
| for (idim = 0; idim < newdim; ++idim) { |
| (*newshape)[idim] = 0; |
| (*newstride)[idim] = 0; |
| } |
| |
| /* Copy the dimensions and strides, except when collapsing */ |
| icombine = 0; |
| for (idim = 0; idim < ndim; ++idim) { |
| /* |
| * The char type may be either signed or unsigned, we |
| * need it to be signed here. |
| */ |
| int label = (signed char)labels[idim]; |
| /* If this label says to merge axes, get the actual label */ |
| if (label < 0) { |
| combineoffset = label; |
| label = labels[idim + label]; |
| } else { |
| combineoffset = 0; |
| if (icombine != idim) { |
| labels[icombine] = labels[idim]; |
| } |
| icombinemap[idim] = icombine; |
| } |
| /* If the label is 0, it's an unlabeled broadcast dimension */ |
| if (label == 0) { |
| (*newshape)[icombine] = shape[idim]; |
| (*newstride)[icombine] = stride[idim]; |
| } else { |
| /* Update the combined axis dimensions and strides */ |
| int i = icombinemap[idim + combineoffset]; |
| CHECK(!(combineoffset < 0 && (*newshape)[i] != 0 && (*newshape)[i] != shape[idim])) |
| << "dimensions in operand " << iop << " for collapsing index '" << label |
| << "' don't match (" << static_cast<int>((*newshape)[i]) << " != " << shape[idim] << ")"; |
| (*newshape)[i] = shape[idim]; |
| (*newstride)[i] += stride[idim]; |
| } |
| |
| /* If the label didn't say to combine axes, increment dest i */ |
| if (combineoffset == 0) { |
| icombine++; |
| } |
| } |
| } |
| |
| inline static int prepare_op_axes(int ndim, |
| int iop, |
| char* labels, |
| int* axes, |
| int ndim_iter, |
| char* iter_labels) { |
| using namespace mxnet_op; |
| int i, label, ibroadcast; |
| |
| ibroadcast = ndim - 1; |
| for (i = ndim_iter - 1; i >= 0; --i) { |
| label = iter_labels[i]; |
| /* |
| * If it's an unlabeled broadcast dimension, choose |
| * the next broadcast dimension from the operand. |
| */ |
| if (label == 0) { |
| while (ibroadcast >= 0 && labels[ibroadcast] != 0) { |
| --ibroadcast; |
| } |
| /* |
| * If we used up all the operand broadcast dimensions, |
| * extend it with a "newaxis" |
| */ |
| if (ibroadcast < 0) { |
| axes[i] = -1; |
| } else { |
| /* Otherwise map to the broadcast axis */ |
| axes[i] = ibroadcast; |
| --ibroadcast; |
| } |
| } else { |
| /* It's a labeled dimension, find the matching one */ |
| char* match = reinterpret_cast<char*>(memchr(labels, label, ndim)); |
| /* If the op doesn't have the label, broadcast it */ |
| if (match == nullptr) { |
| axes[i] = -1; |
| } else { |
| /* Otherwise use it */ |
| axes[i] = match - labels; |
| } |
| } |
| } |
| return 0; |
| } |
| |
| struct NumpyEinsumParam : public dmlc::Parameter<NumpyEinsumParam> { |
| int num_args; |
| int optimize; |
| std::string subscripts; |
| DMLC_DECLARE_PARAMETER(NumpyEinsumParam) { |
| DMLC_DECLARE_FIELD(num_args).set_lower_bound(1).describe("Number of input arrays."); |
| DMLC_DECLARE_FIELD(subscripts) |
| .set_default("") |
| .describe( |
| "Specifies the subscripts for summation as comma separated list" |
| " of subscript labels. An implicit (classical Einstein summation) calculation" |
| " is performed unless the explicit indicator '->' is included as well as" |
| " subscript labels of the precise output form."); |
| DMLC_DECLARE_FIELD(optimize).set_default(0); |
| } |
| void SetAttrDict(std::unordered_map<std::string, std::string>* dict) { |
| std::ostringstream num_args_s, optimize_s, subscripts_s; |
| num_args_s << num_args; |
| optimize_s << optimize; |
| subscripts_s << subscripts; |
| (*dict)["num_args"] = num_args_s.str(); |
| (*dict)["optimize"] = optimize_s.str(); |
| (*dict)["subscripts"] = subscripts_s.str(); |
| } |
| }; |
| |
| class EinsumOp { |
| public: |
| int num_args; |
| int optimize; |
| std::string subscripts; |
| std::shared_ptr<NDArray> tempspace; |
| std::vector<Step> paths; |
| explicit EinsumOp(int num_args, int optimize, std::string subscripts) { |
| this->num_args = num_args; |
| this->optimize = optimize; |
| this->subscripts = subscripts; |
| } |
| bool operator==(const EinsumOp& other) const { |
| return this->num_args == other.num_args && !this->subscripts.compare(other.subscripts) && |
| this->optimize == other.optimize; |
| } |
| }; // class EinsumOp |
| |
| template <int dimension, int req, bool back, typename AType> |
| struct numpy_einsum { |
| template <typename DType> |
| MSHADOW_XINLINE static void Map( |
| index_t i, |
| DType* out, |
| common::StaticArray<DType*, NPY_MAXARGS> op, |
| mshadow::Shape<dimension> oshape, |
| common::StaticArray<mshadow::Shape<dimension>, NPY_MAXARGS> ostride, |
| mshadow::Shape<dimension> reduceshape, |
| common::StaticArray<mshadow::Shape<dimension>, NPY_MAXARGS> rstride, |
| int nop, |
| int iop0, |
| const DType* out_grad) { |
| using namespace mxnet_op; |
| mshadow::Shape<dimension> oidx = unravel(i, oshape); |
| i = back ? dot(oidx, ostride[iop0]) : i; |
| if (req == kWriteTo) { |
| out[i] = (DType)0; |
| } |
| for (int rdim = 0; rdim < dimension; ++rdim) { |
| if (reduceshape[rdim] == 0) { |
| return; |
| } |
| } |
| mshadow::Shape<dimension> ridx = unravel(0, reduceshape); |
| AType sum = 0; |
| do { |
| AType tmp = |
| back ? static_cast<AType>(out_grad[dot(oidx, ostride[nop]) + dot(ridx, rstride[nop])]) |
| : (AType)1; |
| for (int iop = 0; iop < nop; ++iop) { |
| if (iop != iop0) { |
| index_t k = dot(oidx, ostride[iop]) + dot(ridx, rstride[iop]); |
| tmp = tmp * static_cast<AType>(op[iop][k]); |
| } |
| } |
| sum = sum + tmp; |
| } while (inc(&ridx, reduceshape)); |
| out[i] = out[i] + static_cast<DType>(sum); |
| } |
| }; |
| |
| template <typename xpu, bool back> |
| inline void NumpyEinsumProcess(const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs, |
| const char* subscripts, |
| int nop, |
| const OpContext& ctx) { |
| using namespace mxnet_op; |
| |
| /* nop+1 (+1 is for the output) must fit in NPY_MAXARGS */ |
| CHECK(nop < NPY_MAXARGS) << "too many operands provided to einstein sum function"; |
| CHECK(nop >= 1) << "not enough operands provided to einstein sum function"; |
| |
| /* Step 1: Parse the subscripts string into label_counts and op_labels */ |
| int iop, idim, min_label = 127, max_label = 0; |
| char label_counts[128], op_labels[NPY_MAXARGS][NPY_MAXDIMS]; |
| memset(label_counts, 0, sizeof(label_counts)); |
| for (iop = 0; iop < nop; ++iop) { |
| int length = static_cast<int>(strcspn(subscripts, ",-")); |
| |
| CHECK(!(iop == nop - 1 && subscripts[length] == ',')) |
| << "more operands provided to einstein sum function " |
| << "than specified in the subscripts string"; |
| CHECK(!(iop < nop - 1 && subscripts[length] != ',')) |
| << "fewer operands provided to einstein sum function " |
| << "than specified in the subscripts string"; |
| CHECK_GE(parse_operand_subscripts(subscripts, |
| length, |
| inputs[iop + back].shape_.ndim(), |
| iop, |
| op_labels[iop], |
| label_counts, |
| &min_label, |
| &max_label), |
| 0); |
| |
| /* Move subscripts to the start of the labels for the next op */ |
| subscripts += length; |
| if (iop < nop - 1) { |
| subscripts++; |
| } |
| } |
| |
| /* |
| * Find the number of broadcast dimensions, which is the maximum |
| * number of labels == 0 in an op_labels array. |
| */ |
| int ndim_broadcast = 0; |
| for (iop = 0; iop < nop; ++iop) { |
| int count_zeros = 0; |
| int ndim; |
| char* labels = op_labels[iop]; |
| |
| ndim = inputs[iop + back].shape_.ndim(); |
| for (idim = 0; idim < ndim; ++idim) { |
| if (labels[idim] == 0) { |
| ++count_zeros; |
| } |
| } |
| |
| if (count_zeros > ndim_broadcast) { |
| ndim_broadcast = count_zeros; |
| } |
| } |
| |
| /* |
| * If there is no output signature, fill output_labels and ndim_output |
| * using each label that appeared once, in alphabetical order. |
| */ |
| int label, ndim_output; |
| char output_labels[NPY_MAXDIMS]; |
| if (subscripts[0] == '\0') { |
| /* If no output was specified, always broadcast left, as usual. */ |
| for (ndim_output = 0; ndim_output < ndim_broadcast; ++ndim_output) { |
| output_labels[ndim_output] = 0; |
| } |
| for (label = min_label; label <= max_label; ++label) { |
| if (label_counts[label] == 1) { |
| CHECK(ndim_output < NPY_MAXDIMS) << "einstein sum subscript string has too many " |
| << "distinct labels"; |
| output_labels[ndim_output++] = label; |
| } |
| } |
| } else { |
| CHECK(subscripts[0] == '-' && subscripts[1] == '>') << "einstein sum subscript string does not " |
| << "contain proper '->' output specified"; |
| subscripts += 2; |
| |
| /* Parse the output subscript string. */ |
| ndim_output = parse_output_subscripts( |
| subscripts, strlen(subscripts), ndim_broadcast, label_counts, output_labels); |
| CHECK_GE(ndim_output, 0); |
| } |
| |
| /* |
| * Step 2: |
| * Process all the input ops, combining dimensions into their |
| * diagonal where specified. |
| */ |
| std::vector<TShape> opshape(nop), opstride_true(nop); |
| for (iop = 0; iop < nop; ++iop) { |
| char* labels = op_labels[iop]; |
| int combine, ndim; |
| |
| ndim = inputs[iop + back].shape_.ndim(); |
| |
| /* |
| * Check whether any dimensions need to be combined |
| * |
| * The char type may be either signed or unsigned, we |
| * need it to be signed here. |
| */ |
| combine = 0; |
| for (idim = 0; idim < ndim; ++idim) { |
| if ((signed char)labels[idim] < 0) { |
| combine++; |
| } |
| } |
| |
| /* If any dimensions are combined, create a view which combines them */ |
| if (combine) { |
| TShape tshape(ndim - combine, -1); |
| TShape tstride(ndim - combine, -1); |
| get_combined_dims_view(inputs[iop + back], iop, labels, &tshape, &tstride); |
| opshape[iop] = tshape; |
| opstride_true[iop] = tstride; |
| } else { |
| /* No combining needed */ |
| opshape[iop] = inputs[iop + back].shape_; |
| opstride_true[iop] = get_stride(opshape[iop]); |
| } |
| } |
| |
| /* |
| * Step 3: |
| * Set up the labels for the iterator (output + combined labels). |
| * Can just share the output_labels memory, because iter_labels |
| * is output_labels with some more labels appended. |
| */ |
| char* iter_labels = output_labels; |
| int ndim_iter = ndim_output; |
| for (label = min_label; label <= max_label; ++label) { |
| if (label_counts[label] > 0 && memchr(output_labels, label, ndim_output) == nullptr) { |
| CHECK(ndim_iter < NPY_MAXDIMS) << "too many subscripts in einsum"; |
| iter_labels[ndim_iter++] = label; |
| } |
| } |
| |
| /* Step 4: Set up the op_axes for the iterator */ |
| TShape itershape(ndim_iter, -1); |
| std::vector<TShape> iterstride(nop + 1, TShape(ndim_iter, 0)); |
| TShape oshape = back ? inputs[0].shape_ : outputs[0].shape_; |
| TShape ostride_true = get_stride(oshape); |
| TShape reduceshape; |
| std::vector<TShape> remainshape(nop); |
| int op_axes_arrays[NPY_MAXARGS][NPY_MAXDIMS]; |
| int* op_axes[NPY_MAXARGS]; |
| |
| for (iop = 0; iop < nop; ++iop) { |
| op_axes[iop] = op_axes_arrays[iop]; |
| CHECK_GE(prepare_op_axes( |
| opshape[iop].ndim(), iop, op_labels[iop], op_axes[iop], ndim_iter, iter_labels), |
| 0); |
| for (idim = 0; idim < ndim_iter; idim++) { |
| if (op_axes[iop][idim] != -1) { |
| iterstride[iop][idim] = opstride_true[iop][op_axes[iop][idim]]; |
| if (itershape[idim] != -1) { |
| if (itershape[idim] == 1) { |
| itershape[idim] = opshape[iop][op_axes[iop][idim]]; |
| } |
| } else { |
| itershape[idim] = opshape[iop][op_axes[iop][idim]]; |
| } |
| } |
| } |
| } |
| for (idim = 0; idim < ndim_output; ++idim) { |
| iterstride[nop][idim] = ostride_true[idim]; |
| } |
| reduceshape = TShape(ndim_iter - ndim_output, 0); |
| for (idim = ndim_output; idim < ndim_iter; ++idim) { |
| reduceshape[idim - ndim_output] = itershape[idim]; |
| } |
| for (iop = 0; iop < nop; iop++) { |
| std::vector<size_t> rsh; |
| for (idim = 0; idim < ndim_iter; idim++) { |
| if (op_axes_arrays[iop][idim] == -1 || |
| itershape[idim] != opshape[iop][op_axes_arrays[iop][idim]]) { |
| rsh.push_back(itershape[idim]); |
| } |
| } |
| remainshape[iop] = TShape(rsh.begin(), rsh.end()); |
| } |
| |
| // exclude the 0-dim case |
| if (ndim_iter == 0) { |
| ndim_iter = 1; |
| } |
| itershape = pad(itershape, ndim_iter); |
| for (iop = 0; iop <= nop; ++iop) { |
| iterstride[iop] = pad(iterstride[iop], ndim_iter); |
| } |
| oshape = pad(oshape, ndim_iter); |
| reduceshape = pad(reduceshape, ndim_iter); |
| for (iop = 0; iop < nop; ++iop) { |
| opshape[iop] = pad(opshape[iop], ndim_iter); |
| remainshape[iop] = pad(remainshape[iop], ndim_iter); |
| } |
| |
| if (!back) { |
| if (oshape.Size() == 0) { |
| return; |
| } |
| const TBlob& out_data = outputs[0]; |
| MXNET_ACC_TYPE_SWITCH(out_data.type_flag_, DType, AType, { |
| mxnet::common::StaticArray<DType*, NPY_MAXARGS> op; |
| for (iop = 0; iop < nop; ++iop) { |
| op[iop] = inputs[iop].dptr<DType>(); |
| } |
| MXNET_ASSIGN_REQ_SWITCH( |
| req[0], req_type, {MXNET_NDIM_SWITCH_EX(ndim_iter, dimension, { |
| mxnet::common::StaticArray<mshadow::Shape<dimension>, NPY_MAXARGS> ostride_arr; |
| mxnet::common::StaticArray<mshadow::Shape<dimension>, NPY_MAXARGS> rstride_arr; |
| for (iop = 0; iop < nop; ++iop) { |
| mshadow::Shape<dimension> otmp, rtmp; |
| for (idim = 0; idim < dimension; ++idim) { |
| otmp[idim] = idim < ndim_output ? iterstride[iop][idim] : 1; |
| rtmp[idim] = |
| idim < dimension - ndim_output ? iterstride[iop][idim + ndim_output] : 1; |
| } |
| ostride_arr[iop] = otmp; |
| rstride_arr[iop] = rtmp; |
| } |
| Kernel<numpy_einsum<dimension, req_type, 0, AType>, xpu>::Launch( |
| ctx.get_stream<xpu>(), |
| oshape.Size(), |
| out_data.dptr<DType>(), |
| op, |
| oshape.get<dimension>(), |
| ostride_arr, |
| reduceshape.get<dimension>(), |
| rstride_arr, |
| nop, |
| -1, |
| reinterpret_cast<DType*>(NULL)); |
| })}) |
| }) |
| } else { |
| if (oshape.Size() == 0) { |
| for (iop = 0; iop < nop; ++iop) { |
| const TBlob& out_data = outputs[iop]; |
| if (opshape[iop].Size() > 0) { |
| MSHADOW_TYPE_SWITCH( |
| out_data.type_flag_, DType, {MXNET_ASSIGN_REQ_SWITCH(req[iop], req_type, { |
| if (req_type == kWriteTo) { |
| out_data.FlatTo1D<xpu, DType>(ctx.get_stream<xpu>()) = 0; |
| } |
| })}) |
| } |
| } |
| return; |
| } |
| for (int i = 0; i < nop; ++i) { |
| const TBlob& out_data = outputs[i]; |
| const TBlob& out_grad = inputs[0]; |
| std::vector<TShape> opstride(nop + 1, TShape(ndim_iter, 0)); |
| std::vector<TShape> remainstride(nop + 1, TShape(ndim_iter, 0)); |
| for (iop = 0; iop <= nop; ++iop) { |
| int j = 0; |
| for (idim = 0; idim < ndim_iter; ++idim) { |
| if (op_axes_arrays[i][idim] == -1 || |
| (iop != nop && opshape[i][op_axes_arrays[i][idim]] == 1 && |
| op_axes_arrays[iop][idim] != -1 && opshape[iop][op_axes_arrays[iop][idim]] != 1)) { |
| remainstride[iop][j++] = iterstride[iop][idim]; |
| } else { |
| opstride[iop][op_axes_arrays[i][idim]] = iterstride[iop][idim]; |
| } |
| } |
| } |
| MXNET_ACC_TYPE_SWITCH(out_data.type_flag_, DType, AType, { |
| mxnet::common::StaticArray<DType*, NPY_MAXARGS> op; |
| for (iop = 0; iop < nop; ++iop) { |
| op[iop] = inputs[iop + back].dptr<DType>(); |
| } |
| MXNET_ASSIGN_REQ_SWITCH( |
| req[i], req_type, {MXNET_NDIM_SWITCH_EX(ndim_iter, dimension, { |
| mxnet::common::StaticArray<mshadow::Shape<dimension>, NPY_MAXARGS> opstride_arr; |
| mxnet::common::StaticArray<mshadow::Shape<dimension>, NPY_MAXARGS> remainstride_arr; |
| for (iop = 0; iop <= nop; ++iop) { |
| opstride_arr[iop] = opstride[iop].get<dimension>(); |
| remainstride_arr[iop] = remainstride[iop].get<dimension>(); |
| } |
| Kernel<numpy_einsum<dimension, req_type, 1, AType>, xpu>::Launch( |
| ctx.get_stream<xpu>(), |
| opshape[i].Size(), |
| out_data.dptr<DType>(), |
| op, |
| opshape[i].get<dimension>(), |
| opstride_arr, |
| remainshape[i].get<dimension>(), |
| remainstride_arr, |
| nop, |
| i, |
| out_grad.dptr<DType>()); |
| })}) |
| }) |
| } |
| } |
| } |
| |
| template <typename xpu> |
| inline void NumpyEinsumForward(const OpStatePtr& state_ptr, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| using namespace mshadow; |
| using namespace mxnet_op; |
| EinsumOp& state = state_ptr.get_state<EinsumOp>(); |
| int num_args = state.num_args; |
| int optimize = state.optimize; |
| const char* subscripts = state.subscripts.c_str(); |
| Stream<xpu>* s = ctx.get_stream<xpu>(); |
| CHECK_EQ(inputs.size(), num_args); |
| CHECK_EQ(outputs.size(), 1U); |
| if (optimize == 0) { |
| NumpyEinsumProcess<xpu, 0>(inputs, req, outputs, subscripts, num_args, ctx); |
| return; |
| } |
| std::vector<Step>& paths = state.paths; |
| std::vector<std::vector<int> > pos; |
| std::string string_repr; |
| paths = einsum_path(state.subscripts, inputs, true, ctx.run_ctx, &pos, &string_repr); |
| int paths_len = paths.size(); |
| size_t temp_space_size = 0, max_temp_space_size = 0; |
| std::vector<TBlob> operands(inputs), tmp_operands, temp_space_vec(paths_len - 1); |
| for (int i = 0; i + 1 < paths_len; ++i) { |
| temp_space_size += paths[i].oshape.Size(); |
| } |
| for (int i = 0; i < paths_len; ++i) { |
| max_temp_space_size = std::max(max_temp_space_size, paths[i].oshape.Size()); |
| } |
| temp_space_size += max_temp_space_size; |
| MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, { |
| state.tempspace.reset<NDArray>(new NDArray( |
| TShape(Shape1(temp_space_size)), ctx.run_ctx.ctx, false, outputs[0].type_flag_)); |
| Tensor<xpu, 1, DType> temp_space = state.tempspace->data().FlatTo1D<xpu, DType>(); |
| size_t begin = max_temp_space_size; |
| for (int i = 0; i < paths_len - 1; ++i) { |
| TBlob tblob = TBlob(temp_space.Slice(begin, begin + paths[i].oshape.Size())); |
| temp_space_vec[i] = tblob.reshape(paths[i].oshape); |
| begin = begin + paths[i].oshape.Size(); |
| } |
| for (int i = 0; i < paths_len; ++i) { |
| tmp_operands.clear(); |
| |
| // We remove inds from right to left |
| for (const int& p : paths[i].contract_inds) { |
| tmp_operands.push_back(operands[p]); |
| operands.erase(operands.begin() + p); |
| } |
| bool handle_out = (i == paths_len - 1); |
| // Call tensordot if still possible |
| if (paths[i].do_blas) { |
| // Contract! |
| if (paths[i].do_einsum || handle_out) { |
| TBlob max_temp_space = TBlob(temp_space.Slice(0, paths[i].tshape.Size())); |
| max_temp_space.FlatTo1D<xpu, DType>(s) = 0; |
| max_temp_space = max_temp_space.reshape(paths[i].tshape); |
| size_t tensordot_tempspace_size = |
| TensordotWorkspaceSize<xpu>(paths[i].left_pos, |
| paths[i].right_pos, |
| tmp_operands[0], |
| tmp_operands[1], |
| max_temp_space, |
| std::vector<OpReqType>{OpReqType::kWriteTo}); |
| Tensor<xpu, 1, char> tensordot_tempspace = |
| ctx.requested[0].get_space_typed<xpu, 1, char>(Shape1(tensordot_tempspace_size), s); |
| TensordotImpl<xpu>(paths[i].left_pos, |
| paths[i].right_pos, |
| ctx, |
| tmp_operands[0], |
| tmp_operands[1], |
| max_temp_space, |
| std::vector<OpReqType>{OpReqType::kWriteTo}, |
| tensordot_tempspace); |
| NumpyEinsumProcess<xpu, 0>(std::vector<TBlob>{max_temp_space}, |
| handle_out ? req : std::vector<OpReqType>{OpReqType::kWriteTo}, |
| handle_out ? outputs : std::vector<TBlob>{temp_space_vec[i]}, |
| paths[i].blas2einsum_str.c_str(), |
| 1, |
| ctx); |
| } else { |
| size_t tensordot_tempspace_size = |
| TensordotWorkspaceSize<xpu>(paths[i].left_pos, |
| paths[i].right_pos, |
| tmp_operands[0], |
| tmp_operands[1], |
| temp_space_vec[i], |
| std::vector<OpReqType>{OpReqType::kWriteTo}); |
| Tensor<xpu, 1, char> tensordot_tempspace = |
| ctx.requested[0].get_space_typed<xpu, 1, char>(Shape1(tensordot_tempspace_size), s); |
| TensordotImpl<xpu>(paths[i].left_pos, |
| paths[i].right_pos, |
| ctx, |
| tmp_operands[0], |
| tmp_operands[1], |
| temp_space_vec[i], |
| std::vector<OpReqType>{OpReqType::kWriteTo}, |
| tensordot_tempspace); |
| } |
| } else { |
| NumpyEinsumProcess<xpu, 0>(tmp_operands, |
| handle_out ? req : std::vector<OpReqType>{OpReqType::kWriteTo}, |
| handle_out ? outputs : std::vector<TBlob>{temp_space_vec[i]}, |
| paths[i].einsum_str.c_str(), |
| tmp_operands.size(), |
| ctx); |
| } |
| if (!handle_out) { |
| operands.push_back(temp_space_vec[i]); |
| } |
| } |
| }); |
| } |
| |
| template <typename xpu> |
| inline void NumpyEinsumBackward(const OpStatePtr& state_ptr, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| using namespace mshadow; |
| using namespace mshadow_op; |
| const EinsumOp& state = state_ptr.get_state<EinsumOp>(); |
| int num_args = state.num_args; |
| int optimize = state.optimize; |
| const char* subscripts = state.subscripts.c_str(); |
| Stream<xpu>* s = ctx.get_stream<xpu>(); |
| CHECK_EQ(inputs.size(), 1 + num_args); |
| CHECK_EQ(outputs.size(), num_args); |
| if (optimize == 0) { |
| NumpyEinsumProcess<xpu, 1>(inputs, req, outputs, subscripts, num_args, ctx); |
| return; |
| } |
| // calculate temporary space size for temp_grad |
| const std::vector<Step>& paths = state.paths; |
| int paths_len = paths.size(); |
| size_t temp_space_size = 0, max_temp_space_size = 0; |
| for (int i = 0; i < paths_len - 1; ++i) { |
| temp_space_size += paths[i].oshape.Size(); |
| } |
| for (int i = 0; i < paths_len; ++i) { |
| max_temp_space_size = std::max(max_temp_space_size, paths[i].oshape.Size()); |
| } |
| temp_space_size += max_temp_space_size; |
| // replay the forward process |
| std::vector<std::vector<int> > op_idx(paths_len + 1); |
| for (int i = 0; i <= paths_len; ++i) { |
| if (i == 0) { |
| op_idx[i].reserve(num_args); |
| for (int j = 0; j < num_args; ++j) { |
| op_idx[i].push_back(j + 1); |
| } |
| } else { |
| op_idx[i] = op_idx[i - 1]; |
| // We remove inds from right to left |
| for (const int& p : paths[i - 1].contract_inds) { |
| op_idx[i].erase(op_idx[i].begin() + p); |
| } |
| op_idx[i].push_back(-static_cast<int>(i - 1)); |
| } |
| } |
| // calculate temporary space size for tensordot |
| size_t tensordot_max_tempspace_size = 0; |
| size_t begin_tensordot_tempspace = 0; |
| std::vector<TBlob> temp_inputs, temp_outputs; |
| std::vector<OpReqType> temp_req; |
| std::vector<size_t> tensordot_tempspace_size; |
| MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, { |
| for (int i = 0; i < paths_len; ++i) { |
| temp_inputs.clear(); |
| temp_outputs.clear(); |
| temp_req.clear(); |
| bool handle_out = (i == paths_len - 1); |
| |
| if (handle_out) { |
| temp_inputs.push_back(inputs[0]); |
| } else { |
| temp_inputs.push_back( |
| TBlob(reinterpret_cast<DType*>(NULL), paths[i].oshape, xpu::kDevMask)); |
| } |
| for (auto p : paths[i].contract_inds) { |
| int idx = op_idx[i][p]; |
| if (idx >= 1) { |
| temp_inputs.push_back(inputs[idx]); |
| temp_outputs.push_back(outputs[idx - 1]); |
| temp_req.push_back(req[idx - 1]); |
| } else { |
| temp_inputs.push_back( |
| TBlob(reinterpret_cast<DType*>(NULL), paths[-idx].oshape, xpu::kDevMask)); |
| temp_outputs.push_back( |
| TBlob(reinterpret_cast<DType*>(NULL), paths[-idx].oshape, xpu::kDevMask)); |
| temp_req.push_back(OpReqType::kWriteTo); |
| } |
| } |
| size_t cur_tensordot_tempspace_size = 0; |
| if (paths[i].do_blas) { |
| if (paths[i].do_einsum) { |
| cur_tensordot_tempspace_size = TensordotBackwardWorkspaceSize<xpu>( |
| paths[i].left_pos, |
| paths[i].right_pos, |
| TBlob(reinterpret_cast<DType*>(NULL), paths[i].tshape, xpu::kDevMask), |
| temp_inputs[1], |
| temp_inputs[2], |
| temp_outputs[0], |
| temp_outputs[1], |
| temp_req); |
| } else { |
| cur_tensordot_tempspace_size = TensordotBackwardWorkspaceSize<xpu>(paths[i].left_pos, |
| paths[i].right_pos, |
| temp_inputs[0], |
| temp_inputs[1], |
| temp_inputs[2], |
| temp_outputs[0], |
| temp_outputs[1], |
| temp_req); |
| } |
| } |
| tensordot_tempspace_size.push_back(cur_tensordot_tempspace_size); |
| tensordot_max_tempspace_size = |
| std::max(tensordot_max_tempspace_size, cur_tensordot_tempspace_size); |
| } |
| begin_tensordot_tempspace = temp_space_size; |
| temp_space_size += (tensordot_max_tempspace_size + sizeof(DType) - 1) / sizeof(DType); |
| }); |
| // allocate temporary space and propagate |
| std::vector<TBlob> temp_grad(paths_len - 1), temp_data(paths_len - 1); |
| MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, { |
| // allocate temporary space for gradients of intermediate results |
| Tensor<xpu, 1, DType> temp_space = |
| ctx.requested[0].get_space_typed<xpu, 1, DType>(Shape1(temp_space_size), s); |
| size_t begin = max_temp_space_size; |
| for (int i = 0; i + 1 < paths_len; ++i) { |
| TBlob tblob = TBlob(temp_space.Slice(begin, begin + paths[i].oshape.Size())); |
| temp_grad[i] = tblob.reshape(paths[i].oshape); |
| begin = begin + paths[i].oshape.Size(); |
| } |
| |
| // reinterprete ndarray for intermediate results |
| Tensor<xpu, 1, DType> ndarray_space = state.tempspace->data().FlatTo1D<xpu, DType>(); |
| begin = max_temp_space_size; |
| for (int i = 0; i + 1 < paths_len; ++i) { |
| TBlob tblob = TBlob(ndarray_space.Slice(begin, begin + paths[i].oshape.Size())); |
| temp_data[i] = tblob.reshape(paths[i].oshape); |
| begin = begin + paths[i].oshape.Size(); |
| } |
| |
| // go through the paths in the reversed order |
| for (int i = paths_len - 1; i >= 0; i--) { |
| temp_inputs.clear(); |
| temp_outputs.clear(); |
| temp_req.clear(); |
| bool handle_out = (i == paths_len - 1); |
| |
| if (handle_out) { |
| temp_inputs.push_back(inputs[0]); |
| } else { |
| temp_inputs.push_back(temp_grad[i]); |
| } |
| for (auto p : paths[i].contract_inds) { |
| int idx = op_idx[i][p]; |
| if (idx >= 1) { |
| temp_inputs.push_back(inputs[idx]); |
| temp_outputs.push_back(outputs[idx - 1]); |
| temp_req.push_back(req[idx - 1]); |
| } else { |
| temp_inputs.push_back(temp_data[-idx]); |
| temp_outputs.push_back(temp_grad[-idx]); |
| temp_req.push_back(OpReqType::kWriteTo); |
| } |
| } |
| if (paths[i].do_blas) { |
| CHECK_EQ(temp_inputs.size(), 3U); |
| CHECK_EQ(temp_outputs.size(), 2U); |
| CHECK_EQ(temp_req.size(), 2U); |
| Tensor<xpu, 1, DType> tensordot_tempspace = |
| temp_space.Slice(begin_tensordot_tempspace, temp_space_size); |
| Tensor<xpu, 1, char> char_tempspace = |
| Tensor<xpu, 1, char>(reinterpret_cast<char*>(tensordot_tempspace.dptr_), |
| Shape1(tensordot_tempspace_size[i]), |
| tensordot_tempspace.stream_); |
| if (paths[i].do_einsum) { |
| TBlob max_temp_space = TBlob(temp_space.Slice(0, paths[i].tshape.Size())); |
| max_temp_space = max_temp_space.reshape(paths[i].tshape); |
| NumpyEinsumProcess<xpu, 0>(std::vector<TBlob>{temp_inputs[0]}, |
| std::vector<OpReqType>{kWriteTo}, |
| std::vector<TBlob>{max_temp_space}, |
| paths[i].einsum2blas_str.c_str(), |
| 1, |
| ctx); |
| TensordotBackwardImpl<xpu>(paths[i].left_pos, |
| paths[i].right_pos, |
| ctx, |
| max_temp_space, |
| temp_inputs[1], |
| temp_inputs[2], |
| temp_outputs[0], |
| temp_outputs[1], |
| temp_req, |
| char_tempspace); |
| } else { |
| TensordotBackwardImpl<xpu>(paths[i].left_pos, |
| paths[i].right_pos, |
| ctx, |
| temp_inputs[0], |
| temp_inputs[1], |
| temp_inputs[2], |
| temp_outputs[0], |
| temp_outputs[1], |
| temp_req, |
| char_tempspace); |
| } |
| } else { |
| NumpyEinsumProcess<xpu, 1>(temp_inputs, |
| temp_req, |
| temp_outputs, |
| paths[i].einsum_str.c_str(), |
| temp_outputs.size(), |
| ctx); |
| } |
| } |
| }); |
| } |
| |
| } // namespace op |
| } // namespace mxnet |
| |
| #endif // MXNET_OPERATOR_NUMPY_NP_EINSUM_OP_INL_H_ |
