| /*! |
| * Copyright (c) 2017 by Contributors |
| * \file la_op.h |
| * \brief Operators for advanced linear algebra. |
| */ |
| #ifndef MXNET_OPERATOR_TENSOR_LA_OP_H_ |
| #define MXNET_OPERATOR_TENSOR_LA_OP_H_ |
| |
| #include <mxnet/operator_util.h> |
| #include <vector> |
| #include <algorithm> |
| #include "../mshadow_op.h" |
| #include "../mxnet_op.h" |
| #include "../operator_common.h" |
| #include "../elemwise_op_common.h" |
| |
| namespace mxnet { |
| namespace op { |
| |
| // Parameters for general matrix-matrix multiply-accumulate (mac) |
| struct LaMatrixMacParam : public dmlc::Parameter<LaMatrixMacParam> { |
| bool transpose_a, transpose_b; |
| double alpha, beta; |
| DMLC_DECLARE_PARAMETER(LaMatrixMacParam) { |
| DMLC_DECLARE_FIELD(transpose_a) |
| .set_default(false) |
| .describe("Multiply with transposed of first input (A)."); |
| DMLC_DECLARE_FIELD(transpose_b) |
| .set_default(false) |
| .describe("Multiply with transposed of second input (B)."); |
| DMLC_DECLARE_FIELD(alpha) |
| .set_default(1.0) |
| .describe("Scalar factor multiplied with A*B."); |
| DMLC_DECLARE_FIELD(beta) |
| .set_default(1.0) |
| .describe("Scalar factor multiplied with C."); |
| } |
| }; |
| |
| // Parameters for general matrix-matrix multiply |
| struct LaMatrixMultParam : public dmlc::Parameter<LaMatrixMultParam> { |
| bool transpose_a, transpose_b; |
| double alpha; |
| DMLC_DECLARE_PARAMETER(LaMatrixMultParam) { |
| DMLC_DECLARE_FIELD(transpose_a) |
| .set_default(false) |
| .describe("Multiply with transposed of first input (A)."); |
| DMLC_DECLARE_FIELD(transpose_b) |
| .set_default(false) |
| .describe("Multiply with transposed of second input (B)."); |
| DMLC_DECLARE_FIELD(alpha) |
| .set_default(1.0) |
| .describe("Scalar factor multiplied with A*B."); |
| } |
| }; |
| |
| // Parameters for matrix-matrix multiplication where one is a triangular matrix. |
| struct LaTriangMatrixMultParam : public dmlc::Parameter<LaTriangMatrixMultParam> { |
| bool transpose; |
| bool rightside; |
| double alpha; |
| DMLC_DECLARE_PARAMETER(LaTriangMatrixMultParam) { |
| DMLC_DECLARE_FIELD(transpose) |
| .set_default(false) |
| .describe("Use transposed of the triangular matrix"); |
| DMLC_DECLARE_FIELD(rightside) |
| .set_default(false) |
| .describe("Multiply triangular matrix from the right to non-triangular one."); |
| DMLC_DECLARE_FIELD(alpha) |
| .set_default(1.0) |
| .describe("Scalar factor to be applied to the result."); |
| } |
| }; |
| |
| // Common function for shape inference for matrix mult and matrix mac. |
| bool LaMatrixMultMacOpShape(const nnvm::NodeAttrs& attrs, |
| std::vector<TShape>* in_attrs, |
| std::vector<TShape>* out_attrs) { |
| CHECK_GE(in_attrs->size(), 2); |
| CHECK_EQ(out_attrs->size(), 1); |
| bool transpose_a(false), transpose_b(false); |
| if ( in_attrs->size() == 2 ) { |
| // Matrix-Matrix mult |
| transpose_a = nnvm::get<LaMatrixMultParam>(attrs.parsed).transpose_a; |
| transpose_b = nnvm::get<LaMatrixMultParam>(attrs.parsed).transpose_b; |
| } else { |
| // Matrix-Matrix mac |
| transpose_a = nnvm::get<LaMatrixMacParam>(attrs.parsed).transpose_a; |
| transpose_b = nnvm::get<LaMatrixMacParam>(attrs.parsed).transpose_b; |
| } |
| if ( (*in_attrs)[0].ndim() >= 2 && (*in_attrs)[0].ndim() == (*in_attrs)[1].ndim() ) { |
| // Forward shape inference. |
| const int ndim((*in_attrs)[0].ndim()); |
| std::vector<int> oshape(ndim); |
| for ( int i = 0; i < ndim-2; ++i ) { |
| // Both inputs must have same shape except for last two dimensions. |
| if ( (*in_attrs)[0][i] != (*in_attrs)[1][i] ) return false; |
| oshape[i] = (*in_attrs)[0][i]; |
| } |
| CHECK_EQ((transpose_a ? (*in_attrs)[0][ndim-2] : (*in_attrs)[0][ndim-1]), |
| (transpose_b ? (*in_attrs)[1][ndim-1] : (*in_attrs)[1][ndim-2])) |
| << "Incompatible matrix dimensions for multiplication"; |
| oshape[ndim-2] = (transpose_a ? (*in_attrs)[0][ndim-1] : (*in_attrs)[0][ndim-2]); |
| oshape[ndim-1] = (transpose_b ? (*in_attrs)[1][ndim-2] : (*in_attrs)[1][ndim-1]); |
| TShape tshape(oshape.begin(), oshape.end()); |
| SHAPE_ASSIGN_CHECK(*out_attrs, 0, tshape); |
| if ( in_attrs->size() > 2 ) { |
| // Infer/check shape of third operand of a mac. |
| SHAPE_ASSIGN_CHECK(*in_attrs, 2, tshape); |
| } |
| return true; |
| } |
| // Can't do backward inference of shapes for this operator. |
| return false; |
| } |
| |
| bool LaTriangMatrixMultOpShape(const nnvm::NodeAttrs& attrs, |
| std::vector<TShape>* in_attrs, |
| std::vector<TShape>* out_attrs) { |
| const LaTriangMatrixMultParam& param = nnvm::get<LaTriangMatrixMultParam>(attrs.parsed); |
| CHECK_EQ(in_attrs->size(), 2); |
| CHECK_EQ(out_attrs->size(), 1); |
| if ( (*in_attrs)[0].ndim() >= 2 && (*in_attrs)[0].ndim() == (*in_attrs)[1].ndim() ) { |
| // Forward shape inference. |
| const int ndim((*in_attrs)[0].ndim()); |
| CHECK_EQ((*in_attrs)[0][ndim-2], (*in_attrs)[0][ndim-1]) |
| << "First operand must be a tensor of square matrices"; |
| std::vector<int> oshape(ndim); |
| for ( int i = 0; i < ndim-2; ++i ) { |
| // Must have same shape except for last two dimensions. |
| if ( (*in_attrs)[0][i] != (*in_attrs)[1][i] ) return false; |
| oshape[i] = (*in_attrs)[0][i]; |
| } |
| if ( param.rightside ) { |
| // We compute B * A where A is the first and B the second input. |
| CHECK_EQ((*in_attrs)[0][ndim-2], (*in_attrs)[1][ndim-1]) |
| << "Incompatible matrix dimensions for multiplication"; |
| oshape[ndim-2] = (*in_attrs)[1][ndim-2]; |
| oshape[ndim-1] = (param.transpose ? (*in_attrs)[0][ndim-2] : (*in_attrs)[0][ndim-1]); |
| } else { |
| // We compute A * B where A is the first and B the second input. |
| CHECK_EQ((*in_attrs)[1][ndim-2], (*in_attrs)[0][ndim-1]) |
| << "Incompatible matrix dimensions for multiplication"; |
| oshape[ndim-2] = (param.transpose ? (*in_attrs)[0][ndim-1] : (*in_attrs)[0][ndim-2]); |
| oshape[ndim-1] = (*in_attrs)[1][ndim-1]; |
| } |
| TShape tshape(oshape.begin(), oshape.end()); |
| SHAPE_ASSIGN_CHECK(*out_attrs, 0, tshape); |
| return true; |
| } |
| if ( (*out_attrs)[0].ndim() >= 2 ) { |
| // Backward shape inference. |
| const int odim((*out_attrs)[0].ndim()); |
| std::vector<int> ishape1(odim), ishape2(odim); |
| for ( int i = 0; i < odim-2; ++i ) { |
| ishape1[i] = ishape2[i] = (*out_attrs)[0][i]; |
| } |
| if ( param.rightside ) { |
| // We compute B * A where A is the first and B the second input. |
| ishape2[odim-2] = (*out_attrs)[0][odim-2]; |
| ishape1[odim-2] = ishape1[odim-1] = ishape2[odim-1] = (*out_attrs)[0][odim-1]; |
| } else { |
| // We compute A * B where A is the first and B the second input. |
| ishape2[odim-1] = (*out_attrs)[0][odim-1]; |
| ishape1[odim-2] = ishape1[odim-1] = ishape2[odim-2] = (*out_attrs)[0][odim-2]; |
| } |
| TShape tshape1(ishape1.begin(), ishape1.end()); |
| SHAPE_ASSIGN_CHECK(*in_attrs, 0, tshape1); |
| TShape tshape2(ishape2.begin(), ishape2.end()); |
| SHAPE_ASSIGN_CHECK(*in_attrs, 1, tshape2); |
| return true; |
| } |
| return false; |
| } |
| |
| template<int dim> |
| bool LaReduceShape(const nnvm::NodeAttrs& attrs, |
| std::vector<TShape>* in_attrs, |
| std::vector<TShape>* out_attrs) { |
| // Shape for reduction of the dim lowest dimensions to a scalar. |
| // Can only deduct in forward direction. |
| CHECK_EQ(in_attrs->size(), 1); |
| CHECK_EQ(out_attrs->size(), 1); |
| const int ndim((*in_attrs)[0].ndim()); |
| if ( ndim < dim ) { |
| return false; |
| } |
| std::vector<int> oshape(std::max(1, ndim-dim), 1); |
| for ( int i = 0; i < ndim - dim; ++i ) { |
| oshape[i] = (*in_attrs)[0][i]; |
| } |
| // Will reduce all matrices/vectors to a scalar. |
| TShape tshape(oshape.begin(), oshape.end()); |
| SHAPE_ASSIGN_CHECK(*out_attrs, 0, tshape); |
| return true; |
| } |
| |
| // Adapters for calling the various operators with appropriate signatures. |
| template<typename xpu, typename DType, int idim, int odim, int inum, int onum, typename laop> |
| struct LaOpCaller { |
| static void op(const std::vector<TBlob>& inputs, |
| const std::vector<TBlob>& outputs, |
| const int index, |
| const nnvm::NodeAttrs& attrs, |
| mshadow::Stream<xpu> *s) { |
| CHECK(false) << "no specialized LaOpCaller defined for template parameters"; |
| } |
| }; |
| template<typename xpu, typename DType, int idim, int odim, typename laop> |
| struct LaOpCaller<xpu, DType, idim, odim, 1, 1, laop> { |
| static void op(const std::vector<TBlob>& inputs, |
| const std::vector<TBlob>& outputs, |
| const int index, |
| const nnvm::NodeAttrs& attrs, |
| mshadow::Stream<xpu> *s) { |
| laop::op(inputs[0].FlatToKD<xpu, idim+1, DType>(s)[index], |
| outputs[0].FlatToKD<xpu, odim+1, DType>(s)[index], attrs); |
| } |
| }; |
| template<typename xpu, typename DType, int idim, int odim, typename laop> |
| struct LaOpCaller<xpu, DType, idim, odim, 2, 1, laop> { |
| static void op(const std::vector<TBlob>& inputs, |
| const std::vector<TBlob>& outputs, |
| const int index, |
| const nnvm::NodeAttrs& attrs, |
| mshadow::Stream<xpu> *s) { |
| laop::op(inputs[0].FlatToKD<xpu, idim+1, DType>(s)[index], |
| inputs[1].FlatToKD<xpu, idim+1, DType>(s)[index], |
| outputs[0].FlatToKD<xpu, odim+1, DType>(s)[index], |
| attrs); |
| } |
| }; |
| template<typename xpu, typename DType, int idim, int odim, typename laop> |
| struct LaOpCaller<xpu, DType, idim, odim, 3, 1, laop> { |
| static void op(const std::vector<TBlob>& inputs, |
| const std::vector<TBlob>& outputs, |
| const int index, |
| const nnvm::NodeAttrs& attrs, |
| mshadow::Stream<xpu> *s) { |
| laop::op(inputs[0].FlatToKD<xpu, idim+1, DType>(s)[index], |
| inputs[1].FlatToKD<xpu, idim+1, DType>(s)[index], |
| inputs[2].FlatToKD<xpu, idim+1, DType>(s)[index], |
| outputs[0].FlatToKD<xpu, odim+1, DType>(s)[index], |
| attrs); |
| } |
| }; |
| template<typename xpu, typename DType, int idim, int odim, typename laop> |
| struct LaOpCaller<xpu, DType, idim, odim, 3, 2, laop> { |
| static void op(const std::vector<TBlob>& inputs, |
| const std::vector<TBlob>& outputs, |
| const int index, |
| const nnvm::NodeAttrs& attrs, |
| mshadow::Stream<xpu> *s) { |
| laop::op(inputs[0].FlatToKD<xpu, idim+1, DType>(s)[index], |
| inputs[1].FlatToKD<xpu, idim+1, DType>(s)[index], |
| inputs[2].FlatToKD<xpu, idim+1, DType>(s)[index], |
| outputs[0].FlatToKD<xpu, odim+1, DType>(s)[index], |
| outputs[1].FlatToKD<xpu, odim+1, DType>(s)[index], |
| attrs); |
| } |
| }; |
| template<typename xpu, typename DType, int idim, int odim, typename laop> |
| struct LaOpCaller<xpu, DType, idim, odim, 4, 2, laop> { |
| static void op(const std::vector<TBlob>& inputs, |
| const std::vector<TBlob>& outputs, |
| const int index, |
| const nnvm::NodeAttrs& attrs, |
| mshadow::Stream<xpu> *s) { |
| laop::op(inputs[0].FlatToKD<xpu, idim+1, DType>(s)[index], |
| inputs[1].FlatToKD<xpu, idim+1, DType>(s)[index], |
| inputs[2].FlatToKD<xpu, idim+1, DType>(s)[index], |
| inputs[3].FlatToKD<xpu, idim+1, DType>(s)[index], |
| outputs[0].FlatToKD<xpu, odim+1, DType>(s)[index], |
| outputs[1].FlatToKD<xpu, odim+1, DType>(s)[index], |
| attrs); |
| } |
| }; |
| template<typename xpu, typename DType, int idim, int odim, typename laop> |
| struct LaOpCaller<xpu, DType, idim, odim, 4, 3, laop> { |
| static void op(const std::vector<TBlob>& inputs, |
| const std::vector<TBlob>& outputs, |
| const int index, |
| const nnvm::NodeAttrs& attrs, |
| mshadow::Stream<xpu> *s) { |
| laop::op(inputs[0].FlatToKD<xpu, idim+1, DType>(s)[index], |
| inputs[1].FlatToKD<xpu, idim+1, DType>(s)[index], |
| inputs[2].FlatToKD<xpu, idim+1, DType>(s)[index], |
| inputs[3].FlatToKD<xpu, idim+1, DType>(s)[index], |
| outputs[0].FlatToKD<xpu, odim+1, DType>(s)[index], |
| outputs[1].FlatToKD<xpu, odim+1, DType>(s)[index], |
| outputs[2].FlatToKD<xpu, odim+1, DType>(s)[index], |
| attrs); |
| } |
| }; |
| |
| |
| template<typename xpu, int idim, int odim, int inum, int onum, typename laop> |
| void LaOpForward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| using namespace mshadow; |
| Stream<xpu> *s = ctx.get_stream<xpu>(); |
| CHECK_EQ(inputs.size(), inum); |
| CHECK_EQ(outputs.size(), onum); |
| MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, { |
| int N(-1); |
| for ( int i = 0; i < inum; ++i ) { |
| CHECK_EQ(inputs[i].CheckContiguous(), true); |
| const int M(inputs[i].FlatToKD<xpu, idim+1, OType>(s).size(0)); |
| CHECK_EQ((N == -1 || N == M), true); |
| N = M; |
| } |
| for ( int i = 0; i < onum; ++i ) { |
| CHECK_EQ(outputs[i].CheckContiguous(), true); |
| CHECK_EQ((req[i] == kWriteTo || req[i] == kWriteInplace), true); |
| const int M(outputs[i].FlatToKD<xpu, odim+1, OType>(s).size(0)); |
| CHECK_EQ((N == -1 || N == M), true); |
| N = M; |
| } |
| for ( int i = 0; i < N; ++i ) { |
| LaOpCaller<xpu, OType, idim, odim, inum, onum, laop>::op(inputs, outputs, i, attrs, s); |
| } |
| }); |
| } |
| |
| |
| template<typename xpu, int idim, int odim, int inum, int onum, typename laop> |
| void LaOpBackward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| using namespace mshadow; |
| Stream<xpu> *s = ctx.get_stream<xpu>(); |
| CHECK_EQ(inputs.size(), inum); |
| CHECK_EQ(outputs.size(), onum); |
| MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, { |
| int N(-1); |
| for ( int i = 0; i < inum; ++i ) { |
| CHECK_EQ(inputs[i].CheckContiguous(), true); |
| const int M(inputs[i].FlatToKD<xpu, idim+1, OType>(s).size(0)); |
| CHECK_EQ((N == -1 || N == M), true); |
| N = M; |
| } |
| std::vector<TBlob> tspace(outputs); |
| for ( int i = 0; i < onum; ++i ) { |
| CHECK_EQ(outputs[i].CheckContiguous(), true); |
| const int M(outputs[i].FlatToKD<xpu, odim+1, OType>(s).size(0)); |
| CHECK_EQ((N == -1 || N == M), true); |
| N = M; |
| if ( req[i] == kAddTo ) { |
| tspace[i].dptr_ = ctx.requested[ResourceRequest::kTempSpace] |
| .get_space_typed<xpu, 1, OType>(Shape1(outputs[i].Size()), s).dptr_; |
| } |
| } |
| for ( int i = 0; i < N; ++i ) { |
| LaOpCaller<xpu, OType, idim, odim, inum, onum, laop>::op(inputs, tspace, i, attrs, s); |
| } |
| for ( int i = 0; i < onum; ++i ) { |
| if ( req[i] == kAddTo ) { |
| Tensor<xpu, 1, OType> out = outputs[i].FlatTo1D<xpu, OType>(s); |
| out += tspace[i].FlatTo1D<xpu, OType>(s); |
| } |
| } |
| }); |
| } |
| |
| template<typename xpu, int idim, typename laop> |
| void LaReduceForward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| using namespace mshadow; |
| Stream<xpu> *s = ctx.get_stream<xpu>(); |
| CHECK_EQ(inputs.size(), 1); |
| CHECK_EQ(outputs.size(), 1); |
| CHECK_EQ(inputs[0].CheckContiguous(), true); |
| CHECK_EQ(outputs[0].CheckContiguous(), true); |
| MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, { |
| Tensor<xpu, idim+1, OType> in(inputs[0].FlatToKD<xpu, idim+1, OType>(s)); |
| Tensor<xpu, 1, OType> out(outputs[0].FlatTo1D<xpu, OType>(s)); |
| const int N(outputs[0].Size()); |
| CHECK_EQ(in.size(0), N); |
| for ( int i = 0; i < N; ++i ) { |
| laop::op(in[i], out[i], attrs); |
| } |
| }); |
| } |
| |
| template<typename xpu, int idim, typename laop> |
| void LaReduceBackward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| using namespace mshadow; |
| Stream<xpu> *s = ctx.get_stream<xpu>(); |
| CHECK_EQ(inputs.size(), 2); |
| CHECK_EQ(outputs.size(), 1); |
| CHECK_EQ(inputs[0].CheckContiguous(), true); |
| CHECK_EQ(inputs[1].CheckContiguous(), true); |
| CHECK_EQ(outputs[0].CheckContiguous(), true); |
| MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, { |
| const int N(inputs[0].Size()); |
| Tensor<xpu, 1, OType> in0(inputs[0].FlatTo1D<xpu, OType>(s)); |
| Tensor<xpu, idim+1, OType> in1(inputs[1].FlatToKD<xpu, idim+1, OType>(s)); |
| Tensor<xpu, idim+1, OType> out(outputs[0].FlatToKD<xpu, idim+1, OType>(s)); |
| for ( int i = 0; i < N; ++i ) { |
| laop::op(in0[i], in1[i], out[i], attrs, (req[i] == kAddTo)); |
| } |
| }); |
| } |
| |
| } // namespace op |
| } // namespace mxnet |
| |
| #endif // MXNET_OPERATOR_TENSOR_LA_OP_H_ |
